Spaces:
Build error
chore: renormalize line endings to LF
Browse files`.gitattributes` declares `*.py text eol=lf` and `*.sh text eol=lf`, but
74 files (73 .py, 1 .sh) are stored in the index with CRLF line endings,
violating that contract. Every macOS/Linux clone reports these files as
"modified" on fresh checkout because git's diff engine sees the stored
bytes don't match the attribute contract, even though the working tree
and index match byte-for-byte.
Running `git add --renormalize .` rewrites each affected blob so the
stored form matches the attribute declaration. No semantic changes —
every affected file's diff is "N insertions, N deletions" with inserts
and deletes being the same lines modulo line endings.
Follow-up commit adds `.git-blame-ignore-revs` so `git blame` / GitHub
blame skip this mechanical commit.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- headroom/cli/wrap.py +0 -0
- headroom/compress.py +347 -347
- headroom/copilot_auth.py +444 -444
- headroom/install/health.py +28 -28
- headroom/install/providers.py +174 -174
- headroom/install/runtime.py +279 -279
- headroom/providers/aider/install.py +12 -12
- headroom/providers/claude/install.py +63 -63
- headroom/providers/codex/install.py +68 -68
- headroom/providers/copilot/install.py +25 -25
- headroom/providers/cursor/install.py +15 -15
- headroom/providers/install_registry.py +86 -86
- headroom/providers/openclaw/install.py +50 -50
- headroom/proxy/handlers/openai.py +0 -0
- headroom/proxy/server.py +0 -0
- headroom/release_version.py +310 -310
- headroom/subscription/__init__.py +72 -72
- headroom/subscription/base.py +230 -230
- headroom/subscription/client.py +131 -131
- headroom/subscription/codex_rate_limits.py +247 -247
- headroom/subscription/copilot_quota.py +366 -366
- headroom/subscription/models.py +395 -395
- headroom/subscription/session_tracking.py +189 -189
- headroom/subscription/tracker.py +464 -464
- headroom/transforms/content_router.py +0 -0
- scripts/changelog-gen.py +203 -203
- scripts/sync-plugin-versions.py +55 -55
- scripts/tests/test_changelog_gen.py +302 -302
- scripts/tests/test_sync_plugin_versions.py +68 -68
- tests/test_backend_anyllm.py +384 -384
- tests/test_ccr_batch_store.py +126 -126
- tests/test_ccr_response_handler_extra.py +372 -372
- tests/test_cli/test_wrap_copilot.py +335 -335
- tests/test_cli_learn.py +266 -266
- tests/test_codex_rate_limits.py +227 -227
- tests/test_compress_api.py +274 -274
- tests/test_compress_failure.py +39 -39
- tests/test_copilot_quota.py +331 -331
- tests/test_evals_datasets.py +538 -538
- tests/test_evals_metrics.py +132 -132
- tests/test_exceptions.py +40 -40
- tests/test_graph.py +352 -352
- tests/test_install/test_paths.py +68 -68
- tests/test_install/test_runtime.py +468 -468
- tests/test_install/test_supervisors.py +471 -471
- tests/test_plugin_manifests.py +55 -55
- tests/test_pricing.py +130 -130
- tests/test_pricing_litellm.py +97 -97
- tests/test_provider_aider.py +33 -33
- tests/test_provider_claude.py +9 -9
|
The diff for this file is too large to render.
See raw diff
|
|
|
|
@@ -1,347 +1,347 @@
|
|
| 1 |
-
"""One-function compression API for Headroom.
|
| 2 |
-
|
| 3 |
-
The simplest way to use Headroom — no proxy, no config, just compress:
|
| 4 |
-
|
| 5 |
-
from headroom import compress
|
| 6 |
-
|
| 7 |
-
result = compress(messages, model="claude-sonnet-4-5-20250929")
|
| 8 |
-
result.messages # Compressed messages (same format, fewer tokens)
|
| 9 |
-
result.tokens_saved # Tokens saved
|
| 10 |
-
result.compression_ratio # e.g., 0.35 means 65% saved
|
| 11 |
-
|
| 12 |
-
Works with any LLM client, any proxy, any framework. Just compress
|
| 13 |
-
the messages before sending them.
|
| 14 |
-
|
| 15 |
-
Examples:
|
| 16 |
-
|
| 17 |
-
# With Anthropic SDK
|
| 18 |
-
from anthropic import Anthropic
|
| 19 |
-
from headroom import compress
|
| 20 |
-
|
| 21 |
-
client = Anthropic()
|
| 22 |
-
messages = [{"role": "user", "content": huge_tool_output}]
|
| 23 |
-
compressed = compress(messages, model="claude-sonnet-4-5-20250929")
|
| 24 |
-
response = client.messages.create(
|
| 25 |
-
model="claude-sonnet-4-5-20250929",
|
| 26 |
-
messages=compressed.messages,
|
| 27 |
-
)
|
| 28 |
-
|
| 29 |
-
# With OpenAI SDK
|
| 30 |
-
from openai import OpenAI
|
| 31 |
-
from headroom import compress
|
| 32 |
-
|
| 33 |
-
client = OpenAI()
|
| 34 |
-
messages = [{"role": "user", "content": "analyze this"}, {"role": "tool", "content": big_data}]
|
| 35 |
-
compressed = compress(messages, model="gpt-4o")
|
| 36 |
-
response = client.chat.completions.create(model="gpt-4o", messages=compressed.messages)
|
| 37 |
-
|
| 38 |
-
# With LiteLLM
|
| 39 |
-
import litellm
|
| 40 |
-
from headroom import compress
|
| 41 |
-
|
| 42 |
-
messages = [...]
|
| 43 |
-
compressed = compress(messages, model="bedrock/claude-sonnet")
|
| 44 |
-
response = litellm.completion(model="bedrock/claude-sonnet", messages=compressed.messages)
|
| 45 |
-
|
| 46 |
-
# With any HTTP client
|
| 47 |
-
import httpx
|
| 48 |
-
from headroom import compress
|
| 49 |
-
|
| 50 |
-
compressed = compress(messages, model="claude-sonnet-4-5-20250929")
|
| 51 |
-
httpx.post("https://api.anthropic.com/v1/messages", json={
|
| 52 |
-
"model": "claude-sonnet-4-5-20250929",
|
| 53 |
-
"messages": compressed.messages,
|
| 54 |
-
})
|
| 55 |
-
"""
|
| 56 |
-
|
| 57 |
-
from __future__ import annotations
|
| 58 |
-
|
| 59 |
-
import logging
|
| 60 |
-
import threading
|
| 61 |
-
from dataclasses import dataclass, field
|
| 62 |
-
from typing import Any
|
| 63 |
-
|
| 64 |
-
from .observability import get_otel_metrics
|
| 65 |
-
from .pipeline import PipelineExtensionManager, PipelineStage, summarize_routing_markers
|
| 66 |
-
from .utils import extract_user_query as _extract_user_query
|
| 67 |
-
|
| 68 |
-
logger = logging.getLogger(__name__)
|
| 69 |
-
|
| 70 |
-
|
| 71 |
-
# Lazy-initialized singleton pipeline
|
| 72 |
-
_pipeline = None
|
| 73 |
-
_pipeline_lock = threading.Lock()
|
| 74 |
-
|
| 75 |
-
|
| 76 |
-
@dataclass
|
| 77 |
-
class CompressConfig:
|
| 78 |
-
"""User-facing compression options.
|
| 79 |
-
|
| 80 |
-
Controls what gets compressed, how aggressively, and with which model.
|
| 81 |
-
Pass to ``compress()`` or any integration that uses headroom.
|
| 82 |
-
|
| 83 |
-
Examples::
|
| 84 |
-
|
| 85 |
-
# Coding agent (default — skip user messages, protect recent)
|
| 86 |
-
compress(messages, model="gpt-4o")
|
| 87 |
-
|
| 88 |
-
# Financial document (compress everything, keep 50%)
|
| 89 |
-
compress(messages, model="claude-opus-4-20250514",
|
| 90 |
-
compress_user_messages=True,
|
| 91 |
-
target_ratio=0.5,
|
| 92 |
-
protect_recent=0,
|
| 93 |
-
)
|
| 94 |
-
|
| 95 |
-
# Aggressive (logs, search results)
|
| 96 |
-
compress(messages, model="gpt-4o", target_ratio=0.2)
|
| 97 |
-
"""
|
| 98 |
-
|
| 99 |
-
# What to compress
|
| 100 |
-
compress_user_messages: bool = False
|
| 101 |
-
"""Compress user messages too (default: skip them for coding agents).
|
| 102 |
-
Set True for document compression, RAG pipelines, or when user messages
|
| 103 |
-
contain large tool outputs."""
|
| 104 |
-
|
| 105 |
-
compress_system_messages: bool = True
|
| 106 |
-
"""Compress system messages (default: True).
|
| 107 |
-
Set False to preserve system prompts exactly as-is. Useful for voice
|
| 108 |
-
agents where tool definitions and instructions must not be altered."""
|
| 109 |
-
|
| 110 |
-
protect_recent: int = 4
|
| 111 |
-
"""Don't compress the last N messages (they're the active conversation).
|
| 112 |
-
Set 0 to compress everything."""
|
| 113 |
-
|
| 114 |
-
protect_analysis_context: bool = True
|
| 115 |
-
"""Detect 'analyze'/'review' intent and protect code from compression."""
|
| 116 |
-
|
| 117 |
-
# How aggressive
|
| 118 |
-
target_ratio: float | None = None
|
| 119 |
-
"""Keep ratio for Kompress. None = model decides (~15% kept, aggressive).
|
| 120 |
-
0.5 = keep 50% (safe for documents). 0.7 = keep 70% (conservative).
|
| 121 |
-
Only affects Kompress (text compression). SmartCrusher (JSON) has its
|
| 122 |
-
own logic based on array dedup."""
|
| 123 |
-
|
| 124 |
-
min_tokens_to_compress: int = 250
|
| 125 |
-
"""Minimum token count for a message to be compressed.
|
| 126 |
-
Messages shorter than this are left unchanged. Default 250.
|
| 127 |
-
Set lower for voice agents where turns are short."""
|
| 128 |
-
|
| 129 |
-
# Model variant
|
| 130 |
-
kompress_model: str | None = None
|
| 131 |
-
"""Kompress model ID. None = default (chopratejas/kompress-base).
|
| 132 |
-
Set to a HuggingFace model ID for domain-specific compression.
|
| 133 |
-
Set to 'disabled' to skip ML compression entirely
|
| 134 |
-
(only SmartCrusher + CacheAligner will run)."""
|
| 135 |
-
|
| 136 |
-
|
| 137 |
-
@dataclass
|
| 138 |
-
class CompressResult:
|
| 139 |
-
"""Result of compressing messages.
|
| 140 |
-
|
| 141 |
-
Attributes:
|
| 142 |
-
messages: The compressed messages (same format as input).
|
| 143 |
-
tokens_before: Token count before compression.
|
| 144 |
-
tokens_after: Token count after compression.
|
| 145 |
-
tokens_saved: Tokens removed by compression.
|
| 146 |
-
compression_ratio: Ratio of tokens saved (0.0 = no savings, 1.0 = 100% removed).
|
| 147 |
-
transforms_applied: List of transforms that were applied.
|
| 148 |
-
"""
|
| 149 |
-
|
| 150 |
-
messages: list[dict[str, Any]]
|
| 151 |
-
tokens_before: int = 0
|
| 152 |
-
tokens_after: int = 0
|
| 153 |
-
tokens_saved: int = 0
|
| 154 |
-
compression_ratio: float = 0.0
|
| 155 |
-
transforms_applied: list[str] = field(default_factory=list)
|
| 156 |
-
|
| 157 |
-
|
| 158 |
-
def compress(
|
| 159 |
-
messages: list[dict[str, Any]],
|
| 160 |
-
model: str = "claude-sonnet-4-5-20250929",
|
| 161 |
-
model_limit: int = 200000,
|
| 162 |
-
optimize: bool = True,
|
| 163 |
-
hooks: Any = None,
|
| 164 |
-
config: CompressConfig | None = None,
|
| 165 |
-
**kwargs: Any,
|
| 166 |
-
) -> CompressResult:
|
| 167 |
-
"""Compress messages using Headroom's full compression pipeline.
|
| 168 |
-
|
| 169 |
-
This is the simplest way to use Headroom. No proxy, no config needed.
|
| 170 |
-
Just pass messages and get compressed messages back.
|
| 171 |
-
|
| 172 |
-
Args:
|
| 173 |
-
messages: List of messages in Anthropic or OpenAI format.
|
| 174 |
-
model: Model name (used for token counting and context limit).
|
| 175 |
-
model_limit: Model's context window size in tokens.
|
| 176 |
-
optimize: Whether to actually compress (False = passthrough for A/B testing).
|
| 177 |
-
hooks: Optional CompressionHooks instance for custom behavior.
|
| 178 |
-
config: Compression options (CompressConfig). Overrides defaults.
|
| 179 |
-
**kwargs: Shorthand for CompressConfig fields. These override config:
|
| 180 |
-
compress_user_messages, target_ratio, protect_recent,
|
| 181 |
-
protect_analysis_context, kompress_model.
|
| 182 |
-
|
| 183 |
-
Returns:
|
| 184 |
-
CompressResult with compressed messages and metrics.
|
| 185 |
-
|
| 186 |
-
Examples::
|
| 187 |
-
|
| 188 |
-
# Default (coding agent)
|
| 189 |
-
result = compress(messages, model="gpt-4o")
|
| 190 |
-
|
| 191 |
-
# Financial document (keep 50%, compress everything)
|
| 192 |
-
result = compress(messages, model="claude-opus-4-20250514",
|
| 193 |
-
compress_user_messages=True,
|
| 194 |
-
target_ratio=0.5,
|
| 195 |
-
protect_recent=0,
|
| 196 |
-
)
|
| 197 |
-
"""
|
| 198 |
-
if not messages or not optimize:
|
| 199 |
-
return CompressResult(messages=messages)
|
| 200 |
-
|
| 201 |
-
# Build config from explicit config + kwargs
|
| 202 |
-
cfg = config or CompressConfig()
|
| 203 |
-
config_fields = {f.name for f in cfg.__dataclass_fields__.values()}
|
| 204 |
-
for key, value in kwargs.items():
|
| 205 |
-
if key in config_fields:
|
| 206 |
-
setattr(cfg, key, value)
|
| 207 |
-
|
| 208 |
-
pipeline = _get_pipeline()
|
| 209 |
-
pipeline_extensions = PipelineExtensionManager(hooks=hooks, discover=False)
|
| 210 |
-
|
| 211 |
-
try:
|
| 212 |
-
# Compute biases from hooks if provided
|
| 213 |
-
biases = None
|
| 214 |
-
if hooks:
|
| 215 |
-
from headroom.hooks import CompressContext
|
| 216 |
-
|
| 217 |
-
ctx = CompressContext(model=model)
|
| 218 |
-
messages = hooks.pre_compress(messages, ctx)
|
| 219 |
-
biases = hooks.compute_biases(messages, ctx)
|
| 220 |
-
|
| 221 |
-
received_event = pipeline_extensions.emit(
|
| 222 |
-
PipelineStage.INPUT_RECEIVED,
|
| 223 |
-
operation="compress",
|
| 224 |
-
model=model,
|
| 225 |
-
messages=messages,
|
| 226 |
-
)
|
| 227 |
-
if received_event.messages is not None:
|
| 228 |
-
messages = received_event.messages
|
| 229 |
-
|
| 230 |
-
# Extract user query from messages so transforms can score by
|
| 231 |
-
# relevance. Without this, SmartCrusher selects items by statistics
|
| 232 |
-
# alone (position, anomaly) and may drop relevant content.
|
| 233 |
-
context = _extract_user_query(messages)
|
| 234 |
-
|
| 235 |
-
result = pipeline.apply(
|
| 236 |
-
messages=messages,
|
| 237 |
-
model=model,
|
| 238 |
-
model_limit=model_limit,
|
| 239 |
-
context=context,
|
| 240 |
-
biases=biases,
|
| 241 |
-
# Pass CompressConfig options through to transforms
|
| 242 |
-
compress_user_messages=cfg.compress_user_messages,
|
| 243 |
-
compress_system_messages=cfg.compress_system_messages,
|
| 244 |
-
target_ratio=cfg.target_ratio,
|
| 245 |
-
protect_recent=cfg.protect_recent,
|
| 246 |
-
protect_analysis_context=cfg.protect_analysis_context,
|
| 247 |
-
min_tokens_to_compress=cfg.min_tokens_to_compress,
|
| 248 |
-
kompress_model=cfg.kompress_model,
|
| 249 |
-
)
|
| 250 |
-
|
| 251 |
-
tokens_before = result.tokens_before
|
| 252 |
-
tokens_after = result.tokens_after
|
| 253 |
-
compressed_messages = result.messages
|
| 254 |
-
|
| 255 |
-
routing_markers = summarize_routing_markers(result.transforms_applied)
|
| 256 |
-
if routing_markers:
|
| 257 |
-
routed_event = pipeline_extensions.emit(
|
| 258 |
-
PipelineStage.INPUT_ROUTED,
|
| 259 |
-
operation="compress",
|
| 260 |
-
model=model,
|
| 261 |
-
messages=compressed_messages,
|
| 262 |
-
metadata={
|
| 263 |
-
"routing_markers": routing_markers,
|
| 264 |
-
"transforms_applied": result.transforms_applied,
|
| 265 |
-
},
|
| 266 |
-
)
|
| 267 |
-
if routed_event.messages is not None:
|
| 268 |
-
compressed_messages = routed_event.messages
|
| 269 |
-
|
| 270 |
-
compressed_event = pipeline_extensions.emit(
|
| 271 |
-
PipelineStage.INPUT_COMPRESSED,
|
| 272 |
-
operation="compress",
|
| 273 |
-
model=model,
|
| 274 |
-
messages=compressed_messages,
|
| 275 |
-
metadata={
|
| 276 |
-
"tokens_before": tokens_before,
|
| 277 |
-
"tokens_after": tokens_after,
|
| 278 |
-
"transforms_applied": result.transforms_applied,
|
| 279 |
-
},
|
| 280 |
-
)
|
| 281 |
-
if compressed_event.messages is not None:
|
| 282 |
-
compressed_messages = compressed_event.messages
|
| 283 |
-
|
| 284 |
-
tokens_saved = tokens_before - tokens_after
|
| 285 |
-
ratio = tokens_saved / tokens_before if tokens_before > 0 else 0.0
|
| 286 |
-
|
| 287 |
-
# Post-compress hook
|
| 288 |
-
if hooks and tokens_saved > 0:
|
| 289 |
-
from headroom.hooks import CompressEvent
|
| 290 |
-
|
| 291 |
-
hooks.post_compress(
|
| 292 |
-
CompressEvent(
|
| 293 |
-
tokens_before=tokens_before,
|
| 294 |
-
tokens_after=tokens_after,
|
| 295 |
-
tokens_saved=tokens_saved,
|
| 296 |
-
compression_ratio=ratio,
|
| 297 |
-
transforms_applied=result.transforms_applied,
|
| 298 |
-
model=model,
|
| 299 |
-
)
|
| 300 |
-
)
|
| 301 |
-
|
| 302 |
-
return CompressResult(
|
| 303 |
-
messages=compressed_messages,
|
| 304 |
-
tokens_before=tokens_before,
|
| 305 |
-
tokens_after=tokens_after,
|
| 306 |
-
tokens_saved=tokens_saved,
|
| 307 |
-
compression_ratio=ratio,
|
| 308 |
-
transforms_applied=result.transforms_applied,
|
| 309 |
-
)
|
| 310 |
-
|
| 311 |
-
except Exception as e:
|
| 312 |
-
get_otel_metrics().record_compression_failure(
|
| 313 |
-
model=model,
|
| 314 |
-
operation="compress",
|
| 315 |
-
error_type=type(e).__name__,
|
| 316 |
-
)
|
| 317 |
-
logger.warning("Compression failed, returning original messages: %s", e)
|
| 318 |
-
return CompressResult(
|
| 319 |
-
messages=messages,
|
| 320 |
-
tokens_before=0,
|
| 321 |
-
tokens_after=0,
|
| 322 |
-
tokens_saved=0,
|
| 323 |
-
compression_ratio=0.0,
|
| 324 |
-
)
|
| 325 |
-
|
| 326 |
-
|
| 327 |
-
def _get_pipeline() -> Any:
|
| 328 |
-
"""Get or create the singleton compression pipeline."""
|
| 329 |
-
global _pipeline
|
| 330 |
-
|
| 331 |
-
if _pipeline is not None:
|
| 332 |
-
return _pipeline
|
| 333 |
-
|
| 334 |
-
with _pipeline_lock:
|
| 335 |
-
if _pipeline is not None:
|
| 336 |
-
return _pipeline
|
| 337 |
-
|
| 338 |
-
from headroom.transforms import TransformPipeline
|
| 339 |
-
|
| 340 |
-
# Default pipeline: CacheAligner → ContentRouter → IntelligentContext
|
| 341 |
-
# CacheAligner: stabilizes prefix for provider KV cache hits
|
| 342 |
-
# ContentRouter: routes to the right compressor per content type
|
| 343 |
-
# (SmartCrusher for JSON, CodeCompressor for code, Kompress for text)
|
| 344 |
-
# IntelligentContext: enforces token limits with score-based dropping
|
| 345 |
-
_pipeline = TransformPipeline()
|
| 346 |
-
logger.debug("Headroom compression pipeline initialized")
|
| 347 |
-
return _pipeline
|
|
|
|
| 1 |
+
"""One-function compression API for Headroom.
|
| 2 |
+
|
| 3 |
+
The simplest way to use Headroom — no proxy, no config, just compress:
|
| 4 |
+
|
| 5 |
+
from headroom import compress
|
| 6 |
+
|
| 7 |
+
result = compress(messages, model="claude-sonnet-4-5-20250929")
|
| 8 |
+
result.messages # Compressed messages (same format, fewer tokens)
|
| 9 |
+
result.tokens_saved # Tokens saved
|
| 10 |
+
result.compression_ratio # e.g., 0.35 means 65% saved
|
| 11 |
+
|
| 12 |
+
Works with any LLM client, any proxy, any framework. Just compress
|
| 13 |
+
the messages before sending them.
|
| 14 |
+
|
| 15 |
+
Examples:
|
| 16 |
+
|
| 17 |
+
# With Anthropic SDK
|
| 18 |
+
from anthropic import Anthropic
|
| 19 |
+
from headroom import compress
|
| 20 |
+
|
| 21 |
+
client = Anthropic()
|
| 22 |
+
messages = [{"role": "user", "content": huge_tool_output}]
|
| 23 |
+
compressed = compress(messages, model="claude-sonnet-4-5-20250929")
|
| 24 |
+
response = client.messages.create(
|
| 25 |
+
model="claude-sonnet-4-5-20250929",
|
| 26 |
+
messages=compressed.messages,
|
| 27 |
+
)
|
| 28 |
+
|
| 29 |
+
# With OpenAI SDK
|
| 30 |
+
from openai import OpenAI
|
| 31 |
+
from headroom import compress
|
| 32 |
+
|
| 33 |
+
client = OpenAI()
|
| 34 |
+
messages = [{"role": "user", "content": "analyze this"}, {"role": "tool", "content": big_data}]
|
| 35 |
+
compressed = compress(messages, model="gpt-4o")
|
| 36 |
+
response = client.chat.completions.create(model="gpt-4o", messages=compressed.messages)
|
| 37 |
+
|
| 38 |
+
# With LiteLLM
|
| 39 |
+
import litellm
|
| 40 |
+
from headroom import compress
|
| 41 |
+
|
| 42 |
+
messages = [...]
|
| 43 |
+
compressed = compress(messages, model="bedrock/claude-sonnet")
|
| 44 |
+
response = litellm.completion(model="bedrock/claude-sonnet", messages=compressed.messages)
|
| 45 |
+
|
| 46 |
+
# With any HTTP client
|
| 47 |
+
import httpx
|
| 48 |
+
from headroom import compress
|
| 49 |
+
|
| 50 |
+
compressed = compress(messages, model="claude-sonnet-4-5-20250929")
|
| 51 |
+
httpx.post("https://api.anthropic.com/v1/messages", json={
|
| 52 |
+
"model": "claude-sonnet-4-5-20250929",
|
| 53 |
+
"messages": compressed.messages,
|
| 54 |
+
})
|
| 55 |
+
"""
|
| 56 |
+
|
| 57 |
+
from __future__ import annotations
|
| 58 |
+
|
| 59 |
+
import logging
|
| 60 |
+
import threading
|
| 61 |
+
from dataclasses import dataclass, field
|
| 62 |
+
from typing import Any
|
| 63 |
+
|
| 64 |
+
from .observability import get_otel_metrics
|
| 65 |
+
from .pipeline import PipelineExtensionManager, PipelineStage, summarize_routing_markers
|
| 66 |
+
from .utils import extract_user_query as _extract_user_query
|
| 67 |
+
|
| 68 |
+
logger = logging.getLogger(__name__)
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
# Lazy-initialized singleton pipeline
|
| 72 |
+
_pipeline = None
|
| 73 |
+
_pipeline_lock = threading.Lock()
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
@dataclass
|
| 77 |
+
class CompressConfig:
|
| 78 |
+
"""User-facing compression options.
|
| 79 |
+
|
| 80 |
+
Controls what gets compressed, how aggressively, and with which model.
|
| 81 |
+
Pass to ``compress()`` or any integration that uses headroom.
|
| 82 |
+
|
| 83 |
+
Examples::
|
| 84 |
+
|
| 85 |
+
# Coding agent (default — skip user messages, protect recent)
|
| 86 |
+
compress(messages, model="gpt-4o")
|
| 87 |
+
|
| 88 |
+
# Financial document (compress everything, keep 50%)
|
| 89 |
+
compress(messages, model="claude-opus-4-20250514",
|
| 90 |
+
compress_user_messages=True,
|
| 91 |
+
target_ratio=0.5,
|
| 92 |
+
protect_recent=0,
|
| 93 |
+
)
|
| 94 |
+
|
| 95 |
+
# Aggressive (logs, search results)
|
| 96 |
+
compress(messages, model="gpt-4o", target_ratio=0.2)
|
| 97 |
+
"""
|
| 98 |
+
|
| 99 |
+
# What to compress
|
| 100 |
+
compress_user_messages: bool = False
|
| 101 |
+
"""Compress user messages too (default: skip them for coding agents).
|
| 102 |
+
Set True for document compression, RAG pipelines, or when user messages
|
| 103 |
+
contain large tool outputs."""
|
| 104 |
+
|
| 105 |
+
compress_system_messages: bool = True
|
| 106 |
+
"""Compress system messages (default: True).
|
| 107 |
+
Set False to preserve system prompts exactly as-is. Useful for voice
|
| 108 |
+
agents where tool definitions and instructions must not be altered."""
|
| 109 |
+
|
| 110 |
+
protect_recent: int = 4
|
| 111 |
+
"""Don't compress the last N messages (they're the active conversation).
|
| 112 |
+
Set 0 to compress everything."""
|
| 113 |
+
|
| 114 |
+
protect_analysis_context: bool = True
|
| 115 |
+
"""Detect 'analyze'/'review' intent and protect code from compression."""
|
| 116 |
+
|
| 117 |
+
# How aggressive
|
| 118 |
+
target_ratio: float | None = None
|
| 119 |
+
"""Keep ratio for Kompress. None = model decides (~15% kept, aggressive).
|
| 120 |
+
0.5 = keep 50% (safe for documents). 0.7 = keep 70% (conservative).
|
| 121 |
+
Only affects Kompress (text compression). SmartCrusher (JSON) has its
|
| 122 |
+
own logic based on array dedup."""
|
| 123 |
+
|
| 124 |
+
min_tokens_to_compress: int = 250
|
| 125 |
+
"""Minimum token count for a message to be compressed.
|
| 126 |
+
Messages shorter than this are left unchanged. Default 250.
|
| 127 |
+
Set lower for voice agents where turns are short."""
|
| 128 |
+
|
| 129 |
+
# Model variant
|
| 130 |
+
kompress_model: str | None = None
|
| 131 |
+
"""Kompress model ID. None = default (chopratejas/kompress-base).
|
| 132 |
+
Set to a HuggingFace model ID for domain-specific compression.
|
| 133 |
+
Set to 'disabled' to skip ML compression entirely
|
| 134 |
+
(only SmartCrusher + CacheAligner will run)."""
|
| 135 |
+
|
| 136 |
+
|
| 137 |
+
@dataclass
|
| 138 |
+
class CompressResult:
|
| 139 |
+
"""Result of compressing messages.
|
| 140 |
+
|
| 141 |
+
Attributes:
|
| 142 |
+
messages: The compressed messages (same format as input).
|
| 143 |
+
tokens_before: Token count before compression.
|
| 144 |
+
tokens_after: Token count after compression.
|
| 145 |
+
tokens_saved: Tokens removed by compression.
|
| 146 |
+
compression_ratio: Ratio of tokens saved (0.0 = no savings, 1.0 = 100% removed).
|
| 147 |
+
transforms_applied: List of transforms that were applied.
|
| 148 |
+
"""
|
| 149 |
+
|
| 150 |
+
messages: list[dict[str, Any]]
|
| 151 |
+
tokens_before: int = 0
|
| 152 |
+
tokens_after: int = 0
|
| 153 |
+
tokens_saved: int = 0
|
| 154 |
+
compression_ratio: float = 0.0
|
| 155 |
+
transforms_applied: list[str] = field(default_factory=list)
|
| 156 |
+
|
| 157 |
+
|
| 158 |
+
def compress(
|
| 159 |
+
messages: list[dict[str, Any]],
|
| 160 |
+
model: str = "claude-sonnet-4-5-20250929",
|
| 161 |
+
model_limit: int = 200000,
|
| 162 |
+
optimize: bool = True,
|
| 163 |
+
hooks: Any = None,
|
| 164 |
+
config: CompressConfig | None = None,
|
| 165 |
+
**kwargs: Any,
|
| 166 |
+
) -> CompressResult:
|
| 167 |
+
"""Compress messages using Headroom's full compression pipeline.
|
| 168 |
+
|
| 169 |
+
This is the simplest way to use Headroom. No proxy, no config needed.
|
| 170 |
+
Just pass messages and get compressed messages back.
|
| 171 |
+
|
| 172 |
+
Args:
|
| 173 |
+
messages: List of messages in Anthropic or OpenAI format.
|
| 174 |
+
model: Model name (used for token counting and context limit).
|
| 175 |
+
model_limit: Model's context window size in tokens.
|
| 176 |
+
optimize: Whether to actually compress (False = passthrough for A/B testing).
|
| 177 |
+
hooks: Optional CompressionHooks instance for custom behavior.
|
| 178 |
+
config: Compression options (CompressConfig). Overrides defaults.
|
| 179 |
+
**kwargs: Shorthand for CompressConfig fields. These override config:
|
| 180 |
+
compress_user_messages, target_ratio, protect_recent,
|
| 181 |
+
protect_analysis_context, kompress_model.
|
| 182 |
+
|
| 183 |
+
Returns:
|
| 184 |
+
CompressResult with compressed messages and metrics.
|
| 185 |
+
|
| 186 |
+
Examples::
|
| 187 |
+
|
| 188 |
+
# Default (coding agent)
|
| 189 |
+
result = compress(messages, model="gpt-4o")
|
| 190 |
+
|
| 191 |
+
# Financial document (keep 50%, compress everything)
|
| 192 |
+
result = compress(messages, model="claude-opus-4-20250514",
|
| 193 |
+
compress_user_messages=True,
|
| 194 |
+
target_ratio=0.5,
|
| 195 |
+
protect_recent=0,
|
| 196 |
+
)
|
| 197 |
+
"""
|
| 198 |
+
if not messages or not optimize:
|
| 199 |
+
return CompressResult(messages=messages)
|
| 200 |
+
|
| 201 |
+
# Build config from explicit config + kwargs
|
| 202 |
+
cfg = config or CompressConfig()
|
| 203 |
+
config_fields = {f.name for f in cfg.__dataclass_fields__.values()}
|
| 204 |
+
for key, value in kwargs.items():
|
| 205 |
+
if key in config_fields:
|
| 206 |
+
setattr(cfg, key, value)
|
| 207 |
+
|
| 208 |
+
pipeline = _get_pipeline()
|
| 209 |
+
pipeline_extensions = PipelineExtensionManager(hooks=hooks, discover=False)
|
| 210 |
+
|
| 211 |
+
try:
|
| 212 |
+
# Compute biases from hooks if provided
|
| 213 |
+
biases = None
|
| 214 |
+
if hooks:
|
| 215 |
+
from headroom.hooks import CompressContext
|
| 216 |
+
|
| 217 |
+
ctx = CompressContext(model=model)
|
| 218 |
+
messages = hooks.pre_compress(messages, ctx)
|
| 219 |
+
biases = hooks.compute_biases(messages, ctx)
|
| 220 |
+
|
| 221 |
+
received_event = pipeline_extensions.emit(
|
| 222 |
+
PipelineStage.INPUT_RECEIVED,
|
| 223 |
+
operation="compress",
|
| 224 |
+
model=model,
|
| 225 |
+
messages=messages,
|
| 226 |
+
)
|
| 227 |
+
if received_event.messages is not None:
|
| 228 |
+
messages = received_event.messages
|
| 229 |
+
|
| 230 |
+
# Extract user query from messages so transforms can score by
|
| 231 |
+
# relevance. Without this, SmartCrusher selects items by statistics
|
| 232 |
+
# alone (position, anomaly) and may drop relevant content.
|
| 233 |
+
context = _extract_user_query(messages)
|
| 234 |
+
|
| 235 |
+
result = pipeline.apply(
|
| 236 |
+
messages=messages,
|
| 237 |
+
model=model,
|
| 238 |
+
model_limit=model_limit,
|
| 239 |
+
context=context,
|
| 240 |
+
biases=biases,
|
| 241 |
+
# Pass CompressConfig options through to transforms
|
| 242 |
+
compress_user_messages=cfg.compress_user_messages,
|
| 243 |
+
compress_system_messages=cfg.compress_system_messages,
|
| 244 |
+
target_ratio=cfg.target_ratio,
|
| 245 |
+
protect_recent=cfg.protect_recent,
|
| 246 |
+
protect_analysis_context=cfg.protect_analysis_context,
|
| 247 |
+
min_tokens_to_compress=cfg.min_tokens_to_compress,
|
| 248 |
+
kompress_model=cfg.kompress_model,
|
| 249 |
+
)
|
| 250 |
+
|
| 251 |
+
tokens_before = result.tokens_before
|
| 252 |
+
tokens_after = result.tokens_after
|
| 253 |
+
compressed_messages = result.messages
|
| 254 |
+
|
| 255 |
+
routing_markers = summarize_routing_markers(result.transforms_applied)
|
| 256 |
+
if routing_markers:
|
| 257 |
+
routed_event = pipeline_extensions.emit(
|
| 258 |
+
PipelineStage.INPUT_ROUTED,
|
| 259 |
+
operation="compress",
|
| 260 |
+
model=model,
|
| 261 |
+
messages=compressed_messages,
|
| 262 |
+
metadata={
|
| 263 |
+
"routing_markers": routing_markers,
|
| 264 |
+
"transforms_applied": result.transforms_applied,
|
| 265 |
+
},
|
| 266 |
+
)
|
| 267 |
+
if routed_event.messages is not None:
|
| 268 |
+
compressed_messages = routed_event.messages
|
| 269 |
+
|
| 270 |
+
compressed_event = pipeline_extensions.emit(
|
| 271 |
+
PipelineStage.INPUT_COMPRESSED,
|
| 272 |
+
operation="compress",
|
| 273 |
+
model=model,
|
| 274 |
+
messages=compressed_messages,
|
| 275 |
+
metadata={
|
| 276 |
+
"tokens_before": tokens_before,
|
| 277 |
+
"tokens_after": tokens_after,
|
| 278 |
+
"transforms_applied": result.transforms_applied,
|
| 279 |
+
},
|
| 280 |
+
)
|
| 281 |
+
if compressed_event.messages is not None:
|
| 282 |
+
compressed_messages = compressed_event.messages
|
| 283 |
+
|
| 284 |
+
tokens_saved = tokens_before - tokens_after
|
| 285 |
+
ratio = tokens_saved / tokens_before if tokens_before > 0 else 0.0
|
| 286 |
+
|
| 287 |
+
# Post-compress hook
|
| 288 |
+
if hooks and tokens_saved > 0:
|
| 289 |
+
from headroom.hooks import CompressEvent
|
| 290 |
+
|
| 291 |
+
hooks.post_compress(
|
| 292 |
+
CompressEvent(
|
| 293 |
+
tokens_before=tokens_before,
|
| 294 |
+
tokens_after=tokens_after,
|
| 295 |
+
tokens_saved=tokens_saved,
|
| 296 |
+
compression_ratio=ratio,
|
| 297 |
+
transforms_applied=result.transforms_applied,
|
| 298 |
+
model=model,
|
| 299 |
+
)
|
| 300 |
+
)
|
| 301 |
+
|
| 302 |
+
return CompressResult(
|
| 303 |
+
messages=compressed_messages,
|
| 304 |
+
tokens_before=tokens_before,
|
| 305 |
+
tokens_after=tokens_after,
|
| 306 |
+
tokens_saved=tokens_saved,
|
| 307 |
+
compression_ratio=ratio,
|
| 308 |
+
transforms_applied=result.transforms_applied,
|
| 309 |
+
)
|
| 310 |
+
|
| 311 |
+
except Exception as e:
|
| 312 |
+
get_otel_metrics().record_compression_failure(
|
| 313 |
+
model=model,
|
| 314 |
+
operation="compress",
|
| 315 |
+
error_type=type(e).__name__,
|
| 316 |
+
)
|
| 317 |
+
logger.warning("Compression failed, returning original messages: %s", e)
|
| 318 |
+
return CompressResult(
|
| 319 |
+
messages=messages,
|
| 320 |
+
tokens_before=0,
|
| 321 |
+
tokens_after=0,
|
| 322 |
+
tokens_saved=0,
|
| 323 |
+
compression_ratio=0.0,
|
| 324 |
+
)
|
| 325 |
+
|
| 326 |
+
|
| 327 |
+
def _get_pipeline() -> Any:
|
| 328 |
+
"""Get or create the singleton compression pipeline."""
|
| 329 |
+
global _pipeline
|
| 330 |
+
|
| 331 |
+
if _pipeline is not None:
|
| 332 |
+
return _pipeline
|
| 333 |
+
|
| 334 |
+
with _pipeline_lock:
|
| 335 |
+
if _pipeline is not None:
|
| 336 |
+
return _pipeline
|
| 337 |
+
|
| 338 |
+
from headroom.transforms import TransformPipeline
|
| 339 |
+
|
| 340 |
+
# Default pipeline: CacheAligner → ContentRouter → IntelligentContext
|
| 341 |
+
# CacheAligner: stabilizes prefix for provider KV cache hits
|
| 342 |
+
# ContentRouter: routes to the right compressor per content type
|
| 343 |
+
# (SmartCrusher for JSON, CodeCompressor for code, Kompress for text)
|
| 344 |
+
# IntelligentContext: enforces token limits with score-based dropping
|
| 345 |
+
_pipeline = TransformPipeline()
|
| 346 |
+
logger.debug("Headroom compression pipeline initialized")
|
| 347 |
+
return _pipeline
|
|
@@ -1,444 +1,444 @@
|
|
| 1 |
-
"""GitHub Copilot OAuth discovery and API-token exchange helpers."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import asyncio
|
| 6 |
-
import ctypes
|
| 7 |
-
import json
|
| 8 |
-
import logging
|
| 9 |
-
import os
|
| 10 |
-
import subprocess
|
| 11 |
-
import time
|
| 12 |
-
from ctypes import wintypes
|
| 13 |
-
from dataclasses import dataclass
|
| 14 |
-
from datetime import datetime
|
| 15 |
-
from pathlib import Path
|
| 16 |
-
from typing import Any
|
| 17 |
-
from urllib import error as urllib_error
|
| 18 |
-
from urllib import request as urllib_request
|
| 19 |
-
from urllib.parse import urlparse
|
| 20 |
-
|
| 21 |
-
logger = logging.getLogger(__name__)
|
| 22 |
-
|
| 23 |
-
DEFAULT_API_URL = "https://api.githubcopilot.com"
|
| 24 |
-
DEFAULT_TOKEN_EXCHANGE_URL = "https://api.github.com/copilot_internal/v2/token"
|
| 25 |
-
DEFAULT_GITHUB_HOST = "github.com"
|
| 26 |
-
_TOKEN_EXPIRY_BUFFER_S = 60
|
| 27 |
-
_DEFAULT_EDITOR_VERSION = "vscode/1.104.1"
|
| 28 |
-
_DEFAULT_USER_AGENT = "GitHubCopilotChat/0.1"
|
| 29 |
-
|
| 30 |
-
_API_TOKEN_ENV_VARS = (
|
| 31 |
-
"GITHUB_COPILOT_API_TOKEN",
|
| 32 |
-
"COPILOT_PROVIDER_BEARER_TOKEN",
|
| 33 |
-
)
|
| 34 |
-
_OAUTH_TOKEN_ENV_VARS = (
|
| 35 |
-
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 36 |
-
"GITHUB_COPILOT_TOKEN",
|
| 37 |
-
"GITHUB_TOKEN",
|
| 38 |
-
"COPILOT_GITHUB_TOKEN",
|
| 39 |
-
)
|
| 40 |
-
_OAUTH_TOKEN_KEYS = (
|
| 41 |
-
"oauth_token",
|
| 42 |
-
"oauthToken",
|
| 43 |
-
"token",
|
| 44 |
-
"access_token",
|
| 45 |
-
"accessToken",
|
| 46 |
-
)
|
| 47 |
-
_EXPIRY_KEYS = ("expires_at", "expiresAt", "expiry", "expires")
|
| 48 |
-
|
| 49 |
-
|
| 50 |
-
@dataclass(frozen=True)
|
| 51 |
-
class CopilotAPIToken:
|
| 52 |
-
"""Short-lived API token exchanged from a GitHub OAuth token."""
|
| 53 |
-
|
| 54 |
-
token: str
|
| 55 |
-
expires_at: float
|
| 56 |
-
api_url: str = DEFAULT_API_URL
|
| 57 |
-
refresh_in: int | None = None
|
| 58 |
-
sku: str | None = None
|
| 59 |
-
|
| 60 |
-
@property
|
| 61 |
-
def is_valid(self) -> bool:
|
| 62 |
-
return time.time() < (self.expires_at - _TOKEN_EXPIRY_BUFFER_S)
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
def _github_host() -> str:
|
| 66 |
-
return (os.environ.get("GITHUB_COPILOT_HOST") or DEFAULT_GITHUB_HOST).strip().lower()
|
| 67 |
-
|
| 68 |
-
|
| 69 |
-
def _token_exchange_url() -> str:
|
| 70 |
-
return os.environ.get("GITHUB_COPILOT_TOKEN_EXCHANGE_URL", DEFAULT_TOKEN_EXCHANGE_URL).strip()
|
| 71 |
-
|
| 72 |
-
|
| 73 |
-
def _should_exchange_oauth_token() -> bool:
|
| 74 |
-
raw = os.environ.get("GITHUB_COPILOT_USE_TOKEN_EXCHANGE", "").strip().lower()
|
| 75 |
-
return raw in {"1", "true", "yes", "on"}
|
| 76 |
-
|
| 77 |
-
|
| 78 |
-
def _resolve_token_file_paths() -> list[Path]:
|
| 79 |
-
override = os.environ.get("GITHUB_COPILOT_TOKEN_FILE", "").strip()
|
| 80 |
-
if override:
|
| 81 |
-
return [Path(override).expanduser()]
|
| 82 |
-
|
| 83 |
-
paths: list[Path] = []
|
| 84 |
-
local_appdata = os.environ.get("LOCALAPPDATA", "").strip()
|
| 85 |
-
if local_appdata:
|
| 86 |
-
base = Path(local_appdata) / "github-copilot"
|
| 87 |
-
paths.extend([base / "apps.json", base / "hosts.json"])
|
| 88 |
-
|
| 89 |
-
config_base = Path.home() / ".config" / "github-copilot"
|
| 90 |
-
paths.extend([config_base / "apps.json", config_base / "hosts.json"])
|
| 91 |
-
return paths
|
| 92 |
-
|
| 93 |
-
|
| 94 |
-
def _read_gh_cli_oauth_token() -> str | None:
|
| 95 |
-
gh_bin = os.environ.get("GH_PATH", "").strip() or "gh"
|
| 96 |
-
command = [gh_bin, "auth", "token"]
|
| 97 |
-
host = _github_host()
|
| 98 |
-
if host and host != DEFAULT_GITHUB_HOST:
|
| 99 |
-
command.extend(["--hostname", host])
|
| 100 |
-
|
| 101 |
-
try:
|
| 102 |
-
result = subprocess.run(
|
| 103 |
-
command,
|
| 104 |
-
capture_output=True,
|
| 105 |
-
text=True,
|
| 106 |
-
encoding="utf-8",
|
| 107 |
-
errors="replace",
|
| 108 |
-
check=False,
|
| 109 |
-
)
|
| 110 |
-
except OSError as exc:
|
| 111 |
-
logger.debug("Unable to invoke GitHub CLI for Copilot auth discovery: %s", exc)
|
| 112 |
-
return None
|
| 113 |
-
|
| 114 |
-
if result.returncode != 0:
|
| 115 |
-
logger.debug("GitHub CLI auth token lookup failed with exit code %s", result.returncode)
|
| 116 |
-
return None
|
| 117 |
-
|
| 118 |
-
token = result.stdout.strip()
|
| 119 |
-
return token or None
|
| 120 |
-
|
| 121 |
-
|
| 122 |
-
def _read_windows_copilot_cli_oauth_token() -> str | None:
|
| 123 |
-
if os.name != "nt":
|
| 124 |
-
return None
|
| 125 |
-
|
| 126 |
-
class FILETIME(ctypes.Structure):
|
| 127 |
-
_fields_ = [
|
| 128 |
-
("dwLowDateTime", wintypes.DWORD),
|
| 129 |
-
("dwHighDateTime", wintypes.DWORD),
|
| 130 |
-
]
|
| 131 |
-
|
| 132 |
-
class CREDENTIAL(ctypes.Structure):
|
| 133 |
-
_fields_ = [
|
| 134 |
-
("Flags", wintypes.DWORD),
|
| 135 |
-
("Type", wintypes.DWORD),
|
| 136 |
-
("TargetName", wintypes.LPWSTR),
|
| 137 |
-
("Comment", wintypes.LPWSTR),
|
| 138 |
-
("LastWritten", FILETIME),
|
| 139 |
-
("CredentialBlobSize", wintypes.DWORD),
|
| 140 |
-
("CredentialBlob", ctypes.POINTER(ctypes.c_ubyte)),
|
| 141 |
-
("Persist", wintypes.DWORD),
|
| 142 |
-
("AttributeCount", wintypes.DWORD),
|
| 143 |
-
("Attributes", wintypes.LPVOID),
|
| 144 |
-
("TargetAlias", wintypes.LPWSTR),
|
| 145 |
-
("UserName", wintypes.LPWSTR),
|
| 146 |
-
]
|
| 147 |
-
|
| 148 |
-
cred_ptr = ctypes.POINTER(CREDENTIAL)
|
| 149 |
-
credentials = ctypes.POINTER(cred_ptr)()
|
| 150 |
-
count = wintypes.DWORD()
|
| 151 |
-
win_dll = getattr(ctypes, "WinDLL", None)
|
| 152 |
-
if win_dll is None:
|
| 153 |
-
return None
|
| 154 |
-
|
| 155 |
-
advapi32 = win_dll("Advapi32.dll")
|
| 156 |
-
advapi32.CredEnumerateW.argtypes = [
|
| 157 |
-
wintypes.LPCWSTR,
|
| 158 |
-
wintypes.DWORD,
|
| 159 |
-
ctypes.POINTER(wintypes.DWORD),
|
| 160 |
-
ctypes.POINTER(ctypes.POINTER(cred_ptr)),
|
| 161 |
-
]
|
| 162 |
-
advapi32.CredEnumerateW.restype = wintypes.BOOL
|
| 163 |
-
advapi32.CredFree.argtypes = [wintypes.LPVOID]
|
| 164 |
-
|
| 165 |
-
try:
|
| 166 |
-
if not advapi32.CredEnumerateW(None, 0, ctypes.byref(count), ctypes.byref(credentials)):
|
| 167 |
-
return None
|
| 168 |
-
except OSError as exc:
|
| 169 |
-
logger.debug("Unable to enumerate Windows credentials for Copilot auth discovery: %s", exc)
|
| 170 |
-
return None
|
| 171 |
-
|
| 172 |
-
host = _github_host().lower()
|
| 173 |
-
service_prefixes = [f"copilot-cli/{host}:"]
|
| 174 |
-
if "://" not in host:
|
| 175 |
-
service_prefixes.append(f"copilot-cli/https://{host}:")
|
| 176 |
-
|
| 177 |
-
try:
|
| 178 |
-
for idx in range(count.value):
|
| 179 |
-
credential = credentials[idx].contents
|
| 180 |
-
target = (credential.TargetName or "").strip().lower()
|
| 181 |
-
if not any(target.startswith(prefix) for prefix in service_prefixes):
|
| 182 |
-
continue
|
| 183 |
-
if credential.CredentialBlobSize <= 0 or not credential.CredentialBlob:
|
| 184 |
-
continue
|
| 185 |
-
blob = ctypes.string_at(credential.CredentialBlob, credential.CredentialBlobSize)
|
| 186 |
-
token = blob.decode("utf-8", errors="replace").strip()
|
| 187 |
-
if token:
|
| 188 |
-
return token
|
| 189 |
-
finally:
|
| 190 |
-
if credentials:
|
| 191 |
-
advapi32.CredFree(credentials)
|
| 192 |
-
|
| 193 |
-
return None
|
| 194 |
-
|
| 195 |
-
|
| 196 |
-
def _parse_expiry(value: Any) -> float | None:
|
| 197 |
-
if value in (None, ""):
|
| 198 |
-
return None
|
| 199 |
-
|
| 200 |
-
if isinstance(value, int | float):
|
| 201 |
-
number = float(value)
|
| 202 |
-
if number > 10_000_000_000:
|
| 203 |
-
return number / 1000.0
|
| 204 |
-
return number
|
| 205 |
-
|
| 206 |
-
if isinstance(value, str):
|
| 207 |
-
raw = value.strip()
|
| 208 |
-
if not raw:
|
| 209 |
-
return None
|
| 210 |
-
if raw.isdigit():
|
| 211 |
-
return _parse_expiry(int(raw))
|
| 212 |
-
try:
|
| 213 |
-
normalized = raw.replace("Z", "+00:00")
|
| 214 |
-
return datetime.fromisoformat(normalized).timestamp()
|
| 215 |
-
except ValueError:
|
| 216 |
-
return None
|
| 217 |
-
|
| 218 |
-
return None
|
| 219 |
-
|
| 220 |
-
|
| 221 |
-
def _entry_expired(entry: dict[str, Any]) -> bool:
|
| 222 |
-
for key in _EXPIRY_KEYS:
|
| 223 |
-
expiry = _parse_expiry(entry.get(key))
|
| 224 |
-
if expiry is None:
|
| 225 |
-
continue
|
| 226 |
-
return time.time() >= (expiry - _TOKEN_EXPIRY_BUFFER_S)
|
| 227 |
-
return False
|
| 228 |
-
|
| 229 |
-
|
| 230 |
-
def _extract_oauth_token(entry: dict[str, Any]) -> str | None:
|
| 231 |
-
if _entry_expired(entry):
|
| 232 |
-
return None
|
| 233 |
-
|
| 234 |
-
for key in _OAUTH_TOKEN_KEYS:
|
| 235 |
-
value = entry.get(key)
|
| 236 |
-
if isinstance(value, str) and value.strip():
|
| 237 |
-
return value.strip()
|
| 238 |
-
|
| 239 |
-
for value in entry.values():
|
| 240 |
-
if isinstance(value, dict):
|
| 241 |
-
nested = _extract_oauth_token(value)
|
| 242 |
-
if nested:
|
| 243 |
-
return nested
|
| 244 |
-
|
| 245 |
-
return None
|
| 246 |
-
|
| 247 |
-
|
| 248 |
-
def _iter_file_entries(payload: Any) -> list[tuple[str, dict[str, Any]]]:
|
| 249 |
-
entries: list[tuple[str, dict[str, Any]]] = []
|
| 250 |
-
if isinstance(payload, dict):
|
| 251 |
-
for key, value in payload.items():
|
| 252 |
-
if isinstance(value, dict):
|
| 253 |
-
entries.append((str(key), value))
|
| 254 |
-
elif isinstance(payload, list):
|
| 255 |
-
for idx, value in enumerate(payload):
|
| 256 |
-
if isinstance(value, dict):
|
| 257 |
-
key = str(value.get("host") or value.get("githubHost") or idx)
|
| 258 |
-
entries.append((key, value))
|
| 259 |
-
return entries
|
| 260 |
-
|
| 261 |
-
|
| 262 |
-
def read_cached_oauth_token() -> str | None:
|
| 263 |
-
"""Return a GitHub OAuth token for Copilot, if one is available."""
|
| 264 |
-
|
| 265 |
-
for env_var in _OAUTH_TOKEN_ENV_VARS:
|
| 266 |
-
token = os.environ.get(env_var, "").strip()
|
| 267 |
-
if token:
|
| 268 |
-
return token
|
| 269 |
-
|
| 270 |
-
windows_copilot_token = _read_windows_copilot_cli_oauth_token()
|
| 271 |
-
if windows_copilot_token:
|
| 272 |
-
return windows_copilot_token
|
| 273 |
-
|
| 274 |
-
gh_token = _read_gh_cli_oauth_token()
|
| 275 |
-
if gh_token:
|
| 276 |
-
return gh_token
|
| 277 |
-
|
| 278 |
-
host = _github_host()
|
| 279 |
-
for path in _resolve_token_file_paths():
|
| 280 |
-
try:
|
| 281 |
-
payload = json.loads(path.read_text(encoding="utf-8"))
|
| 282 |
-
except FileNotFoundError:
|
| 283 |
-
continue
|
| 284 |
-
except Exception as exc:
|
| 285 |
-
logger.debug("Unable to read Copilot credentials file %s: %s", path, exc)
|
| 286 |
-
continue
|
| 287 |
-
|
| 288 |
-
for key, entry in _iter_file_entries(payload):
|
| 289 |
-
if host not in key.lower():
|
| 290 |
-
continue
|
| 291 |
-
cached_token = _extract_oauth_token(entry)
|
| 292 |
-
if cached_token:
|
| 293 |
-
return cached_token
|
| 294 |
-
|
| 295 |
-
return None
|
| 296 |
-
|
| 297 |
-
|
| 298 |
-
def resolve_client_bearer_token() -> str | None:
|
| 299 |
-
"""Return a bearer token suitable for satisfying Copilot provider auth checks."""
|
| 300 |
-
|
| 301 |
-
for env_var in _API_TOKEN_ENV_VARS:
|
| 302 |
-
token = os.environ.get(env_var, "").strip()
|
| 303 |
-
if token:
|
| 304 |
-
return token
|
| 305 |
-
return read_cached_oauth_token()
|
| 306 |
-
|
| 307 |
-
|
| 308 |
-
def has_oauth_auth() -> bool:
|
| 309 |
-
"""Return True when existing Copilot auth can be reused."""
|
| 310 |
-
|
| 311 |
-
return resolve_client_bearer_token() is not None
|
| 312 |
-
|
| 313 |
-
|
| 314 |
-
def is_copilot_api_url(url: str | None) -> bool:
|
| 315 |
-
"""Return True when the upstream URL points at GitHub Copilot."""
|
| 316 |
-
|
| 317 |
-
if not url:
|
| 318 |
-
return False
|
| 319 |
-
parsed = urlparse(url)
|
| 320 |
-
host = parsed.netloc.lower() or parsed.path.lower()
|
| 321 |
-
return "githubcopilot.com" in host
|
| 322 |
-
|
| 323 |
-
|
| 324 |
-
def build_copilot_upstream_url(base_url: str, path: str) -> str:
|
| 325 |
-
"""Build an upstream URL, normalizing GitHub Copilot's non-/v1 path layout."""
|
| 326 |
-
|
| 327 |
-
normalized_base = base_url.rstrip("/")
|
| 328 |
-
normalized_path = path if path.startswith("/") else f"/{path}"
|
| 329 |
-
if is_copilot_api_url(normalized_base) and normalized_path.startswith("/v1/"):
|
| 330 |
-
normalized_path = normalized_path[3:]
|
| 331 |
-
return f"{normalized_base}{normalized_path}"
|
| 332 |
-
|
| 333 |
-
|
| 334 |
-
class CopilotTokenProvider:
|
| 335 |
-
"""Resolve and cache short-lived Copilot API tokens."""
|
| 336 |
-
|
| 337 |
-
def __init__(self) -> None:
|
| 338 |
-
self._lock = asyncio.Lock()
|
| 339 |
-
self._cached: CopilotAPIToken | None = None
|
| 340 |
-
|
| 341 |
-
async def get_api_token(self) -> CopilotAPIToken:
|
| 342 |
-
explicit_api_token = os.environ.get("GITHUB_COPILOT_API_TOKEN", "").strip()
|
| 343 |
-
if explicit_api_token:
|
| 344 |
-
return CopilotAPIToken(
|
| 345 |
-
token=explicit_api_token,
|
| 346 |
-
expires_at=time.time() + 3600,
|
| 347 |
-
api_url=os.environ.get("GITHUB_COPILOT_API_URL", DEFAULT_API_URL).strip()
|
| 348 |
-
or DEFAULT_API_URL,
|
| 349 |
-
)
|
| 350 |
-
|
| 351 |
-
cached = self._cached
|
| 352 |
-
if cached is not None and cached.is_valid:
|
| 353 |
-
return cached
|
| 354 |
-
|
| 355 |
-
async with self._lock:
|
| 356 |
-
cached = self._cached
|
| 357 |
-
if cached is not None and cached.is_valid:
|
| 358 |
-
return cached
|
| 359 |
-
|
| 360 |
-
oauth_token = read_cached_oauth_token()
|
| 361 |
-
if not oauth_token:
|
| 362 |
-
raise RuntimeError("No GitHub Copilot OAuth token is available.")
|
| 363 |
-
|
| 364 |
-
if not _should_exchange_oauth_token():
|
| 365 |
-
direct_token = CopilotAPIToken(
|
| 366 |
-
token=oauth_token,
|
| 367 |
-
expires_at=time.time() + 3600,
|
| 368 |
-
api_url=os.environ.get("GITHUB_COPILOT_API_URL", DEFAULT_API_URL).strip()
|
| 369 |
-
or DEFAULT_API_URL,
|
| 370 |
-
)
|
| 371 |
-
self._cached = direct_token
|
| 372 |
-
return direct_token
|
| 373 |
-
|
| 374 |
-
exchanged = await self._exchange_token(oauth_token)
|
| 375 |
-
self._cached = exchanged
|
| 376 |
-
return exchanged
|
| 377 |
-
|
| 378 |
-
async def _exchange_token(self, oauth_token: str) -> CopilotAPIToken:
|
| 379 |
-
headers = {
|
| 380 |
-
"Authorization": f"token {oauth_token}",
|
| 381 |
-
"Accept": "application/json",
|
| 382 |
-
"Editor-Version": os.environ.get(
|
| 383 |
-
"GITHUB_COPILOT_EDITOR_VERSION", _DEFAULT_EDITOR_VERSION
|
| 384 |
-
),
|
| 385 |
-
"User-Agent": _DEFAULT_USER_AGENT,
|
| 386 |
-
}
|
| 387 |
-
payload = await asyncio.to_thread(self._exchange_token_sync, headers)
|
| 388 |
-
token = str(payload.get("token") or "").strip()
|
| 389 |
-
if not token:
|
| 390 |
-
raise RuntimeError("Copilot token exchange returned an empty token.")
|
| 391 |
-
|
| 392 |
-
expires_at = _parse_expiry(payload.get("expires_at")) or (time.time() + 1800)
|
| 393 |
-
raw_endpoints = payload.get("endpoints")
|
| 394 |
-
endpoints: dict[str, Any] = raw_endpoints if isinstance(raw_endpoints, dict) else {}
|
| 395 |
-
api_url = str(endpoints.get("api") or DEFAULT_API_URL).strip() or DEFAULT_API_URL
|
| 396 |
-
refresh_in = payload.get("refresh_in")
|
| 397 |
-
sku = payload.get("sku")
|
| 398 |
-
return CopilotAPIToken(
|
| 399 |
-
token=token,
|
| 400 |
-
expires_at=expires_at,
|
| 401 |
-
api_url=api_url,
|
| 402 |
-
refresh_in=int(refresh_in) if isinstance(refresh_in, int | float) else None,
|
| 403 |
-
sku=str(sku) if isinstance(sku, str) and sku.strip() else None,
|
| 404 |
-
)
|
| 405 |
-
|
| 406 |
-
@staticmethod
|
| 407 |
-
def _exchange_token_sync(headers: dict[str, str]) -> dict[str, Any]:
|
| 408 |
-
request = urllib_request.Request(_token_exchange_url(), headers=headers, method="GET")
|
| 409 |
-
try:
|
| 410 |
-
with urllib_request.urlopen(request, timeout=10.0) as response:
|
| 411 |
-
payload = json.loads(response.read().decode("utf-8"))
|
| 412 |
-
return payload if isinstance(payload, dict) else {}
|
| 413 |
-
except urllib_error.HTTPError as exc:
|
| 414 |
-
body = exc.read().decode("utf-8", errors="replace")
|
| 415 |
-
raise RuntimeError(
|
| 416 |
-
f"Copilot token exchange failed with HTTP {exc.code}: {body}"
|
| 417 |
-
) from exc
|
| 418 |
-
|
| 419 |
-
|
| 420 |
-
_provider: CopilotTokenProvider | None = None
|
| 421 |
-
|
| 422 |
-
|
| 423 |
-
def get_copilot_token_provider() -> CopilotTokenProvider:
|
| 424 |
-
"""Return the shared Copilot token provider."""
|
| 425 |
-
|
| 426 |
-
global _provider
|
| 427 |
-
if _provider is None:
|
| 428 |
-
_provider = CopilotTokenProvider()
|
| 429 |
-
return _provider
|
| 430 |
-
|
| 431 |
-
|
| 432 |
-
async def apply_copilot_api_auth(headers: dict[str, str], *, url: str) -> dict[str, str]:
|
| 433 |
-
"""Replace Authorization with a fresh Copilot API token when targeting Copilot."""
|
| 434 |
-
|
| 435 |
-
resolved = dict(headers)
|
| 436 |
-
if not is_copilot_api_url(url):
|
| 437 |
-
return resolved
|
| 438 |
-
|
| 439 |
-
token = await get_copilot_token_provider().get_api_token()
|
| 440 |
-
for key in list(resolved):
|
| 441 |
-
if key.lower() == "authorization":
|
| 442 |
-
resolved.pop(key)
|
| 443 |
-
resolved["Authorization"] = f"Bearer {token.token}"
|
| 444 |
-
return resolved
|
|
|
|
| 1 |
+
"""GitHub Copilot OAuth discovery and API-token exchange helpers."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import asyncio
|
| 6 |
+
import ctypes
|
| 7 |
+
import json
|
| 8 |
+
import logging
|
| 9 |
+
import os
|
| 10 |
+
import subprocess
|
| 11 |
+
import time
|
| 12 |
+
from ctypes import wintypes
|
| 13 |
+
from dataclasses import dataclass
|
| 14 |
+
from datetime import datetime
|
| 15 |
+
from pathlib import Path
|
| 16 |
+
from typing import Any
|
| 17 |
+
from urllib import error as urllib_error
|
| 18 |
+
from urllib import request as urllib_request
|
| 19 |
+
from urllib.parse import urlparse
|
| 20 |
+
|
| 21 |
+
logger = logging.getLogger(__name__)
|
| 22 |
+
|
| 23 |
+
DEFAULT_API_URL = "https://api.githubcopilot.com"
|
| 24 |
+
DEFAULT_TOKEN_EXCHANGE_URL = "https://api.github.com/copilot_internal/v2/token"
|
| 25 |
+
DEFAULT_GITHUB_HOST = "github.com"
|
| 26 |
+
_TOKEN_EXPIRY_BUFFER_S = 60
|
| 27 |
+
_DEFAULT_EDITOR_VERSION = "vscode/1.104.1"
|
| 28 |
+
_DEFAULT_USER_AGENT = "GitHubCopilotChat/0.1"
|
| 29 |
+
|
| 30 |
+
_API_TOKEN_ENV_VARS = (
|
| 31 |
+
"GITHUB_COPILOT_API_TOKEN",
|
| 32 |
+
"COPILOT_PROVIDER_BEARER_TOKEN",
|
| 33 |
+
)
|
| 34 |
+
_OAUTH_TOKEN_ENV_VARS = (
|
| 35 |
+
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 36 |
+
"GITHUB_COPILOT_TOKEN",
|
| 37 |
+
"GITHUB_TOKEN",
|
| 38 |
+
"COPILOT_GITHUB_TOKEN",
|
| 39 |
+
)
|
| 40 |
+
_OAUTH_TOKEN_KEYS = (
|
| 41 |
+
"oauth_token",
|
| 42 |
+
"oauthToken",
|
| 43 |
+
"token",
|
| 44 |
+
"access_token",
|
| 45 |
+
"accessToken",
|
| 46 |
+
)
|
| 47 |
+
_EXPIRY_KEYS = ("expires_at", "expiresAt", "expiry", "expires")
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
@dataclass(frozen=True)
|
| 51 |
+
class CopilotAPIToken:
|
| 52 |
+
"""Short-lived API token exchanged from a GitHub OAuth token."""
|
| 53 |
+
|
| 54 |
+
token: str
|
| 55 |
+
expires_at: float
|
| 56 |
+
api_url: str = DEFAULT_API_URL
|
| 57 |
+
refresh_in: int | None = None
|
| 58 |
+
sku: str | None = None
|
| 59 |
+
|
| 60 |
+
@property
|
| 61 |
+
def is_valid(self) -> bool:
|
| 62 |
+
return time.time() < (self.expires_at - _TOKEN_EXPIRY_BUFFER_S)
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def _github_host() -> str:
|
| 66 |
+
return (os.environ.get("GITHUB_COPILOT_HOST") or DEFAULT_GITHUB_HOST).strip().lower()
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
def _token_exchange_url() -> str:
|
| 70 |
+
return os.environ.get("GITHUB_COPILOT_TOKEN_EXCHANGE_URL", DEFAULT_TOKEN_EXCHANGE_URL).strip()
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def _should_exchange_oauth_token() -> bool:
|
| 74 |
+
raw = os.environ.get("GITHUB_COPILOT_USE_TOKEN_EXCHANGE", "").strip().lower()
|
| 75 |
+
return raw in {"1", "true", "yes", "on"}
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def _resolve_token_file_paths() -> list[Path]:
|
| 79 |
+
override = os.environ.get("GITHUB_COPILOT_TOKEN_FILE", "").strip()
|
| 80 |
+
if override:
|
| 81 |
+
return [Path(override).expanduser()]
|
| 82 |
+
|
| 83 |
+
paths: list[Path] = []
|
| 84 |
+
local_appdata = os.environ.get("LOCALAPPDATA", "").strip()
|
| 85 |
+
if local_appdata:
|
| 86 |
+
base = Path(local_appdata) / "github-copilot"
|
| 87 |
+
paths.extend([base / "apps.json", base / "hosts.json"])
|
| 88 |
+
|
| 89 |
+
config_base = Path.home() / ".config" / "github-copilot"
|
| 90 |
+
paths.extend([config_base / "apps.json", config_base / "hosts.json"])
|
| 91 |
+
return paths
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def _read_gh_cli_oauth_token() -> str | None:
|
| 95 |
+
gh_bin = os.environ.get("GH_PATH", "").strip() or "gh"
|
| 96 |
+
command = [gh_bin, "auth", "token"]
|
| 97 |
+
host = _github_host()
|
| 98 |
+
if host and host != DEFAULT_GITHUB_HOST:
|
| 99 |
+
command.extend(["--hostname", host])
|
| 100 |
+
|
| 101 |
+
try:
|
| 102 |
+
result = subprocess.run(
|
| 103 |
+
command,
|
| 104 |
+
capture_output=True,
|
| 105 |
+
text=True,
|
| 106 |
+
encoding="utf-8",
|
| 107 |
+
errors="replace",
|
| 108 |
+
check=False,
|
| 109 |
+
)
|
| 110 |
+
except OSError as exc:
|
| 111 |
+
logger.debug("Unable to invoke GitHub CLI for Copilot auth discovery: %s", exc)
|
| 112 |
+
return None
|
| 113 |
+
|
| 114 |
+
if result.returncode != 0:
|
| 115 |
+
logger.debug("GitHub CLI auth token lookup failed with exit code %s", result.returncode)
|
| 116 |
+
return None
|
| 117 |
+
|
| 118 |
+
token = result.stdout.strip()
|
| 119 |
+
return token or None
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
def _read_windows_copilot_cli_oauth_token() -> str | None:
|
| 123 |
+
if os.name != "nt":
|
| 124 |
+
return None
|
| 125 |
+
|
| 126 |
+
class FILETIME(ctypes.Structure):
|
| 127 |
+
_fields_ = [
|
| 128 |
+
("dwLowDateTime", wintypes.DWORD),
|
| 129 |
+
("dwHighDateTime", wintypes.DWORD),
|
| 130 |
+
]
|
| 131 |
+
|
| 132 |
+
class CREDENTIAL(ctypes.Structure):
|
| 133 |
+
_fields_ = [
|
| 134 |
+
("Flags", wintypes.DWORD),
|
| 135 |
+
("Type", wintypes.DWORD),
|
| 136 |
+
("TargetName", wintypes.LPWSTR),
|
| 137 |
+
("Comment", wintypes.LPWSTR),
|
| 138 |
+
("LastWritten", FILETIME),
|
| 139 |
+
("CredentialBlobSize", wintypes.DWORD),
|
| 140 |
+
("CredentialBlob", ctypes.POINTER(ctypes.c_ubyte)),
|
| 141 |
+
("Persist", wintypes.DWORD),
|
| 142 |
+
("AttributeCount", wintypes.DWORD),
|
| 143 |
+
("Attributes", wintypes.LPVOID),
|
| 144 |
+
("TargetAlias", wintypes.LPWSTR),
|
| 145 |
+
("UserName", wintypes.LPWSTR),
|
| 146 |
+
]
|
| 147 |
+
|
| 148 |
+
cred_ptr = ctypes.POINTER(CREDENTIAL)
|
| 149 |
+
credentials = ctypes.POINTER(cred_ptr)()
|
| 150 |
+
count = wintypes.DWORD()
|
| 151 |
+
win_dll = getattr(ctypes, "WinDLL", None)
|
| 152 |
+
if win_dll is None:
|
| 153 |
+
return None
|
| 154 |
+
|
| 155 |
+
advapi32 = win_dll("Advapi32.dll")
|
| 156 |
+
advapi32.CredEnumerateW.argtypes = [
|
| 157 |
+
wintypes.LPCWSTR,
|
| 158 |
+
wintypes.DWORD,
|
| 159 |
+
ctypes.POINTER(wintypes.DWORD),
|
| 160 |
+
ctypes.POINTER(ctypes.POINTER(cred_ptr)),
|
| 161 |
+
]
|
| 162 |
+
advapi32.CredEnumerateW.restype = wintypes.BOOL
|
| 163 |
+
advapi32.CredFree.argtypes = [wintypes.LPVOID]
|
| 164 |
+
|
| 165 |
+
try:
|
| 166 |
+
if not advapi32.CredEnumerateW(None, 0, ctypes.byref(count), ctypes.byref(credentials)):
|
| 167 |
+
return None
|
| 168 |
+
except OSError as exc:
|
| 169 |
+
logger.debug("Unable to enumerate Windows credentials for Copilot auth discovery: %s", exc)
|
| 170 |
+
return None
|
| 171 |
+
|
| 172 |
+
host = _github_host().lower()
|
| 173 |
+
service_prefixes = [f"copilot-cli/{host}:"]
|
| 174 |
+
if "://" not in host:
|
| 175 |
+
service_prefixes.append(f"copilot-cli/https://{host}:")
|
| 176 |
+
|
| 177 |
+
try:
|
| 178 |
+
for idx in range(count.value):
|
| 179 |
+
credential = credentials[idx].contents
|
| 180 |
+
target = (credential.TargetName or "").strip().lower()
|
| 181 |
+
if not any(target.startswith(prefix) for prefix in service_prefixes):
|
| 182 |
+
continue
|
| 183 |
+
if credential.CredentialBlobSize <= 0 or not credential.CredentialBlob:
|
| 184 |
+
continue
|
| 185 |
+
blob = ctypes.string_at(credential.CredentialBlob, credential.CredentialBlobSize)
|
| 186 |
+
token = blob.decode("utf-8", errors="replace").strip()
|
| 187 |
+
if token:
|
| 188 |
+
return token
|
| 189 |
+
finally:
|
| 190 |
+
if credentials:
|
| 191 |
+
advapi32.CredFree(credentials)
|
| 192 |
+
|
| 193 |
+
return None
|
| 194 |
+
|
| 195 |
+
|
| 196 |
+
def _parse_expiry(value: Any) -> float | None:
|
| 197 |
+
if value in (None, ""):
|
| 198 |
+
return None
|
| 199 |
+
|
| 200 |
+
if isinstance(value, int | float):
|
| 201 |
+
number = float(value)
|
| 202 |
+
if number > 10_000_000_000:
|
| 203 |
+
return number / 1000.0
|
| 204 |
+
return number
|
| 205 |
+
|
| 206 |
+
if isinstance(value, str):
|
| 207 |
+
raw = value.strip()
|
| 208 |
+
if not raw:
|
| 209 |
+
return None
|
| 210 |
+
if raw.isdigit():
|
| 211 |
+
return _parse_expiry(int(raw))
|
| 212 |
+
try:
|
| 213 |
+
normalized = raw.replace("Z", "+00:00")
|
| 214 |
+
return datetime.fromisoformat(normalized).timestamp()
|
| 215 |
+
except ValueError:
|
| 216 |
+
return None
|
| 217 |
+
|
| 218 |
+
return None
|
| 219 |
+
|
| 220 |
+
|
| 221 |
+
def _entry_expired(entry: dict[str, Any]) -> bool:
|
| 222 |
+
for key in _EXPIRY_KEYS:
|
| 223 |
+
expiry = _parse_expiry(entry.get(key))
|
| 224 |
+
if expiry is None:
|
| 225 |
+
continue
|
| 226 |
+
return time.time() >= (expiry - _TOKEN_EXPIRY_BUFFER_S)
|
| 227 |
+
return False
|
| 228 |
+
|
| 229 |
+
|
| 230 |
+
def _extract_oauth_token(entry: dict[str, Any]) -> str | None:
|
| 231 |
+
if _entry_expired(entry):
|
| 232 |
+
return None
|
| 233 |
+
|
| 234 |
+
for key in _OAUTH_TOKEN_KEYS:
|
| 235 |
+
value = entry.get(key)
|
| 236 |
+
if isinstance(value, str) and value.strip():
|
| 237 |
+
return value.strip()
|
| 238 |
+
|
| 239 |
+
for value in entry.values():
|
| 240 |
+
if isinstance(value, dict):
|
| 241 |
+
nested = _extract_oauth_token(value)
|
| 242 |
+
if nested:
|
| 243 |
+
return nested
|
| 244 |
+
|
| 245 |
+
return None
|
| 246 |
+
|
| 247 |
+
|
| 248 |
+
def _iter_file_entries(payload: Any) -> list[tuple[str, dict[str, Any]]]:
|
| 249 |
+
entries: list[tuple[str, dict[str, Any]]] = []
|
| 250 |
+
if isinstance(payload, dict):
|
| 251 |
+
for key, value in payload.items():
|
| 252 |
+
if isinstance(value, dict):
|
| 253 |
+
entries.append((str(key), value))
|
| 254 |
+
elif isinstance(payload, list):
|
| 255 |
+
for idx, value in enumerate(payload):
|
| 256 |
+
if isinstance(value, dict):
|
| 257 |
+
key = str(value.get("host") or value.get("githubHost") or idx)
|
| 258 |
+
entries.append((key, value))
|
| 259 |
+
return entries
|
| 260 |
+
|
| 261 |
+
|
| 262 |
+
def read_cached_oauth_token() -> str | None:
|
| 263 |
+
"""Return a GitHub OAuth token for Copilot, if one is available."""
|
| 264 |
+
|
| 265 |
+
for env_var in _OAUTH_TOKEN_ENV_VARS:
|
| 266 |
+
token = os.environ.get(env_var, "").strip()
|
| 267 |
+
if token:
|
| 268 |
+
return token
|
| 269 |
+
|
| 270 |
+
windows_copilot_token = _read_windows_copilot_cli_oauth_token()
|
| 271 |
+
if windows_copilot_token:
|
| 272 |
+
return windows_copilot_token
|
| 273 |
+
|
| 274 |
+
gh_token = _read_gh_cli_oauth_token()
|
| 275 |
+
if gh_token:
|
| 276 |
+
return gh_token
|
| 277 |
+
|
| 278 |
+
host = _github_host()
|
| 279 |
+
for path in _resolve_token_file_paths():
|
| 280 |
+
try:
|
| 281 |
+
payload = json.loads(path.read_text(encoding="utf-8"))
|
| 282 |
+
except FileNotFoundError:
|
| 283 |
+
continue
|
| 284 |
+
except Exception as exc:
|
| 285 |
+
logger.debug("Unable to read Copilot credentials file %s: %s", path, exc)
|
| 286 |
+
continue
|
| 287 |
+
|
| 288 |
+
for key, entry in _iter_file_entries(payload):
|
| 289 |
+
if host not in key.lower():
|
| 290 |
+
continue
|
| 291 |
+
cached_token = _extract_oauth_token(entry)
|
| 292 |
+
if cached_token:
|
| 293 |
+
return cached_token
|
| 294 |
+
|
| 295 |
+
return None
|
| 296 |
+
|
| 297 |
+
|
| 298 |
+
def resolve_client_bearer_token() -> str | None:
|
| 299 |
+
"""Return a bearer token suitable for satisfying Copilot provider auth checks."""
|
| 300 |
+
|
| 301 |
+
for env_var in _API_TOKEN_ENV_VARS:
|
| 302 |
+
token = os.environ.get(env_var, "").strip()
|
| 303 |
+
if token:
|
| 304 |
+
return token
|
| 305 |
+
return read_cached_oauth_token()
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
def has_oauth_auth() -> bool:
|
| 309 |
+
"""Return True when existing Copilot auth can be reused."""
|
| 310 |
+
|
| 311 |
+
return resolve_client_bearer_token() is not None
|
| 312 |
+
|
| 313 |
+
|
| 314 |
+
def is_copilot_api_url(url: str | None) -> bool:
|
| 315 |
+
"""Return True when the upstream URL points at GitHub Copilot."""
|
| 316 |
+
|
| 317 |
+
if not url:
|
| 318 |
+
return False
|
| 319 |
+
parsed = urlparse(url)
|
| 320 |
+
host = parsed.netloc.lower() or parsed.path.lower()
|
| 321 |
+
return "githubcopilot.com" in host
|
| 322 |
+
|
| 323 |
+
|
| 324 |
+
def build_copilot_upstream_url(base_url: str, path: str) -> str:
|
| 325 |
+
"""Build an upstream URL, normalizing GitHub Copilot's non-/v1 path layout."""
|
| 326 |
+
|
| 327 |
+
normalized_base = base_url.rstrip("/")
|
| 328 |
+
normalized_path = path if path.startswith("/") else f"/{path}"
|
| 329 |
+
if is_copilot_api_url(normalized_base) and normalized_path.startswith("/v1/"):
|
| 330 |
+
normalized_path = normalized_path[3:]
|
| 331 |
+
return f"{normalized_base}{normalized_path}"
|
| 332 |
+
|
| 333 |
+
|
| 334 |
+
class CopilotTokenProvider:
|
| 335 |
+
"""Resolve and cache short-lived Copilot API tokens."""
|
| 336 |
+
|
| 337 |
+
def __init__(self) -> None:
|
| 338 |
+
self._lock = asyncio.Lock()
|
| 339 |
+
self._cached: CopilotAPIToken | None = None
|
| 340 |
+
|
| 341 |
+
async def get_api_token(self) -> CopilotAPIToken:
|
| 342 |
+
explicit_api_token = os.environ.get("GITHUB_COPILOT_API_TOKEN", "").strip()
|
| 343 |
+
if explicit_api_token:
|
| 344 |
+
return CopilotAPIToken(
|
| 345 |
+
token=explicit_api_token,
|
| 346 |
+
expires_at=time.time() + 3600,
|
| 347 |
+
api_url=os.environ.get("GITHUB_COPILOT_API_URL", DEFAULT_API_URL).strip()
|
| 348 |
+
or DEFAULT_API_URL,
|
| 349 |
+
)
|
| 350 |
+
|
| 351 |
+
cached = self._cached
|
| 352 |
+
if cached is not None and cached.is_valid:
|
| 353 |
+
return cached
|
| 354 |
+
|
| 355 |
+
async with self._lock:
|
| 356 |
+
cached = self._cached
|
| 357 |
+
if cached is not None and cached.is_valid:
|
| 358 |
+
return cached
|
| 359 |
+
|
| 360 |
+
oauth_token = read_cached_oauth_token()
|
| 361 |
+
if not oauth_token:
|
| 362 |
+
raise RuntimeError("No GitHub Copilot OAuth token is available.")
|
| 363 |
+
|
| 364 |
+
if not _should_exchange_oauth_token():
|
| 365 |
+
direct_token = CopilotAPIToken(
|
| 366 |
+
token=oauth_token,
|
| 367 |
+
expires_at=time.time() + 3600,
|
| 368 |
+
api_url=os.environ.get("GITHUB_COPILOT_API_URL", DEFAULT_API_URL).strip()
|
| 369 |
+
or DEFAULT_API_URL,
|
| 370 |
+
)
|
| 371 |
+
self._cached = direct_token
|
| 372 |
+
return direct_token
|
| 373 |
+
|
| 374 |
+
exchanged = await self._exchange_token(oauth_token)
|
| 375 |
+
self._cached = exchanged
|
| 376 |
+
return exchanged
|
| 377 |
+
|
| 378 |
+
async def _exchange_token(self, oauth_token: str) -> CopilotAPIToken:
|
| 379 |
+
headers = {
|
| 380 |
+
"Authorization": f"token {oauth_token}",
|
| 381 |
+
"Accept": "application/json",
|
| 382 |
+
"Editor-Version": os.environ.get(
|
| 383 |
+
"GITHUB_COPILOT_EDITOR_VERSION", _DEFAULT_EDITOR_VERSION
|
| 384 |
+
),
|
| 385 |
+
"User-Agent": _DEFAULT_USER_AGENT,
|
| 386 |
+
}
|
| 387 |
+
payload = await asyncio.to_thread(self._exchange_token_sync, headers)
|
| 388 |
+
token = str(payload.get("token") or "").strip()
|
| 389 |
+
if not token:
|
| 390 |
+
raise RuntimeError("Copilot token exchange returned an empty token.")
|
| 391 |
+
|
| 392 |
+
expires_at = _parse_expiry(payload.get("expires_at")) or (time.time() + 1800)
|
| 393 |
+
raw_endpoints = payload.get("endpoints")
|
| 394 |
+
endpoints: dict[str, Any] = raw_endpoints if isinstance(raw_endpoints, dict) else {}
|
| 395 |
+
api_url = str(endpoints.get("api") or DEFAULT_API_URL).strip() or DEFAULT_API_URL
|
| 396 |
+
refresh_in = payload.get("refresh_in")
|
| 397 |
+
sku = payload.get("sku")
|
| 398 |
+
return CopilotAPIToken(
|
| 399 |
+
token=token,
|
| 400 |
+
expires_at=expires_at,
|
| 401 |
+
api_url=api_url,
|
| 402 |
+
refresh_in=int(refresh_in) if isinstance(refresh_in, int | float) else None,
|
| 403 |
+
sku=str(sku) if isinstance(sku, str) and sku.strip() else None,
|
| 404 |
+
)
|
| 405 |
+
|
| 406 |
+
@staticmethod
|
| 407 |
+
def _exchange_token_sync(headers: dict[str, str]) -> dict[str, Any]:
|
| 408 |
+
request = urllib_request.Request(_token_exchange_url(), headers=headers, method="GET")
|
| 409 |
+
try:
|
| 410 |
+
with urllib_request.urlopen(request, timeout=10.0) as response:
|
| 411 |
+
payload = json.loads(response.read().decode("utf-8"))
|
| 412 |
+
return payload if isinstance(payload, dict) else {}
|
| 413 |
+
except urllib_error.HTTPError as exc:
|
| 414 |
+
body = exc.read().decode("utf-8", errors="replace")
|
| 415 |
+
raise RuntimeError(
|
| 416 |
+
f"Copilot token exchange failed with HTTP {exc.code}: {body}"
|
| 417 |
+
) from exc
|
| 418 |
+
|
| 419 |
+
|
| 420 |
+
_provider: CopilotTokenProvider | None = None
|
| 421 |
+
|
| 422 |
+
|
| 423 |
+
def get_copilot_token_provider() -> CopilotTokenProvider:
|
| 424 |
+
"""Return the shared Copilot token provider."""
|
| 425 |
+
|
| 426 |
+
global _provider
|
| 427 |
+
if _provider is None:
|
| 428 |
+
_provider = CopilotTokenProvider()
|
| 429 |
+
return _provider
|
| 430 |
+
|
| 431 |
+
|
| 432 |
+
async def apply_copilot_api_auth(headers: dict[str, str], *, url: str) -> dict[str, str]:
|
| 433 |
+
"""Replace Authorization with a fresh Copilot API token when targeting Copilot."""
|
| 434 |
+
|
| 435 |
+
resolved = dict(headers)
|
| 436 |
+
if not is_copilot_api_url(url):
|
| 437 |
+
return resolved
|
| 438 |
+
|
| 439 |
+
token = await get_copilot_token_provider().get_api_token()
|
| 440 |
+
for key in list(resolved):
|
| 441 |
+
if key.lower() == "authorization":
|
| 442 |
+
resolved.pop(key)
|
| 443 |
+
resolved["Authorization"] = f"Bearer {token.token}"
|
| 444 |
+
return resolved
|
|
@@ -1,28 +1,28 @@
|
|
| 1 |
-
"""Health helpers for persistent deployments."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import json
|
| 6 |
-
import urllib.error
|
| 7 |
-
import urllib.request
|
| 8 |
-
from typing import Any
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
def probe_json(url: str, timeout: float = 2.0) -> dict[str, Any] | None:
|
| 12 |
-
"""Return a JSON payload from the URL when reachable."""
|
| 13 |
-
|
| 14 |
-
try:
|
| 15 |
-
with urllib.request.urlopen(url, timeout=timeout) as response:
|
| 16 |
-
payload = json.loads(response.read().decode("utf-8"))
|
| 17 |
-
except (OSError, urllib.error.URLError, ValueError, json.JSONDecodeError):
|
| 18 |
-
return None
|
| 19 |
-
return payload if isinstance(payload, dict) else None
|
| 20 |
-
|
| 21 |
-
|
| 22 |
-
def probe_ready(url: str, timeout: float = 2.0) -> bool:
|
| 23 |
-
"""Return True when the ready endpoint reports readiness."""
|
| 24 |
-
|
| 25 |
-
payload = probe_json(url, timeout=timeout)
|
| 26 |
-
if not isinstance(payload, dict):
|
| 27 |
-
return False
|
| 28 |
-
return bool(payload.get("ready", False) or payload.get("status") == "healthy")
|
|
|
|
| 1 |
+
"""Health helpers for persistent deployments."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import json
|
| 6 |
+
import urllib.error
|
| 7 |
+
import urllib.request
|
| 8 |
+
from typing import Any
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def probe_json(url: str, timeout: float = 2.0) -> dict[str, Any] | None:
|
| 12 |
+
"""Return a JSON payload from the URL when reachable."""
|
| 13 |
+
|
| 14 |
+
try:
|
| 15 |
+
with urllib.request.urlopen(url, timeout=timeout) as response:
|
| 16 |
+
payload = json.loads(response.read().decode("utf-8"))
|
| 17 |
+
except (OSError, urllib.error.URLError, ValueError, json.JSONDecodeError):
|
| 18 |
+
return None
|
| 19 |
+
return payload if isinstance(payload, dict) else None
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def probe_ready(url: str, timeout: float = 2.0) -> bool:
|
| 23 |
+
"""Return True when the ready endpoint reports readiness."""
|
| 24 |
+
|
| 25 |
+
payload = probe_json(url, timeout=timeout)
|
| 26 |
+
if not isinstance(payload, dict):
|
| 27 |
+
return False
|
| 28 |
+
return bool(payload.get("ready", False) or payload.get("status") == "healthy")
|
|
@@ -1,174 +1,174 @@
|
|
| 1 |
-
"""Tool-target configuration for persistent deployments."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import os
|
| 6 |
-
import re
|
| 7 |
-
import subprocess
|
| 8 |
-
from pathlib import Path
|
| 9 |
-
|
| 10 |
-
from headroom.providers.install_registry import (
|
| 11 |
-
apply_provider_scope_mutations,
|
| 12 |
-
revert_provider_scope_mutation,
|
| 13 |
-
)
|
| 14 |
-
|
| 15 |
-
from .models import ConfigScope, DeploymentManifest, ManagedMutation
|
| 16 |
-
from .paths import (
|
| 17 |
-
unix_system_env_targets,
|
| 18 |
-
unix_user_env_targets,
|
| 19 |
-
)
|
| 20 |
-
|
| 21 |
-
_ENV_MARKER_START = "# >>> headroom persistent env >>>"
|
| 22 |
-
_ENV_MARKER_END = "# <<< headroom persistent env <<<"
|
| 23 |
-
_ENV_PATTERN = re.compile(
|
| 24 |
-
re.escape(_ENV_MARKER_START) + r".*?" + re.escape(_ENV_MARKER_END),
|
| 25 |
-
re.DOTALL,
|
| 26 |
-
)
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
def _merge_marker_block(file_path: Path, block: str, pattern: re.Pattern[str], marker: str) -> str:
|
| 30 |
-
if file_path.exists():
|
| 31 |
-
existing = file_path.read_text()
|
| 32 |
-
if marker in existing:
|
| 33 |
-
return pattern.sub(block, existing)
|
| 34 |
-
return existing.rstrip() + "\n\n" + block + "\n"
|
| 35 |
-
return block + "\n"
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
def _env_block(values: dict[str, str]) -> str:
|
| 39 |
-
lines = [_ENV_MARKER_START]
|
| 40 |
-
for name, value in values.items():
|
| 41 |
-
lines.append(f'export {name}="{value}"')
|
| 42 |
-
lines.append(_ENV_MARKER_END)
|
| 43 |
-
return "\n".join(lines)
|
| 44 |
-
|
| 45 |
-
|
| 46 |
-
def _powershell_literal(value: str) -> str:
|
| 47 |
-
return "'" + value.replace("'", "''") + "'"
|
| 48 |
-
|
| 49 |
-
|
| 50 |
-
def _unix_scope_values(manifest: DeploymentManifest) -> dict[str, str]:
|
| 51 |
-
merged = dict(manifest.base_env)
|
| 52 |
-
for env_map in manifest.tool_envs.values():
|
| 53 |
-
merged.update(env_map)
|
| 54 |
-
return merged
|
| 55 |
-
|
| 56 |
-
|
| 57 |
-
def _apply_unix_env_scope(manifest: DeploymentManifest) -> list[ManagedMutation]:
|
| 58 |
-
values = _unix_scope_values(manifest)
|
| 59 |
-
block = _env_block(values)
|
| 60 |
-
if manifest.scope == ConfigScope.USER.value:
|
| 61 |
-
targets = unix_user_env_targets()
|
| 62 |
-
else:
|
| 63 |
-
targets = unix_system_env_targets()
|
| 64 |
-
mutations: list[ManagedMutation] = []
|
| 65 |
-
for path in targets:
|
| 66 |
-
path.parent.mkdir(parents=True, exist_ok=True)
|
| 67 |
-
merged = _merge_marker_block(path, block, _ENV_PATTERN, _ENV_MARKER_START)
|
| 68 |
-
path.write_text(merged)
|
| 69 |
-
mutations.append(ManagedMutation(target="env", kind="shell-block", path=str(path)))
|
| 70 |
-
return mutations
|
| 71 |
-
|
| 72 |
-
|
| 73 |
-
def _remove_unix_env_scope(mutations: list[ManagedMutation]) -> None:
|
| 74 |
-
for mutation in mutations:
|
| 75 |
-
if mutation.kind != "shell-block" or not mutation.path:
|
| 76 |
-
continue
|
| 77 |
-
path = Path(mutation.path)
|
| 78 |
-
if not path.exists():
|
| 79 |
-
continue
|
| 80 |
-
content = path.read_text()
|
| 81 |
-
if _ENV_MARKER_START not in content:
|
| 82 |
-
continue
|
| 83 |
-
path.write_text(_ENV_PATTERN.sub("", content).strip() + "\n")
|
| 84 |
-
|
| 85 |
-
|
| 86 |
-
def _apply_windows_env_scope(manifest: DeploymentManifest) -> list[ManagedMutation]:
|
| 87 |
-
scope_name = "Machine" if manifest.scope == ConfigScope.SYSTEM.value else "User"
|
| 88 |
-
merged = _unix_scope_values(manifest)
|
| 89 |
-
mutations: list[ManagedMutation] = []
|
| 90 |
-
for name, value in merged.items():
|
| 91 |
-
previous = subprocess.run(
|
| 92 |
-
[
|
| 93 |
-
"powershell",
|
| 94 |
-
"-NoProfile",
|
| 95 |
-
"-Command",
|
| 96 |
-
f"$value = [Environment]::GetEnvironmentVariable({_powershell_literal(name)},{_powershell_literal(scope_name)}); "
|
| 97 |
-
"if ($null -eq $value) { '__HEADROOM_UNSET__' } else { $value }",
|
| 98 |
-
],
|
| 99 |
-
capture_output=True,
|
| 100 |
-
text=True,
|
| 101 |
-
check=True,
|
| 102 |
-
).stdout.strip()
|
| 103 |
-
command = [
|
| 104 |
-
"powershell",
|
| 105 |
-
"-NoProfile",
|
| 106 |
-
"-Command",
|
| 107 |
-
f"[Environment]::SetEnvironmentVariable({_powershell_literal(name)},{_powershell_literal(value)},{_powershell_literal(scope_name)})",
|
| 108 |
-
]
|
| 109 |
-
subprocess.run(command, check=True)
|
| 110 |
-
mutations.append(
|
| 111 |
-
ManagedMutation(
|
| 112 |
-
target="env",
|
| 113 |
-
kind="windows-env",
|
| 114 |
-
data={
|
| 115 |
-
"name": name,
|
| 116 |
-
"scope": scope_name,
|
| 117 |
-
"previous": None if previous == "__HEADROOM_UNSET__" else previous,
|
| 118 |
-
},
|
| 119 |
-
)
|
| 120 |
-
)
|
| 121 |
-
return mutations
|
| 122 |
-
|
| 123 |
-
|
| 124 |
-
def _remove_windows_env_scope(mutations: list[ManagedMutation]) -> None:
|
| 125 |
-
for mutation in mutations:
|
| 126 |
-
if mutation.kind != "windows-env":
|
| 127 |
-
continue
|
| 128 |
-
name = mutation.data.get("name")
|
| 129 |
-
if not isinstance(name, str):
|
| 130 |
-
raise ValueError("Windows environment mutation is missing a variable name")
|
| 131 |
-
scope_name = mutation.data.get("scope", "User")
|
| 132 |
-
if not isinstance(scope_name, str):
|
| 133 |
-
raise ValueError("Windows environment mutation is missing a valid scope")
|
| 134 |
-
previous = mutation.data.get("previous")
|
| 135 |
-
if previous is None:
|
| 136 |
-
value_literal = "$null"
|
| 137 |
-
else:
|
| 138 |
-
value_literal = _powershell_literal(previous)
|
| 139 |
-
command = [
|
| 140 |
-
"powershell",
|
| 141 |
-
"-NoProfile",
|
| 142 |
-
"-Command",
|
| 143 |
-
f"[Environment]::SetEnvironmentVariable({_powershell_literal(name)},{value_literal},{_powershell_literal(scope_name)})",
|
| 144 |
-
]
|
| 145 |
-
subprocess.run(command, check=True)
|
| 146 |
-
|
| 147 |
-
|
| 148 |
-
def apply_mutations(manifest: DeploymentManifest) -> list[ManagedMutation]:
|
| 149 |
-
"""Apply provider/user/system configuration for a deployment."""
|
| 150 |
-
|
| 151 |
-
mutations: list[ManagedMutation] = []
|
| 152 |
-
if manifest.scope in {ConfigScope.USER.value, ConfigScope.SYSTEM.value}:
|
| 153 |
-
if os.name == "nt":
|
| 154 |
-
mutations.extend(_apply_windows_env_scope(manifest))
|
| 155 |
-
else:
|
| 156 |
-
mutations.extend(_apply_unix_env_scope(manifest))
|
| 157 |
-
mutations.extend(apply_provider_scope_mutations(manifest))
|
| 158 |
-
return mutations
|
| 159 |
-
|
| 160 |
-
return [*mutations, *apply_provider_scope_mutations(manifest)]
|
| 161 |
-
|
| 162 |
-
|
| 163 |
-
def revert_mutations(manifest: DeploymentManifest) -> None:
|
| 164 |
-
"""Undo the stored mutations for a deployment."""
|
| 165 |
-
|
| 166 |
-
if manifest.scope in {ConfigScope.USER.value, ConfigScope.SYSTEM.value}:
|
| 167 |
-
shell_mutations = [m for m in manifest.mutations if m.target == "env"]
|
| 168 |
-
if os.name == "nt":
|
| 169 |
-
_remove_windows_env_scope(shell_mutations)
|
| 170 |
-
else:
|
| 171 |
-
_remove_unix_env_scope(shell_mutations)
|
| 172 |
-
|
| 173 |
-
for mutation in manifest.mutations:
|
| 174 |
-
revert_provider_scope_mutation(manifest, mutation)
|
|
|
|
| 1 |
+
"""Tool-target configuration for persistent deployments."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import os
|
| 6 |
+
import re
|
| 7 |
+
import subprocess
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
|
| 10 |
+
from headroom.providers.install_registry import (
|
| 11 |
+
apply_provider_scope_mutations,
|
| 12 |
+
revert_provider_scope_mutation,
|
| 13 |
+
)
|
| 14 |
+
|
| 15 |
+
from .models import ConfigScope, DeploymentManifest, ManagedMutation
|
| 16 |
+
from .paths import (
|
| 17 |
+
unix_system_env_targets,
|
| 18 |
+
unix_user_env_targets,
|
| 19 |
+
)
|
| 20 |
+
|
| 21 |
+
_ENV_MARKER_START = "# >>> headroom persistent env >>>"
|
| 22 |
+
_ENV_MARKER_END = "# <<< headroom persistent env <<<"
|
| 23 |
+
_ENV_PATTERN = re.compile(
|
| 24 |
+
re.escape(_ENV_MARKER_START) + r".*?" + re.escape(_ENV_MARKER_END),
|
| 25 |
+
re.DOTALL,
|
| 26 |
+
)
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def _merge_marker_block(file_path: Path, block: str, pattern: re.Pattern[str], marker: str) -> str:
|
| 30 |
+
if file_path.exists():
|
| 31 |
+
existing = file_path.read_text()
|
| 32 |
+
if marker in existing:
|
| 33 |
+
return pattern.sub(block, existing)
|
| 34 |
+
return existing.rstrip() + "\n\n" + block + "\n"
|
| 35 |
+
return block + "\n"
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def _env_block(values: dict[str, str]) -> str:
|
| 39 |
+
lines = [_ENV_MARKER_START]
|
| 40 |
+
for name, value in values.items():
|
| 41 |
+
lines.append(f'export {name}="{value}"')
|
| 42 |
+
lines.append(_ENV_MARKER_END)
|
| 43 |
+
return "\n".join(lines)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def _powershell_literal(value: str) -> str:
|
| 47 |
+
return "'" + value.replace("'", "''") + "'"
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def _unix_scope_values(manifest: DeploymentManifest) -> dict[str, str]:
|
| 51 |
+
merged = dict(manifest.base_env)
|
| 52 |
+
for env_map in manifest.tool_envs.values():
|
| 53 |
+
merged.update(env_map)
|
| 54 |
+
return merged
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
def _apply_unix_env_scope(manifest: DeploymentManifest) -> list[ManagedMutation]:
|
| 58 |
+
values = _unix_scope_values(manifest)
|
| 59 |
+
block = _env_block(values)
|
| 60 |
+
if manifest.scope == ConfigScope.USER.value:
|
| 61 |
+
targets = unix_user_env_targets()
|
| 62 |
+
else:
|
| 63 |
+
targets = unix_system_env_targets()
|
| 64 |
+
mutations: list[ManagedMutation] = []
|
| 65 |
+
for path in targets:
|
| 66 |
+
path.parent.mkdir(parents=True, exist_ok=True)
|
| 67 |
+
merged = _merge_marker_block(path, block, _ENV_PATTERN, _ENV_MARKER_START)
|
| 68 |
+
path.write_text(merged)
|
| 69 |
+
mutations.append(ManagedMutation(target="env", kind="shell-block", path=str(path)))
|
| 70 |
+
return mutations
|
| 71 |
+
|
| 72 |
+
|
| 73 |
+
def _remove_unix_env_scope(mutations: list[ManagedMutation]) -> None:
|
| 74 |
+
for mutation in mutations:
|
| 75 |
+
if mutation.kind != "shell-block" or not mutation.path:
|
| 76 |
+
continue
|
| 77 |
+
path = Path(mutation.path)
|
| 78 |
+
if not path.exists():
|
| 79 |
+
continue
|
| 80 |
+
content = path.read_text()
|
| 81 |
+
if _ENV_MARKER_START not in content:
|
| 82 |
+
continue
|
| 83 |
+
path.write_text(_ENV_PATTERN.sub("", content).strip() + "\n")
|
| 84 |
+
|
| 85 |
+
|
| 86 |
+
def _apply_windows_env_scope(manifest: DeploymentManifest) -> list[ManagedMutation]:
|
| 87 |
+
scope_name = "Machine" if manifest.scope == ConfigScope.SYSTEM.value else "User"
|
| 88 |
+
merged = _unix_scope_values(manifest)
|
| 89 |
+
mutations: list[ManagedMutation] = []
|
| 90 |
+
for name, value in merged.items():
|
| 91 |
+
previous = subprocess.run(
|
| 92 |
+
[
|
| 93 |
+
"powershell",
|
| 94 |
+
"-NoProfile",
|
| 95 |
+
"-Command",
|
| 96 |
+
f"$value = [Environment]::GetEnvironmentVariable({_powershell_literal(name)},{_powershell_literal(scope_name)}); "
|
| 97 |
+
"if ($null -eq $value) { '__HEADROOM_UNSET__' } else { $value }",
|
| 98 |
+
],
|
| 99 |
+
capture_output=True,
|
| 100 |
+
text=True,
|
| 101 |
+
check=True,
|
| 102 |
+
).stdout.strip()
|
| 103 |
+
command = [
|
| 104 |
+
"powershell",
|
| 105 |
+
"-NoProfile",
|
| 106 |
+
"-Command",
|
| 107 |
+
f"[Environment]::SetEnvironmentVariable({_powershell_literal(name)},{_powershell_literal(value)},{_powershell_literal(scope_name)})",
|
| 108 |
+
]
|
| 109 |
+
subprocess.run(command, check=True)
|
| 110 |
+
mutations.append(
|
| 111 |
+
ManagedMutation(
|
| 112 |
+
target="env",
|
| 113 |
+
kind="windows-env",
|
| 114 |
+
data={
|
| 115 |
+
"name": name,
|
| 116 |
+
"scope": scope_name,
|
| 117 |
+
"previous": None if previous == "__HEADROOM_UNSET__" else previous,
|
| 118 |
+
},
|
| 119 |
+
)
|
| 120 |
+
)
|
| 121 |
+
return mutations
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
def _remove_windows_env_scope(mutations: list[ManagedMutation]) -> None:
|
| 125 |
+
for mutation in mutations:
|
| 126 |
+
if mutation.kind != "windows-env":
|
| 127 |
+
continue
|
| 128 |
+
name = mutation.data.get("name")
|
| 129 |
+
if not isinstance(name, str):
|
| 130 |
+
raise ValueError("Windows environment mutation is missing a variable name")
|
| 131 |
+
scope_name = mutation.data.get("scope", "User")
|
| 132 |
+
if not isinstance(scope_name, str):
|
| 133 |
+
raise ValueError("Windows environment mutation is missing a valid scope")
|
| 134 |
+
previous = mutation.data.get("previous")
|
| 135 |
+
if previous is None:
|
| 136 |
+
value_literal = "$null"
|
| 137 |
+
else:
|
| 138 |
+
value_literal = _powershell_literal(previous)
|
| 139 |
+
command = [
|
| 140 |
+
"powershell",
|
| 141 |
+
"-NoProfile",
|
| 142 |
+
"-Command",
|
| 143 |
+
f"[Environment]::SetEnvironmentVariable({_powershell_literal(name)},{value_literal},{_powershell_literal(scope_name)})",
|
| 144 |
+
]
|
| 145 |
+
subprocess.run(command, check=True)
|
| 146 |
+
|
| 147 |
+
|
| 148 |
+
def apply_mutations(manifest: DeploymentManifest) -> list[ManagedMutation]:
|
| 149 |
+
"""Apply provider/user/system configuration for a deployment."""
|
| 150 |
+
|
| 151 |
+
mutations: list[ManagedMutation] = []
|
| 152 |
+
if manifest.scope in {ConfigScope.USER.value, ConfigScope.SYSTEM.value}:
|
| 153 |
+
if os.name == "nt":
|
| 154 |
+
mutations.extend(_apply_windows_env_scope(manifest))
|
| 155 |
+
else:
|
| 156 |
+
mutations.extend(_apply_unix_env_scope(manifest))
|
| 157 |
+
mutations.extend(apply_provider_scope_mutations(manifest))
|
| 158 |
+
return mutations
|
| 159 |
+
|
| 160 |
+
return [*mutations, *apply_provider_scope_mutations(manifest)]
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
def revert_mutations(manifest: DeploymentManifest) -> None:
|
| 164 |
+
"""Undo the stored mutations for a deployment."""
|
| 165 |
+
|
| 166 |
+
if manifest.scope in {ConfigScope.USER.value, ConfigScope.SYSTEM.value}:
|
| 167 |
+
shell_mutations = [m for m in manifest.mutations if m.target == "env"]
|
| 168 |
+
if os.name == "nt":
|
| 169 |
+
_remove_windows_env_scope(shell_mutations)
|
| 170 |
+
else:
|
| 171 |
+
_remove_unix_env_scope(shell_mutations)
|
| 172 |
+
|
| 173 |
+
for mutation in manifest.mutations:
|
| 174 |
+
revert_provider_scope_mutation(manifest, mutation)
|
|
@@ -1,279 +1,279 @@
|
|
| 1 |
-
"""Runtime helpers for persistent deployments."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import os
|
| 6 |
-
import shutil
|
| 7 |
-
import signal
|
| 8 |
-
import subprocess
|
| 9 |
-
import sys
|
| 10 |
-
import time
|
| 11 |
-
from pathlib import Path
|
| 12 |
-
from typing import Any
|
| 13 |
-
|
| 14 |
-
from .health import probe_ready
|
| 15 |
-
from .models import DeploymentManifest, InstallPreset, RuntimeKind
|
| 16 |
-
from .paths import log_path, pid_path
|
| 17 |
-
|
| 18 |
-
PASSTHROUGH_ENV_PREFIXES = (
|
| 19 |
-
"HEADROOM_",
|
| 20 |
-
"ANTHROPIC_",
|
| 21 |
-
"OPENAI_",
|
| 22 |
-
"GEMINI_",
|
| 23 |
-
"AWS_",
|
| 24 |
-
"AZURE_",
|
| 25 |
-
"VERTEX_",
|
| 26 |
-
"GOOGLE_",
|
| 27 |
-
"GOOGLE_CLOUD_",
|
| 28 |
-
"MISTRAL_",
|
| 29 |
-
"GROQ_",
|
| 30 |
-
"OPENROUTER_",
|
| 31 |
-
"XAI_",
|
| 32 |
-
"TOGETHER_",
|
| 33 |
-
"COHERE_",
|
| 34 |
-
"OLLAMA_",
|
| 35 |
-
"LITELLM_",
|
| 36 |
-
"OTEL_",
|
| 37 |
-
"SUPABASE_",
|
| 38 |
-
"QDRANT_",
|
| 39 |
-
"NEO4J_",
|
| 40 |
-
"LANGSMITH_",
|
| 41 |
-
)
|
| 42 |
-
|
| 43 |
-
|
| 44 |
-
def _is_windows() -> bool:
|
| 45 |
-
return sys.platform.startswith("win")
|
| 46 |
-
|
| 47 |
-
|
| 48 |
-
def _deployment_env(manifest: DeploymentManifest) -> dict[str, str]:
|
| 49 |
-
return {
|
| 50 |
-
"HEADROOM_DEPLOYMENT_PROFILE": manifest.profile,
|
| 51 |
-
"HEADROOM_DEPLOYMENT_PRESET": manifest.preset,
|
| 52 |
-
"HEADROOM_DEPLOYMENT_RUNTIME": manifest.runtime_kind,
|
| 53 |
-
"HEADROOM_DEPLOYMENT_SUPERVISOR": manifest.supervisor_kind,
|
| 54 |
-
"HEADROOM_DEPLOYMENT_SCOPE": manifest.scope,
|
| 55 |
-
}
|
| 56 |
-
|
| 57 |
-
|
| 58 |
-
def resolve_headroom_command() -> list[str]:
|
| 59 |
-
"""Resolve the most reliable command to invoke headroom."""
|
| 60 |
-
|
| 61 |
-
headroom_bin = shutil.which("headroom")
|
| 62 |
-
if headroom_bin:
|
| 63 |
-
return [headroom_bin]
|
| 64 |
-
return [sys.executable, "-m", "headroom.cli"]
|
| 65 |
-
|
| 66 |
-
|
| 67 |
-
def _runtime_env(manifest: DeploymentManifest) -> dict[str, str]:
|
| 68 |
-
env = os.environ.copy()
|
| 69 |
-
env.update(manifest.base_env)
|
| 70 |
-
env.update(_deployment_env(manifest))
|
| 71 |
-
return env
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
def _ensure_host_dirs() -> None:
|
| 75 |
-
for subdir in (".headroom", ".claude", ".codex", ".gemini"):
|
| 76 |
-
(Path.home() / subdir).mkdir(parents=True, exist_ok=True)
|
| 77 |
-
|
| 78 |
-
|
| 79 |
-
def _mount_source(home: str, subdir: str) -> str:
|
| 80 |
-
if _is_windows():
|
| 81 |
-
return f"{home}\\{subdir}"
|
| 82 |
-
return f"{home}/{subdir}"
|
| 83 |
-
|
| 84 |
-
|
| 85 |
-
def build_runtime_command(manifest: DeploymentManifest) -> list[str]:
|
| 86 |
-
"""Build the raw foreground command that runs the proxy."""
|
| 87 |
-
|
| 88 |
-
if manifest.runtime_kind == RuntimeKind.PYTHON.value:
|
| 89 |
-
return [sys.executable, "-m", "headroom.cli", "proxy", *manifest.proxy_args]
|
| 90 |
-
|
| 91 |
-
_ensure_host_dirs()
|
| 92 |
-
home = str(Path.home())
|
| 93 |
-
container_home = "/tmp/headroom-home"
|
| 94 |
-
command = [
|
| 95 |
-
"docker",
|
| 96 |
-
"run",
|
| 97 |
-
"--rm",
|
| 98 |
-
"--name",
|
| 99 |
-
manifest.container_name,
|
| 100 |
-
"-p",
|
| 101 |
-
f"127.0.0.1:{manifest.port}:{manifest.port}",
|
| 102 |
-
"--workdir",
|
| 103 |
-
container_home,
|
| 104 |
-
"--env",
|
| 105 |
-
f"HOME={container_home}",
|
| 106 |
-
"--env",
|
| 107 |
-
"PYTHONUNBUFFERED=1",
|
| 108 |
-
# Canonical Headroom filesystem contract (issue #175).
|
| 109 |
-
"--env",
|
| 110 |
-
f"HEADROOM_WORKSPACE_DIR={container_home}/.headroom",
|
| 111 |
-
"--env",
|
| 112 |
-
f"HEADROOM_CONFIG_DIR={container_home}/.headroom/config",
|
| 113 |
-
"--volume",
|
| 114 |
-
f"{_mount_source(home, '.headroom')}:{container_home}/.headroom",
|
| 115 |
-
"--volume",
|
| 116 |
-
f"{_mount_source(home, '.claude')}:{container_home}/.claude",
|
| 117 |
-
"--volume",
|
| 118 |
-
f"{_mount_source(home, '.codex')}:{container_home}/.codex",
|
| 119 |
-
"--volume",
|
| 120 |
-
f"{_mount_source(home, '.gemini')}:{container_home}/.gemini",
|
| 121 |
-
]
|
| 122 |
-
if not _is_windows():
|
| 123 |
-
getuid = getattr(os, "getuid", None)
|
| 124 |
-
getgid = getattr(os, "getgid", None)
|
| 125 |
-
if callable(getuid) and callable(getgid):
|
| 126 |
-
command.extend(["--user", f"{getuid()}:{getgid()}"])
|
| 127 |
-
runtime_env = {**manifest.base_env, **_deployment_env(manifest)}
|
| 128 |
-
for name, value in runtime_env.items():
|
| 129 |
-
command.extend(["--env", f"{name}={value}"])
|
| 130 |
-
for name in sorted(os.environ):
|
| 131 |
-
if name.startswith(PASSTHROUGH_ENV_PREFIXES):
|
| 132 |
-
command.extend(["--env", name])
|
| 133 |
-
command.extend(
|
| 134 |
-
[
|
| 135 |
-
manifest.image,
|
| 136 |
-
"headroom",
|
| 137 |
-
"proxy",
|
| 138 |
-
"--host",
|
| 139 |
-
"0.0.0.0",
|
| 140 |
-
*manifest.proxy_args[2:],
|
| 141 |
-
]
|
| 142 |
-
)
|
| 143 |
-
return command
|
| 144 |
-
|
| 145 |
-
|
| 146 |
-
def _write_pid(profile: str, pid: int) -> None:
|
| 147 |
-
path = pid_path(profile)
|
| 148 |
-
path.parent.mkdir(parents=True, exist_ok=True)
|
| 149 |
-
path.write_text(str(pid))
|
| 150 |
-
|
| 151 |
-
|
| 152 |
-
def _read_pid(profile: str) -> int | None:
|
| 153 |
-
path = pid_path(profile)
|
| 154 |
-
if not path.exists():
|
| 155 |
-
return None
|
| 156 |
-
try:
|
| 157 |
-
return int(path.read_text().strip())
|
| 158 |
-
except ValueError:
|
| 159 |
-
return None
|
| 160 |
-
|
| 161 |
-
|
| 162 |
-
def _clear_pid(profile: str) -> None:
|
| 163 |
-
path = pid_path(profile)
|
| 164 |
-
if path.exists():
|
| 165 |
-
path.unlink()
|
| 166 |
-
|
| 167 |
-
|
| 168 |
-
def run_foreground(manifest: DeploymentManifest) -> int:
|
| 169 |
-
"""Run the raw runtime command in the foreground."""
|
| 170 |
-
|
| 171 |
-
command = build_runtime_command(manifest)
|
| 172 |
-
env = _runtime_env(manifest)
|
| 173 |
-
log_file_path = log_path(manifest.profile)
|
| 174 |
-
log_file_path.parent.mkdir(parents=True, exist_ok=True)
|
| 175 |
-
|
| 176 |
-
with open(log_file_path, "a", encoding="utf-8", errors="replace") as log_file:
|
| 177 |
-
proc = subprocess.Popen(command, env=env, stdout=log_file, stderr=log_file)
|
| 178 |
-
_write_pid(manifest.profile, proc.pid)
|
| 179 |
-
|
| 180 |
-
def _cleanup(signum: int | None = None, frame: Any = None) -> None:
|
| 181 |
-
if proc.poll() is None:
|
| 182 |
-
proc.terminate()
|
| 183 |
-
try:
|
| 184 |
-
proc.wait(timeout=10)
|
| 185 |
-
except subprocess.TimeoutExpired:
|
| 186 |
-
proc.kill()
|
| 187 |
-
|
| 188 |
-
signal.signal(signal.SIGINT, _cleanup)
|
| 189 |
-
signal.signal(signal.SIGTERM, _cleanup)
|
| 190 |
-
try:
|
| 191 |
-
return proc.wait()
|
| 192 |
-
finally:
|
| 193 |
-
_clear_pid(manifest.profile)
|
| 194 |
-
|
| 195 |
-
|
| 196 |
-
def start_detached_agent(profile: str) -> subprocess.Popen[str]:
|
| 197 |
-
"""Start `headroom install agent run` detached for the given profile."""
|
| 198 |
-
|
| 199 |
-
command = [*resolve_headroom_command(), "install", "agent", "run", "--profile", profile]
|
| 200 |
-
log_file_path = log_path(profile)
|
| 201 |
-
log_file_path.parent.mkdir(parents=True, exist_ok=True)
|
| 202 |
-
log_file = open(log_file_path, "a", encoding="utf-8", errors="replace") # noqa: SIM115
|
| 203 |
-
|
| 204 |
-
kwargs: dict[str, Any] = {"stdout": log_file, "stderr": log_file}
|
| 205 |
-
if _is_windows():
|
| 206 |
-
kwargs["creationflags"] = getattr(subprocess, "DETACHED_PROCESS", 0) | getattr(
|
| 207 |
-
subprocess, "CREATE_NEW_PROCESS_GROUP", 0
|
| 208 |
-
)
|
| 209 |
-
else:
|
| 210 |
-
kwargs["start_new_session"] = True
|
| 211 |
-
return subprocess.Popen(command, **kwargs)
|
| 212 |
-
|
| 213 |
-
|
| 214 |
-
def start_persistent_docker(manifest: DeploymentManifest) -> None:
|
| 215 |
-
"""Start a persistent Docker container with restart policy."""
|
| 216 |
-
|
| 217 |
-
command = build_runtime_command(manifest)
|
| 218 |
-
docker_cmd = [
|
| 219 |
-
"docker",
|
| 220 |
-
"run",
|
| 221 |
-
"-d",
|
| 222 |
-
"--restart",
|
| 223 |
-
"unless-stopped",
|
| 224 |
-
"--name",
|
| 225 |
-
manifest.container_name,
|
| 226 |
-
*command[5:], # drop initial `docker run --rm --name ...`
|
| 227 |
-
]
|
| 228 |
-
subprocess.run(["docker", "rm", "-f", manifest.container_name], capture_output=True, text=True)
|
| 229 |
-
subprocess.run(docker_cmd, check=True)
|
| 230 |
-
|
| 231 |
-
|
| 232 |
-
def stop_runtime(manifest: DeploymentManifest) -> None:
|
| 233 |
-
"""Stop the raw runtime for the deployment."""
|
| 234 |
-
|
| 235 |
-
if manifest.preset == InstallPreset.PERSISTENT_DOCKER.value:
|
| 236 |
-
subprocess.run(["docker", "stop", manifest.container_name], capture_output=True, text=True)
|
| 237 |
-
subprocess.run(
|
| 238 |
-
["docker", "rm", "-f", manifest.container_name], capture_output=True, text=True
|
| 239 |
-
)
|
| 240 |
-
return
|
| 241 |
-
|
| 242 |
-
pid = _read_pid(manifest.profile)
|
| 243 |
-
if pid is None:
|
| 244 |
-
return
|
| 245 |
-
try:
|
| 246 |
-
os.kill(pid, signal.SIGTERM)
|
| 247 |
-
except OSError:
|
| 248 |
-
pass
|
| 249 |
-
_clear_pid(manifest.profile)
|
| 250 |
-
|
| 251 |
-
|
| 252 |
-
def wait_ready(manifest: DeploymentManifest, timeout_seconds: int = 30) -> bool:
|
| 253 |
-
"""Wait for the deployment to report ready."""
|
| 254 |
-
|
| 255 |
-
for _ in range(timeout_seconds):
|
| 256 |
-
if probe_ready(manifest.health_url):
|
| 257 |
-
return True
|
| 258 |
-
time.sleep(1)
|
| 259 |
-
return False
|
| 260 |
-
|
| 261 |
-
|
| 262 |
-
def runtime_status(manifest: DeploymentManifest) -> str:
|
| 263 |
-
"""Return a short status string for the deployment runtime."""
|
| 264 |
-
|
| 265 |
-
if manifest.preset == InstallPreset.PERSISTENT_DOCKER.value:
|
| 266 |
-
result = subprocess.run(
|
| 267 |
-
["docker", "ps", "--format", "{{.Names}}"], capture_output=True, text=True
|
| 268 |
-
)
|
| 269 |
-
if manifest.container_name in result.stdout.splitlines():
|
| 270 |
-
return "running"
|
| 271 |
-
return "stopped"
|
| 272 |
-
pid = _read_pid(manifest.profile)
|
| 273 |
-
if pid is None:
|
| 274 |
-
return "stopped"
|
| 275 |
-
try:
|
| 276 |
-
os.kill(pid, 0)
|
| 277 |
-
except OSError:
|
| 278 |
-
return "stopped"
|
| 279 |
-
return "running"
|
|
|
|
| 1 |
+
"""Runtime helpers for persistent deployments."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import os
|
| 6 |
+
import shutil
|
| 7 |
+
import signal
|
| 8 |
+
import subprocess
|
| 9 |
+
import sys
|
| 10 |
+
import time
|
| 11 |
+
from pathlib import Path
|
| 12 |
+
from typing import Any
|
| 13 |
+
|
| 14 |
+
from .health import probe_ready
|
| 15 |
+
from .models import DeploymentManifest, InstallPreset, RuntimeKind
|
| 16 |
+
from .paths import log_path, pid_path
|
| 17 |
+
|
| 18 |
+
PASSTHROUGH_ENV_PREFIXES = (
|
| 19 |
+
"HEADROOM_",
|
| 20 |
+
"ANTHROPIC_",
|
| 21 |
+
"OPENAI_",
|
| 22 |
+
"GEMINI_",
|
| 23 |
+
"AWS_",
|
| 24 |
+
"AZURE_",
|
| 25 |
+
"VERTEX_",
|
| 26 |
+
"GOOGLE_",
|
| 27 |
+
"GOOGLE_CLOUD_",
|
| 28 |
+
"MISTRAL_",
|
| 29 |
+
"GROQ_",
|
| 30 |
+
"OPENROUTER_",
|
| 31 |
+
"XAI_",
|
| 32 |
+
"TOGETHER_",
|
| 33 |
+
"COHERE_",
|
| 34 |
+
"OLLAMA_",
|
| 35 |
+
"LITELLM_",
|
| 36 |
+
"OTEL_",
|
| 37 |
+
"SUPABASE_",
|
| 38 |
+
"QDRANT_",
|
| 39 |
+
"NEO4J_",
|
| 40 |
+
"LANGSMITH_",
|
| 41 |
+
)
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def _is_windows() -> bool:
|
| 45 |
+
return sys.platform.startswith("win")
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def _deployment_env(manifest: DeploymentManifest) -> dict[str, str]:
|
| 49 |
+
return {
|
| 50 |
+
"HEADROOM_DEPLOYMENT_PROFILE": manifest.profile,
|
| 51 |
+
"HEADROOM_DEPLOYMENT_PRESET": manifest.preset,
|
| 52 |
+
"HEADROOM_DEPLOYMENT_RUNTIME": manifest.runtime_kind,
|
| 53 |
+
"HEADROOM_DEPLOYMENT_SUPERVISOR": manifest.supervisor_kind,
|
| 54 |
+
"HEADROOM_DEPLOYMENT_SCOPE": manifest.scope,
|
| 55 |
+
}
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
def resolve_headroom_command() -> list[str]:
|
| 59 |
+
"""Resolve the most reliable command to invoke headroom."""
|
| 60 |
+
|
| 61 |
+
headroom_bin = shutil.which("headroom")
|
| 62 |
+
if headroom_bin:
|
| 63 |
+
return [headroom_bin]
|
| 64 |
+
return [sys.executable, "-m", "headroom.cli"]
|
| 65 |
+
|
| 66 |
+
|
| 67 |
+
def _runtime_env(manifest: DeploymentManifest) -> dict[str, str]:
|
| 68 |
+
env = os.environ.copy()
|
| 69 |
+
env.update(manifest.base_env)
|
| 70 |
+
env.update(_deployment_env(manifest))
|
| 71 |
+
return env
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def _ensure_host_dirs() -> None:
|
| 75 |
+
for subdir in (".headroom", ".claude", ".codex", ".gemini"):
|
| 76 |
+
(Path.home() / subdir).mkdir(parents=True, exist_ok=True)
|
| 77 |
+
|
| 78 |
+
|
| 79 |
+
def _mount_source(home: str, subdir: str) -> str:
|
| 80 |
+
if _is_windows():
|
| 81 |
+
return f"{home}\\{subdir}"
|
| 82 |
+
return f"{home}/{subdir}"
|
| 83 |
+
|
| 84 |
+
|
| 85 |
+
def build_runtime_command(manifest: DeploymentManifest) -> list[str]:
|
| 86 |
+
"""Build the raw foreground command that runs the proxy."""
|
| 87 |
+
|
| 88 |
+
if manifest.runtime_kind == RuntimeKind.PYTHON.value:
|
| 89 |
+
return [sys.executable, "-m", "headroom.cli", "proxy", *manifest.proxy_args]
|
| 90 |
+
|
| 91 |
+
_ensure_host_dirs()
|
| 92 |
+
home = str(Path.home())
|
| 93 |
+
container_home = "/tmp/headroom-home"
|
| 94 |
+
command = [
|
| 95 |
+
"docker",
|
| 96 |
+
"run",
|
| 97 |
+
"--rm",
|
| 98 |
+
"--name",
|
| 99 |
+
manifest.container_name,
|
| 100 |
+
"-p",
|
| 101 |
+
f"127.0.0.1:{manifest.port}:{manifest.port}",
|
| 102 |
+
"--workdir",
|
| 103 |
+
container_home,
|
| 104 |
+
"--env",
|
| 105 |
+
f"HOME={container_home}",
|
| 106 |
+
"--env",
|
| 107 |
+
"PYTHONUNBUFFERED=1",
|
| 108 |
+
# Canonical Headroom filesystem contract (issue #175).
|
| 109 |
+
"--env",
|
| 110 |
+
f"HEADROOM_WORKSPACE_DIR={container_home}/.headroom",
|
| 111 |
+
"--env",
|
| 112 |
+
f"HEADROOM_CONFIG_DIR={container_home}/.headroom/config",
|
| 113 |
+
"--volume",
|
| 114 |
+
f"{_mount_source(home, '.headroom')}:{container_home}/.headroom",
|
| 115 |
+
"--volume",
|
| 116 |
+
f"{_mount_source(home, '.claude')}:{container_home}/.claude",
|
| 117 |
+
"--volume",
|
| 118 |
+
f"{_mount_source(home, '.codex')}:{container_home}/.codex",
|
| 119 |
+
"--volume",
|
| 120 |
+
f"{_mount_source(home, '.gemini')}:{container_home}/.gemini",
|
| 121 |
+
]
|
| 122 |
+
if not _is_windows():
|
| 123 |
+
getuid = getattr(os, "getuid", None)
|
| 124 |
+
getgid = getattr(os, "getgid", None)
|
| 125 |
+
if callable(getuid) and callable(getgid):
|
| 126 |
+
command.extend(["--user", f"{getuid()}:{getgid()}"])
|
| 127 |
+
runtime_env = {**manifest.base_env, **_deployment_env(manifest)}
|
| 128 |
+
for name, value in runtime_env.items():
|
| 129 |
+
command.extend(["--env", f"{name}={value}"])
|
| 130 |
+
for name in sorted(os.environ):
|
| 131 |
+
if name.startswith(PASSTHROUGH_ENV_PREFIXES):
|
| 132 |
+
command.extend(["--env", name])
|
| 133 |
+
command.extend(
|
| 134 |
+
[
|
| 135 |
+
manifest.image,
|
| 136 |
+
"headroom",
|
| 137 |
+
"proxy",
|
| 138 |
+
"--host",
|
| 139 |
+
"0.0.0.0",
|
| 140 |
+
*manifest.proxy_args[2:],
|
| 141 |
+
]
|
| 142 |
+
)
|
| 143 |
+
return command
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def _write_pid(profile: str, pid: int) -> None:
|
| 147 |
+
path = pid_path(profile)
|
| 148 |
+
path.parent.mkdir(parents=True, exist_ok=True)
|
| 149 |
+
path.write_text(str(pid))
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
def _read_pid(profile: str) -> int | None:
|
| 153 |
+
path = pid_path(profile)
|
| 154 |
+
if not path.exists():
|
| 155 |
+
return None
|
| 156 |
+
try:
|
| 157 |
+
return int(path.read_text().strip())
|
| 158 |
+
except ValueError:
|
| 159 |
+
return None
|
| 160 |
+
|
| 161 |
+
|
| 162 |
+
def _clear_pid(profile: str) -> None:
|
| 163 |
+
path = pid_path(profile)
|
| 164 |
+
if path.exists():
|
| 165 |
+
path.unlink()
|
| 166 |
+
|
| 167 |
+
|
| 168 |
+
def run_foreground(manifest: DeploymentManifest) -> int:
|
| 169 |
+
"""Run the raw runtime command in the foreground."""
|
| 170 |
+
|
| 171 |
+
command = build_runtime_command(manifest)
|
| 172 |
+
env = _runtime_env(manifest)
|
| 173 |
+
log_file_path = log_path(manifest.profile)
|
| 174 |
+
log_file_path.parent.mkdir(parents=True, exist_ok=True)
|
| 175 |
+
|
| 176 |
+
with open(log_file_path, "a", encoding="utf-8", errors="replace") as log_file:
|
| 177 |
+
proc = subprocess.Popen(command, env=env, stdout=log_file, stderr=log_file)
|
| 178 |
+
_write_pid(manifest.profile, proc.pid)
|
| 179 |
+
|
| 180 |
+
def _cleanup(signum: int | None = None, frame: Any = None) -> None:
|
| 181 |
+
if proc.poll() is None:
|
| 182 |
+
proc.terminate()
|
| 183 |
+
try:
|
| 184 |
+
proc.wait(timeout=10)
|
| 185 |
+
except subprocess.TimeoutExpired:
|
| 186 |
+
proc.kill()
|
| 187 |
+
|
| 188 |
+
signal.signal(signal.SIGINT, _cleanup)
|
| 189 |
+
signal.signal(signal.SIGTERM, _cleanup)
|
| 190 |
+
try:
|
| 191 |
+
return proc.wait()
|
| 192 |
+
finally:
|
| 193 |
+
_clear_pid(manifest.profile)
|
| 194 |
+
|
| 195 |
+
|
| 196 |
+
def start_detached_agent(profile: str) -> subprocess.Popen[str]:
|
| 197 |
+
"""Start `headroom install agent run` detached for the given profile."""
|
| 198 |
+
|
| 199 |
+
command = [*resolve_headroom_command(), "install", "agent", "run", "--profile", profile]
|
| 200 |
+
log_file_path = log_path(profile)
|
| 201 |
+
log_file_path.parent.mkdir(parents=True, exist_ok=True)
|
| 202 |
+
log_file = open(log_file_path, "a", encoding="utf-8", errors="replace") # noqa: SIM115
|
| 203 |
+
|
| 204 |
+
kwargs: dict[str, Any] = {"stdout": log_file, "stderr": log_file}
|
| 205 |
+
if _is_windows():
|
| 206 |
+
kwargs["creationflags"] = getattr(subprocess, "DETACHED_PROCESS", 0) | getattr(
|
| 207 |
+
subprocess, "CREATE_NEW_PROCESS_GROUP", 0
|
| 208 |
+
)
|
| 209 |
+
else:
|
| 210 |
+
kwargs["start_new_session"] = True
|
| 211 |
+
return subprocess.Popen(command, **kwargs)
|
| 212 |
+
|
| 213 |
+
|
| 214 |
+
def start_persistent_docker(manifest: DeploymentManifest) -> None:
|
| 215 |
+
"""Start a persistent Docker container with restart policy."""
|
| 216 |
+
|
| 217 |
+
command = build_runtime_command(manifest)
|
| 218 |
+
docker_cmd = [
|
| 219 |
+
"docker",
|
| 220 |
+
"run",
|
| 221 |
+
"-d",
|
| 222 |
+
"--restart",
|
| 223 |
+
"unless-stopped",
|
| 224 |
+
"--name",
|
| 225 |
+
manifest.container_name,
|
| 226 |
+
*command[5:], # drop initial `docker run --rm --name ...`
|
| 227 |
+
]
|
| 228 |
+
subprocess.run(["docker", "rm", "-f", manifest.container_name], capture_output=True, text=True)
|
| 229 |
+
subprocess.run(docker_cmd, check=True)
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
def stop_runtime(manifest: DeploymentManifest) -> None:
|
| 233 |
+
"""Stop the raw runtime for the deployment."""
|
| 234 |
+
|
| 235 |
+
if manifest.preset == InstallPreset.PERSISTENT_DOCKER.value:
|
| 236 |
+
subprocess.run(["docker", "stop", manifest.container_name], capture_output=True, text=True)
|
| 237 |
+
subprocess.run(
|
| 238 |
+
["docker", "rm", "-f", manifest.container_name], capture_output=True, text=True
|
| 239 |
+
)
|
| 240 |
+
return
|
| 241 |
+
|
| 242 |
+
pid = _read_pid(manifest.profile)
|
| 243 |
+
if pid is None:
|
| 244 |
+
return
|
| 245 |
+
try:
|
| 246 |
+
os.kill(pid, signal.SIGTERM)
|
| 247 |
+
except OSError:
|
| 248 |
+
pass
|
| 249 |
+
_clear_pid(manifest.profile)
|
| 250 |
+
|
| 251 |
+
|
| 252 |
+
def wait_ready(manifest: DeploymentManifest, timeout_seconds: int = 30) -> bool:
|
| 253 |
+
"""Wait for the deployment to report ready."""
|
| 254 |
+
|
| 255 |
+
for _ in range(timeout_seconds):
|
| 256 |
+
if probe_ready(manifest.health_url):
|
| 257 |
+
return True
|
| 258 |
+
time.sleep(1)
|
| 259 |
+
return False
|
| 260 |
+
|
| 261 |
+
|
| 262 |
+
def runtime_status(manifest: DeploymentManifest) -> str:
|
| 263 |
+
"""Return a short status string for the deployment runtime."""
|
| 264 |
+
|
| 265 |
+
if manifest.preset == InstallPreset.PERSISTENT_DOCKER.value:
|
| 266 |
+
result = subprocess.run(
|
| 267 |
+
["docker", "ps", "--format", "{{.Names}}"], capture_output=True, text=True
|
| 268 |
+
)
|
| 269 |
+
if manifest.container_name in result.stdout.splitlines():
|
| 270 |
+
return "running"
|
| 271 |
+
return "stopped"
|
| 272 |
+
pid = _read_pid(manifest.profile)
|
| 273 |
+
if pid is None:
|
| 274 |
+
return "stopped"
|
| 275 |
+
try:
|
| 276 |
+
os.kill(pid, 0)
|
| 277 |
+
except OSError:
|
| 278 |
+
return "stopped"
|
| 279 |
+
return "running"
|
|
@@ -1,12 +1,12 @@
|
|
| 1 |
-
"""Aider install-time helpers."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
from .runtime import build_launch_env
|
| 6 |
-
|
| 7 |
-
|
| 8 |
-
def build_install_env(*, port: int, backend: str) -> dict[str, str]:
|
| 9 |
-
"""Build the persistent install environment for Aider."""
|
| 10 |
-
del backend
|
| 11 |
-
env, _lines = build_launch_env(port=port, environ={})
|
| 12 |
-
return {key: env[key] for key in ("OPENAI_API_BASE", "ANTHROPIC_BASE_URL")}
|
|
|
|
| 1 |
+
"""Aider install-time helpers."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from .runtime import build_launch_env
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def build_install_env(*, port: int, backend: str) -> dict[str, str]:
|
| 9 |
+
"""Build the persistent install environment for Aider."""
|
| 10 |
+
del backend
|
| 11 |
+
env, _lines = build_launch_env(port=port, environ={})
|
| 12 |
+
return {key: env[key] for key in ("OPENAI_API_BASE", "ANTHROPIC_BASE_URL")}
|
|
@@ -1,63 +1,63 @@
|
|
| 1 |
-
"""Claude install-time helpers."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import json
|
| 6 |
-
from pathlib import Path
|
| 7 |
-
|
| 8 |
-
from headroom.install.models import ConfigScope, DeploymentManifest, ManagedMutation, ToolTarget
|
| 9 |
-
from headroom.install.paths import claude_settings_path
|
| 10 |
-
|
| 11 |
-
from .runtime import proxy_base_url
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
def build_install_env(*, port: int, backend: str) -> dict[str, str]:
|
| 15 |
-
"""Build the persistent install environment for Claude."""
|
| 16 |
-
del backend
|
| 17 |
-
return {"ANTHROPIC_BASE_URL": proxy_base_url(port)}
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
def apply_provider_scope(manifest: DeploymentManifest) -> ManagedMutation | None:
|
| 21 |
-
"""Apply Claude provider-scope configuration when requested."""
|
| 22 |
-
if manifest.scope != ConfigScope.PROVIDER.value:
|
| 23 |
-
return None
|
| 24 |
-
|
| 25 |
-
path = claude_settings_path()
|
| 26 |
-
path.parent.mkdir(parents=True, exist_ok=True)
|
| 27 |
-
payload: dict[str, object] = {}
|
| 28 |
-
if path.exists():
|
| 29 |
-
payload = json.loads(path.read_text())
|
| 30 |
-
env = payload.get("env")
|
| 31 |
-
env_map = dict(env) if isinstance(env, dict) else {}
|
| 32 |
-
values = manifest.tool_envs.get(ToolTarget.CLAUDE.value, {})
|
| 33 |
-
previous = {name: env_map.get(name) for name in values}
|
| 34 |
-
env_map.update(values)
|
| 35 |
-
payload["env"] = env_map
|
| 36 |
-
path.write_text(json.dumps(payload, indent=2) + "\n")
|
| 37 |
-
return ManagedMutation(
|
| 38 |
-
target=ToolTarget.CLAUDE.value,
|
| 39 |
-
kind="json-env",
|
| 40 |
-
path=str(path),
|
| 41 |
-
data={"previous": previous},
|
| 42 |
-
)
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
def revert_provider_scope(mutation: ManagedMutation, manifest: DeploymentManifest) -> None:
|
| 46 |
-
"""Revert Claude provider-scope configuration."""
|
| 47 |
-
if not mutation.path:
|
| 48 |
-
return
|
| 49 |
-
path = Path(mutation.path)
|
| 50 |
-
if not path.exists():
|
| 51 |
-
return
|
| 52 |
-
payload = json.loads(path.read_text())
|
| 53 |
-
env = payload.get("env")
|
| 54 |
-
env_map = dict(env) if isinstance(env, dict) else {}
|
| 55 |
-
previous: dict[str, object] = mutation.data.get("previous", {})
|
| 56 |
-
values = manifest.tool_envs.get(ToolTarget.CLAUDE.value, {})
|
| 57 |
-
for name in values:
|
| 58 |
-
if previous.get(name) is None:
|
| 59 |
-
env_map.pop(name, None)
|
| 60 |
-
else:
|
| 61 |
-
env_map[name] = previous[name]
|
| 62 |
-
payload["env"] = env_map
|
| 63 |
-
path.write_text(json.dumps(payload, indent=2) + "\n")
|
|
|
|
| 1 |
+
"""Claude install-time helpers."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import json
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
from headroom.install.models import ConfigScope, DeploymentManifest, ManagedMutation, ToolTarget
|
| 9 |
+
from headroom.install.paths import claude_settings_path
|
| 10 |
+
|
| 11 |
+
from .runtime import proxy_base_url
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def build_install_env(*, port: int, backend: str) -> dict[str, str]:
|
| 15 |
+
"""Build the persistent install environment for Claude."""
|
| 16 |
+
del backend
|
| 17 |
+
return {"ANTHROPIC_BASE_URL": proxy_base_url(port)}
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
def apply_provider_scope(manifest: DeploymentManifest) -> ManagedMutation | None:
|
| 21 |
+
"""Apply Claude provider-scope configuration when requested."""
|
| 22 |
+
if manifest.scope != ConfigScope.PROVIDER.value:
|
| 23 |
+
return None
|
| 24 |
+
|
| 25 |
+
path = claude_settings_path()
|
| 26 |
+
path.parent.mkdir(parents=True, exist_ok=True)
|
| 27 |
+
payload: dict[str, object] = {}
|
| 28 |
+
if path.exists():
|
| 29 |
+
payload = json.loads(path.read_text())
|
| 30 |
+
env = payload.get("env")
|
| 31 |
+
env_map = dict(env) if isinstance(env, dict) else {}
|
| 32 |
+
values = manifest.tool_envs.get(ToolTarget.CLAUDE.value, {})
|
| 33 |
+
previous = {name: env_map.get(name) for name in values}
|
| 34 |
+
env_map.update(values)
|
| 35 |
+
payload["env"] = env_map
|
| 36 |
+
path.write_text(json.dumps(payload, indent=2) + "\n")
|
| 37 |
+
return ManagedMutation(
|
| 38 |
+
target=ToolTarget.CLAUDE.value,
|
| 39 |
+
kind="json-env",
|
| 40 |
+
path=str(path),
|
| 41 |
+
data={"previous": previous},
|
| 42 |
+
)
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def revert_provider_scope(mutation: ManagedMutation, manifest: DeploymentManifest) -> None:
|
| 46 |
+
"""Revert Claude provider-scope configuration."""
|
| 47 |
+
if not mutation.path:
|
| 48 |
+
return
|
| 49 |
+
path = Path(mutation.path)
|
| 50 |
+
if not path.exists():
|
| 51 |
+
return
|
| 52 |
+
payload = json.loads(path.read_text())
|
| 53 |
+
env = payload.get("env")
|
| 54 |
+
env_map = dict(env) if isinstance(env, dict) else {}
|
| 55 |
+
previous: dict[str, object] = mutation.data.get("previous", {})
|
| 56 |
+
values = manifest.tool_envs.get(ToolTarget.CLAUDE.value, {})
|
| 57 |
+
for name in values:
|
| 58 |
+
if previous.get(name) is None:
|
| 59 |
+
env_map.pop(name, None)
|
| 60 |
+
else:
|
| 61 |
+
env_map[name] = previous[name]
|
| 62 |
+
payload["env"] = env_map
|
| 63 |
+
path.write_text(json.dumps(payload, indent=2) + "\n")
|
|
@@ -1,68 +1,68 @@
|
|
| 1 |
-
"""Codex install-time helpers."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import re
|
| 6 |
-
from pathlib import Path
|
| 7 |
-
|
| 8 |
-
from headroom.install.models import ConfigScope, DeploymentManifest, ManagedMutation, ToolTarget
|
| 9 |
-
from headroom.install.paths import codex_config_path
|
| 10 |
-
|
| 11 |
-
from .runtime import proxy_base_url
|
| 12 |
-
|
| 13 |
-
_CODEX_MARKER_START = "# --- Headroom persistent provider ---"
|
| 14 |
-
_CODEX_MARKER_END = "# --- end Headroom persistent provider ---"
|
| 15 |
-
_CODEX_PATTERN = re.compile(
|
| 16 |
-
re.escape(_CODEX_MARKER_START) + r".*?" + re.escape(_CODEX_MARKER_END),
|
| 17 |
-
re.DOTALL,
|
| 18 |
-
)
|
| 19 |
-
|
| 20 |
-
|
| 21 |
-
def build_install_env(*, port: int, backend: str) -> dict[str, str]:
|
| 22 |
-
"""Build the persistent install environment for Codex."""
|
| 23 |
-
del backend
|
| 24 |
-
return {"OPENAI_BASE_URL": proxy_base_url(port)}
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
def apply_provider_scope(manifest: DeploymentManifest) -> ManagedMutation | None:
|
| 28 |
-
"""Apply Codex provider-scope configuration when requested."""
|
| 29 |
-
if manifest.scope != ConfigScope.PROVIDER.value:
|
| 30 |
-
return None
|
| 31 |
-
|
| 32 |
-
path = codex_config_path()
|
| 33 |
-
path.parent.mkdir(parents=True, exist_ok=True)
|
| 34 |
-
section = (
|
| 35 |
-
f"{_CODEX_MARKER_START}\n"
|
| 36 |
-
'model_provider = "headroom"\n\n'
|
| 37 |
-
"[model_providers.headroom]\n"
|
| 38 |
-
'name = "Headroom persistent proxy"\n'
|
| 39 |
-
f'base_url = "{proxy_base_url(manifest.port)}"\n'
|
| 40 |
-
'env_key = "OPENAI_API_KEY"\n'
|
| 41 |
-
"requires_openai_auth = true\n"
|
| 42 |
-
"supports_websockets = true\n"
|
| 43 |
-
f"{_CODEX_MARKER_END}\n"
|
| 44 |
-
)
|
| 45 |
-
if path.exists():
|
| 46 |
-
existing = path.read_text()
|
| 47 |
-
if _CODEX_MARKER_START in existing:
|
| 48 |
-
merged = _CODEX_PATTERN.sub(section, existing)
|
| 49 |
-
else:
|
| 50 |
-
merged = existing.rstrip() + "\n\n" + section + "\n"
|
| 51 |
-
else:
|
| 52 |
-
merged = section + "\n"
|
| 53 |
-
path.write_text(merged)
|
| 54 |
-
return ManagedMutation(target=ToolTarget.CODEX.value, kind="toml-block", path=str(path))
|
| 55 |
-
|
| 56 |
-
|
| 57 |
-
def revert_provider_scope(mutation: ManagedMutation, manifest: DeploymentManifest) -> None:
|
| 58 |
-
"""Revert Codex provider-scope configuration."""
|
| 59 |
-
del manifest
|
| 60 |
-
if not mutation.path:
|
| 61 |
-
return
|
| 62 |
-
path = Path(mutation.path)
|
| 63 |
-
if not path.exists():
|
| 64 |
-
return
|
| 65 |
-
content = path.read_text()
|
| 66 |
-
if _CODEX_MARKER_START not in content:
|
| 67 |
-
return
|
| 68 |
-
path.write_text(_CODEX_PATTERN.sub("", content).strip() + "\n")
|
|
|
|
| 1 |
+
"""Codex install-time helpers."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import re
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
from headroom.install.models import ConfigScope, DeploymentManifest, ManagedMutation, ToolTarget
|
| 9 |
+
from headroom.install.paths import codex_config_path
|
| 10 |
+
|
| 11 |
+
from .runtime import proxy_base_url
|
| 12 |
+
|
| 13 |
+
_CODEX_MARKER_START = "# --- Headroom persistent provider ---"
|
| 14 |
+
_CODEX_MARKER_END = "# --- end Headroom persistent provider ---"
|
| 15 |
+
_CODEX_PATTERN = re.compile(
|
| 16 |
+
re.escape(_CODEX_MARKER_START) + r".*?" + re.escape(_CODEX_MARKER_END),
|
| 17 |
+
re.DOTALL,
|
| 18 |
+
)
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
def build_install_env(*, port: int, backend: str) -> dict[str, str]:
|
| 22 |
+
"""Build the persistent install environment for Codex."""
|
| 23 |
+
del backend
|
| 24 |
+
return {"OPENAI_BASE_URL": proxy_base_url(port)}
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def apply_provider_scope(manifest: DeploymentManifest) -> ManagedMutation | None:
|
| 28 |
+
"""Apply Codex provider-scope configuration when requested."""
|
| 29 |
+
if manifest.scope != ConfigScope.PROVIDER.value:
|
| 30 |
+
return None
|
| 31 |
+
|
| 32 |
+
path = codex_config_path()
|
| 33 |
+
path.parent.mkdir(parents=True, exist_ok=True)
|
| 34 |
+
section = (
|
| 35 |
+
f"{_CODEX_MARKER_START}\n"
|
| 36 |
+
'model_provider = "headroom"\n\n'
|
| 37 |
+
"[model_providers.headroom]\n"
|
| 38 |
+
'name = "Headroom persistent proxy"\n'
|
| 39 |
+
f'base_url = "{proxy_base_url(manifest.port)}"\n'
|
| 40 |
+
'env_key = "OPENAI_API_KEY"\n'
|
| 41 |
+
"requires_openai_auth = true\n"
|
| 42 |
+
"supports_websockets = true\n"
|
| 43 |
+
f"{_CODEX_MARKER_END}\n"
|
| 44 |
+
)
|
| 45 |
+
if path.exists():
|
| 46 |
+
existing = path.read_text()
|
| 47 |
+
if _CODEX_MARKER_START in existing:
|
| 48 |
+
merged = _CODEX_PATTERN.sub(section, existing)
|
| 49 |
+
else:
|
| 50 |
+
merged = existing.rstrip() + "\n\n" + section + "\n"
|
| 51 |
+
else:
|
| 52 |
+
merged = section + "\n"
|
| 53 |
+
path.write_text(merged)
|
| 54 |
+
return ManagedMutation(target=ToolTarget.CODEX.value, kind="toml-block", path=str(path))
|
| 55 |
+
|
| 56 |
+
|
| 57 |
+
def revert_provider_scope(mutation: ManagedMutation, manifest: DeploymentManifest) -> None:
|
| 58 |
+
"""Revert Codex provider-scope configuration."""
|
| 59 |
+
del manifest
|
| 60 |
+
if not mutation.path:
|
| 61 |
+
return
|
| 62 |
+
path = Path(mutation.path)
|
| 63 |
+
if not path.exists():
|
| 64 |
+
return
|
| 65 |
+
content = path.read_text()
|
| 66 |
+
if _CODEX_MARKER_START not in content:
|
| 67 |
+
return
|
| 68 |
+
path.write_text(_CODEX_PATTERN.sub("", content).strip() + "\n")
|
|
@@ -1,25 +1,25 @@
|
|
| 1 |
-
"""Copilot install-time helpers."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
from .wrap import build_launch_env, resolve_provider_type
|
| 6 |
-
|
| 7 |
-
|
| 8 |
-
def build_install_env(*, port: int, backend: str) -> dict[str, str]:
|
| 9 |
-
"""Build the persistent install environment for Copilot."""
|
| 10 |
-
provider_type = resolve_provider_type(backend, "auto", {"HEADROOM_BACKEND": backend})
|
| 11 |
-
env, _lines = build_launch_env(
|
| 12 |
-
port=port,
|
| 13 |
-
provider_type=provider_type,
|
| 14 |
-
wire_api=None,
|
| 15 |
-
environ={},
|
| 16 |
-
)
|
| 17 |
-
return {
|
| 18 |
-
key: env[key]
|
| 19 |
-
for key in (
|
| 20 |
-
"COPILOT_PROVIDER_TYPE",
|
| 21 |
-
"COPILOT_PROVIDER_BASE_URL",
|
| 22 |
-
"COPILOT_PROVIDER_WIRE_API",
|
| 23 |
-
)
|
| 24 |
-
if key in env
|
| 25 |
-
}
|
|
|
|
| 1 |
+
"""Copilot install-time helpers."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from .wrap import build_launch_env, resolve_provider_type
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def build_install_env(*, port: int, backend: str) -> dict[str, str]:
|
| 9 |
+
"""Build the persistent install environment for Copilot."""
|
| 10 |
+
provider_type = resolve_provider_type(backend, "auto", {"HEADROOM_BACKEND": backend})
|
| 11 |
+
env, _lines = build_launch_env(
|
| 12 |
+
port=port,
|
| 13 |
+
provider_type=provider_type,
|
| 14 |
+
wire_api=None,
|
| 15 |
+
environ={},
|
| 16 |
+
)
|
| 17 |
+
return {
|
| 18 |
+
key: env[key]
|
| 19 |
+
for key in (
|
| 20 |
+
"COPILOT_PROVIDER_TYPE",
|
| 21 |
+
"COPILOT_PROVIDER_BASE_URL",
|
| 22 |
+
"COPILOT_PROVIDER_WIRE_API",
|
| 23 |
+
)
|
| 24 |
+
if key in env
|
| 25 |
+
}
|
|
@@ -1,15 +1,15 @@
|
|
| 1 |
-
"""Cursor install-time helpers."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
from .runtime import build_proxy_targets
|
| 6 |
-
|
| 7 |
-
|
| 8 |
-
def build_install_env(*, port: int, backend: str) -> dict[str, str]:
|
| 9 |
-
"""Build the persistent install environment for Cursor."""
|
| 10 |
-
del backend
|
| 11 |
-
targets = build_proxy_targets(port)
|
| 12 |
-
return {
|
| 13 |
-
"OPENAI_BASE_URL": targets.openai_base_url,
|
| 14 |
-
"ANTHROPIC_BASE_URL": targets.anthropic_base_url,
|
| 15 |
-
}
|
|
|
|
| 1 |
+
"""Cursor install-time helpers."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from .runtime import build_proxy_targets
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def build_install_env(*, port: int, backend: str) -> dict[str, str]:
|
| 9 |
+
"""Build the persistent install environment for Cursor."""
|
| 10 |
+
del backend
|
| 11 |
+
targets = build_proxy_targets(port)
|
| 12 |
+
return {
|
| 13 |
+
"OPENAI_BASE_URL": targets.openai_base_url,
|
| 14 |
+
"ANTHROPIC_BASE_URL": targets.anthropic_base_url,
|
| 15 |
+
}
|
|
@@ -1,86 +1,86 @@
|
|
| 1 |
-
"""Install-time provider registry helpers."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
from collections.abc import Callable
|
| 6 |
-
|
| 7 |
-
from headroom.install.models import DeploymentManifest, ManagedMutation
|
| 8 |
-
from headroom.providers.aider.install import build_install_env as _build_aider_install_env
|
| 9 |
-
from headroom.providers.claude.install import (
|
| 10 |
-
apply_provider_scope as _apply_claude_provider_scope,
|
| 11 |
-
)
|
| 12 |
-
from headroom.providers.claude.install import (
|
| 13 |
-
build_install_env as _build_claude_install_env,
|
| 14 |
-
)
|
| 15 |
-
from headroom.providers.claude.install import (
|
| 16 |
-
revert_provider_scope as _revert_claude_provider_scope,
|
| 17 |
-
)
|
| 18 |
-
from headroom.providers.codex.install import (
|
| 19 |
-
apply_provider_scope as _apply_codex_provider_scope,
|
| 20 |
-
)
|
| 21 |
-
from headroom.providers.codex.install import build_install_env as _build_codex_install_env
|
| 22 |
-
from headroom.providers.codex.install import (
|
| 23 |
-
revert_provider_scope as _revert_codex_provider_scope,
|
| 24 |
-
)
|
| 25 |
-
from headroom.providers.copilot.install import (
|
| 26 |
-
build_install_env as _build_copilot_install_env,
|
| 27 |
-
)
|
| 28 |
-
from headroom.providers.cursor.install import build_install_env as _build_cursor_install_env
|
| 29 |
-
from headroom.providers.openclaw.install import (
|
| 30 |
-
apply_provider_scope as _apply_openclaw_provider_scope,
|
| 31 |
-
)
|
| 32 |
-
from headroom.providers.openclaw.install import (
|
| 33 |
-
revert_provider_scope as _revert_openclaw_provider_scope,
|
| 34 |
-
)
|
| 35 |
-
|
| 36 |
-
_InstallEnvBuilder = Callable[..., dict[str, str]]
|
| 37 |
-
_ProviderScopeApplier = Callable[[DeploymentManifest], ManagedMutation | None]
|
| 38 |
-
_ProviderScopeReverter = Callable[[ManagedMutation, DeploymentManifest], None]
|
| 39 |
-
|
| 40 |
-
_ENV_BUILDERS: dict[str, _InstallEnvBuilder] = {
|
| 41 |
-
"claude": _build_claude_install_env,
|
| 42 |
-
"copilot": _build_copilot_install_env,
|
| 43 |
-
"codex": _build_codex_install_env,
|
| 44 |
-
"aider": _build_aider_install_env,
|
| 45 |
-
"cursor": _build_cursor_install_env,
|
| 46 |
-
}
|
| 47 |
-
|
| 48 |
-
_PROVIDER_SCOPE_HANDLERS: dict[str, tuple[_ProviderScopeApplier, _ProviderScopeReverter]] = {
|
| 49 |
-
"claude": (_apply_claude_provider_scope, _revert_claude_provider_scope),
|
| 50 |
-
"codex": (_apply_codex_provider_scope, _revert_codex_provider_scope),
|
| 51 |
-
"openclaw": (_apply_openclaw_provider_scope, _revert_openclaw_provider_scope),
|
| 52 |
-
}
|
| 53 |
-
|
| 54 |
-
|
| 55 |
-
def build_install_target_envs(
|
| 56 |
-
port: int, backend: str, targets: list[str]
|
| 57 |
-
) -> dict[str, dict[str, str]]:
|
| 58 |
-
"""Build per-target install environment values via provider slices."""
|
| 59 |
-
target_envs: dict[str, dict[str, str]] = {}
|
| 60 |
-
for target in targets:
|
| 61 |
-
builder = _ENV_BUILDERS.get(target)
|
| 62 |
-
if builder is None:
|
| 63 |
-
continue
|
| 64 |
-
target_envs[target] = builder(port=port, backend=backend)
|
| 65 |
-
return target_envs
|
| 66 |
-
|
| 67 |
-
|
| 68 |
-
def apply_provider_scope_mutations(manifest: DeploymentManifest) -> list[ManagedMutation]:
|
| 69 |
-
"""Apply provider-scope mutations owned by provider slices."""
|
| 70 |
-
mutations: list[ManagedMutation] = []
|
| 71 |
-
for target in manifest.targets:
|
| 72 |
-
handlers = _PROVIDER_SCOPE_HANDLERS.get(target)
|
| 73 |
-
if handlers is None:
|
| 74 |
-
continue
|
| 75 |
-
mutation = handlers[0](manifest)
|
| 76 |
-
if mutation is not None:
|
| 77 |
-
mutations.append(mutation)
|
| 78 |
-
return mutations
|
| 79 |
-
|
| 80 |
-
|
| 81 |
-
def revert_provider_scope_mutation(manifest: DeploymentManifest, mutation: ManagedMutation) -> None:
|
| 82 |
-
"""Revert a provider-scope mutation via the owning provider slice."""
|
| 83 |
-
handlers = _PROVIDER_SCOPE_HANDLERS.get(mutation.target)
|
| 84 |
-
if handlers is None:
|
| 85 |
-
return
|
| 86 |
-
handlers[1](mutation, manifest)
|
|
|
|
| 1 |
+
"""Install-time provider registry helpers."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
from collections.abc import Callable
|
| 6 |
+
|
| 7 |
+
from headroom.install.models import DeploymentManifest, ManagedMutation
|
| 8 |
+
from headroom.providers.aider.install import build_install_env as _build_aider_install_env
|
| 9 |
+
from headroom.providers.claude.install import (
|
| 10 |
+
apply_provider_scope as _apply_claude_provider_scope,
|
| 11 |
+
)
|
| 12 |
+
from headroom.providers.claude.install import (
|
| 13 |
+
build_install_env as _build_claude_install_env,
|
| 14 |
+
)
|
| 15 |
+
from headroom.providers.claude.install import (
|
| 16 |
+
revert_provider_scope as _revert_claude_provider_scope,
|
| 17 |
+
)
|
| 18 |
+
from headroom.providers.codex.install import (
|
| 19 |
+
apply_provider_scope as _apply_codex_provider_scope,
|
| 20 |
+
)
|
| 21 |
+
from headroom.providers.codex.install import build_install_env as _build_codex_install_env
|
| 22 |
+
from headroom.providers.codex.install import (
|
| 23 |
+
revert_provider_scope as _revert_codex_provider_scope,
|
| 24 |
+
)
|
| 25 |
+
from headroom.providers.copilot.install import (
|
| 26 |
+
build_install_env as _build_copilot_install_env,
|
| 27 |
+
)
|
| 28 |
+
from headroom.providers.cursor.install import build_install_env as _build_cursor_install_env
|
| 29 |
+
from headroom.providers.openclaw.install import (
|
| 30 |
+
apply_provider_scope as _apply_openclaw_provider_scope,
|
| 31 |
+
)
|
| 32 |
+
from headroom.providers.openclaw.install import (
|
| 33 |
+
revert_provider_scope as _revert_openclaw_provider_scope,
|
| 34 |
+
)
|
| 35 |
+
|
| 36 |
+
_InstallEnvBuilder = Callable[..., dict[str, str]]
|
| 37 |
+
_ProviderScopeApplier = Callable[[DeploymentManifest], ManagedMutation | None]
|
| 38 |
+
_ProviderScopeReverter = Callable[[ManagedMutation, DeploymentManifest], None]
|
| 39 |
+
|
| 40 |
+
_ENV_BUILDERS: dict[str, _InstallEnvBuilder] = {
|
| 41 |
+
"claude": _build_claude_install_env,
|
| 42 |
+
"copilot": _build_copilot_install_env,
|
| 43 |
+
"codex": _build_codex_install_env,
|
| 44 |
+
"aider": _build_aider_install_env,
|
| 45 |
+
"cursor": _build_cursor_install_env,
|
| 46 |
+
}
|
| 47 |
+
|
| 48 |
+
_PROVIDER_SCOPE_HANDLERS: dict[str, tuple[_ProviderScopeApplier, _ProviderScopeReverter]] = {
|
| 49 |
+
"claude": (_apply_claude_provider_scope, _revert_claude_provider_scope),
|
| 50 |
+
"codex": (_apply_codex_provider_scope, _revert_codex_provider_scope),
|
| 51 |
+
"openclaw": (_apply_openclaw_provider_scope, _revert_openclaw_provider_scope),
|
| 52 |
+
}
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def build_install_target_envs(
|
| 56 |
+
port: int, backend: str, targets: list[str]
|
| 57 |
+
) -> dict[str, dict[str, str]]:
|
| 58 |
+
"""Build per-target install environment values via provider slices."""
|
| 59 |
+
target_envs: dict[str, dict[str, str]] = {}
|
| 60 |
+
for target in targets:
|
| 61 |
+
builder = _ENV_BUILDERS.get(target)
|
| 62 |
+
if builder is None:
|
| 63 |
+
continue
|
| 64 |
+
target_envs[target] = builder(port=port, backend=backend)
|
| 65 |
+
return target_envs
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def apply_provider_scope_mutations(manifest: DeploymentManifest) -> list[ManagedMutation]:
|
| 69 |
+
"""Apply provider-scope mutations owned by provider slices."""
|
| 70 |
+
mutations: list[ManagedMutation] = []
|
| 71 |
+
for target in manifest.targets:
|
| 72 |
+
handlers = _PROVIDER_SCOPE_HANDLERS.get(target)
|
| 73 |
+
if handlers is None:
|
| 74 |
+
continue
|
| 75 |
+
mutation = handlers[0](manifest)
|
| 76 |
+
if mutation is not None:
|
| 77 |
+
mutations.append(mutation)
|
| 78 |
+
return mutations
|
| 79 |
+
|
| 80 |
+
|
| 81 |
+
def revert_provider_scope_mutation(manifest: DeploymentManifest, mutation: ManagedMutation) -> None:
|
| 82 |
+
"""Revert a provider-scope mutation via the owning provider slice."""
|
| 83 |
+
handlers = _PROVIDER_SCOPE_HANDLERS.get(mutation.target)
|
| 84 |
+
if handlers is None:
|
| 85 |
+
return
|
| 86 |
+
handlers[1](mutation, manifest)
|
|
@@ -1,50 +1,50 @@
|
|
| 1 |
-
"""OpenClaw install-time helpers."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import click
|
| 6 |
-
|
| 7 |
-
from headroom.install.models import DeploymentManifest, ManagedMutation, ToolTarget
|
| 8 |
-
from headroom.install.paths import openclaw_config_path
|
| 9 |
-
from headroom.install.runtime import resolve_headroom_command
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
def shutil_which(name: str) -> str | None:
|
| 13 |
-
from shutil import which
|
| 14 |
-
|
| 15 |
-
return which(name)
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
def _invoke_openclaw(command: list[str]) -> None:
|
| 19 |
-
import subprocess
|
| 20 |
-
|
| 21 |
-
subprocess.run(command, check=True)
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
def apply_provider_scope(manifest: DeploymentManifest) -> ManagedMutation:
|
| 25 |
-
"""Configure OpenClaw to route through the persistent proxy."""
|
| 26 |
-
if not shutil_which("openclaw"):
|
| 27 |
-
raise click.ClickException("openclaw not found in PATH; cannot apply provider scope.")
|
| 28 |
-
command = [
|
| 29 |
-
*resolve_headroom_command(),
|
| 30 |
-
"wrap",
|
| 31 |
-
"openclaw",
|
| 32 |
-
"--no-auto-start",
|
| 33 |
-
"--proxy-port",
|
| 34 |
-
str(manifest.port),
|
| 35 |
-
]
|
| 36 |
-
_invoke_openclaw(command)
|
| 37 |
-
return ManagedMutation(
|
| 38 |
-
target=ToolTarget.OPENCLAW.value,
|
| 39 |
-
kind="openclaw-wrap",
|
| 40 |
-
path=str(openclaw_config_path()),
|
| 41 |
-
)
|
| 42 |
-
|
| 43 |
-
|
| 44 |
-
def revert_provider_scope(mutation: ManagedMutation, manifest: DeploymentManifest) -> None:
|
| 45 |
-
"""Undo OpenClaw persistent proxy configuration."""
|
| 46 |
-
del mutation, manifest
|
| 47 |
-
if not shutil_which("openclaw"):
|
| 48 |
-
return
|
| 49 |
-
command = [*resolve_headroom_command(), "unwrap", "openclaw"]
|
| 50 |
-
_invoke_openclaw(command)
|
|
|
|
| 1 |
+
"""OpenClaw install-time helpers."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import click
|
| 6 |
+
|
| 7 |
+
from headroom.install.models import DeploymentManifest, ManagedMutation, ToolTarget
|
| 8 |
+
from headroom.install.paths import openclaw_config_path
|
| 9 |
+
from headroom.install.runtime import resolve_headroom_command
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
def shutil_which(name: str) -> str | None:
|
| 13 |
+
from shutil import which
|
| 14 |
+
|
| 15 |
+
return which(name)
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def _invoke_openclaw(command: list[str]) -> None:
|
| 19 |
+
import subprocess
|
| 20 |
+
|
| 21 |
+
subprocess.run(command, check=True)
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def apply_provider_scope(manifest: DeploymentManifest) -> ManagedMutation:
|
| 25 |
+
"""Configure OpenClaw to route through the persistent proxy."""
|
| 26 |
+
if not shutil_which("openclaw"):
|
| 27 |
+
raise click.ClickException("openclaw not found in PATH; cannot apply provider scope.")
|
| 28 |
+
command = [
|
| 29 |
+
*resolve_headroom_command(),
|
| 30 |
+
"wrap",
|
| 31 |
+
"openclaw",
|
| 32 |
+
"--no-auto-start",
|
| 33 |
+
"--proxy-port",
|
| 34 |
+
str(manifest.port),
|
| 35 |
+
]
|
| 36 |
+
_invoke_openclaw(command)
|
| 37 |
+
return ManagedMutation(
|
| 38 |
+
target=ToolTarget.OPENCLAW.value,
|
| 39 |
+
kind="openclaw-wrap",
|
| 40 |
+
path=str(openclaw_config_path()),
|
| 41 |
+
)
|
| 42 |
+
|
| 43 |
+
|
| 44 |
+
def revert_provider_scope(mutation: ManagedMutation, manifest: DeploymentManifest) -> None:
|
| 45 |
+
"""Undo OpenClaw persistent proxy configuration."""
|
| 46 |
+
del mutation, manifest
|
| 47 |
+
if not shutil_which("openclaw"):
|
| 48 |
+
return
|
| 49 |
+
command = [*resolve_headroom_command(), "unwrap", "openclaw"]
|
| 50 |
+
_invoke_openclaw(command)
|
|
The diff for this file is too large to render.
See raw diff
|
|
|
|
The diff for this file is too large to render.
See raw diff
|
|
|
|
@@ -1,310 +1,310 @@
|
|
| 1 |
-
"""Release version helpers for the GitHub Actions release workflow."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import os
|
| 6 |
-
import re
|
| 7 |
-
import subprocess
|
| 8 |
-
from collections.abc import Sequence
|
| 9 |
-
from dataclasses import dataclass, replace
|
| 10 |
-
from pathlib import Path
|
| 11 |
-
|
| 12 |
-
SEMVER_RE = re.compile(r"^(\d+)\.(\d+)\.(\d+)$")
|
| 13 |
-
RELEASE_TAG_RE = re.compile(r"^v(\d+)\.(\d+)\.(\d+)(?:\.(\d+))?$")
|
| 14 |
-
CONVENTIONAL_COMMIT_RE = re.compile(
|
| 15 |
-
r"^(feat|fix|ci|chore|perf|refactor|docs|style|test)(\(.+\))?(!)?:\s*(.+)$"
|
| 16 |
-
)
|
| 17 |
-
BREAKING_CHANGE_RE = re.compile(r"^BREAKING CHANGE:\s*(.+)$", re.MULTILINE)
|
| 18 |
-
FIELD_SEP = "\x1f"
|
| 19 |
-
RECORD_SEP = "\x1e"
|
| 20 |
-
GIT_LOG_FORMAT = "%s%x1f%b%x1e"
|
| 21 |
-
BUMP_PRIORITY = {"patch": 0, "minor": 1, "major": 2}
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
@dataclass(frozen=True, order=True)
|
| 25 |
-
class SemVer:
|
| 26 |
-
"""Semantic version tuple with simple bump helpers."""
|
| 27 |
-
|
| 28 |
-
major: int
|
| 29 |
-
minor: int
|
| 30 |
-
patch: int
|
| 31 |
-
|
| 32 |
-
@classmethod
|
| 33 |
-
def parse(cls, value: str) -> SemVer:
|
| 34 |
-
match = SEMVER_RE.match(value)
|
| 35 |
-
if not match:
|
| 36 |
-
raise ValueError(f"Invalid semantic version: {value}")
|
| 37 |
-
return cls(*(int(part) for part in match.groups()))
|
| 38 |
-
|
| 39 |
-
def bump(self, level: str) -> SemVer:
|
| 40 |
-
if level == "major":
|
| 41 |
-
return SemVer(self.major + 1, 0, 0)
|
| 42 |
-
if level == "minor":
|
| 43 |
-
return SemVer(self.major, self.minor + 1, 0)
|
| 44 |
-
if level == "patch":
|
| 45 |
-
return SemVer(self.major, self.minor, self.patch + 1)
|
| 46 |
-
raise ValueError(f"Unsupported bump level: {level}")
|
| 47 |
-
|
| 48 |
-
def __str__(self) -> str:
|
| 49 |
-
return f"{self.major}.{self.minor}.{self.patch}"
|
| 50 |
-
|
| 51 |
-
|
| 52 |
-
@dataclass(frozen=True)
|
| 53 |
-
class ReleaseVersionInfo:
|
| 54 |
-
"""Workflow outputs for release version calculation."""
|
| 55 |
-
|
| 56 |
-
version: str
|
| 57 |
-
npm_version: str
|
| 58 |
-
canonical: str
|
| 59 |
-
height: str
|
| 60 |
-
bump: str
|
| 61 |
-
previous_tag: str
|
| 62 |
-
|
| 63 |
-
def as_outputs(self) -> dict[str, str]:
|
| 64 |
-
return {
|
| 65 |
-
"version": self.version,
|
| 66 |
-
"npm_version": self.npm_version,
|
| 67 |
-
"canonical": self.canonical,
|
| 68 |
-
"height": self.height,
|
| 69 |
-
"bump": self.bump,
|
| 70 |
-
"previous_tag": self.previous_tag,
|
| 71 |
-
}
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
@dataclass(frozen=True, order=True)
|
| 75 |
-
class ReleaseTag:
|
| 76 |
-
"""Parsed release tag metadata used for sorting and normalization."""
|
| 77 |
-
|
| 78 |
-
version: SemVer
|
| 79 |
-
legacy_height: int = -1
|
| 80 |
-
raw: str = ""
|
| 81 |
-
|
| 82 |
-
|
| 83 |
-
@dataclass(frozen=True)
|
| 84 |
-
class CommitInfo:
|
| 85 |
-
"""Commit subject/body pair used for bump detection."""
|
| 86 |
-
|
| 87 |
-
subject: str
|
| 88 |
-
body: str = ""
|
| 89 |
-
|
| 90 |
-
|
| 91 |
-
def parse_release_tag(tag: str) -> ReleaseTag:
|
| 92 |
-
"""Parse a release tag, preserving legacy fourth-component ordering."""
|
| 93 |
-
|
| 94 |
-
match = RELEASE_TAG_RE.match(tag)
|
| 95 |
-
if not match:
|
| 96 |
-
raise ValueError(f"Invalid release tag: {tag}")
|
| 97 |
-
major, minor, patch, extra = match.groups()
|
| 98 |
-
return ReleaseTag(
|
| 99 |
-
version=SemVer(int(major), int(minor), int(patch)),
|
| 100 |
-
legacy_height=int(extra) if extra is not None else -1,
|
| 101 |
-
raw=tag,
|
| 102 |
-
)
|
| 103 |
-
|
| 104 |
-
|
| 105 |
-
def normalize_release_tag(tag: str) -> SemVer:
|
| 106 |
-
"""Collapse historic 4-part release tags into their base semantic version."""
|
| 107 |
-
|
| 108 |
-
return parse_release_tag(tag).version
|
| 109 |
-
|
| 110 |
-
|
| 111 |
-
def find_latest_release_tag(tags: Sequence[str]) -> str | None:
|
| 112 |
-
"""Return the latest release tag after normalizing legacy 4-part tags."""
|
| 113 |
-
|
| 114 |
-
candidates: list[ReleaseTag] = []
|
| 115 |
-
for tag in tags:
|
| 116 |
-
if RELEASE_TAG_RE.match(tag):
|
| 117 |
-
candidates.append(parse_release_tag(tag))
|
| 118 |
-
if not candidates:
|
| 119 |
-
return None
|
| 120 |
-
candidates.sort(reverse=True)
|
| 121 |
-
return candidates[0].raw
|
| 122 |
-
|
| 123 |
-
|
| 124 |
-
def _merge_summary(subject: str, body: str) -> str:
|
| 125 |
-
"""Return the first meaningful body line for merge commits."""
|
| 126 |
-
|
| 127 |
-
if not subject.startswith("Merge "):
|
| 128 |
-
return ""
|
| 129 |
-
for line in body.splitlines():
|
| 130 |
-
stripped = line.strip()
|
| 131 |
-
if stripped:
|
| 132 |
-
return stripped
|
| 133 |
-
return ""
|
| 134 |
-
|
| 135 |
-
|
| 136 |
-
def classify_commit_bump(commit: CommitInfo) -> str:
|
| 137 |
-
"""Classify one commit using conventional commit semantics."""
|
| 138 |
-
|
| 139 |
-
merge_summary = _merge_summary(commit.subject, commit.body)
|
| 140 |
-
candidates = [commit.subject]
|
| 141 |
-
if merge_summary:
|
| 142 |
-
candidates.insert(0, merge_summary)
|
| 143 |
-
|
| 144 |
-
has_breaking_change = bool(BREAKING_CHANGE_RE.search(commit.body))
|
| 145 |
-
for candidate in candidates:
|
| 146 |
-
match = CONVENTIONAL_COMMIT_RE.match(candidate)
|
| 147 |
-
if not match:
|
| 148 |
-
continue
|
| 149 |
-
if has_breaking_change or bool(match.group(3)):
|
| 150 |
-
return "major"
|
| 151 |
-
if match.group(1) == "feat":
|
| 152 |
-
return "minor"
|
| 153 |
-
return "patch"
|
| 154 |
-
|
| 155 |
-
if has_breaking_change:
|
| 156 |
-
return "major"
|
| 157 |
-
return "patch"
|
| 158 |
-
|
| 159 |
-
|
| 160 |
-
def determine_bump_level(commits: Sequence[CommitInfo]) -> str:
|
| 161 |
-
"""Return the highest required bump across a commit range."""
|
| 162 |
-
|
| 163 |
-
level = "patch"
|
| 164 |
-
for commit in commits:
|
| 165 |
-
candidate = classify_commit_bump(commit)
|
| 166 |
-
if BUMP_PRIORITY[candidate] > BUMP_PRIORITY[level]:
|
| 167 |
-
level = candidate
|
| 168 |
-
return level
|
| 169 |
-
|
| 170 |
-
|
| 171 |
-
def compute_release_version(
|
| 172 |
-
canonical_version: str,
|
| 173 |
-
level: str,
|
| 174 |
-
tags: Sequence[str],
|
| 175 |
-
manual_version: str = "",
|
| 176 |
-
) -> ReleaseVersionInfo:
|
| 177 |
-
"""Compute the next release version from the canonical version and existing tags."""
|
| 178 |
-
|
| 179 |
-
if manual_version:
|
| 180 |
-
manual = str(SemVer.parse(manual_version))
|
| 181 |
-
return ReleaseVersionInfo(
|
| 182 |
-
version=manual,
|
| 183 |
-
npm_version=manual,
|
| 184 |
-
canonical=canonical_version,
|
| 185 |
-
height="0",
|
| 186 |
-
bump="manual",
|
| 187 |
-
previous_tag="",
|
| 188 |
-
)
|
| 189 |
-
|
| 190 |
-
canonical = SemVer.parse(canonical_version)
|
| 191 |
-
previous_tag = find_latest_release_tag(tags)
|
| 192 |
-
current = canonical
|
| 193 |
-
if previous_tag is not None:
|
| 194 |
-
current = max(current, normalize_release_tag(previous_tag))
|
| 195 |
-
|
| 196 |
-
next_version = str(current.bump(level))
|
| 197 |
-
return ReleaseVersionInfo(
|
| 198 |
-
version=next_version,
|
| 199 |
-
npm_version=next_version,
|
| 200 |
-
canonical=canonical_version,
|
| 201 |
-
height="0",
|
| 202 |
-
bump=level,
|
| 203 |
-
previous_tag=previous_tag or "",
|
| 204 |
-
)
|
| 205 |
-
|
| 206 |
-
|
| 207 |
-
def get_canonical_version(root: Path) -> str:
|
| 208 |
-
"""Read the canonical project version from pyproject.toml."""
|
| 209 |
-
|
| 210 |
-
try:
|
| 211 |
-
import tomllib
|
| 212 |
-
except ModuleNotFoundError: # pragma: no cover - Python 3.10 compatibility
|
| 213 |
-
import tomli as tomllib
|
| 214 |
-
|
| 215 |
-
with open(root / "pyproject.toml", "rb") as file:
|
| 216 |
-
project = tomllib.load(file)["project"]
|
| 217 |
-
return str(project["version"])
|
| 218 |
-
|
| 219 |
-
|
| 220 |
-
def list_release_tags(root: Path) -> list[str]:
|
| 221 |
-
"""List release tags from the local Git checkout."""
|
| 222 |
-
|
| 223 |
-
result = subprocess.run(
|
| 224 |
-
["git", "tag", "-l", "v*"],
|
| 225 |
-
cwd=root,
|
| 226 |
-
check=True,
|
| 227 |
-
capture_output=True,
|
| 228 |
-
text=True,
|
| 229 |
-
)
|
| 230 |
-
return [tag.strip() for tag in result.stdout.splitlines() if tag.strip()]
|
| 231 |
-
|
| 232 |
-
|
| 233 |
-
def list_release_commits(root: Path, previous_tag: str) -> list[CommitInfo]:
|
| 234 |
-
"""List commit subject/body pairs since the previous release tag."""
|
| 235 |
-
|
| 236 |
-
cmd = ["git", "log", "--first-parent", f"--pretty=format:{GIT_LOG_FORMAT}"]
|
| 237 |
-
if previous_tag:
|
| 238 |
-
cmd.append(f"{previous_tag}..HEAD")
|
| 239 |
-
else:
|
| 240 |
-
cmd.append("HEAD")
|
| 241 |
-
|
| 242 |
-
result = subprocess.run(
|
| 243 |
-
cmd,
|
| 244 |
-
cwd=root,
|
| 245 |
-
check=True,
|
| 246 |
-
capture_output=True,
|
| 247 |
-
text=True,
|
| 248 |
-
)
|
| 249 |
-
|
| 250 |
-
commits: list[CommitInfo] = []
|
| 251 |
-
for raw_entry in result.stdout.split(RECORD_SEP):
|
| 252 |
-
if not raw_entry or FIELD_SEP not in raw_entry:
|
| 253 |
-
continue
|
| 254 |
-
subject, body = raw_entry.split(FIELD_SEP, 1)
|
| 255 |
-
commits.append(CommitInfo(subject=subject.strip(), body=body.strip()))
|
| 256 |
-
return commits
|
| 257 |
-
|
| 258 |
-
|
| 259 |
-
def commit_height_since(root: Path, previous_tag: str) -> str:
|
| 260 |
-
"""Count commits since the previous release tag for changelog/debug outputs."""
|
| 261 |
-
|
| 262 |
-
if not previous_tag:
|
| 263 |
-
return "0"
|
| 264 |
-
|
| 265 |
-
result = subprocess.run(
|
| 266 |
-
["git", "rev-list", f"{previous_tag}..HEAD", "--count"],
|
| 267 |
-
cwd=root,
|
| 268 |
-
check=True,
|
| 269 |
-
capture_output=True,
|
| 270 |
-
text=True,
|
| 271 |
-
)
|
| 272 |
-
return result.stdout.strip() or "0"
|
| 273 |
-
|
| 274 |
-
|
| 275 |
-
def write_github_outputs(info: ReleaseVersionInfo, output_path: str) -> None:
|
| 276 |
-
"""Append workflow outputs to the GitHub Actions output file."""
|
| 277 |
-
|
| 278 |
-
with open(output_path, "a", encoding="utf-8") as output_file:
|
| 279 |
-
for key, value in info.as_outputs().items():
|
| 280 |
-
output_file.write(f"{key}={value}\n")
|
| 281 |
-
|
| 282 |
-
|
| 283 |
-
def main() -> None:
|
| 284 |
-
root = Path.cwd()
|
| 285 |
-
manual_version = os.environ.get("MANUAL_VER", "").strip()
|
| 286 |
-
tags = list_release_tags(root)
|
| 287 |
-
previous_tag = find_latest_release_tag(tags) or ""
|
| 288 |
-
level = os.environ.get("LEVEL", "").strip()
|
| 289 |
-
if not level:
|
| 290 |
-
level = determine_bump_level(list_release_commits(root, previous_tag))
|
| 291 |
-
|
| 292 |
-
info = compute_release_version(
|
| 293 |
-
canonical_version=get_canonical_version(root),
|
| 294 |
-
level=level,
|
| 295 |
-
tags=tags,
|
| 296 |
-
manual_version=manual_version,
|
| 297 |
-
)
|
| 298 |
-
info = replace(info, height=commit_height_since(root, info.previous_tag))
|
| 299 |
-
|
| 300 |
-
output_path = os.environ.get("GITHUB_OUTPUT", "").strip()
|
| 301 |
-
if output_path:
|
| 302 |
-
write_github_outputs(info, output_path)
|
| 303 |
-
return
|
| 304 |
-
|
| 305 |
-
for key, value in info.as_outputs().items():
|
| 306 |
-
print(f"{key}={value}")
|
| 307 |
-
|
| 308 |
-
|
| 309 |
-
if __name__ == "__main__":
|
| 310 |
-
main()
|
|
|
|
| 1 |
+
"""Release version helpers for the GitHub Actions release workflow."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import os
|
| 6 |
+
import re
|
| 7 |
+
import subprocess
|
| 8 |
+
from collections.abc import Sequence
|
| 9 |
+
from dataclasses import dataclass, replace
|
| 10 |
+
from pathlib import Path
|
| 11 |
+
|
| 12 |
+
SEMVER_RE = re.compile(r"^(\d+)\.(\d+)\.(\d+)$")
|
| 13 |
+
RELEASE_TAG_RE = re.compile(r"^v(\d+)\.(\d+)\.(\d+)(?:\.(\d+))?$")
|
| 14 |
+
CONVENTIONAL_COMMIT_RE = re.compile(
|
| 15 |
+
r"^(feat|fix|ci|chore|perf|refactor|docs|style|test)(\(.+\))?(!)?:\s*(.+)$"
|
| 16 |
+
)
|
| 17 |
+
BREAKING_CHANGE_RE = re.compile(r"^BREAKING CHANGE:\s*(.+)$", re.MULTILINE)
|
| 18 |
+
FIELD_SEP = "\x1f"
|
| 19 |
+
RECORD_SEP = "\x1e"
|
| 20 |
+
GIT_LOG_FORMAT = "%s%x1f%b%x1e"
|
| 21 |
+
BUMP_PRIORITY = {"patch": 0, "minor": 1, "major": 2}
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
@dataclass(frozen=True, order=True)
|
| 25 |
+
class SemVer:
|
| 26 |
+
"""Semantic version tuple with simple bump helpers."""
|
| 27 |
+
|
| 28 |
+
major: int
|
| 29 |
+
minor: int
|
| 30 |
+
patch: int
|
| 31 |
+
|
| 32 |
+
@classmethod
|
| 33 |
+
def parse(cls, value: str) -> SemVer:
|
| 34 |
+
match = SEMVER_RE.match(value)
|
| 35 |
+
if not match:
|
| 36 |
+
raise ValueError(f"Invalid semantic version: {value}")
|
| 37 |
+
return cls(*(int(part) for part in match.groups()))
|
| 38 |
+
|
| 39 |
+
def bump(self, level: str) -> SemVer:
|
| 40 |
+
if level == "major":
|
| 41 |
+
return SemVer(self.major + 1, 0, 0)
|
| 42 |
+
if level == "minor":
|
| 43 |
+
return SemVer(self.major, self.minor + 1, 0)
|
| 44 |
+
if level == "patch":
|
| 45 |
+
return SemVer(self.major, self.minor, self.patch + 1)
|
| 46 |
+
raise ValueError(f"Unsupported bump level: {level}")
|
| 47 |
+
|
| 48 |
+
def __str__(self) -> str:
|
| 49 |
+
return f"{self.major}.{self.minor}.{self.patch}"
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
@dataclass(frozen=True)
|
| 53 |
+
class ReleaseVersionInfo:
|
| 54 |
+
"""Workflow outputs for release version calculation."""
|
| 55 |
+
|
| 56 |
+
version: str
|
| 57 |
+
npm_version: str
|
| 58 |
+
canonical: str
|
| 59 |
+
height: str
|
| 60 |
+
bump: str
|
| 61 |
+
previous_tag: str
|
| 62 |
+
|
| 63 |
+
def as_outputs(self) -> dict[str, str]:
|
| 64 |
+
return {
|
| 65 |
+
"version": self.version,
|
| 66 |
+
"npm_version": self.npm_version,
|
| 67 |
+
"canonical": self.canonical,
|
| 68 |
+
"height": self.height,
|
| 69 |
+
"bump": self.bump,
|
| 70 |
+
"previous_tag": self.previous_tag,
|
| 71 |
+
}
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
@dataclass(frozen=True, order=True)
|
| 75 |
+
class ReleaseTag:
|
| 76 |
+
"""Parsed release tag metadata used for sorting and normalization."""
|
| 77 |
+
|
| 78 |
+
version: SemVer
|
| 79 |
+
legacy_height: int = -1
|
| 80 |
+
raw: str = ""
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
@dataclass(frozen=True)
|
| 84 |
+
class CommitInfo:
|
| 85 |
+
"""Commit subject/body pair used for bump detection."""
|
| 86 |
+
|
| 87 |
+
subject: str
|
| 88 |
+
body: str = ""
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
def parse_release_tag(tag: str) -> ReleaseTag:
|
| 92 |
+
"""Parse a release tag, preserving legacy fourth-component ordering."""
|
| 93 |
+
|
| 94 |
+
match = RELEASE_TAG_RE.match(tag)
|
| 95 |
+
if not match:
|
| 96 |
+
raise ValueError(f"Invalid release tag: {tag}")
|
| 97 |
+
major, minor, patch, extra = match.groups()
|
| 98 |
+
return ReleaseTag(
|
| 99 |
+
version=SemVer(int(major), int(minor), int(patch)),
|
| 100 |
+
legacy_height=int(extra) if extra is not None else -1,
|
| 101 |
+
raw=tag,
|
| 102 |
+
)
|
| 103 |
+
|
| 104 |
+
|
| 105 |
+
def normalize_release_tag(tag: str) -> SemVer:
|
| 106 |
+
"""Collapse historic 4-part release tags into their base semantic version."""
|
| 107 |
+
|
| 108 |
+
return parse_release_tag(tag).version
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
def find_latest_release_tag(tags: Sequence[str]) -> str | None:
|
| 112 |
+
"""Return the latest release tag after normalizing legacy 4-part tags."""
|
| 113 |
+
|
| 114 |
+
candidates: list[ReleaseTag] = []
|
| 115 |
+
for tag in tags:
|
| 116 |
+
if RELEASE_TAG_RE.match(tag):
|
| 117 |
+
candidates.append(parse_release_tag(tag))
|
| 118 |
+
if not candidates:
|
| 119 |
+
return None
|
| 120 |
+
candidates.sort(reverse=True)
|
| 121 |
+
return candidates[0].raw
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
def _merge_summary(subject: str, body: str) -> str:
|
| 125 |
+
"""Return the first meaningful body line for merge commits."""
|
| 126 |
+
|
| 127 |
+
if not subject.startswith("Merge "):
|
| 128 |
+
return ""
|
| 129 |
+
for line in body.splitlines():
|
| 130 |
+
stripped = line.strip()
|
| 131 |
+
if stripped:
|
| 132 |
+
return stripped
|
| 133 |
+
return ""
|
| 134 |
+
|
| 135 |
+
|
| 136 |
+
def classify_commit_bump(commit: CommitInfo) -> str:
|
| 137 |
+
"""Classify one commit using conventional commit semantics."""
|
| 138 |
+
|
| 139 |
+
merge_summary = _merge_summary(commit.subject, commit.body)
|
| 140 |
+
candidates = [commit.subject]
|
| 141 |
+
if merge_summary:
|
| 142 |
+
candidates.insert(0, merge_summary)
|
| 143 |
+
|
| 144 |
+
has_breaking_change = bool(BREAKING_CHANGE_RE.search(commit.body))
|
| 145 |
+
for candidate in candidates:
|
| 146 |
+
match = CONVENTIONAL_COMMIT_RE.match(candidate)
|
| 147 |
+
if not match:
|
| 148 |
+
continue
|
| 149 |
+
if has_breaking_change or bool(match.group(3)):
|
| 150 |
+
return "major"
|
| 151 |
+
if match.group(1) == "feat":
|
| 152 |
+
return "minor"
|
| 153 |
+
return "patch"
|
| 154 |
+
|
| 155 |
+
if has_breaking_change:
|
| 156 |
+
return "major"
|
| 157 |
+
return "patch"
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
def determine_bump_level(commits: Sequence[CommitInfo]) -> str:
|
| 161 |
+
"""Return the highest required bump across a commit range."""
|
| 162 |
+
|
| 163 |
+
level = "patch"
|
| 164 |
+
for commit in commits:
|
| 165 |
+
candidate = classify_commit_bump(commit)
|
| 166 |
+
if BUMP_PRIORITY[candidate] > BUMP_PRIORITY[level]:
|
| 167 |
+
level = candidate
|
| 168 |
+
return level
|
| 169 |
+
|
| 170 |
+
|
| 171 |
+
def compute_release_version(
|
| 172 |
+
canonical_version: str,
|
| 173 |
+
level: str,
|
| 174 |
+
tags: Sequence[str],
|
| 175 |
+
manual_version: str = "",
|
| 176 |
+
) -> ReleaseVersionInfo:
|
| 177 |
+
"""Compute the next release version from the canonical version and existing tags."""
|
| 178 |
+
|
| 179 |
+
if manual_version:
|
| 180 |
+
manual = str(SemVer.parse(manual_version))
|
| 181 |
+
return ReleaseVersionInfo(
|
| 182 |
+
version=manual,
|
| 183 |
+
npm_version=manual,
|
| 184 |
+
canonical=canonical_version,
|
| 185 |
+
height="0",
|
| 186 |
+
bump="manual",
|
| 187 |
+
previous_tag="",
|
| 188 |
+
)
|
| 189 |
+
|
| 190 |
+
canonical = SemVer.parse(canonical_version)
|
| 191 |
+
previous_tag = find_latest_release_tag(tags)
|
| 192 |
+
current = canonical
|
| 193 |
+
if previous_tag is not None:
|
| 194 |
+
current = max(current, normalize_release_tag(previous_tag))
|
| 195 |
+
|
| 196 |
+
next_version = str(current.bump(level))
|
| 197 |
+
return ReleaseVersionInfo(
|
| 198 |
+
version=next_version,
|
| 199 |
+
npm_version=next_version,
|
| 200 |
+
canonical=canonical_version,
|
| 201 |
+
height="0",
|
| 202 |
+
bump=level,
|
| 203 |
+
previous_tag=previous_tag or "",
|
| 204 |
+
)
|
| 205 |
+
|
| 206 |
+
|
| 207 |
+
def get_canonical_version(root: Path) -> str:
|
| 208 |
+
"""Read the canonical project version from pyproject.toml."""
|
| 209 |
+
|
| 210 |
+
try:
|
| 211 |
+
import tomllib
|
| 212 |
+
except ModuleNotFoundError: # pragma: no cover - Python 3.10 compatibility
|
| 213 |
+
import tomli as tomllib
|
| 214 |
+
|
| 215 |
+
with open(root / "pyproject.toml", "rb") as file:
|
| 216 |
+
project = tomllib.load(file)["project"]
|
| 217 |
+
return str(project["version"])
|
| 218 |
+
|
| 219 |
+
|
| 220 |
+
def list_release_tags(root: Path) -> list[str]:
|
| 221 |
+
"""List release tags from the local Git checkout."""
|
| 222 |
+
|
| 223 |
+
result = subprocess.run(
|
| 224 |
+
["git", "tag", "-l", "v*"],
|
| 225 |
+
cwd=root,
|
| 226 |
+
check=True,
|
| 227 |
+
capture_output=True,
|
| 228 |
+
text=True,
|
| 229 |
+
)
|
| 230 |
+
return [tag.strip() for tag in result.stdout.splitlines() if tag.strip()]
|
| 231 |
+
|
| 232 |
+
|
| 233 |
+
def list_release_commits(root: Path, previous_tag: str) -> list[CommitInfo]:
|
| 234 |
+
"""List commit subject/body pairs since the previous release tag."""
|
| 235 |
+
|
| 236 |
+
cmd = ["git", "log", "--first-parent", f"--pretty=format:{GIT_LOG_FORMAT}"]
|
| 237 |
+
if previous_tag:
|
| 238 |
+
cmd.append(f"{previous_tag}..HEAD")
|
| 239 |
+
else:
|
| 240 |
+
cmd.append("HEAD")
|
| 241 |
+
|
| 242 |
+
result = subprocess.run(
|
| 243 |
+
cmd,
|
| 244 |
+
cwd=root,
|
| 245 |
+
check=True,
|
| 246 |
+
capture_output=True,
|
| 247 |
+
text=True,
|
| 248 |
+
)
|
| 249 |
+
|
| 250 |
+
commits: list[CommitInfo] = []
|
| 251 |
+
for raw_entry in result.stdout.split(RECORD_SEP):
|
| 252 |
+
if not raw_entry or FIELD_SEP not in raw_entry:
|
| 253 |
+
continue
|
| 254 |
+
subject, body = raw_entry.split(FIELD_SEP, 1)
|
| 255 |
+
commits.append(CommitInfo(subject=subject.strip(), body=body.strip()))
|
| 256 |
+
return commits
|
| 257 |
+
|
| 258 |
+
|
| 259 |
+
def commit_height_since(root: Path, previous_tag: str) -> str:
|
| 260 |
+
"""Count commits since the previous release tag for changelog/debug outputs."""
|
| 261 |
+
|
| 262 |
+
if not previous_tag:
|
| 263 |
+
return "0"
|
| 264 |
+
|
| 265 |
+
result = subprocess.run(
|
| 266 |
+
["git", "rev-list", f"{previous_tag}..HEAD", "--count"],
|
| 267 |
+
cwd=root,
|
| 268 |
+
check=True,
|
| 269 |
+
capture_output=True,
|
| 270 |
+
text=True,
|
| 271 |
+
)
|
| 272 |
+
return result.stdout.strip() or "0"
|
| 273 |
+
|
| 274 |
+
|
| 275 |
+
def write_github_outputs(info: ReleaseVersionInfo, output_path: str) -> None:
|
| 276 |
+
"""Append workflow outputs to the GitHub Actions output file."""
|
| 277 |
+
|
| 278 |
+
with open(output_path, "a", encoding="utf-8") as output_file:
|
| 279 |
+
for key, value in info.as_outputs().items():
|
| 280 |
+
output_file.write(f"{key}={value}\n")
|
| 281 |
+
|
| 282 |
+
|
| 283 |
+
def main() -> None:
|
| 284 |
+
root = Path.cwd()
|
| 285 |
+
manual_version = os.environ.get("MANUAL_VER", "").strip()
|
| 286 |
+
tags = list_release_tags(root)
|
| 287 |
+
previous_tag = find_latest_release_tag(tags) or ""
|
| 288 |
+
level = os.environ.get("LEVEL", "").strip()
|
| 289 |
+
if not level:
|
| 290 |
+
level = determine_bump_level(list_release_commits(root, previous_tag))
|
| 291 |
+
|
| 292 |
+
info = compute_release_version(
|
| 293 |
+
canonical_version=get_canonical_version(root),
|
| 294 |
+
level=level,
|
| 295 |
+
tags=tags,
|
| 296 |
+
manual_version=manual_version,
|
| 297 |
+
)
|
| 298 |
+
info = replace(info, height=commit_height_since(root, info.previous_tag))
|
| 299 |
+
|
| 300 |
+
output_path = os.environ.get("GITHUB_OUTPUT", "").strip()
|
| 301 |
+
if output_path:
|
| 302 |
+
write_github_outputs(info, output_path)
|
| 303 |
+
return
|
| 304 |
+
|
| 305 |
+
for key, value in info.as_outputs().items():
|
| 306 |
+
print(f"{key}={value}")
|
| 307 |
+
|
| 308 |
+
|
| 309 |
+
if __name__ == "__main__":
|
| 310 |
+
main()
|
|
@@ -1,72 +1,72 @@
|
|
| 1 |
-
"""Subscription window tracking for Anthropic Claude Code accounts and Codex rate limits."""
|
| 2 |
-
|
| 3 |
-
from headroom.subscription.base import (
|
| 4 |
-
QuotaTracker,
|
| 5 |
-
QuotaTrackerRegistry,
|
| 6 |
-
get_quota_registry,
|
| 7 |
-
reset_quota_registry,
|
| 8 |
-
)
|
| 9 |
-
from headroom.subscription.client import SubscriptionClient, read_cached_oauth_token
|
| 10 |
-
from headroom.subscription.codex_rate_limits import (
|
| 11 |
-
CodexCreditsSnapshot,
|
| 12 |
-
CodexRateLimitSnapshot,
|
| 13 |
-
CodexRateLimitState,
|
| 14 |
-
CodexRateLimitWindow,
|
| 15 |
-
get_codex_rate_limit_state,
|
| 16 |
-
parse_codex_rate_limits,
|
| 17 |
-
)
|
| 18 |
-
from headroom.subscription.copilot_quota import (
|
| 19 |
-
CopilotQuotaCategory,
|
| 20 |
-
CopilotQuotaSnapshot,
|
| 21 |
-
CopilotQuotaState,
|
| 22 |
-
discover_github_token,
|
| 23 |
-
get_copilot_quota_tracker,
|
| 24 |
-
parse_copilot_quota,
|
| 25 |
-
)
|
| 26 |
-
from headroom.subscription.models import (
|
| 27 |
-
ExtraUsage,
|
| 28 |
-
HeadroomContribution,
|
| 29 |
-
RateLimitWindow,
|
| 30 |
-
SubscriptionSnapshot,
|
| 31 |
-
SubscriptionState,
|
| 32 |
-
WindowDiscrepancy,
|
| 33 |
-
WindowTokens,
|
| 34 |
-
)
|
| 35 |
-
from headroom.subscription.tracker import (
|
| 36 |
-
SubscriptionTracker,
|
| 37 |
-
configure_subscription_tracker,
|
| 38 |
-
get_subscription_tracker,
|
| 39 |
-
shutdown_subscription_tracker,
|
| 40 |
-
)
|
| 41 |
-
|
| 42 |
-
__all__ = [
|
| 43 |
-
"CodexCreditsSnapshot",
|
| 44 |
-
"CodexRateLimitSnapshot",
|
| 45 |
-
"CodexRateLimitState",
|
| 46 |
-
"CodexRateLimitWindow",
|
| 47 |
-
"CopilotQuotaCategory",
|
| 48 |
-
"CopilotQuotaSnapshot",
|
| 49 |
-
"CopilotQuotaState",
|
| 50 |
-
"ExtraUsage",
|
| 51 |
-
"HeadroomContribution",
|
| 52 |
-
"QuotaTracker",
|
| 53 |
-
"QuotaTrackerRegistry",
|
| 54 |
-
"RateLimitWindow",
|
| 55 |
-
"SubscriptionClient",
|
| 56 |
-
"SubscriptionSnapshot",
|
| 57 |
-
"SubscriptionState",
|
| 58 |
-
"SubscriptionTracker",
|
| 59 |
-
"WindowDiscrepancy",
|
| 60 |
-
"WindowTokens",
|
| 61 |
-
"configure_subscription_tracker",
|
| 62 |
-
"discover_github_token",
|
| 63 |
-
"get_codex_rate_limit_state",
|
| 64 |
-
"get_copilot_quota_tracker",
|
| 65 |
-
"get_quota_registry",
|
| 66 |
-
"get_subscription_tracker",
|
| 67 |
-
"parse_codex_rate_limits",
|
| 68 |
-
"parse_copilot_quota",
|
| 69 |
-
"read_cached_oauth_token",
|
| 70 |
-
"reset_quota_registry",
|
| 71 |
-
"shutdown_subscription_tracker",
|
| 72 |
-
]
|
|
|
|
| 1 |
+
"""Subscription window tracking for Anthropic Claude Code accounts and Codex rate limits."""
|
| 2 |
+
|
| 3 |
+
from headroom.subscription.base import (
|
| 4 |
+
QuotaTracker,
|
| 5 |
+
QuotaTrackerRegistry,
|
| 6 |
+
get_quota_registry,
|
| 7 |
+
reset_quota_registry,
|
| 8 |
+
)
|
| 9 |
+
from headroom.subscription.client import SubscriptionClient, read_cached_oauth_token
|
| 10 |
+
from headroom.subscription.codex_rate_limits import (
|
| 11 |
+
CodexCreditsSnapshot,
|
| 12 |
+
CodexRateLimitSnapshot,
|
| 13 |
+
CodexRateLimitState,
|
| 14 |
+
CodexRateLimitWindow,
|
| 15 |
+
get_codex_rate_limit_state,
|
| 16 |
+
parse_codex_rate_limits,
|
| 17 |
+
)
|
| 18 |
+
from headroom.subscription.copilot_quota import (
|
| 19 |
+
CopilotQuotaCategory,
|
| 20 |
+
CopilotQuotaSnapshot,
|
| 21 |
+
CopilotQuotaState,
|
| 22 |
+
discover_github_token,
|
| 23 |
+
get_copilot_quota_tracker,
|
| 24 |
+
parse_copilot_quota,
|
| 25 |
+
)
|
| 26 |
+
from headroom.subscription.models import (
|
| 27 |
+
ExtraUsage,
|
| 28 |
+
HeadroomContribution,
|
| 29 |
+
RateLimitWindow,
|
| 30 |
+
SubscriptionSnapshot,
|
| 31 |
+
SubscriptionState,
|
| 32 |
+
WindowDiscrepancy,
|
| 33 |
+
WindowTokens,
|
| 34 |
+
)
|
| 35 |
+
from headroom.subscription.tracker import (
|
| 36 |
+
SubscriptionTracker,
|
| 37 |
+
configure_subscription_tracker,
|
| 38 |
+
get_subscription_tracker,
|
| 39 |
+
shutdown_subscription_tracker,
|
| 40 |
+
)
|
| 41 |
+
|
| 42 |
+
__all__ = [
|
| 43 |
+
"CodexCreditsSnapshot",
|
| 44 |
+
"CodexRateLimitSnapshot",
|
| 45 |
+
"CodexRateLimitState",
|
| 46 |
+
"CodexRateLimitWindow",
|
| 47 |
+
"CopilotQuotaCategory",
|
| 48 |
+
"CopilotQuotaSnapshot",
|
| 49 |
+
"CopilotQuotaState",
|
| 50 |
+
"ExtraUsage",
|
| 51 |
+
"HeadroomContribution",
|
| 52 |
+
"QuotaTracker",
|
| 53 |
+
"QuotaTrackerRegistry",
|
| 54 |
+
"RateLimitWindow",
|
| 55 |
+
"SubscriptionClient",
|
| 56 |
+
"SubscriptionSnapshot",
|
| 57 |
+
"SubscriptionState",
|
| 58 |
+
"SubscriptionTracker",
|
| 59 |
+
"WindowDiscrepancy",
|
| 60 |
+
"WindowTokens",
|
| 61 |
+
"configure_subscription_tracker",
|
| 62 |
+
"discover_github_token",
|
| 63 |
+
"get_codex_rate_limit_state",
|
| 64 |
+
"get_copilot_quota_tracker",
|
| 65 |
+
"get_quota_registry",
|
| 66 |
+
"get_subscription_tracker",
|
| 67 |
+
"parse_codex_rate_limits",
|
| 68 |
+
"parse_copilot_quota",
|
| 69 |
+
"read_cached_oauth_token",
|
| 70 |
+
"reset_quota_registry",
|
| 71 |
+
"shutdown_subscription_tracker",
|
| 72 |
+
]
|
|
@@ -1,230 +1,230 @@
|
|
| 1 |
-
"""Base abstractions for pluggable AI-tool quota / rate-limit trackers.
|
| 2 |
-
|
| 3 |
-
Every provider tracker (Anthropic, Codex, Copilot, …) inherits from
|
| 4 |
-
:class:`QuotaTracker` and is registered with the process-global
|
| 5 |
-
:class:`QuotaTrackerRegistry`. ``server.py`` only interacts with the
|
| 6 |
-
registry — adding a new provider requires *zero* changes to the server.
|
| 7 |
-
|
| 8 |
-
Quick-start for a new provider::
|
| 9 |
-
|
| 10 |
-
from headroom.subscription.base import QuotaTracker, get_quota_registry
|
| 11 |
-
|
| 12 |
-
class GeminiQuotaTracker(QuotaTracker):
|
| 13 |
-
key = "gemini_quota"
|
| 14 |
-
label = "Google Gemini"
|
| 15 |
-
|
| 16 |
-
def is_available(self) -> bool:
|
| 17 |
-
return bool(os.environ.get("GOOGLE_API_KEY"))
|
| 18 |
-
|
| 19 |
-
async def start(self) -> None: ... # launch background poll
|
| 20 |
-
async def stop(self) -> None: ... # cancel poll task
|
| 21 |
-
|
| 22 |
-
def get_stats(self) -> dict | None:
|
| 23 |
-
return ... # serialisable dict or None if no data yet
|
| 24 |
-
|
| 25 |
-
get_quota_registry().register(GeminiQuotaTracker())
|
| 26 |
-
"""
|
| 27 |
-
|
| 28 |
-
from __future__ import annotations
|
| 29 |
-
|
| 30 |
-
import abc
|
| 31 |
-
import logging
|
| 32 |
-
from threading import Lock
|
| 33 |
-
from typing import Any
|
| 34 |
-
|
| 35 |
-
logger = logging.getLogger(__name__)
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
class QuotaTracker(abc.ABC):
|
| 39 |
-
"""Abstract base for a single AI-tool quota / rate-limit tracker.
|
| 40 |
-
|
| 41 |
-
Subclasses must define :attr:`key`, :attr:`label`, and
|
| 42 |
-
:meth:`get_stats`. All other methods have sensible defaults.
|
| 43 |
-
"""
|
| 44 |
-
|
| 45 |
-
# ------------------------------------------------------------------ #
|
| 46 |
-
# Class-level identity — subclasses should override as class attributes
|
| 47 |
-
# ------------------------------------------------------------------ #
|
| 48 |
-
|
| 49 |
-
@property
|
| 50 |
-
@abc.abstractmethod
|
| 51 |
-
def key(self) -> str:
|
| 52 |
-
"""Stats key used in ``/stats`` and the dashboard.
|
| 53 |
-
|
| 54 |
-
Must be unique across all registered trackers.
|
| 55 |
-
Examples: ``"subscription_window"``, ``"codex_rate_limits"``.
|
| 56 |
-
"""
|
| 57 |
-
|
| 58 |
-
@property
|
| 59 |
-
@abc.abstractmethod
|
| 60 |
-
def label(self) -> str:
|
| 61 |
-
"""Human-readable name for log messages.
|
| 62 |
-
|
| 63 |
-
Example: ``"Anthropic Claude Code"``.
|
| 64 |
-
"""
|
| 65 |
-
|
| 66 |
-
# ------------------------------------------------------------------ #
|
| 67 |
-
# Availability gate
|
| 68 |
-
# ------------------------------------------------------------------ #
|
| 69 |
-
|
| 70 |
-
def is_available(self) -> bool:
|
| 71 |
-
"""Return ``True`` if this tracker should be activated.
|
| 72 |
-
|
| 73 |
-
Override to gate on environment variables, config flags, etc.
|
| 74 |
-
The registry calls this before :meth:`start` and skips trackers
|
| 75 |
-
that return ``False``. Default: always available.
|
| 76 |
-
"""
|
| 77 |
-
return True
|
| 78 |
-
|
| 79 |
-
# ------------------------------------------------------------------ #
|
| 80 |
-
# Lifecycle — default no-ops (suitable for passive/header-based trackers)
|
| 81 |
-
# ------------------------------------------------------------------ #
|
| 82 |
-
|
| 83 |
-
async def start(self) -> None: # noqa: B027
|
| 84 |
-
"""Start background polling. No-op for passive trackers."""
|
| 85 |
-
|
| 86 |
-
async def stop(self) -> None: # noqa: B027
|
| 87 |
-
"""Stop background polling. No-op for passive trackers."""
|
| 88 |
-
|
| 89 |
-
# ------------------------------------------------------------------ #
|
| 90 |
-
# Stats
|
| 91 |
-
# ------------------------------------------------------------------ #
|
| 92 |
-
|
| 93 |
-
@abc.abstractmethod
|
| 94 |
-
def get_stats(self) -> dict[str, Any] | None:
|
| 95 |
-
"""Return the current snapshot as a serialisable dict, or ``None``.
|
| 96 |
-
|
| 97 |
-
``None`` means "no data yet" and causes the key to be omitted from
|
| 98 |
-
``/stats`` rather than appearing as ``null``.
|
| 99 |
-
"""
|
| 100 |
-
|
| 101 |
-
|
| 102 |
-
# --------------------------------------------------------------------------- #
|
| 103 |
-
# Registry
|
| 104 |
-
# --------------------------------------------------------------------------- #
|
| 105 |
-
|
| 106 |
-
|
| 107 |
-
class QuotaTrackerRegistry:
|
| 108 |
-
"""Process-global registry of all :class:`QuotaTracker` instances.
|
| 109 |
-
|
| 110 |
-
Typical usage::
|
| 111 |
-
|
| 112 |
-
registry = get_quota_registry()
|
| 113 |
-
registry.register(SubscriptionTracker(...))
|
| 114 |
-
registry.register(get_codex_rate_limit_state())
|
| 115 |
-
registry.register(get_copilot_quota_tracker())
|
| 116 |
-
|
| 117 |
-
# server startup
|
| 118 |
-
await registry.start_all()
|
| 119 |
-
|
| 120 |
-
# /stats assembly
|
| 121 |
-
stats.update(registry.get_all_stats())
|
| 122 |
-
|
| 123 |
-
# server shutdown
|
| 124 |
-
await registry.stop_all()
|
| 125 |
-
"""
|
| 126 |
-
|
| 127 |
-
def __init__(self) -> None:
|
| 128 |
-
self._trackers: list[QuotaTracker] = []
|
| 129 |
-
self._lock = Lock()
|
| 130 |
-
|
| 131 |
-
# ------------------------------------------------------------------ #
|
| 132 |
-
# Registration
|
| 133 |
-
# ------------------------------------------------------------------ #
|
| 134 |
-
|
| 135 |
-
def register(self, tracker: QuotaTracker) -> None:
|
| 136 |
-
"""Register a tracker. Duplicate keys are rejected."""
|
| 137 |
-
with self._lock:
|
| 138 |
-
existing_keys = {t.key for t in self._trackers}
|
| 139 |
-
if tracker.key in existing_keys:
|
| 140 |
-
raise ValueError(
|
| 141 |
-
f"A tracker with key '{tracker.key}' is already registered. "
|
| 142 |
-
"Each tracker must have a unique key."
|
| 143 |
-
)
|
| 144 |
-
self._trackers.append(tracker)
|
| 145 |
-
|
| 146 |
-
def get(self, key: str) -> QuotaTracker | None:
|
| 147 |
-
"""Return the registered tracker for *key*, or ``None``."""
|
| 148 |
-
with self._lock:
|
| 149 |
-
for t in self._trackers:
|
| 150 |
-
if t.key == key:
|
| 151 |
-
return t
|
| 152 |
-
return None
|
| 153 |
-
|
| 154 |
-
@property
|
| 155 |
-
def trackers(self) -> list[QuotaTracker]:
|
| 156 |
-
"""Read-only snapshot of the registered tracker list."""
|
| 157 |
-
with self._lock:
|
| 158 |
-
return list(self._trackers)
|
| 159 |
-
|
| 160 |
-
# ------------------------------------------------------------------ #
|
| 161 |
-
# Lifecycle
|
| 162 |
-
# ------------------------------------------------------------------ #
|
| 163 |
-
|
| 164 |
-
async def start_all(self) -> None:
|
| 165 |
-
"""Start every available tracker and log its status."""
|
| 166 |
-
for tracker in self.trackers:
|
| 167 |
-
if tracker.is_available():
|
| 168 |
-
await tracker.start()
|
| 169 |
-
logger.info("%s quota tracking: ENABLED", tracker.label)
|
| 170 |
-
else:
|
| 171 |
-
logger.info("%s quota tracking: DISABLED (not available)", tracker.label)
|
| 172 |
-
|
| 173 |
-
async def stop_all(self) -> None:
|
| 174 |
-
"""Stop all registered trackers (regardless of availability)."""
|
| 175 |
-
for tracker in self.trackers:
|
| 176 |
-
try:
|
| 177 |
-
await tracker.stop()
|
| 178 |
-
except Exception as exc: # noqa: BLE001
|
| 179 |
-
logger.warning("Error stopping %s tracker: %s", tracker.label, exc)
|
| 180 |
-
|
| 181 |
-
# ------------------------------------------------------------------ #
|
| 182 |
-
# Stats
|
| 183 |
-
# ------------------------------------------------------------------ #
|
| 184 |
-
|
| 185 |
-
def get_all_stats(self) -> dict[str, dict[str, Any] | None]:
|
| 186 |
-
"""Return ``{key: stats_dict}`` for every available tracker.
|
| 187 |
-
|
| 188 |
-
Trackers that are unavailable or return ``None`` are excluded.
|
| 189 |
-
"""
|
| 190 |
-
result: dict[str, dict[str, Any] | None] = {}
|
| 191 |
-
for tracker in self.trackers:
|
| 192 |
-
if not tracker.is_available():
|
| 193 |
-
continue
|
| 194 |
-
stats = tracker.get_stats()
|
| 195 |
-
if stats is not None:
|
| 196 |
-
result[tracker.key] = stats
|
| 197 |
-
return result
|
| 198 |
-
|
| 199 |
-
def get_stats(self, key: str) -> dict[str, Any] | None:
|
| 200 |
-
"""Return stats for a single tracker by key, or ``None``."""
|
| 201 |
-
tracker = self.get(key)
|
| 202 |
-
return tracker.get_stats() if tracker is not None else None
|
| 203 |
-
|
| 204 |
-
|
| 205 |
-
# --------------------------------------------------------------------------- #
|
| 206 |
-
# Process-global singleton
|
| 207 |
-
# --------------------------------------------------------------------------- #
|
| 208 |
-
|
| 209 |
-
_registry: QuotaTrackerRegistry | None = None
|
| 210 |
-
_registry_lock = Lock()
|
| 211 |
-
|
| 212 |
-
|
| 213 |
-
def get_quota_registry() -> QuotaTrackerRegistry:
|
| 214 |
-
"""Return the process-global :class:`QuotaTrackerRegistry` singleton."""
|
| 215 |
-
global _registry
|
| 216 |
-
if _registry is None:
|
| 217 |
-
with _registry_lock:
|
| 218 |
-
if _registry is None:
|
| 219 |
-
_registry = QuotaTrackerRegistry()
|
| 220 |
-
return _registry
|
| 221 |
-
|
| 222 |
-
|
| 223 |
-
def reset_quota_registry() -> None:
|
| 224 |
-
"""Replace the global registry with a fresh empty instance.
|
| 225 |
-
|
| 226 |
-
Intended for use in tests only.
|
| 227 |
-
"""
|
| 228 |
-
global _registry
|
| 229 |
-
with _registry_lock:
|
| 230 |
-
_registry = QuotaTrackerRegistry()
|
|
|
|
| 1 |
+
"""Base abstractions for pluggable AI-tool quota / rate-limit trackers.
|
| 2 |
+
|
| 3 |
+
Every provider tracker (Anthropic, Codex, Copilot, …) inherits from
|
| 4 |
+
:class:`QuotaTracker` and is registered with the process-global
|
| 5 |
+
:class:`QuotaTrackerRegistry`. ``server.py`` only interacts with the
|
| 6 |
+
registry — adding a new provider requires *zero* changes to the server.
|
| 7 |
+
|
| 8 |
+
Quick-start for a new provider::
|
| 9 |
+
|
| 10 |
+
from headroom.subscription.base import QuotaTracker, get_quota_registry
|
| 11 |
+
|
| 12 |
+
class GeminiQuotaTracker(QuotaTracker):
|
| 13 |
+
key = "gemini_quota"
|
| 14 |
+
label = "Google Gemini"
|
| 15 |
+
|
| 16 |
+
def is_available(self) -> bool:
|
| 17 |
+
return bool(os.environ.get("GOOGLE_API_KEY"))
|
| 18 |
+
|
| 19 |
+
async def start(self) -> None: ... # launch background poll
|
| 20 |
+
async def stop(self) -> None: ... # cancel poll task
|
| 21 |
+
|
| 22 |
+
def get_stats(self) -> dict | None:
|
| 23 |
+
return ... # serialisable dict or None if no data yet
|
| 24 |
+
|
| 25 |
+
get_quota_registry().register(GeminiQuotaTracker())
|
| 26 |
+
"""
|
| 27 |
+
|
| 28 |
+
from __future__ import annotations
|
| 29 |
+
|
| 30 |
+
import abc
|
| 31 |
+
import logging
|
| 32 |
+
from threading import Lock
|
| 33 |
+
from typing import Any
|
| 34 |
+
|
| 35 |
+
logger = logging.getLogger(__name__)
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
class QuotaTracker(abc.ABC):
|
| 39 |
+
"""Abstract base for a single AI-tool quota / rate-limit tracker.
|
| 40 |
+
|
| 41 |
+
Subclasses must define :attr:`key`, :attr:`label`, and
|
| 42 |
+
:meth:`get_stats`. All other methods have sensible defaults.
|
| 43 |
+
"""
|
| 44 |
+
|
| 45 |
+
# ------------------------------------------------------------------ #
|
| 46 |
+
# Class-level identity — subclasses should override as class attributes
|
| 47 |
+
# ------------------------------------------------------------------ #
|
| 48 |
+
|
| 49 |
+
@property
|
| 50 |
+
@abc.abstractmethod
|
| 51 |
+
def key(self) -> str:
|
| 52 |
+
"""Stats key used in ``/stats`` and the dashboard.
|
| 53 |
+
|
| 54 |
+
Must be unique across all registered trackers.
|
| 55 |
+
Examples: ``"subscription_window"``, ``"codex_rate_limits"``.
|
| 56 |
+
"""
|
| 57 |
+
|
| 58 |
+
@property
|
| 59 |
+
@abc.abstractmethod
|
| 60 |
+
def label(self) -> str:
|
| 61 |
+
"""Human-readable name for log messages.
|
| 62 |
+
|
| 63 |
+
Example: ``"Anthropic Claude Code"``.
|
| 64 |
+
"""
|
| 65 |
+
|
| 66 |
+
# ------------------------------------------------------------------ #
|
| 67 |
+
# Availability gate
|
| 68 |
+
# ------------------------------------------------------------------ #
|
| 69 |
+
|
| 70 |
+
def is_available(self) -> bool:
|
| 71 |
+
"""Return ``True`` if this tracker should be activated.
|
| 72 |
+
|
| 73 |
+
Override to gate on environment variables, config flags, etc.
|
| 74 |
+
The registry calls this before :meth:`start` and skips trackers
|
| 75 |
+
that return ``False``. Default: always available.
|
| 76 |
+
"""
|
| 77 |
+
return True
|
| 78 |
+
|
| 79 |
+
# ------------------------------------------------------------------ #
|
| 80 |
+
# Lifecycle — default no-ops (suitable for passive/header-based trackers)
|
| 81 |
+
# ------------------------------------------------------------------ #
|
| 82 |
+
|
| 83 |
+
async def start(self) -> None: # noqa: B027
|
| 84 |
+
"""Start background polling. No-op for passive trackers."""
|
| 85 |
+
|
| 86 |
+
async def stop(self) -> None: # noqa: B027
|
| 87 |
+
"""Stop background polling. No-op for passive trackers."""
|
| 88 |
+
|
| 89 |
+
# ------------------------------------------------------------------ #
|
| 90 |
+
# Stats
|
| 91 |
+
# ------------------------------------------------------------------ #
|
| 92 |
+
|
| 93 |
+
@abc.abstractmethod
|
| 94 |
+
def get_stats(self) -> dict[str, Any] | None:
|
| 95 |
+
"""Return the current snapshot as a serialisable dict, or ``None``.
|
| 96 |
+
|
| 97 |
+
``None`` means "no data yet" and causes the key to be omitted from
|
| 98 |
+
``/stats`` rather than appearing as ``null``.
|
| 99 |
+
"""
|
| 100 |
+
|
| 101 |
+
|
| 102 |
+
# --------------------------------------------------------------------------- #
|
| 103 |
+
# Registry
|
| 104 |
+
# --------------------------------------------------------------------------- #
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
class QuotaTrackerRegistry:
|
| 108 |
+
"""Process-global registry of all :class:`QuotaTracker` instances.
|
| 109 |
+
|
| 110 |
+
Typical usage::
|
| 111 |
+
|
| 112 |
+
registry = get_quota_registry()
|
| 113 |
+
registry.register(SubscriptionTracker(...))
|
| 114 |
+
registry.register(get_codex_rate_limit_state())
|
| 115 |
+
registry.register(get_copilot_quota_tracker())
|
| 116 |
+
|
| 117 |
+
# server startup
|
| 118 |
+
await registry.start_all()
|
| 119 |
+
|
| 120 |
+
# /stats assembly
|
| 121 |
+
stats.update(registry.get_all_stats())
|
| 122 |
+
|
| 123 |
+
# server shutdown
|
| 124 |
+
await registry.stop_all()
|
| 125 |
+
"""
|
| 126 |
+
|
| 127 |
+
def __init__(self) -> None:
|
| 128 |
+
self._trackers: list[QuotaTracker] = []
|
| 129 |
+
self._lock = Lock()
|
| 130 |
+
|
| 131 |
+
# ------------------------------------------------------------------ #
|
| 132 |
+
# Registration
|
| 133 |
+
# ------------------------------------------------------------------ #
|
| 134 |
+
|
| 135 |
+
def register(self, tracker: QuotaTracker) -> None:
|
| 136 |
+
"""Register a tracker. Duplicate keys are rejected."""
|
| 137 |
+
with self._lock:
|
| 138 |
+
existing_keys = {t.key for t in self._trackers}
|
| 139 |
+
if tracker.key in existing_keys:
|
| 140 |
+
raise ValueError(
|
| 141 |
+
f"A tracker with key '{tracker.key}' is already registered. "
|
| 142 |
+
"Each tracker must have a unique key."
|
| 143 |
+
)
|
| 144 |
+
self._trackers.append(tracker)
|
| 145 |
+
|
| 146 |
+
def get(self, key: str) -> QuotaTracker | None:
|
| 147 |
+
"""Return the registered tracker for *key*, or ``None``."""
|
| 148 |
+
with self._lock:
|
| 149 |
+
for t in self._trackers:
|
| 150 |
+
if t.key == key:
|
| 151 |
+
return t
|
| 152 |
+
return None
|
| 153 |
+
|
| 154 |
+
@property
|
| 155 |
+
def trackers(self) -> list[QuotaTracker]:
|
| 156 |
+
"""Read-only snapshot of the registered tracker list."""
|
| 157 |
+
with self._lock:
|
| 158 |
+
return list(self._trackers)
|
| 159 |
+
|
| 160 |
+
# ------------------------------------------------------------------ #
|
| 161 |
+
# Lifecycle
|
| 162 |
+
# ------------------------------------------------------------------ #
|
| 163 |
+
|
| 164 |
+
async def start_all(self) -> None:
|
| 165 |
+
"""Start every available tracker and log its status."""
|
| 166 |
+
for tracker in self.trackers:
|
| 167 |
+
if tracker.is_available():
|
| 168 |
+
await tracker.start()
|
| 169 |
+
logger.info("%s quota tracking: ENABLED", tracker.label)
|
| 170 |
+
else:
|
| 171 |
+
logger.info("%s quota tracking: DISABLED (not available)", tracker.label)
|
| 172 |
+
|
| 173 |
+
async def stop_all(self) -> None:
|
| 174 |
+
"""Stop all registered trackers (regardless of availability)."""
|
| 175 |
+
for tracker in self.trackers:
|
| 176 |
+
try:
|
| 177 |
+
await tracker.stop()
|
| 178 |
+
except Exception as exc: # noqa: BLE001
|
| 179 |
+
logger.warning("Error stopping %s tracker: %s", tracker.label, exc)
|
| 180 |
+
|
| 181 |
+
# ------------------------------------------------------------------ #
|
| 182 |
+
# Stats
|
| 183 |
+
# ------------------------------------------------------------------ #
|
| 184 |
+
|
| 185 |
+
def get_all_stats(self) -> dict[str, dict[str, Any] | None]:
|
| 186 |
+
"""Return ``{key: stats_dict}`` for every available tracker.
|
| 187 |
+
|
| 188 |
+
Trackers that are unavailable or return ``None`` are excluded.
|
| 189 |
+
"""
|
| 190 |
+
result: dict[str, dict[str, Any] | None] = {}
|
| 191 |
+
for tracker in self.trackers:
|
| 192 |
+
if not tracker.is_available():
|
| 193 |
+
continue
|
| 194 |
+
stats = tracker.get_stats()
|
| 195 |
+
if stats is not None:
|
| 196 |
+
result[tracker.key] = stats
|
| 197 |
+
return result
|
| 198 |
+
|
| 199 |
+
def get_stats(self, key: str) -> dict[str, Any] | None:
|
| 200 |
+
"""Return stats for a single tracker by key, or ``None``."""
|
| 201 |
+
tracker = self.get(key)
|
| 202 |
+
return tracker.get_stats() if tracker is not None else None
|
| 203 |
+
|
| 204 |
+
|
| 205 |
+
# --------------------------------------------------------------------------- #
|
| 206 |
+
# Process-global singleton
|
| 207 |
+
# --------------------------------------------------------------------------- #
|
| 208 |
+
|
| 209 |
+
_registry: QuotaTrackerRegistry | None = None
|
| 210 |
+
_registry_lock = Lock()
|
| 211 |
+
|
| 212 |
+
|
| 213 |
+
def get_quota_registry() -> QuotaTrackerRegistry:
|
| 214 |
+
"""Return the process-global :class:`QuotaTrackerRegistry` singleton."""
|
| 215 |
+
global _registry
|
| 216 |
+
if _registry is None:
|
| 217 |
+
with _registry_lock:
|
| 218 |
+
if _registry is None:
|
| 219 |
+
_registry = QuotaTrackerRegistry()
|
| 220 |
+
return _registry
|
| 221 |
+
|
| 222 |
+
|
| 223 |
+
def reset_quota_registry() -> None:
|
| 224 |
+
"""Replace the global registry with a fresh empty instance.
|
| 225 |
+
|
| 226 |
+
Intended for use in tests only.
|
| 227 |
+
"""
|
| 228 |
+
global _registry
|
| 229 |
+
with _registry_lock:
|
| 230 |
+
_registry = QuotaTrackerRegistry()
|
|
@@ -1,131 +1,131 @@
|
|
| 1 |
-
"""Async HTTP client for Anthropic's OAuth usage API.
|
| 2 |
-
|
| 3 |
-
Endpoint: GET https://api.anthropic.com/api/oauth/usage
|
| 4 |
-
Required header: anthropic-beta: oauth-2025-04-20
|
| 5 |
-
Auth: Authorization: Bearer <oauth_access_token>
|
| 6 |
-
|
| 7 |
-
Token resolution order (highest → lowest priority):
|
| 8 |
-
1. Explicit token passed to :meth:`fetch`
|
| 9 |
-
2. ``CLAUDE_CODE_OAUTH_TOKEN`` env-var
|
| 10 |
-
3. ``~/.claude/.credentials.json`` → ``claudeAiOauth.accessToken``
|
| 11 |
-
(respects ``CLAUDE_CONFIG_DIR`` env-var override)
|
| 12 |
-
"""
|
| 13 |
-
|
| 14 |
-
from __future__ import annotations
|
| 15 |
-
|
| 16 |
-
import json
|
| 17 |
-
import logging
|
| 18 |
-
import os
|
| 19 |
-
from pathlib import Path
|
| 20 |
-
from typing import Any
|
| 21 |
-
|
| 22 |
-
import httpx
|
| 23 |
-
|
| 24 |
-
from headroom.subscription.models import SubscriptionSnapshot
|
| 25 |
-
|
| 26 |
-
logger = logging.getLogger(__name__)
|
| 27 |
-
|
| 28 |
-
_USAGE_URL = "https://api.anthropic.com/api/oauth/usage"
|
| 29 |
-
_BETA_HEADER = "oauth-2025-04-20"
|
| 30 |
-
_TOKEN_EXPIRY_BUFFER_S = 60
|
| 31 |
-
|
| 32 |
-
|
| 33 |
-
def _credentials_path() -> Path:
|
| 34 |
-
base = os.environ.get("CLAUDE_CONFIG_DIR") or os.path.expanduser("~/.claude")
|
| 35 |
-
return Path(base) / ".credentials.json"
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
def _load_credentials_file() -> dict[str, Any] | None:
|
| 39 |
-
"""Load raw credentials dict from the Claude Code credentials file."""
|
| 40 |
-
path = _credentials_path()
|
| 41 |
-
try:
|
| 42 |
-
with path.open() as fh:
|
| 43 |
-
return json.load(fh) # type: ignore[no-any-return]
|
| 44 |
-
except FileNotFoundError:
|
| 45 |
-
return None
|
| 46 |
-
except Exception as exc:
|
| 47 |
-
logger.debug("Cannot read credentials file %s: %s", path, exc)
|
| 48 |
-
return None
|
| 49 |
-
|
| 50 |
-
|
| 51 |
-
def read_cached_oauth_token() -> str | None:
|
| 52 |
-
"""Resolve a stored OAuth token for background polling (no request needed).
|
| 53 |
-
|
| 54 |
-
Returns the raw access token string if found and not expired, else None.
|
| 55 |
-
"""
|
| 56 |
-
# 1. Env var
|
| 57 |
-
env_token = os.environ.get("CLAUDE_CODE_OAUTH_TOKEN", "").strip()
|
| 58 |
-
if env_token:
|
| 59 |
-
return env_token
|
| 60 |
-
|
| 61 |
-
# 2. Credentials file
|
| 62 |
-
creds = _load_credentials_file()
|
| 63 |
-
if not creds:
|
| 64 |
-
return None
|
| 65 |
-
oauth = creds.get("claudeAiOauth") or {}
|
| 66 |
-
token = oauth.get("accessToken") or ""
|
| 67 |
-
if not token:
|
| 68 |
-
return None
|
| 69 |
-
|
| 70 |
-
# Check expiry (Anthropic stores timestamp in milliseconds)
|
| 71 |
-
expires_at_ms = oauth.get("expiresAt")
|
| 72 |
-
if expires_at_ms is not None:
|
| 73 |
-
import time
|
| 74 |
-
|
| 75 |
-
now_ms = time.time() * 1000
|
| 76 |
-
if now_ms >= (expires_at_ms - _TOKEN_EXPIRY_BUFFER_S * 1000):
|
| 77 |
-
logger.debug("Cached OAuth token expired; skipping background poll")
|
| 78 |
-
return None
|
| 79 |
-
|
| 80 |
-
return token
|
| 81 |
-
|
| 82 |
-
|
| 83 |
-
class SubscriptionClient:
|
| 84 |
-
"""Thin async wrapper around the Anthropic OAuth usage endpoint."""
|
| 85 |
-
|
| 86 |
-
def __init__(self, timeout: float = 10.0) -> None:
|
| 87 |
-
self._timeout = timeout
|
| 88 |
-
|
| 89 |
-
async def fetch(self, token: str | None = None) -> SubscriptionSnapshot | None:
|
| 90 |
-
"""Fetch current subscription window data.
|
| 91 |
-
|
| 92 |
-
:param token: OAuth access token. When *None*, falls back to
|
| 93 |
-
:func:`read_cached_oauth_token`.
|
| 94 |
-
:returns: :class:`SubscriptionSnapshot` or *None* on auth failure /
|
| 95 |
-
unsupported account.
|
| 96 |
-
"""
|
| 97 |
-
resolved = (token or "").strip() or read_cached_oauth_token()
|
| 98 |
-
if not resolved:
|
| 99 |
-
logger.debug("No OAuth token available for subscription polling")
|
| 100 |
-
return None
|
| 101 |
-
|
| 102 |
-
headers = {
|
| 103 |
-
"Authorization": f"Bearer {resolved}",
|
| 104 |
-
"anthropic-beta": _BETA_HEADER,
|
| 105 |
-
"Content-Type": "application/json",
|
| 106 |
-
}
|
| 107 |
-
|
| 108 |
-
try:
|
| 109 |
-
async with httpx.AsyncClient(timeout=self._timeout) as client:
|
| 110 |
-
resp = await client.get(_USAGE_URL, headers=headers)
|
| 111 |
-
|
| 112 |
-
if resp.status_code == 401:
|
| 113 |
-
logger.debug("OAuth token rejected (401) by Anthropic usage API")
|
| 114 |
-
return None
|
| 115 |
-
if resp.status_code == 404:
|
| 116 |
-
# API key accounts (non-subscription) return 404
|
| 117 |
-
logger.debug("Subscription usage API returned 404; likely API-key account")
|
| 118 |
-
return None
|
| 119 |
-
if resp.status_code != 200:
|
| 120 |
-
logger.warning("Anthropic usage API returned %s", resp.status_code)
|
| 121 |
-
return None
|
| 122 |
-
|
| 123 |
-
data: dict[str, Any] = resp.json()
|
| 124 |
-
return SubscriptionSnapshot.from_api_response(data, token=resolved)
|
| 125 |
-
|
| 126 |
-
except httpx.TimeoutException:
|
| 127 |
-
logger.debug("Timeout fetching Anthropic subscription window")
|
| 128 |
-
return None
|
| 129 |
-
except Exception as exc:
|
| 130 |
-
logger.warning("Error fetching subscription window: %s", exc)
|
| 131 |
-
return None
|
|
|
|
| 1 |
+
"""Async HTTP client for Anthropic's OAuth usage API.
|
| 2 |
+
|
| 3 |
+
Endpoint: GET https://api.anthropic.com/api/oauth/usage
|
| 4 |
+
Required header: anthropic-beta: oauth-2025-04-20
|
| 5 |
+
Auth: Authorization: Bearer <oauth_access_token>
|
| 6 |
+
|
| 7 |
+
Token resolution order (highest → lowest priority):
|
| 8 |
+
1. Explicit token passed to :meth:`fetch`
|
| 9 |
+
2. ``CLAUDE_CODE_OAUTH_TOKEN`` env-var
|
| 10 |
+
3. ``~/.claude/.credentials.json`` → ``claudeAiOauth.accessToken``
|
| 11 |
+
(respects ``CLAUDE_CONFIG_DIR`` env-var override)
|
| 12 |
+
"""
|
| 13 |
+
|
| 14 |
+
from __future__ import annotations
|
| 15 |
+
|
| 16 |
+
import json
|
| 17 |
+
import logging
|
| 18 |
+
import os
|
| 19 |
+
from pathlib import Path
|
| 20 |
+
from typing import Any
|
| 21 |
+
|
| 22 |
+
import httpx
|
| 23 |
+
|
| 24 |
+
from headroom.subscription.models import SubscriptionSnapshot
|
| 25 |
+
|
| 26 |
+
logger = logging.getLogger(__name__)
|
| 27 |
+
|
| 28 |
+
_USAGE_URL = "https://api.anthropic.com/api/oauth/usage"
|
| 29 |
+
_BETA_HEADER = "oauth-2025-04-20"
|
| 30 |
+
_TOKEN_EXPIRY_BUFFER_S = 60
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
def _credentials_path() -> Path:
|
| 34 |
+
base = os.environ.get("CLAUDE_CONFIG_DIR") or os.path.expanduser("~/.claude")
|
| 35 |
+
return Path(base) / ".credentials.json"
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
def _load_credentials_file() -> dict[str, Any] | None:
|
| 39 |
+
"""Load raw credentials dict from the Claude Code credentials file."""
|
| 40 |
+
path = _credentials_path()
|
| 41 |
+
try:
|
| 42 |
+
with path.open() as fh:
|
| 43 |
+
return json.load(fh) # type: ignore[no-any-return]
|
| 44 |
+
except FileNotFoundError:
|
| 45 |
+
return None
|
| 46 |
+
except Exception as exc:
|
| 47 |
+
logger.debug("Cannot read credentials file %s: %s", path, exc)
|
| 48 |
+
return None
|
| 49 |
+
|
| 50 |
+
|
| 51 |
+
def read_cached_oauth_token() -> str | None:
|
| 52 |
+
"""Resolve a stored OAuth token for background polling (no request needed).
|
| 53 |
+
|
| 54 |
+
Returns the raw access token string if found and not expired, else None.
|
| 55 |
+
"""
|
| 56 |
+
# 1. Env var
|
| 57 |
+
env_token = os.environ.get("CLAUDE_CODE_OAUTH_TOKEN", "").strip()
|
| 58 |
+
if env_token:
|
| 59 |
+
return env_token
|
| 60 |
+
|
| 61 |
+
# 2. Credentials file
|
| 62 |
+
creds = _load_credentials_file()
|
| 63 |
+
if not creds:
|
| 64 |
+
return None
|
| 65 |
+
oauth = creds.get("claudeAiOauth") or {}
|
| 66 |
+
token = oauth.get("accessToken") or ""
|
| 67 |
+
if not token:
|
| 68 |
+
return None
|
| 69 |
+
|
| 70 |
+
# Check expiry (Anthropic stores timestamp in milliseconds)
|
| 71 |
+
expires_at_ms = oauth.get("expiresAt")
|
| 72 |
+
if expires_at_ms is not None:
|
| 73 |
+
import time
|
| 74 |
+
|
| 75 |
+
now_ms = time.time() * 1000
|
| 76 |
+
if now_ms >= (expires_at_ms - _TOKEN_EXPIRY_BUFFER_S * 1000):
|
| 77 |
+
logger.debug("Cached OAuth token expired; skipping background poll")
|
| 78 |
+
return None
|
| 79 |
+
|
| 80 |
+
return token
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
class SubscriptionClient:
|
| 84 |
+
"""Thin async wrapper around the Anthropic OAuth usage endpoint."""
|
| 85 |
+
|
| 86 |
+
def __init__(self, timeout: float = 10.0) -> None:
|
| 87 |
+
self._timeout = timeout
|
| 88 |
+
|
| 89 |
+
async def fetch(self, token: str | None = None) -> SubscriptionSnapshot | None:
|
| 90 |
+
"""Fetch current subscription window data.
|
| 91 |
+
|
| 92 |
+
:param token: OAuth access token. When *None*, falls back to
|
| 93 |
+
:func:`read_cached_oauth_token`.
|
| 94 |
+
:returns: :class:`SubscriptionSnapshot` or *None* on auth failure /
|
| 95 |
+
unsupported account.
|
| 96 |
+
"""
|
| 97 |
+
resolved = (token or "").strip() or read_cached_oauth_token()
|
| 98 |
+
if not resolved:
|
| 99 |
+
logger.debug("No OAuth token available for subscription polling")
|
| 100 |
+
return None
|
| 101 |
+
|
| 102 |
+
headers = {
|
| 103 |
+
"Authorization": f"Bearer {resolved}",
|
| 104 |
+
"anthropic-beta": _BETA_HEADER,
|
| 105 |
+
"Content-Type": "application/json",
|
| 106 |
+
}
|
| 107 |
+
|
| 108 |
+
try:
|
| 109 |
+
async with httpx.AsyncClient(timeout=self._timeout) as client:
|
| 110 |
+
resp = await client.get(_USAGE_URL, headers=headers)
|
| 111 |
+
|
| 112 |
+
if resp.status_code == 401:
|
| 113 |
+
logger.debug("OAuth token rejected (401) by Anthropic usage API")
|
| 114 |
+
return None
|
| 115 |
+
if resp.status_code == 404:
|
| 116 |
+
# API key accounts (non-subscription) return 404
|
| 117 |
+
logger.debug("Subscription usage API returned 404; likely API-key account")
|
| 118 |
+
return None
|
| 119 |
+
if resp.status_code != 200:
|
| 120 |
+
logger.warning("Anthropic usage API returned %s", resp.status_code)
|
| 121 |
+
return None
|
| 122 |
+
|
| 123 |
+
data: dict[str, Any] = resp.json()
|
| 124 |
+
return SubscriptionSnapshot.from_api_response(data, token=resolved)
|
| 125 |
+
|
| 126 |
+
except httpx.TimeoutException:
|
| 127 |
+
logger.debug("Timeout fetching Anthropic subscription window")
|
| 128 |
+
return None
|
| 129 |
+
except Exception as exc:
|
| 130 |
+
logger.warning("Error fetching subscription window: %s", exc)
|
| 131 |
+
return None
|
|
@@ -1,247 +1,247 @@
|
|
| 1 |
-
"""Passive tracking of OpenAI Codex rate-limit window data from response headers.
|
| 2 |
-
|
| 3 |
-
Codex (OpenAI) embeds rate-limit data directly in API response headers
|
| 4 |
-
(``x-codex-primary-used-percent``, ``x-codex-primary-window-minutes``, etc.)
|
| 5 |
-
rather than exposing a dedicated usage endpoint. This module captures those
|
| 6 |
-
headers from responses that headroom proxies and makes them available in
|
| 7 |
-
``/stats`` and the dashboard.
|
| 8 |
-
|
| 9 |
-
Header schema (parsed by codex-rs ``rate_limits.rs``):
|
| 10 |
-
x-codex-primary-used-percent float 0-100
|
| 11 |
-
x-codex-primary-window-minutes int window size
|
| 12 |
-
x-codex-primary-reset-at int Unix timestamp (seconds)
|
| 13 |
-
x-codex-secondary-used-percent float 0-100 (optional)
|
| 14 |
-
x-codex-secondary-window-minutes int (optional)
|
| 15 |
-
x-codex-secondary-reset-at int (optional)
|
| 16 |
-
x-codex-credits-has-credits bool
|
| 17 |
-
x-codex-credits-unlimited bool
|
| 18 |
-
x-codex-credits-balance str e.g. "$5.00"
|
| 19 |
-
x-codex-promo-message str server announcement
|
| 20 |
-
x-codex-limit-name str e.g. "gpt-5.2-codex-sonic"
|
| 21 |
-
"""
|
| 22 |
-
|
| 23 |
-
from __future__ import annotations
|
| 24 |
-
|
| 25 |
-
import time
|
| 26 |
-
from dataclasses import dataclass, field
|
| 27 |
-
from threading import Lock
|
| 28 |
-
|
| 29 |
-
from headroom.subscription.base import QuotaTracker
|
| 30 |
-
|
| 31 |
-
|
| 32 |
-
@dataclass
|
| 33 |
-
class CodexRateLimitWindow:
|
| 34 |
-
"""Usage data for a single rolling rate-limit window."""
|
| 35 |
-
|
| 36 |
-
used_percent: float
|
| 37 |
-
window_minutes: int | None = None
|
| 38 |
-
resets_at: int | None = None # Unix timestamp (seconds)
|
| 39 |
-
|
| 40 |
-
@property
|
| 41 |
-
def window_label(self) -> str:
|
| 42 |
-
if self.window_minutes is None:
|
| 43 |
-
return "unknown"
|
| 44 |
-
if self.window_minutes < 60:
|
| 45 |
-
return f"{self.window_minutes}m"
|
| 46 |
-
hours = self.window_minutes // 60
|
| 47 |
-
mins = self.window_minutes % 60
|
| 48 |
-
return f"{hours}h{mins:02d}m" if mins else f"{hours}h"
|
| 49 |
-
|
| 50 |
-
@property
|
| 51 |
-
def seconds_until_reset(self) -> int | None:
|
| 52 |
-
if self.resets_at is None:
|
| 53 |
-
return None
|
| 54 |
-
return max(0, self.resets_at - int(time.time()))
|
| 55 |
-
|
| 56 |
-
def to_dict(self) -> dict:
|
| 57 |
-
return {
|
| 58 |
-
"used_percent": self.used_percent,
|
| 59 |
-
"window_minutes": self.window_minutes,
|
| 60 |
-
"window_label": self.window_label,
|
| 61 |
-
"resets_at": self.resets_at,
|
| 62 |
-
"seconds_until_reset": self.seconds_until_reset,
|
| 63 |
-
}
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
@dataclass
|
| 67 |
-
class CodexCreditsSnapshot:
|
| 68 |
-
"""OpenAI credits balance for Codex."""
|
| 69 |
-
|
| 70 |
-
has_credits: bool
|
| 71 |
-
unlimited: bool
|
| 72 |
-
balance: str | None = None
|
| 73 |
-
|
| 74 |
-
def to_dict(self) -> dict:
|
| 75 |
-
return {
|
| 76 |
-
"has_credits": self.has_credits,
|
| 77 |
-
"unlimited": self.unlimited,
|
| 78 |
-
"balance": self.balance,
|
| 79 |
-
}
|
| 80 |
-
|
| 81 |
-
|
| 82 |
-
@dataclass
|
| 83 |
-
class CodexRateLimitSnapshot:
|
| 84 |
-
"""Full rate-limit snapshot parsed from a single Codex API response."""
|
| 85 |
-
|
| 86 |
-
limit_id: str = "codex"
|
| 87 |
-
limit_name: str | None = None
|
| 88 |
-
primary: CodexRateLimitWindow | None = None
|
| 89 |
-
secondary: CodexRateLimitWindow | None = None
|
| 90 |
-
credits: CodexCreditsSnapshot | None = None
|
| 91 |
-
promo_message: str | None = None
|
| 92 |
-
captured_at: float = field(default_factory=time.time)
|
| 93 |
-
|
| 94 |
-
def to_dict(self) -> dict:
|
| 95 |
-
return {
|
| 96 |
-
"limit_id": self.limit_id,
|
| 97 |
-
"limit_name": self.limit_name,
|
| 98 |
-
"primary": self.primary.to_dict() if self.primary else None,
|
| 99 |
-
"secondary": self.secondary.to_dict() if self.secondary else None,
|
| 100 |
-
"credits": self.credits.to_dict() if self.credits else None,
|
| 101 |
-
"promo_message": self.promo_message,
|
| 102 |
-
"captured_at": self.captured_at,
|
| 103 |
-
}
|
| 104 |
-
|
| 105 |
-
|
| 106 |
-
# ---------------------------------------------------------------------------
|
| 107 |
-
# Header parsing helpers
|
| 108 |
-
# ---------------------------------------------------------------------------
|
| 109 |
-
|
| 110 |
-
|
| 111 |
-
def _parse_float(headers: dict[str, str], name: str) -> float | None:
|
| 112 |
-
raw = headers.get(name)
|
| 113 |
-
if raw is None:
|
| 114 |
-
return None
|
| 115 |
-
try:
|
| 116 |
-
v = float(raw)
|
| 117 |
-
return v if v == v else None # NaN guard
|
| 118 |
-
except (ValueError, TypeError):
|
| 119 |
-
return None
|
| 120 |
-
|
| 121 |
-
|
| 122 |
-
def _parse_int(headers: dict[str, str], name: str) -> int | None:
|
| 123 |
-
raw = headers.get(name)
|
| 124 |
-
if raw is None:
|
| 125 |
-
return None
|
| 126 |
-
try:
|
| 127 |
-
return int(raw)
|
| 128 |
-
except (ValueError, TypeError):
|
| 129 |
-
return None
|
| 130 |
-
|
| 131 |
-
|
| 132 |
-
def _parse_bool(headers: dict[str, str], name: str) -> bool | None:
|
| 133 |
-
raw = headers.get(name)
|
| 134 |
-
if raw is None:
|
| 135 |
-
return None
|
| 136 |
-
if raw.lower() in ("true", "1"):
|
| 137 |
-
return True
|
| 138 |
-
if raw.lower() in ("false", "0"):
|
| 139 |
-
return False
|
| 140 |
-
return None
|
| 141 |
-
|
| 142 |
-
|
| 143 |
-
def _parse_window(headers: dict[str, str], prefix: str, which: str) -> CodexRateLimitWindow | None:
|
| 144 |
-
used_pct = _parse_float(headers, f"{prefix}-{which}-used-percent")
|
| 145 |
-
if used_pct is None:
|
| 146 |
-
return None
|
| 147 |
-
return CodexRateLimitWindow(
|
| 148 |
-
used_percent=used_pct,
|
| 149 |
-
window_minutes=_parse_int(headers, f"{prefix}-{which}-window-minutes"),
|
| 150 |
-
resets_at=_parse_int(headers, f"{prefix}-{which}-reset-at"),
|
| 151 |
-
)
|
| 152 |
-
|
| 153 |
-
|
| 154 |
-
def _parse_credits(headers: dict[str, str]) -> CodexCreditsSnapshot | None:
|
| 155 |
-
has_credits = _parse_bool(headers, "x-codex-credits-has-credits")
|
| 156 |
-
if has_credits is None:
|
| 157 |
-
return None
|
| 158 |
-
unlimited = _parse_bool(headers, "x-codex-credits-unlimited") or False
|
| 159 |
-
raw_balance = headers.get("x-codex-credits-balance", "").strip()
|
| 160 |
-
return CodexCreditsSnapshot(
|
| 161 |
-
has_credits=has_credits,
|
| 162 |
-
unlimited=unlimited,
|
| 163 |
-
balance=raw_balance or None,
|
| 164 |
-
)
|
| 165 |
-
|
| 166 |
-
|
| 167 |
-
def parse_codex_rate_limits(headers: dict[str, str]) -> CodexRateLimitSnapshot | None:
|
| 168 |
-
"""Parse a :class:`CodexRateLimitSnapshot` from a dict of HTTP response headers.
|
| 169 |
-
|
| 170 |
-
Returns ``None`` when no Codex rate-limit headers are present (e.g. the
|
| 171 |
-
response came from a non-Codex OpenAI endpoint or a cached reply).
|
| 172 |
-
"""
|
| 173 |
-
prefix = "x-codex"
|
| 174 |
-
primary = _parse_window(headers, prefix, "primary")
|
| 175 |
-
secondary = _parse_window(headers, prefix, "secondary")
|
| 176 |
-
credits = _parse_credits(headers)
|
| 177 |
-
raw_promo = headers.get("x-codex-promo-message", "").strip()
|
| 178 |
-
promo = raw_promo or None
|
| 179 |
-
raw_limit_name = headers.get("x-codex-limit-name", "").strip()
|
| 180 |
-
limit_name = raw_limit_name or None
|
| 181 |
-
|
| 182 |
-
if primary is None and secondary is None and credits is None and promo is None:
|
| 183 |
-
return None # Not a Codex response with rate-limit headers
|
| 184 |
-
|
| 185 |
-
return CodexRateLimitSnapshot(
|
| 186 |
-
limit_id="codex",
|
| 187 |
-
limit_name=limit_name,
|
| 188 |
-
primary=primary,
|
| 189 |
-
secondary=secondary,
|
| 190 |
-
credits=credits,
|
| 191 |
-
promo_message=promo,
|
| 192 |
-
)
|
| 193 |
-
|
| 194 |
-
|
| 195 |
-
# ---------------------------------------------------------------------------
|
| 196 |
-
# Singleton state store
|
| 197 |
-
# ---------------------------------------------------------------------------
|
| 198 |
-
|
| 199 |
-
|
| 200 |
-
class CodexRateLimitState(QuotaTracker):
|
| 201 |
-
"""Thread-safe store for the latest Codex rate-limit snapshot.
|
| 202 |
-
|
| 203 |
-
Implements :class:`~headroom.subscription.base.QuotaTracker` so it can
|
| 204 |
-
be registered with the :class:`~headroom.subscription.base.QuotaTrackerRegistry`.
|
| 205 |
-
This tracker is *passive* — it is updated by the OpenAI proxy handler
|
| 206 |
-
each time a response containing ``x-codex-*`` headers passes through
|
| 207 |
-
headroom, so :meth:`start` and :meth:`stop` are no-ops.
|
| 208 |
-
"""
|
| 209 |
-
|
| 210 |
-
# QuotaTracker identity
|
| 211 |
-
key = "codex_rate_limits"
|
| 212 |
-
label = "OpenAI Codex"
|
| 213 |
-
|
| 214 |
-
def __init__(self) -> None:
|
| 215 |
-
self._lock = Lock()
|
| 216 |
-
self._latest: CodexRateLimitSnapshot | None = None
|
| 217 |
-
|
| 218 |
-
def update_from_headers(self, headers: dict[str, str]) -> None:
|
| 219 |
-
"""Update state from a response header dict (no-op if no Codex headers)."""
|
| 220 |
-
snapshot = parse_codex_rate_limits(headers)
|
| 221 |
-
if snapshot is None:
|
| 222 |
-
return
|
| 223 |
-
with self._lock:
|
| 224 |
-
self._latest = snapshot
|
| 225 |
-
|
| 226 |
-
@property
|
| 227 |
-
def latest(self) -> CodexRateLimitSnapshot | None:
|
| 228 |
-
with self._lock:
|
| 229 |
-
return self._latest
|
| 230 |
-
|
| 231 |
-
def get_stats(self) -> dict | None:
|
| 232 |
-
snap = self.latest
|
| 233 |
-
return snap.to_dict() if snap is not None else None
|
| 234 |
-
|
| 235 |
-
|
| 236 |
-
_state: CodexRateLimitState | None = None
|
| 237 |
-
_state_lock = Lock()
|
| 238 |
-
|
| 239 |
-
|
| 240 |
-
def get_codex_rate_limit_state() -> CodexRateLimitState:
|
| 241 |
-
"""Return the process-global :class:`CodexRateLimitState` singleton."""
|
| 242 |
-
global _state
|
| 243 |
-
if _state is None:
|
| 244 |
-
with _state_lock:
|
| 245 |
-
if _state is None:
|
| 246 |
-
_state = CodexRateLimitState()
|
| 247 |
-
return _state
|
|
|
|
| 1 |
+
"""Passive tracking of OpenAI Codex rate-limit window data from response headers.
|
| 2 |
+
|
| 3 |
+
Codex (OpenAI) embeds rate-limit data directly in API response headers
|
| 4 |
+
(``x-codex-primary-used-percent``, ``x-codex-primary-window-minutes``, etc.)
|
| 5 |
+
rather than exposing a dedicated usage endpoint. This module captures those
|
| 6 |
+
headers from responses that headroom proxies and makes them available in
|
| 7 |
+
``/stats`` and the dashboard.
|
| 8 |
+
|
| 9 |
+
Header schema (parsed by codex-rs ``rate_limits.rs``):
|
| 10 |
+
x-codex-primary-used-percent float 0-100
|
| 11 |
+
x-codex-primary-window-minutes int window size
|
| 12 |
+
x-codex-primary-reset-at int Unix timestamp (seconds)
|
| 13 |
+
x-codex-secondary-used-percent float 0-100 (optional)
|
| 14 |
+
x-codex-secondary-window-minutes int (optional)
|
| 15 |
+
x-codex-secondary-reset-at int (optional)
|
| 16 |
+
x-codex-credits-has-credits bool
|
| 17 |
+
x-codex-credits-unlimited bool
|
| 18 |
+
x-codex-credits-balance str e.g. "$5.00"
|
| 19 |
+
x-codex-promo-message str server announcement
|
| 20 |
+
x-codex-limit-name str e.g. "gpt-5.2-codex-sonic"
|
| 21 |
+
"""
|
| 22 |
+
|
| 23 |
+
from __future__ import annotations
|
| 24 |
+
|
| 25 |
+
import time
|
| 26 |
+
from dataclasses import dataclass, field
|
| 27 |
+
from threading import Lock
|
| 28 |
+
|
| 29 |
+
from headroom.subscription.base import QuotaTracker
|
| 30 |
+
|
| 31 |
+
|
| 32 |
+
@dataclass
|
| 33 |
+
class CodexRateLimitWindow:
|
| 34 |
+
"""Usage data for a single rolling rate-limit window."""
|
| 35 |
+
|
| 36 |
+
used_percent: float
|
| 37 |
+
window_minutes: int | None = None
|
| 38 |
+
resets_at: int | None = None # Unix timestamp (seconds)
|
| 39 |
+
|
| 40 |
+
@property
|
| 41 |
+
def window_label(self) -> str:
|
| 42 |
+
if self.window_minutes is None:
|
| 43 |
+
return "unknown"
|
| 44 |
+
if self.window_minutes < 60:
|
| 45 |
+
return f"{self.window_minutes}m"
|
| 46 |
+
hours = self.window_minutes // 60
|
| 47 |
+
mins = self.window_minutes % 60
|
| 48 |
+
return f"{hours}h{mins:02d}m" if mins else f"{hours}h"
|
| 49 |
+
|
| 50 |
+
@property
|
| 51 |
+
def seconds_until_reset(self) -> int | None:
|
| 52 |
+
if self.resets_at is None:
|
| 53 |
+
return None
|
| 54 |
+
return max(0, self.resets_at - int(time.time()))
|
| 55 |
+
|
| 56 |
+
def to_dict(self) -> dict:
|
| 57 |
+
return {
|
| 58 |
+
"used_percent": self.used_percent,
|
| 59 |
+
"window_minutes": self.window_minutes,
|
| 60 |
+
"window_label": self.window_label,
|
| 61 |
+
"resets_at": self.resets_at,
|
| 62 |
+
"seconds_until_reset": self.seconds_until_reset,
|
| 63 |
+
}
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
@dataclass
|
| 67 |
+
class CodexCreditsSnapshot:
|
| 68 |
+
"""OpenAI credits balance for Codex."""
|
| 69 |
+
|
| 70 |
+
has_credits: bool
|
| 71 |
+
unlimited: bool
|
| 72 |
+
balance: str | None = None
|
| 73 |
+
|
| 74 |
+
def to_dict(self) -> dict:
|
| 75 |
+
return {
|
| 76 |
+
"has_credits": self.has_credits,
|
| 77 |
+
"unlimited": self.unlimited,
|
| 78 |
+
"balance": self.balance,
|
| 79 |
+
}
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
@dataclass
|
| 83 |
+
class CodexRateLimitSnapshot:
|
| 84 |
+
"""Full rate-limit snapshot parsed from a single Codex API response."""
|
| 85 |
+
|
| 86 |
+
limit_id: str = "codex"
|
| 87 |
+
limit_name: str | None = None
|
| 88 |
+
primary: CodexRateLimitWindow | None = None
|
| 89 |
+
secondary: CodexRateLimitWindow | None = None
|
| 90 |
+
credits: CodexCreditsSnapshot | None = None
|
| 91 |
+
promo_message: str | None = None
|
| 92 |
+
captured_at: float = field(default_factory=time.time)
|
| 93 |
+
|
| 94 |
+
def to_dict(self) -> dict:
|
| 95 |
+
return {
|
| 96 |
+
"limit_id": self.limit_id,
|
| 97 |
+
"limit_name": self.limit_name,
|
| 98 |
+
"primary": self.primary.to_dict() if self.primary else None,
|
| 99 |
+
"secondary": self.secondary.to_dict() if self.secondary else None,
|
| 100 |
+
"credits": self.credits.to_dict() if self.credits else None,
|
| 101 |
+
"promo_message": self.promo_message,
|
| 102 |
+
"captured_at": self.captured_at,
|
| 103 |
+
}
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
# ---------------------------------------------------------------------------
|
| 107 |
+
# Header parsing helpers
|
| 108 |
+
# ---------------------------------------------------------------------------
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
def _parse_float(headers: dict[str, str], name: str) -> float | None:
|
| 112 |
+
raw = headers.get(name)
|
| 113 |
+
if raw is None:
|
| 114 |
+
return None
|
| 115 |
+
try:
|
| 116 |
+
v = float(raw)
|
| 117 |
+
return v if v == v else None # NaN guard
|
| 118 |
+
except (ValueError, TypeError):
|
| 119 |
+
return None
|
| 120 |
+
|
| 121 |
+
|
| 122 |
+
def _parse_int(headers: dict[str, str], name: str) -> int | None:
|
| 123 |
+
raw = headers.get(name)
|
| 124 |
+
if raw is None:
|
| 125 |
+
return None
|
| 126 |
+
try:
|
| 127 |
+
return int(raw)
|
| 128 |
+
except (ValueError, TypeError):
|
| 129 |
+
return None
|
| 130 |
+
|
| 131 |
+
|
| 132 |
+
def _parse_bool(headers: dict[str, str], name: str) -> bool | None:
|
| 133 |
+
raw = headers.get(name)
|
| 134 |
+
if raw is None:
|
| 135 |
+
return None
|
| 136 |
+
if raw.lower() in ("true", "1"):
|
| 137 |
+
return True
|
| 138 |
+
if raw.lower() in ("false", "0"):
|
| 139 |
+
return False
|
| 140 |
+
return None
|
| 141 |
+
|
| 142 |
+
|
| 143 |
+
def _parse_window(headers: dict[str, str], prefix: str, which: str) -> CodexRateLimitWindow | None:
|
| 144 |
+
used_pct = _parse_float(headers, f"{prefix}-{which}-used-percent")
|
| 145 |
+
if used_pct is None:
|
| 146 |
+
return None
|
| 147 |
+
return CodexRateLimitWindow(
|
| 148 |
+
used_percent=used_pct,
|
| 149 |
+
window_minutes=_parse_int(headers, f"{prefix}-{which}-window-minutes"),
|
| 150 |
+
resets_at=_parse_int(headers, f"{prefix}-{which}-reset-at"),
|
| 151 |
+
)
|
| 152 |
+
|
| 153 |
+
|
| 154 |
+
def _parse_credits(headers: dict[str, str]) -> CodexCreditsSnapshot | None:
|
| 155 |
+
has_credits = _parse_bool(headers, "x-codex-credits-has-credits")
|
| 156 |
+
if has_credits is None:
|
| 157 |
+
return None
|
| 158 |
+
unlimited = _parse_bool(headers, "x-codex-credits-unlimited") or False
|
| 159 |
+
raw_balance = headers.get("x-codex-credits-balance", "").strip()
|
| 160 |
+
return CodexCreditsSnapshot(
|
| 161 |
+
has_credits=has_credits,
|
| 162 |
+
unlimited=unlimited,
|
| 163 |
+
balance=raw_balance or None,
|
| 164 |
+
)
|
| 165 |
+
|
| 166 |
+
|
| 167 |
+
def parse_codex_rate_limits(headers: dict[str, str]) -> CodexRateLimitSnapshot | None:
|
| 168 |
+
"""Parse a :class:`CodexRateLimitSnapshot` from a dict of HTTP response headers.
|
| 169 |
+
|
| 170 |
+
Returns ``None`` when no Codex rate-limit headers are present (e.g. the
|
| 171 |
+
response came from a non-Codex OpenAI endpoint or a cached reply).
|
| 172 |
+
"""
|
| 173 |
+
prefix = "x-codex"
|
| 174 |
+
primary = _parse_window(headers, prefix, "primary")
|
| 175 |
+
secondary = _parse_window(headers, prefix, "secondary")
|
| 176 |
+
credits = _parse_credits(headers)
|
| 177 |
+
raw_promo = headers.get("x-codex-promo-message", "").strip()
|
| 178 |
+
promo = raw_promo or None
|
| 179 |
+
raw_limit_name = headers.get("x-codex-limit-name", "").strip()
|
| 180 |
+
limit_name = raw_limit_name or None
|
| 181 |
+
|
| 182 |
+
if primary is None and secondary is None and credits is None and promo is None:
|
| 183 |
+
return None # Not a Codex response with rate-limit headers
|
| 184 |
+
|
| 185 |
+
return CodexRateLimitSnapshot(
|
| 186 |
+
limit_id="codex",
|
| 187 |
+
limit_name=limit_name,
|
| 188 |
+
primary=primary,
|
| 189 |
+
secondary=secondary,
|
| 190 |
+
credits=credits,
|
| 191 |
+
promo_message=promo,
|
| 192 |
+
)
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
# ---------------------------------------------------------------------------
|
| 196 |
+
# Singleton state store
|
| 197 |
+
# ---------------------------------------------------------------------------
|
| 198 |
+
|
| 199 |
+
|
| 200 |
+
class CodexRateLimitState(QuotaTracker):
|
| 201 |
+
"""Thread-safe store for the latest Codex rate-limit snapshot.
|
| 202 |
+
|
| 203 |
+
Implements :class:`~headroom.subscription.base.QuotaTracker` so it can
|
| 204 |
+
be registered with the :class:`~headroom.subscription.base.QuotaTrackerRegistry`.
|
| 205 |
+
This tracker is *passive* — it is updated by the OpenAI proxy handler
|
| 206 |
+
each time a response containing ``x-codex-*`` headers passes through
|
| 207 |
+
headroom, so :meth:`start` and :meth:`stop` are no-ops.
|
| 208 |
+
"""
|
| 209 |
+
|
| 210 |
+
# QuotaTracker identity
|
| 211 |
+
key = "codex_rate_limits"
|
| 212 |
+
label = "OpenAI Codex"
|
| 213 |
+
|
| 214 |
+
def __init__(self) -> None:
|
| 215 |
+
self._lock = Lock()
|
| 216 |
+
self._latest: CodexRateLimitSnapshot | None = None
|
| 217 |
+
|
| 218 |
+
def update_from_headers(self, headers: dict[str, str]) -> None:
|
| 219 |
+
"""Update state from a response header dict (no-op if no Codex headers)."""
|
| 220 |
+
snapshot = parse_codex_rate_limits(headers)
|
| 221 |
+
if snapshot is None:
|
| 222 |
+
return
|
| 223 |
+
with self._lock:
|
| 224 |
+
self._latest = snapshot
|
| 225 |
+
|
| 226 |
+
@property
|
| 227 |
+
def latest(self) -> CodexRateLimitSnapshot | None:
|
| 228 |
+
with self._lock:
|
| 229 |
+
return self._latest
|
| 230 |
+
|
| 231 |
+
def get_stats(self) -> dict | None:
|
| 232 |
+
snap = self.latest
|
| 233 |
+
return snap.to_dict() if snap is not None else None
|
| 234 |
+
|
| 235 |
+
|
| 236 |
+
_state: CodexRateLimitState | None = None
|
| 237 |
+
_state_lock = Lock()
|
| 238 |
+
|
| 239 |
+
|
| 240 |
+
def get_codex_rate_limit_state() -> CodexRateLimitState:
|
| 241 |
+
"""Return the process-global :class:`CodexRateLimitState` singleton."""
|
| 242 |
+
global _state
|
| 243 |
+
if _state is None:
|
| 244 |
+
with _state_lock:
|
| 245 |
+
if _state is None:
|
| 246 |
+
_state = CodexRateLimitState()
|
| 247 |
+
return _state
|
|
@@ -1,366 +1,366 @@
|
|
| 1 |
-
"""GitHub Copilot monthly quota tracking via the copilot_internal/user API.
|
| 2 |
-
|
| 3 |
-
GitHub Copilot exposes per-category monthly quotas (chat, completions,
|
| 4 |
-
premium_interactions) at ``GET https://api.github.com/copilot_internal/user``
|
| 5 |
-
authenticated with a GitHub Bearer token.
|
| 6 |
-
|
| 7 |
-
Token discovery order (first non-empty wins):
|
| 8 |
-
1. GITHUB_COPILOT_GITHUB_TOKEN
|
| 9 |
-
2. GITHUB_TOKEN
|
| 10 |
-
3. COPILOT_GITHUB_TOKEN
|
| 11 |
-
4. GITHUB_COPILOT_API_TOKEN
|
| 12 |
-
|
| 13 |
-
The poll is triggered every ``poll_interval_s`` seconds (default 60) as long as
|
| 14 |
-
a GitHub token is available. No proxy traffic interception is needed — tokens
|
| 15 |
-
come from the environment at headroom start-up.
|
| 16 |
-
|
| 17 |
-
API response schema (relevant fields from ``/copilot_internal/user``):
|
| 18 |
-
login str GitHub username
|
| 19 |
-
copilot_plan str "free" | "individual" | "business" | "enterprise"
|
| 20 |
-
access_type_sku str plan SKU string
|
| 21 |
-
quota_reset_date_utc str ISO-8601 date when monthly quota resets
|
| 22 |
-
quota_snapshots:
|
| 23 |
-
chat / completions / premium_interactions:
|
| 24 |
-
entitlement int total monthly allocation
|
| 25 |
-
remaining int remaining uses this month
|
| 26 |
-
quota_remaining int (alias for remaining)
|
| 27 |
-
percent_remaining float 0-100
|
| 28 |
-
overage_count int uses beyond entitlement
|
| 29 |
-
overage_permitted bool whether overage is allowed
|
| 30 |
-
unlimited bool whether this category is unlimited
|
| 31 |
-
timestamp_utc str when the snapshot was recorded
|
| 32 |
-
"""
|
| 33 |
-
|
| 34 |
-
from __future__ import annotations
|
| 35 |
-
|
| 36 |
-
import asyncio
|
| 37 |
-
import logging
|
| 38 |
-
import os
|
| 39 |
-
import time
|
| 40 |
-
from dataclasses import dataclass, field
|
| 41 |
-
from threading import Lock
|
| 42 |
-
from typing import Any
|
| 43 |
-
|
| 44 |
-
from headroom.subscription.base import QuotaTracker
|
| 45 |
-
|
| 46 |
-
logger = logging.getLogger(__name__)
|
| 47 |
-
|
| 48 |
-
_GITHUB_API_BASE = "https://api.github.com"
|
| 49 |
-
_TOKEN_ENV_VARS = [
|
| 50 |
-
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 51 |
-
"GITHUB_TOKEN",
|
| 52 |
-
"COPILOT_GITHUB_TOKEN",
|
| 53 |
-
"GITHUB_COPILOT_API_TOKEN",
|
| 54 |
-
]
|
| 55 |
-
|
| 56 |
-
# Categories surfaced by the quota_snapshots endpoint
|
| 57 |
-
QUOTA_CATEGORIES = ("chat", "completions", "premium_interactions")
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
# ---------------------------------------------------------------------------
|
| 61 |
-
# Data classes
|
| 62 |
-
# ---------------------------------------------------------------------------
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
@dataclass
|
| 66 |
-
class CopilotQuotaCategory:
|
| 67 |
-
"""Quota data for a single Copilot usage category."""
|
| 68 |
-
|
| 69 |
-
name: str
|
| 70 |
-
entitlement: int | None = None # total monthly allocation
|
| 71 |
-
remaining: int | None = None # remaining uses
|
| 72 |
-
percent_remaining: float | None = None # 0-100
|
| 73 |
-
overage_count: int = 0 # uses beyond entitlement
|
| 74 |
-
overage_permitted: bool = False
|
| 75 |
-
unlimited: bool = False
|
| 76 |
-
timestamp_utc: str | None = None
|
| 77 |
-
|
| 78 |
-
@property
|
| 79 |
-
def used(self) -> int | None:
|
| 80 |
-
if self.entitlement is not None and self.remaining is not None:
|
| 81 |
-
return max(0, self.entitlement - self.remaining)
|
| 82 |
-
return None
|
| 83 |
-
|
| 84 |
-
@property
|
| 85 |
-
def used_percent(self) -> float | None:
|
| 86 |
-
if self.unlimited:
|
| 87 |
-
return 0.0
|
| 88 |
-
if self.percent_remaining is not None:
|
| 89 |
-
return max(0.0, 100.0 - self.percent_remaining)
|
| 90 |
-
if self.entitlement and self.entitlement > 0 and self.remaining is not None:
|
| 91 |
-
return 100.0 * (self.entitlement - self.remaining) / self.entitlement
|
| 92 |
-
return None
|
| 93 |
-
|
| 94 |
-
def to_dict(self) -> dict[str, Any]:
|
| 95 |
-
return {
|
| 96 |
-
"name": self.name,
|
| 97 |
-
"entitlement": self.entitlement,
|
| 98 |
-
"remaining": self.remaining,
|
| 99 |
-
"used": self.used,
|
| 100 |
-
"percent_remaining": self.percent_remaining,
|
| 101 |
-
"used_percent": self.used_percent,
|
| 102 |
-
"overage_count": self.overage_count,
|
| 103 |
-
"overage_permitted": self.overage_permitted,
|
| 104 |
-
"unlimited": self.unlimited,
|
| 105 |
-
"timestamp_utc": self.timestamp_utc,
|
| 106 |
-
}
|
| 107 |
-
|
| 108 |
-
|
| 109 |
-
@dataclass
|
| 110 |
-
class CopilotQuotaSnapshot:
|
| 111 |
-
"""Full quota snapshot from one /copilot_internal/user response."""
|
| 112 |
-
|
| 113 |
-
login: str | None = None
|
| 114 |
-
copilot_plan: str | None = None
|
| 115 |
-
access_type_sku: str | None = None
|
| 116 |
-
quota_reset_date_utc: str | None = None
|
| 117 |
-
categories: dict[str, CopilotQuotaCategory] = field(default_factory=dict)
|
| 118 |
-
fetched_at: float = field(default_factory=time.time)
|
| 119 |
-
|
| 120 |
-
def to_dict(self) -> dict[str, Any]:
|
| 121 |
-
return {
|
| 122 |
-
"login": self.login,
|
| 123 |
-
"copilot_plan": self.copilot_plan,
|
| 124 |
-
"access_type_sku": self.access_type_sku,
|
| 125 |
-
"quota_reset_date_utc": self.quota_reset_date_utc,
|
| 126 |
-
"categories": {k: v.to_dict() for k, v in self.categories.items()},
|
| 127 |
-
"fetched_at": self.fetched_at,
|
| 128 |
-
}
|
| 129 |
-
|
| 130 |
-
|
| 131 |
-
@dataclass
|
| 132 |
-
class CopilotQuotaState:
|
| 133 |
-
"""Thread-safe singleton state for Copilot quota tracking."""
|
| 134 |
-
|
| 135 |
-
latest: CopilotQuotaSnapshot | None = None
|
| 136 |
-
last_error: str | None = None
|
| 137 |
-
last_updated: float | None = None
|
| 138 |
-
|
| 139 |
-
def to_dict(self) -> dict[str, Any]:
|
| 140 |
-
return {
|
| 141 |
-
"latest": self.latest.to_dict() if self.latest else None,
|
| 142 |
-
"last_error": self.last_error,
|
| 143 |
-
"last_updated": self.last_updated,
|
| 144 |
-
}
|
| 145 |
-
|
| 146 |
-
|
| 147 |
-
# ---------------------------------------------------------------------------
|
| 148 |
-
# Parsing
|
| 149 |
-
# ---------------------------------------------------------------------------
|
| 150 |
-
|
| 151 |
-
|
| 152 |
-
def parse_copilot_quota(data: dict[str, Any]) -> CopilotQuotaSnapshot:
|
| 153 |
-
"""Parse a /copilot_internal/user response into a ``CopilotQuotaSnapshot``."""
|
| 154 |
-
snapshot = CopilotQuotaSnapshot(
|
| 155 |
-
login=data.get("login"),
|
| 156 |
-
copilot_plan=data.get("copilot_plan"),
|
| 157 |
-
access_type_sku=data.get("access_type_sku"),
|
| 158 |
-
quota_reset_date_utc=data.get("quota_reset_date_utc") or data.get("quota_reset_date"),
|
| 159 |
-
)
|
| 160 |
-
|
| 161 |
-
raw_qs = data.get("quota_snapshots") or {}
|
| 162 |
-
for cat_name in QUOTA_CATEGORIES:
|
| 163 |
-
raw = raw_qs.get(cat_name)
|
| 164 |
-
if not raw:
|
| 165 |
-
continue
|
| 166 |
-
remaining = raw.get("remaining") or raw.get("quota_remaining")
|
| 167 |
-
cat = CopilotQuotaCategory(
|
| 168 |
-
name=cat_name,
|
| 169 |
-
entitlement=raw.get("entitlement"),
|
| 170 |
-
remaining=remaining,
|
| 171 |
-
percent_remaining=raw.get("percent_remaining"),
|
| 172 |
-
overage_count=raw.get("overage_count") or 0,
|
| 173 |
-
overage_permitted=bool(raw.get("overage_permitted")),
|
| 174 |
-
unlimited=bool(raw.get("unlimited")),
|
| 175 |
-
timestamp_utc=raw.get("timestamp_utc"),
|
| 176 |
-
)
|
| 177 |
-
snapshot.categories[cat_name] = cat
|
| 178 |
-
|
| 179 |
-
return snapshot
|
| 180 |
-
|
| 181 |
-
|
| 182 |
-
# ---------------------------------------------------------------------------
|
| 183 |
-
# Token discovery
|
| 184 |
-
# ---------------------------------------------------------------------------
|
| 185 |
-
|
| 186 |
-
|
| 187 |
-
def discover_github_token() -> str | None:
|
| 188 |
-
"""Return the first GitHub token found from known environment variables."""
|
| 189 |
-
for var in _TOKEN_ENV_VARS:
|
| 190 |
-
val = os.environ.get(var, "").strip()
|
| 191 |
-
if val:
|
| 192 |
-
return val
|
| 193 |
-
return None
|
| 194 |
-
|
| 195 |
-
|
| 196 |
-
# ---------------------------------------------------------------------------
|
| 197 |
-
# Background tracker singleton
|
| 198 |
-
# ---------------------------------------------------------------------------
|
| 199 |
-
|
| 200 |
-
|
| 201 |
-
class _CopilotQuotaTracker(QuotaTracker):
|
| 202 |
-
"""Singleton background poller for GitHub Copilot quota.
|
| 203 |
-
|
| 204 |
-
Implements :class:`~headroom.subscription.base.QuotaTracker` so it can be
|
| 205 |
-
registered with the :class:`~headroom.subscription.base.QuotaTrackerRegistry`.
|
| 206 |
-
Availability is gated on a GitHub token being present in the environment.
|
| 207 |
-
"""
|
| 208 |
-
|
| 209 |
-
# QuotaTracker identity
|
| 210 |
-
key = "copilot_quota"
|
| 211 |
-
label = "GitHub Copilot"
|
| 212 |
-
|
| 213 |
-
def __init__(self, poll_interval_s: float = 60.0) -> None:
|
| 214 |
-
self._poll_interval_s = poll_interval_s
|
| 215 |
-
self._state = CopilotQuotaState()
|
| 216 |
-
self._lock = Lock()
|
| 217 |
-
self._stop_event: asyncio.Event | None = None
|
| 218 |
-
self._task: asyncio.Task | None = None # type: ignore[type-arg]
|
| 219 |
-
|
| 220 |
-
# ------------------------------------------------------------------
|
| 221 |
-
# QuotaTracker interface
|
| 222 |
-
# ------------------------------------------------------------------
|
| 223 |
-
|
| 224 |
-
def is_available(self) -> bool:
|
| 225 |
-
"""Returns ``True`` when a GitHub token is available in the environment."""
|
| 226 |
-
return discover_github_token() is not None
|
| 227 |
-
|
| 228 |
-
def get_stats(self) -> dict[str, Any] | None:
|
| 229 |
-
"""Return the latest quota state dict, or ``None`` if no data yet."""
|
| 230 |
-
data = self.state
|
| 231 |
-
if not data.get("latest"):
|
| 232 |
-
return None
|
| 233 |
-
return data
|
| 234 |
-
|
| 235 |
-
# ------------------------------------------------------------------
|
| 236 |
-
# Lifecycle
|
| 237 |
-
# ------------------------------------------------------------------
|
| 238 |
-
|
| 239 |
-
async def start(self) -> None:
|
| 240 |
-
"""Start the background polling loop."""
|
| 241 |
-
if self._task is not None and not self._task.done():
|
| 242 |
-
return
|
| 243 |
-
self._stop_event = asyncio.Event()
|
| 244 |
-
self._task = asyncio.create_task(self._poll_loop())
|
| 245 |
-
|
| 246 |
-
async def stop(self) -> None:
|
| 247 |
-
"""Stop the polling loop."""
|
| 248 |
-
if self._stop_event:
|
| 249 |
-
self._stop_event.set()
|
| 250 |
-
if self._task:
|
| 251 |
-
try:
|
| 252 |
-
await asyncio.wait_for(self._task, timeout=5.0)
|
| 253 |
-
except (asyncio.TimeoutError, asyncio.CancelledError):
|
| 254 |
-
# Mirror SubscriptionTracker.stop(): on timeout or outer
|
| 255 |
-
# cancellation, cancel the underlying poll task. Without
|
| 256 |
-
# this, a wedged poll task would leak past ``stop()``.
|
| 257 |
-
self._task.cancel()
|
| 258 |
-
|
| 259 |
-
# ------------------------------------------------------------------
|
| 260 |
-
# State
|
| 261 |
-
# ------------------------------------------------------------------
|
| 262 |
-
|
| 263 |
-
@property
|
| 264 |
-
def state(self) -> dict[str, Any]:
|
| 265 |
-
with self._lock:
|
| 266 |
-
return self._state.to_dict()
|
| 267 |
-
|
| 268 |
-
# ------------------------------------------------------------------
|
| 269 |
-
# Poll loop
|
| 270 |
-
# ------------------------------------------------------------------
|
| 271 |
-
|
| 272 |
-
async def _poll_loop(self) -> None:
|
| 273 |
-
assert self._stop_event is not None
|
| 274 |
-
while not self._stop_event.is_set():
|
| 275 |
-
try:
|
| 276 |
-
await self._maybe_poll()
|
| 277 |
-
except Exception as exc:
|
| 278 |
-
logger.warning("Copilot quota poll error: %s", exc)
|
| 279 |
-
try:
|
| 280 |
-
# NOTE: do NOT wrap in asyncio.shield() — shield prevents the
|
| 281 |
-
# inner Event.wait() from being cancelled when wait_for times
|
| 282 |
-
# out, leaking one Task per poll interval. See the matching
|
| 283 |
-
# note in headroom/subscription/tracker.py:_poll_loop.
|
| 284 |
-
await asyncio.wait_for(
|
| 285 |
-
self._stop_event.wait(),
|
| 286 |
-
timeout=self._poll_interval_s,
|
| 287 |
-
)
|
| 288 |
-
break # stop event was set
|
| 289 |
-
except asyncio.TimeoutError:
|
| 290 |
-
pass # normal: poll interval elapsed
|
| 291 |
-
|
| 292 |
-
async def _maybe_poll(self) -> None:
|
| 293 |
-
token = discover_github_token()
|
| 294 |
-
if not token:
|
| 295 |
-
return
|
| 296 |
-
|
| 297 |
-
try:
|
| 298 |
-
import aiohttp
|
| 299 |
-
except ImportError:
|
| 300 |
-
logger.debug("aiohttp not available; skipping Copilot quota poll")
|
| 301 |
-
return
|
| 302 |
-
|
| 303 |
-
url = f"{_GITHUB_API_BASE}/copilot_internal/user"
|
| 304 |
-
headers = {
|
| 305 |
-
"Authorization": f"Bearer {token}",
|
| 306 |
-
"Accept": "application/json",
|
| 307 |
-
"X-GitHub-Api-Version": "2022-11-28",
|
| 308 |
-
}
|
| 309 |
-
|
| 310 |
-
try:
|
| 311 |
-
async with aiohttp.ClientSession() as session:
|
| 312 |
-
async with session.get(
|
| 313 |
-
url, headers=headers, timeout=aiohttp.ClientTimeout(total=10)
|
| 314 |
-
) as resp:
|
| 315 |
-
if resp.status == 401:
|
| 316 |
-
with self._lock:
|
| 317 |
-
self._state.last_error = "unauthorized — check GITHUB_TOKEN"
|
| 318 |
-
return
|
| 319 |
-
if resp.status == 404:
|
| 320 |
-
# API-key-only account or endpoint not available
|
| 321 |
-
with self._lock:
|
| 322 |
-
self._state.last_error = "endpoint not found (non-Copilot account?)"
|
| 323 |
-
return
|
| 324 |
-
if not resp.ok:
|
| 325 |
-
with self._lock:
|
| 326 |
-
self._state.last_error = f"HTTP {resp.status}"
|
| 327 |
-
return
|
| 328 |
-
|
| 329 |
-
data = await resp.json()
|
| 330 |
-
except Exception as exc:
|
| 331 |
-
with self._lock:
|
| 332 |
-
self._state.last_error = str(exc)
|
| 333 |
-
logger.debug("Copilot quota fetch failed: %s", exc)
|
| 334 |
-
return
|
| 335 |
-
|
| 336 |
-
try:
|
| 337 |
-
snapshot = parse_copilot_quota(data)
|
| 338 |
-
except Exception as exc:
|
| 339 |
-
with self._lock:
|
| 340 |
-
self._state.last_error = f"parse error: {exc}"
|
| 341 |
-
return
|
| 342 |
-
|
| 343 |
-
with self._lock:
|
| 344 |
-
self._state.latest = snapshot
|
| 345 |
-
self._state.last_error = None
|
| 346 |
-
self._state.last_updated = time.time()
|
| 347 |
-
|
| 348 |
-
logger.debug(
|
| 349 |
-
"Copilot quota polled: plan=%s categories=%s",
|
| 350 |
-
snapshot.copilot_plan,
|
| 351 |
-
list(snapshot.categories.keys()),
|
| 352 |
-
)
|
| 353 |
-
|
| 354 |
-
|
| 355 |
-
_singleton_lock = Lock()
|
| 356 |
-
_singleton: _CopilotQuotaTracker | None = None
|
| 357 |
-
|
| 358 |
-
|
| 359 |
-
def get_copilot_quota_tracker() -> _CopilotQuotaTracker:
|
| 360 |
-
"""Return the global ``_CopilotQuotaTracker`` singleton."""
|
| 361 |
-
global _singleton
|
| 362 |
-
if _singleton is None:
|
| 363 |
-
with _singleton_lock:
|
| 364 |
-
if _singleton is None:
|
| 365 |
-
_singleton = _CopilotQuotaTracker()
|
| 366 |
-
return _singleton
|
|
|
|
| 1 |
+
"""GitHub Copilot monthly quota tracking via the copilot_internal/user API.
|
| 2 |
+
|
| 3 |
+
GitHub Copilot exposes per-category monthly quotas (chat, completions,
|
| 4 |
+
premium_interactions) at ``GET https://api.github.com/copilot_internal/user``
|
| 5 |
+
authenticated with a GitHub Bearer token.
|
| 6 |
+
|
| 7 |
+
Token discovery order (first non-empty wins):
|
| 8 |
+
1. GITHUB_COPILOT_GITHUB_TOKEN
|
| 9 |
+
2. GITHUB_TOKEN
|
| 10 |
+
3. COPILOT_GITHUB_TOKEN
|
| 11 |
+
4. GITHUB_COPILOT_API_TOKEN
|
| 12 |
+
|
| 13 |
+
The poll is triggered every ``poll_interval_s`` seconds (default 60) as long as
|
| 14 |
+
a GitHub token is available. No proxy traffic interception is needed — tokens
|
| 15 |
+
come from the environment at headroom start-up.
|
| 16 |
+
|
| 17 |
+
API response schema (relevant fields from ``/copilot_internal/user``):
|
| 18 |
+
login str GitHub username
|
| 19 |
+
copilot_plan str "free" | "individual" | "business" | "enterprise"
|
| 20 |
+
access_type_sku str plan SKU string
|
| 21 |
+
quota_reset_date_utc str ISO-8601 date when monthly quota resets
|
| 22 |
+
quota_snapshots:
|
| 23 |
+
chat / completions / premium_interactions:
|
| 24 |
+
entitlement int total monthly allocation
|
| 25 |
+
remaining int remaining uses this month
|
| 26 |
+
quota_remaining int (alias for remaining)
|
| 27 |
+
percent_remaining float 0-100
|
| 28 |
+
overage_count int uses beyond entitlement
|
| 29 |
+
overage_permitted bool whether overage is allowed
|
| 30 |
+
unlimited bool whether this category is unlimited
|
| 31 |
+
timestamp_utc str when the snapshot was recorded
|
| 32 |
+
"""
|
| 33 |
+
|
| 34 |
+
from __future__ import annotations
|
| 35 |
+
|
| 36 |
+
import asyncio
|
| 37 |
+
import logging
|
| 38 |
+
import os
|
| 39 |
+
import time
|
| 40 |
+
from dataclasses import dataclass, field
|
| 41 |
+
from threading import Lock
|
| 42 |
+
from typing import Any
|
| 43 |
+
|
| 44 |
+
from headroom.subscription.base import QuotaTracker
|
| 45 |
+
|
| 46 |
+
logger = logging.getLogger(__name__)
|
| 47 |
+
|
| 48 |
+
_GITHUB_API_BASE = "https://api.github.com"
|
| 49 |
+
_TOKEN_ENV_VARS = [
|
| 50 |
+
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 51 |
+
"GITHUB_TOKEN",
|
| 52 |
+
"COPILOT_GITHUB_TOKEN",
|
| 53 |
+
"GITHUB_COPILOT_API_TOKEN",
|
| 54 |
+
]
|
| 55 |
+
|
| 56 |
+
# Categories surfaced by the quota_snapshots endpoint
|
| 57 |
+
QUOTA_CATEGORIES = ("chat", "completions", "premium_interactions")
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
# ---------------------------------------------------------------------------
|
| 61 |
+
# Data classes
|
| 62 |
+
# ---------------------------------------------------------------------------
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
@dataclass
|
| 66 |
+
class CopilotQuotaCategory:
|
| 67 |
+
"""Quota data for a single Copilot usage category."""
|
| 68 |
+
|
| 69 |
+
name: str
|
| 70 |
+
entitlement: int | None = None # total monthly allocation
|
| 71 |
+
remaining: int | None = None # remaining uses
|
| 72 |
+
percent_remaining: float | None = None # 0-100
|
| 73 |
+
overage_count: int = 0 # uses beyond entitlement
|
| 74 |
+
overage_permitted: bool = False
|
| 75 |
+
unlimited: bool = False
|
| 76 |
+
timestamp_utc: str | None = None
|
| 77 |
+
|
| 78 |
+
@property
|
| 79 |
+
def used(self) -> int | None:
|
| 80 |
+
if self.entitlement is not None and self.remaining is not None:
|
| 81 |
+
return max(0, self.entitlement - self.remaining)
|
| 82 |
+
return None
|
| 83 |
+
|
| 84 |
+
@property
|
| 85 |
+
def used_percent(self) -> float | None:
|
| 86 |
+
if self.unlimited:
|
| 87 |
+
return 0.0
|
| 88 |
+
if self.percent_remaining is not None:
|
| 89 |
+
return max(0.0, 100.0 - self.percent_remaining)
|
| 90 |
+
if self.entitlement and self.entitlement > 0 and self.remaining is not None:
|
| 91 |
+
return 100.0 * (self.entitlement - self.remaining) / self.entitlement
|
| 92 |
+
return None
|
| 93 |
+
|
| 94 |
+
def to_dict(self) -> dict[str, Any]:
|
| 95 |
+
return {
|
| 96 |
+
"name": self.name,
|
| 97 |
+
"entitlement": self.entitlement,
|
| 98 |
+
"remaining": self.remaining,
|
| 99 |
+
"used": self.used,
|
| 100 |
+
"percent_remaining": self.percent_remaining,
|
| 101 |
+
"used_percent": self.used_percent,
|
| 102 |
+
"overage_count": self.overage_count,
|
| 103 |
+
"overage_permitted": self.overage_permitted,
|
| 104 |
+
"unlimited": self.unlimited,
|
| 105 |
+
"timestamp_utc": self.timestamp_utc,
|
| 106 |
+
}
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
@dataclass
|
| 110 |
+
class CopilotQuotaSnapshot:
|
| 111 |
+
"""Full quota snapshot from one /copilot_internal/user response."""
|
| 112 |
+
|
| 113 |
+
login: str | None = None
|
| 114 |
+
copilot_plan: str | None = None
|
| 115 |
+
access_type_sku: str | None = None
|
| 116 |
+
quota_reset_date_utc: str | None = None
|
| 117 |
+
categories: dict[str, CopilotQuotaCategory] = field(default_factory=dict)
|
| 118 |
+
fetched_at: float = field(default_factory=time.time)
|
| 119 |
+
|
| 120 |
+
def to_dict(self) -> dict[str, Any]:
|
| 121 |
+
return {
|
| 122 |
+
"login": self.login,
|
| 123 |
+
"copilot_plan": self.copilot_plan,
|
| 124 |
+
"access_type_sku": self.access_type_sku,
|
| 125 |
+
"quota_reset_date_utc": self.quota_reset_date_utc,
|
| 126 |
+
"categories": {k: v.to_dict() for k, v in self.categories.items()},
|
| 127 |
+
"fetched_at": self.fetched_at,
|
| 128 |
+
}
|
| 129 |
+
|
| 130 |
+
|
| 131 |
+
@dataclass
|
| 132 |
+
class CopilotQuotaState:
|
| 133 |
+
"""Thread-safe singleton state for Copilot quota tracking."""
|
| 134 |
+
|
| 135 |
+
latest: CopilotQuotaSnapshot | None = None
|
| 136 |
+
last_error: str | None = None
|
| 137 |
+
last_updated: float | None = None
|
| 138 |
+
|
| 139 |
+
def to_dict(self) -> dict[str, Any]:
|
| 140 |
+
return {
|
| 141 |
+
"latest": self.latest.to_dict() if self.latest else None,
|
| 142 |
+
"last_error": self.last_error,
|
| 143 |
+
"last_updated": self.last_updated,
|
| 144 |
+
}
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
# ---------------------------------------------------------------------------
|
| 148 |
+
# Parsing
|
| 149 |
+
# ---------------------------------------------------------------------------
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
def parse_copilot_quota(data: dict[str, Any]) -> CopilotQuotaSnapshot:
|
| 153 |
+
"""Parse a /copilot_internal/user response into a ``CopilotQuotaSnapshot``."""
|
| 154 |
+
snapshot = CopilotQuotaSnapshot(
|
| 155 |
+
login=data.get("login"),
|
| 156 |
+
copilot_plan=data.get("copilot_plan"),
|
| 157 |
+
access_type_sku=data.get("access_type_sku"),
|
| 158 |
+
quota_reset_date_utc=data.get("quota_reset_date_utc") or data.get("quota_reset_date"),
|
| 159 |
+
)
|
| 160 |
+
|
| 161 |
+
raw_qs = data.get("quota_snapshots") or {}
|
| 162 |
+
for cat_name in QUOTA_CATEGORIES:
|
| 163 |
+
raw = raw_qs.get(cat_name)
|
| 164 |
+
if not raw:
|
| 165 |
+
continue
|
| 166 |
+
remaining = raw.get("remaining") or raw.get("quota_remaining")
|
| 167 |
+
cat = CopilotQuotaCategory(
|
| 168 |
+
name=cat_name,
|
| 169 |
+
entitlement=raw.get("entitlement"),
|
| 170 |
+
remaining=remaining,
|
| 171 |
+
percent_remaining=raw.get("percent_remaining"),
|
| 172 |
+
overage_count=raw.get("overage_count") or 0,
|
| 173 |
+
overage_permitted=bool(raw.get("overage_permitted")),
|
| 174 |
+
unlimited=bool(raw.get("unlimited")),
|
| 175 |
+
timestamp_utc=raw.get("timestamp_utc"),
|
| 176 |
+
)
|
| 177 |
+
snapshot.categories[cat_name] = cat
|
| 178 |
+
|
| 179 |
+
return snapshot
|
| 180 |
+
|
| 181 |
+
|
| 182 |
+
# ---------------------------------------------------------------------------
|
| 183 |
+
# Token discovery
|
| 184 |
+
# ---------------------------------------------------------------------------
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
def discover_github_token() -> str | None:
|
| 188 |
+
"""Return the first GitHub token found from known environment variables."""
|
| 189 |
+
for var in _TOKEN_ENV_VARS:
|
| 190 |
+
val = os.environ.get(var, "").strip()
|
| 191 |
+
if val:
|
| 192 |
+
return val
|
| 193 |
+
return None
|
| 194 |
+
|
| 195 |
+
|
| 196 |
+
# ---------------------------------------------------------------------------
|
| 197 |
+
# Background tracker singleton
|
| 198 |
+
# ---------------------------------------------------------------------------
|
| 199 |
+
|
| 200 |
+
|
| 201 |
+
class _CopilotQuotaTracker(QuotaTracker):
|
| 202 |
+
"""Singleton background poller for GitHub Copilot quota.
|
| 203 |
+
|
| 204 |
+
Implements :class:`~headroom.subscription.base.QuotaTracker` so it can be
|
| 205 |
+
registered with the :class:`~headroom.subscription.base.QuotaTrackerRegistry`.
|
| 206 |
+
Availability is gated on a GitHub token being present in the environment.
|
| 207 |
+
"""
|
| 208 |
+
|
| 209 |
+
# QuotaTracker identity
|
| 210 |
+
key = "copilot_quota"
|
| 211 |
+
label = "GitHub Copilot"
|
| 212 |
+
|
| 213 |
+
def __init__(self, poll_interval_s: float = 60.0) -> None:
|
| 214 |
+
self._poll_interval_s = poll_interval_s
|
| 215 |
+
self._state = CopilotQuotaState()
|
| 216 |
+
self._lock = Lock()
|
| 217 |
+
self._stop_event: asyncio.Event | None = None
|
| 218 |
+
self._task: asyncio.Task | None = None # type: ignore[type-arg]
|
| 219 |
+
|
| 220 |
+
# ------------------------------------------------------------------
|
| 221 |
+
# QuotaTracker interface
|
| 222 |
+
# ------------------------------------------------------------------
|
| 223 |
+
|
| 224 |
+
def is_available(self) -> bool:
|
| 225 |
+
"""Returns ``True`` when a GitHub token is available in the environment."""
|
| 226 |
+
return discover_github_token() is not None
|
| 227 |
+
|
| 228 |
+
def get_stats(self) -> dict[str, Any] | None:
|
| 229 |
+
"""Return the latest quota state dict, or ``None`` if no data yet."""
|
| 230 |
+
data = self.state
|
| 231 |
+
if not data.get("latest"):
|
| 232 |
+
return None
|
| 233 |
+
return data
|
| 234 |
+
|
| 235 |
+
# ------------------------------------------------------------------
|
| 236 |
+
# Lifecycle
|
| 237 |
+
# ------------------------------------------------------------------
|
| 238 |
+
|
| 239 |
+
async def start(self) -> None:
|
| 240 |
+
"""Start the background polling loop."""
|
| 241 |
+
if self._task is not None and not self._task.done():
|
| 242 |
+
return
|
| 243 |
+
self._stop_event = asyncio.Event()
|
| 244 |
+
self._task = asyncio.create_task(self._poll_loop())
|
| 245 |
+
|
| 246 |
+
async def stop(self) -> None:
|
| 247 |
+
"""Stop the polling loop."""
|
| 248 |
+
if self._stop_event:
|
| 249 |
+
self._stop_event.set()
|
| 250 |
+
if self._task:
|
| 251 |
+
try:
|
| 252 |
+
await asyncio.wait_for(self._task, timeout=5.0)
|
| 253 |
+
except (asyncio.TimeoutError, asyncio.CancelledError):
|
| 254 |
+
# Mirror SubscriptionTracker.stop(): on timeout or outer
|
| 255 |
+
# cancellation, cancel the underlying poll task. Without
|
| 256 |
+
# this, a wedged poll task would leak past ``stop()``.
|
| 257 |
+
self._task.cancel()
|
| 258 |
+
|
| 259 |
+
# ------------------------------------------------------------------
|
| 260 |
+
# State
|
| 261 |
+
# ------------------------------------------------------------------
|
| 262 |
+
|
| 263 |
+
@property
|
| 264 |
+
def state(self) -> dict[str, Any]:
|
| 265 |
+
with self._lock:
|
| 266 |
+
return self._state.to_dict()
|
| 267 |
+
|
| 268 |
+
# ------------------------------------------------------------------
|
| 269 |
+
# Poll loop
|
| 270 |
+
# ------------------------------------------------------------------
|
| 271 |
+
|
| 272 |
+
async def _poll_loop(self) -> None:
|
| 273 |
+
assert self._stop_event is not None
|
| 274 |
+
while not self._stop_event.is_set():
|
| 275 |
+
try:
|
| 276 |
+
await self._maybe_poll()
|
| 277 |
+
except Exception as exc:
|
| 278 |
+
logger.warning("Copilot quota poll error: %s", exc)
|
| 279 |
+
try:
|
| 280 |
+
# NOTE: do NOT wrap in asyncio.shield() — shield prevents the
|
| 281 |
+
# inner Event.wait() from being cancelled when wait_for times
|
| 282 |
+
# out, leaking one Task per poll interval. See the matching
|
| 283 |
+
# note in headroom/subscription/tracker.py:_poll_loop.
|
| 284 |
+
await asyncio.wait_for(
|
| 285 |
+
self._stop_event.wait(),
|
| 286 |
+
timeout=self._poll_interval_s,
|
| 287 |
+
)
|
| 288 |
+
break # stop event was set
|
| 289 |
+
except asyncio.TimeoutError:
|
| 290 |
+
pass # normal: poll interval elapsed
|
| 291 |
+
|
| 292 |
+
async def _maybe_poll(self) -> None:
|
| 293 |
+
token = discover_github_token()
|
| 294 |
+
if not token:
|
| 295 |
+
return
|
| 296 |
+
|
| 297 |
+
try:
|
| 298 |
+
import aiohttp
|
| 299 |
+
except ImportError:
|
| 300 |
+
logger.debug("aiohttp not available; skipping Copilot quota poll")
|
| 301 |
+
return
|
| 302 |
+
|
| 303 |
+
url = f"{_GITHUB_API_BASE}/copilot_internal/user"
|
| 304 |
+
headers = {
|
| 305 |
+
"Authorization": f"Bearer {token}",
|
| 306 |
+
"Accept": "application/json",
|
| 307 |
+
"X-GitHub-Api-Version": "2022-11-28",
|
| 308 |
+
}
|
| 309 |
+
|
| 310 |
+
try:
|
| 311 |
+
async with aiohttp.ClientSession() as session:
|
| 312 |
+
async with session.get(
|
| 313 |
+
url, headers=headers, timeout=aiohttp.ClientTimeout(total=10)
|
| 314 |
+
) as resp:
|
| 315 |
+
if resp.status == 401:
|
| 316 |
+
with self._lock:
|
| 317 |
+
self._state.last_error = "unauthorized — check GITHUB_TOKEN"
|
| 318 |
+
return
|
| 319 |
+
if resp.status == 404:
|
| 320 |
+
# API-key-only account or endpoint not available
|
| 321 |
+
with self._lock:
|
| 322 |
+
self._state.last_error = "endpoint not found (non-Copilot account?)"
|
| 323 |
+
return
|
| 324 |
+
if not resp.ok:
|
| 325 |
+
with self._lock:
|
| 326 |
+
self._state.last_error = f"HTTP {resp.status}"
|
| 327 |
+
return
|
| 328 |
+
|
| 329 |
+
data = await resp.json()
|
| 330 |
+
except Exception as exc:
|
| 331 |
+
with self._lock:
|
| 332 |
+
self._state.last_error = str(exc)
|
| 333 |
+
logger.debug("Copilot quota fetch failed: %s", exc)
|
| 334 |
+
return
|
| 335 |
+
|
| 336 |
+
try:
|
| 337 |
+
snapshot = parse_copilot_quota(data)
|
| 338 |
+
except Exception as exc:
|
| 339 |
+
with self._lock:
|
| 340 |
+
self._state.last_error = f"parse error: {exc}"
|
| 341 |
+
return
|
| 342 |
+
|
| 343 |
+
with self._lock:
|
| 344 |
+
self._state.latest = snapshot
|
| 345 |
+
self._state.last_error = None
|
| 346 |
+
self._state.last_updated = time.time()
|
| 347 |
+
|
| 348 |
+
logger.debug(
|
| 349 |
+
"Copilot quota polled: plan=%s categories=%s",
|
| 350 |
+
snapshot.copilot_plan,
|
| 351 |
+
list(snapshot.categories.keys()),
|
| 352 |
+
)
|
| 353 |
+
|
| 354 |
+
|
| 355 |
+
_singleton_lock = Lock()
|
| 356 |
+
_singleton: _CopilotQuotaTracker | None = None
|
| 357 |
+
|
| 358 |
+
|
| 359 |
+
def get_copilot_quota_tracker() -> _CopilotQuotaTracker:
|
| 360 |
+
"""Return the global ``_CopilotQuotaTracker`` singleton."""
|
| 361 |
+
global _singleton
|
| 362 |
+
if _singleton is None:
|
| 363 |
+
with _singleton_lock:
|
| 364 |
+
if _singleton is None:
|
| 365 |
+
_singleton = _CopilotQuotaTracker()
|
| 366 |
+
return _singleton
|
|
@@ -1,395 +1,395 @@
|
|
| 1 |
-
"""Data models for Anthropic subscription window tracking.
|
| 2 |
-
|
| 3 |
-
Mirrors the Anthropic OAuth usage API response exactly, including:
|
| 4 |
-
- five_hour / seven_day rolling windows (utilization + reset times)
|
| 5 |
-
- seven_day_opus / seven_day_sonnet per-model 7-day windows
|
| 6 |
-
- extra_usage overage block (credits stored in cents by Anthropic)
|
| 7 |
-
- Headroom contribution: tokens conserved by compression, rtk, cache
|
| 8 |
-
- Window discrepancy detection (surge pricing, cache-miss anomalies)
|
| 9 |
-
"""
|
| 10 |
-
|
| 11 |
-
from __future__ import annotations
|
| 12 |
-
|
| 13 |
-
from dataclasses import dataclass, field
|
| 14 |
-
from datetime import datetime, timezone
|
| 15 |
-
from typing import Any
|
| 16 |
-
|
| 17 |
-
# ---------------------------------------------------------------------------
|
| 18 |
-
# Helpers
|
| 19 |
-
# ---------------------------------------------------------------------------
|
| 20 |
-
|
| 21 |
-
|
| 22 |
-
def _utc_now() -> datetime:
|
| 23 |
-
return datetime.now(timezone.utc)
|
| 24 |
-
|
| 25 |
-
|
| 26 |
-
def _to_utc_iso(dt: datetime) -> str:
|
| 27 |
-
return dt.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
def _parse_timestamp(value: Any) -> datetime | None:
|
| 31 |
-
if not isinstance(value, str) or not value:
|
| 32 |
-
return None
|
| 33 |
-
normalized = value.replace("Z", "+00:00")
|
| 34 |
-
try:
|
| 35 |
-
dt = datetime.fromisoformat(normalized)
|
| 36 |
-
except ValueError:
|
| 37 |
-
return None
|
| 38 |
-
if dt.tzinfo is None:
|
| 39 |
-
dt = dt.replace(tzinfo=timezone.utc)
|
| 40 |
-
return dt.astimezone(timezone.utc)
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
def _safe_float(value: Any) -> float | None:
|
| 44 |
-
if value is None:
|
| 45 |
-
return None
|
| 46 |
-
try:
|
| 47 |
-
return float(value)
|
| 48 |
-
except (TypeError, ValueError):
|
| 49 |
-
return None
|
| 50 |
-
|
| 51 |
-
|
| 52 |
-
def _safe_int(value: Any) -> int | None:
|
| 53 |
-
if value is None:
|
| 54 |
-
return None
|
| 55 |
-
try:
|
| 56 |
-
return int(value)
|
| 57 |
-
except (TypeError, ValueError):
|
| 58 |
-
return None
|
| 59 |
-
|
| 60 |
-
|
| 61 |
-
# ---------------------------------------------------------------------------
|
| 62 |
-
# Rate-limit window (five_hour / seven_day / seven_day_opus / seven_day_sonnet)
|
| 63 |
-
# ---------------------------------------------------------------------------
|
| 64 |
-
|
| 65 |
-
|
| 66 |
-
@dataclass
|
| 67 |
-
class RateLimitWindow:
|
| 68 |
-
"""A single rolling rate-limit window returned by the Anthropic usage API.
|
| 69 |
-
|
| 70 |
-
``used`` and ``limit`` are in Anthropic's internal token-equivalent units
|
| 71 |
-
(not raw tokens; Anthropic weights tokens differently per model family).
|
| 72 |
-
``utilization_pct`` is the authoritative 0–100 % figure from the API.
|
| 73 |
-
"""
|
| 74 |
-
|
| 75 |
-
used: int = 0
|
| 76 |
-
limit: int = 0
|
| 77 |
-
utilization_pct: float = 0.0
|
| 78 |
-
resets_at: datetime | None = None
|
| 79 |
-
|
| 80 |
-
@classmethod
|
| 81 |
-
def from_api_dict(cls, data: dict[str, Any]) -> RateLimitWindow:
|
| 82 |
-
return cls(
|
| 83 |
-
used=int(data.get("used") or 0),
|
| 84 |
-
limit=int(data.get("limit") or 0),
|
| 85 |
-
utilization_pct=float(data.get("utilization") or 0.0),
|
| 86 |
-
resets_at=_parse_timestamp(data.get("resets_at")),
|
| 87 |
-
)
|
| 88 |
-
|
| 89 |
-
def seconds_to_reset(self, *, now: datetime | None = None) -> float | None:
|
| 90 |
-
if self.resets_at is None:
|
| 91 |
-
return None
|
| 92 |
-
return max((self.resets_at - (now or _utc_now())).total_seconds(), 0.0)
|
| 93 |
-
|
| 94 |
-
def to_dict(self) -> dict[str, Any]:
|
| 95 |
-
return {
|
| 96 |
-
"used": self.used,
|
| 97 |
-
"limit": self.limit,
|
| 98 |
-
"utilization_pct": round(self.utilization_pct, 2),
|
| 99 |
-
"resets_at": _to_utc_iso(self.resets_at) if self.resets_at else None,
|
| 100 |
-
"seconds_to_reset": self.seconds_to_reset(),
|
| 101 |
-
}
|
| 102 |
-
|
| 103 |
-
|
| 104 |
-
# ---------------------------------------------------------------------------
|
| 105 |
-
# Extra-usage / overage block
|
| 106 |
-
# ---------------------------------------------------------------------------
|
| 107 |
-
|
| 108 |
-
|
| 109 |
-
@dataclass
|
| 110 |
-
class ExtraUsage:
|
| 111 |
-
"""Overage / extra-usage block from the Anthropic usage API.
|
| 112 |
-
|
| 113 |
-
``monthly_limit_cents`` and ``used_credits_cents`` are in US cents as
|
| 114 |
-
returned by the API (divide by 100 for USD).
|
| 115 |
-
"""
|
| 116 |
-
|
| 117 |
-
is_enabled: bool = False
|
| 118 |
-
monthly_limit_cents: int | None = None
|
| 119 |
-
used_credits_cents: int | None = None
|
| 120 |
-
utilization_pct: float | None = None
|
| 121 |
-
|
| 122 |
-
@classmethod
|
| 123 |
-
def from_api_dict(cls, data: dict[str, Any]) -> ExtraUsage:
|
| 124 |
-
return cls(
|
| 125 |
-
is_enabled=bool(data.get("is_enabled", False)),
|
| 126 |
-
monthly_limit_cents=_safe_int(data.get("monthly_limit")),
|
| 127 |
-
used_credits_cents=_safe_int(data.get("used_credits")),
|
| 128 |
-
utilization_pct=_safe_float(data.get("utilization")),
|
| 129 |
-
)
|
| 130 |
-
|
| 131 |
-
@property
|
| 132 |
-
def monthly_limit_usd(self) -> float | None:
|
| 133 |
-
if self.monthly_limit_cents is None:
|
| 134 |
-
return None
|
| 135 |
-
return self.monthly_limit_cents / 100.0
|
| 136 |
-
|
| 137 |
-
@property
|
| 138 |
-
def used_credits_usd(self) -> float | None:
|
| 139 |
-
if self.used_credits_cents is None:
|
| 140 |
-
return None
|
| 141 |
-
return self.used_credits_cents / 100.0
|
| 142 |
-
|
| 143 |
-
def to_dict(self) -> dict[str, Any]:
|
| 144 |
-
return {
|
| 145 |
-
"is_enabled": self.is_enabled,
|
| 146 |
-
"monthly_limit_usd": round(self.monthly_limit_usd, 2)
|
| 147 |
-
if self.monthly_limit_usd is not None
|
| 148 |
-
else None,
|
| 149 |
-
"used_credits_usd": round(self.used_credits_usd, 4)
|
| 150 |
-
if self.used_credits_usd is not None
|
| 151 |
-
else None,
|
| 152 |
-
"utilization_pct": round(self.utilization_pct, 2)
|
| 153 |
-
if self.utilization_pct is not None
|
| 154 |
-
else None,
|
| 155 |
-
}
|
| 156 |
-
|
| 157 |
-
|
| 158 |
-
# ---------------------------------------------------------------------------
|
| 159 |
-
# Full snapshot from one API poll
|
| 160 |
-
# ---------------------------------------------------------------------------
|
| 161 |
-
|
| 162 |
-
|
| 163 |
-
@dataclass
|
| 164 |
-
class SubscriptionSnapshot:
|
| 165 |
-
"""One complete poll of GET /api/oauth/usage."""
|
| 166 |
-
|
| 167 |
-
five_hour: RateLimitWindow = field(default_factory=RateLimitWindow)
|
| 168 |
-
seven_day: RateLimitWindow = field(default_factory=RateLimitWindow)
|
| 169 |
-
seven_day_opus: RateLimitWindow | None = None
|
| 170 |
-
seven_day_sonnet: RateLimitWindow | None = None
|
| 171 |
-
extra_usage: ExtraUsage = field(default_factory=ExtraUsage)
|
| 172 |
-
polled_at: datetime = field(default_factory=_utc_now)
|
| 173 |
-
token_prefix: str = ""
|
| 174 |
-
"""First 8 chars of the OAuth token (for multi-account detection)."""
|
| 175 |
-
|
| 176 |
-
@classmethod
|
| 177 |
-
def from_api_response(cls, data: dict[str, Any], *, token: str = "") -> SubscriptionSnapshot:
|
| 178 |
-
snap = cls(token_prefix=token[:8] if token else "")
|
| 179 |
-
if "five_hour" in data and data["five_hour"]:
|
| 180 |
-
snap.five_hour = RateLimitWindow.from_api_dict(data["five_hour"])
|
| 181 |
-
if "seven_day" in data and data["seven_day"]:
|
| 182 |
-
snap.seven_day = RateLimitWindow.from_api_dict(data["seven_day"])
|
| 183 |
-
if "seven_day_opus" in data and data["seven_day_opus"]:
|
| 184 |
-
snap.seven_day_opus = RateLimitWindow.from_api_dict(data["seven_day_opus"])
|
| 185 |
-
if "seven_day_sonnet" in data and data["seven_day_sonnet"]:
|
| 186 |
-
snap.seven_day_sonnet = RateLimitWindow.from_api_dict(data["seven_day_sonnet"])
|
| 187 |
-
if "extra_usage" in data and data["extra_usage"]:
|
| 188 |
-
snap.extra_usage = ExtraUsage.from_api_dict(data["extra_usage"])
|
| 189 |
-
return snap
|
| 190 |
-
|
| 191 |
-
def to_dict(self) -> dict[str, Any]:
|
| 192 |
-
d: dict[str, Any] = {
|
| 193 |
-
"five_hour": self.five_hour.to_dict(),
|
| 194 |
-
"seven_day": self.seven_day.to_dict(),
|
| 195 |
-
"extra_usage": self.extra_usage.to_dict(),
|
| 196 |
-
"polled_at": _to_utc_iso(self.polled_at),
|
| 197 |
-
"token_prefix": self.token_prefix,
|
| 198 |
-
}
|
| 199 |
-
if self.seven_day_opus:
|
| 200 |
-
d["seven_day_opus"] = self.seven_day_opus.to_dict()
|
| 201 |
-
if self.seven_day_sonnet:
|
| 202 |
-
d["seven_day_sonnet"] = self.seven_day_sonnet.to_dict()
|
| 203 |
-
return d
|
| 204 |
-
|
| 205 |
-
|
| 206 |
-
# ---------------------------------------------------------------------------
|
| 207 |
-
# Transcript-based window token breakdown
|
| 208 |
-
# ---------------------------------------------------------------------------
|
| 209 |
-
|
| 210 |
-
|
| 211 |
-
@dataclass
|
| 212 |
-
class WindowTokens:
|
| 213 |
-
"""Token breakdown from Claude transcript JSONL files for one time window."""
|
| 214 |
-
|
| 215 |
-
input: int = 0
|
| 216 |
-
output: int = 0
|
| 217 |
-
cache_reads: int = 0
|
| 218 |
-
cache_writes_5m: int = 0
|
| 219 |
-
cache_writes_1h: int = 0
|
| 220 |
-
cache_writes_total: int = 0
|
| 221 |
-
by_model: dict[str, dict[str, int]] = field(default_factory=dict)
|
| 222 |
-
weighted_token_equivalent: float = 0.0
|
| 223 |
-
"""Sonnet-normalised weighted total (opus×2, sonnet×1, haiku×0.5)."""
|
| 224 |
-
|
| 225 |
-
def total_raw(self) -> int:
|
| 226 |
-
return self.input + self.output + self.cache_reads + self.cache_writes_total
|
| 227 |
-
|
| 228 |
-
def to_dict(self) -> dict[str, Any]:
|
| 229 |
-
return {
|
| 230 |
-
"input": self.input,
|
| 231 |
-
"output": self.output,
|
| 232 |
-
"cache_reads": self.cache_reads,
|
| 233 |
-
"cache_writes_5m": self.cache_writes_5m,
|
| 234 |
-
"cache_writes_1h": self.cache_writes_1h,
|
| 235 |
-
"cache_writes_total": self.cache_writes_total,
|
| 236 |
-
"total_raw": self.total_raw(),
|
| 237 |
-
"weighted_token_equivalent": round(self.weighted_token_equivalent, 1),
|
| 238 |
-
"by_model": self.by_model,
|
| 239 |
-
}
|
| 240 |
-
|
| 241 |
-
|
| 242 |
-
# ---------------------------------------------------------------------------
|
| 243 |
-
# Headroom contribution estimate
|
| 244 |
-
# ---------------------------------------------------------------------------
|
| 245 |
-
|
| 246 |
-
|
| 247 |
-
@dataclass
|
| 248 |
-
class HeadroomContribution:
|
| 249 |
-
"""Tokens conserved within the current 5h window by Headroom's layers.
|
| 250 |
-
|
| 251 |
-
These are cumulative counters reset when the 5h window rolls over.
|
| 252 |
-
"""
|
| 253 |
-
|
| 254 |
-
tokens_submitted: int = 0
|
| 255 |
-
"""Raw input tokens actually forwarded to Anthropic by the proxy."""
|
| 256 |
-
|
| 257 |
-
tokens_saved_compression: int = 0
|
| 258 |
-
"""Input tokens removed by proxy compression."""
|
| 259 |
-
|
| 260 |
-
tokens_saved_rtk: int = 0
|
| 261 |
-
"""Tokens avoided by CLI filtering (rtk) before reaching context."""
|
| 262 |
-
|
| 263 |
-
tokens_saved_cache_reads: int = 0
|
| 264 |
-
"""Input tokens served from Anthropic prefix-cache (discounted reads)."""
|
| 265 |
-
|
| 266 |
-
compression_savings_usd: float = 0.0
|
| 267 |
-
cache_savings_usd: float = 0.0
|
| 268 |
-
|
| 269 |
-
def total_saved(self) -> int:
|
| 270 |
-
return self.tokens_saved_compression + self.tokens_saved_rtk + self.tokens_saved_cache_reads
|
| 271 |
-
|
| 272 |
-
def total_savings_usd(self) -> float:
|
| 273 |
-
return self.compression_savings_usd + self.cache_savings_usd
|
| 274 |
-
|
| 275 |
-
def raw_without_headroom(self) -> int:
|
| 276 |
-
return self.tokens_submitted + self.tokens_saved_compression + self.tokens_saved_rtk
|
| 277 |
-
|
| 278 |
-
def efficiency_pct(self) -> float:
|
| 279 |
-
raw = self.raw_without_headroom()
|
| 280 |
-
if raw == 0:
|
| 281 |
-
return 0.0
|
| 282 |
-
return round(self.total_saved() / raw * 100, 1)
|
| 283 |
-
|
| 284 |
-
def to_dict(self) -> dict[str, Any]:
|
| 285 |
-
return {
|
| 286 |
-
"tokens_submitted": self.tokens_submitted,
|
| 287 |
-
"tokens_saved": {
|
| 288 |
-
"compression": self.tokens_saved_compression,
|
| 289 |
-
"rtk": self.tokens_saved_rtk,
|
| 290 |
-
"cache_reads": self.tokens_saved_cache_reads,
|
| 291 |
-
"total": self.total_saved(),
|
| 292 |
-
},
|
| 293 |
-
"raw_without_headroom": self.raw_without_headroom(),
|
| 294 |
-
"efficiency_pct": self.efficiency_pct(),
|
| 295 |
-
"savings_usd": {
|
| 296 |
-
"compression": round(self.compression_savings_usd, 4),
|
| 297 |
-
"cache": round(self.cache_savings_usd, 4),
|
| 298 |
-
"total": round(self.total_savings_usd(), 4),
|
| 299 |
-
},
|
| 300 |
-
}
|
| 301 |
-
|
| 302 |
-
|
| 303 |
-
# ---------------------------------------------------------------------------
|
| 304 |
-
# Anomaly / discrepancy record
|
| 305 |
-
# ---------------------------------------------------------------------------
|
| 306 |
-
|
| 307 |
-
|
| 308 |
-
@dataclass
|
| 309 |
-
class WindowDiscrepancy:
|
| 310 |
-
"""Detected anomaly between expected and API-reported utilization."""
|
| 311 |
-
|
| 312 |
-
kind: str
|
| 313 |
-
"""'surge_pricing' | 'cache_miss' | 'none'"""
|
| 314 |
-
|
| 315 |
-
description: str = ""
|
| 316 |
-
severity: str = "info"
|
| 317 |
-
"""'info' | 'warning' | 'alert'"""
|
| 318 |
-
|
| 319 |
-
expected_utilization_pct: float | None = None
|
| 320 |
-
actual_utilization_pct: float | None = None
|
| 321 |
-
delta_pct: float | None = None
|
| 322 |
-
|
| 323 |
-
def to_dict(self) -> dict[str, Any]:
|
| 324 |
-
return {
|
| 325 |
-
"kind": self.kind,
|
| 326 |
-
"description": self.description,
|
| 327 |
-
"severity": self.severity,
|
| 328 |
-
"expected_utilization_pct": self.expected_utilization_pct,
|
| 329 |
-
"actual_utilization_pct": self.actual_utilization_pct,
|
| 330 |
-
"delta_pct": self.delta_pct,
|
| 331 |
-
}
|
| 332 |
-
|
| 333 |
-
|
| 334 |
-
# ---------------------------------------------------------------------------
|
| 335 |
-
# Full tracker state
|
| 336 |
-
# ---------------------------------------------------------------------------
|
| 337 |
-
|
| 338 |
-
|
| 339 |
-
@dataclass
|
| 340 |
-
class SubscriptionState:
|
| 341 |
-
"""Persistent state for the subscription tracker."""
|
| 342 |
-
|
| 343 |
-
latest: SubscriptionSnapshot | None = None
|
| 344 |
-
window_tokens: WindowTokens | None = None
|
| 345 |
-
"""Transcript-derived token breakdown for the current 5h window."""
|
| 346 |
-
|
| 347 |
-
contribution: HeadroomContribution = field(default_factory=HeadroomContribution)
|
| 348 |
-
discrepancies: list[WindowDiscrepancy] = field(default_factory=list)
|
| 349 |
-
history: list[SubscriptionSnapshot] = field(default_factory=list)
|
| 350 |
-
|
| 351 |
-
poll_count: int = 0
|
| 352 |
-
poll_errors: int = 0
|
| 353 |
-
last_error: str | None = None
|
| 354 |
-
last_active_at: datetime | None = None
|
| 355 |
-
|
| 356 |
-
_MAX_HISTORY: int = field(default=100, init=False, repr=False)
|
| 357 |
-
_MAX_DISCREPANCIES: int = field(default=20, init=False, repr=False)
|
| 358 |
-
|
| 359 |
-
def add_snapshot(self, snapshot: SubscriptionSnapshot) -> None:
|
| 360 |
-
self.latest = snapshot
|
| 361 |
-
self.history.append(snapshot)
|
| 362 |
-
if len(self.history) > self._MAX_HISTORY:
|
| 363 |
-
self.history = self.history[-self._MAX_HISTORY :]
|
| 364 |
-
self.poll_count += 1
|
| 365 |
-
|
| 366 |
-
def mark_error(self, msg: str) -> None:
|
| 367 |
-
self.poll_errors += 1
|
| 368 |
-
self.last_error = msg
|
| 369 |
-
|
| 370 |
-
def add_discrepancy(self, d: WindowDiscrepancy) -> None:
|
| 371 |
-
self.discrepancies.append(d)
|
| 372 |
-
if len(self.discrepancies) > self._MAX_DISCREPANCIES:
|
| 373 |
-
self.discrepancies = self.discrepancies[-self._MAX_DISCREPANCIES :]
|
| 374 |
-
|
| 375 |
-
def is_active(self, *, active_window_s: float = 60.0) -> bool:
|
| 376 |
-
if self.last_active_at is None:
|
| 377 |
-
return False
|
| 378 |
-
return (_utc_now() - self.last_active_at).total_seconds() <= active_window_s
|
| 379 |
-
|
| 380 |
-
def to_dict(self) -> dict[str, Any]:
|
| 381 |
-
return {
|
| 382 |
-
"latest": self.latest.to_dict() if self.latest else None,
|
| 383 |
-
"window_tokens": self.window_tokens.to_dict() if self.window_tokens else None,
|
| 384 |
-
"contribution": self.contribution.to_dict(),
|
| 385 |
-
"discrepancies": [d.to_dict() for d in self.discrepancies[-5:]],
|
| 386 |
-
"poll_count": self.poll_count,
|
| 387 |
-
"poll_errors": self.poll_errors,
|
| 388 |
-
"last_error": self.last_error,
|
| 389 |
-
"last_active_at": _to_utc_iso(self.last_active_at) if self.last_active_at else None,
|
| 390 |
-
}
|
| 391 |
-
|
| 392 |
-
def to_persist_dict(self) -> dict[str, Any]:
|
| 393 |
-
d = self.to_dict()
|
| 394 |
-
d["history"] = [s.to_dict() for s in self.history[-20:]]
|
| 395 |
-
return d
|
|
|
|
| 1 |
+
"""Data models for Anthropic subscription window tracking.
|
| 2 |
+
|
| 3 |
+
Mirrors the Anthropic OAuth usage API response exactly, including:
|
| 4 |
+
- five_hour / seven_day rolling windows (utilization + reset times)
|
| 5 |
+
- seven_day_opus / seven_day_sonnet per-model 7-day windows
|
| 6 |
+
- extra_usage overage block (credits stored in cents by Anthropic)
|
| 7 |
+
- Headroom contribution: tokens conserved by compression, rtk, cache
|
| 8 |
+
- Window discrepancy detection (surge pricing, cache-miss anomalies)
|
| 9 |
+
"""
|
| 10 |
+
|
| 11 |
+
from __future__ import annotations
|
| 12 |
+
|
| 13 |
+
from dataclasses import dataclass, field
|
| 14 |
+
from datetime import datetime, timezone
|
| 15 |
+
from typing import Any
|
| 16 |
+
|
| 17 |
+
# ---------------------------------------------------------------------------
|
| 18 |
+
# Helpers
|
| 19 |
+
# ---------------------------------------------------------------------------
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def _utc_now() -> datetime:
|
| 23 |
+
return datetime.now(timezone.utc)
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def _to_utc_iso(dt: datetime) -> str:
|
| 27 |
+
return dt.astimezone(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def _parse_timestamp(value: Any) -> datetime | None:
|
| 31 |
+
if not isinstance(value, str) or not value:
|
| 32 |
+
return None
|
| 33 |
+
normalized = value.replace("Z", "+00:00")
|
| 34 |
+
try:
|
| 35 |
+
dt = datetime.fromisoformat(normalized)
|
| 36 |
+
except ValueError:
|
| 37 |
+
return None
|
| 38 |
+
if dt.tzinfo is None:
|
| 39 |
+
dt = dt.replace(tzinfo=timezone.utc)
|
| 40 |
+
return dt.astimezone(timezone.utc)
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def _safe_float(value: Any) -> float | None:
|
| 44 |
+
if value is None:
|
| 45 |
+
return None
|
| 46 |
+
try:
|
| 47 |
+
return float(value)
|
| 48 |
+
except (TypeError, ValueError):
|
| 49 |
+
return None
|
| 50 |
+
|
| 51 |
+
|
| 52 |
+
def _safe_int(value: Any) -> int | None:
|
| 53 |
+
if value is None:
|
| 54 |
+
return None
|
| 55 |
+
try:
|
| 56 |
+
return int(value)
|
| 57 |
+
except (TypeError, ValueError):
|
| 58 |
+
return None
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
# ---------------------------------------------------------------------------
|
| 62 |
+
# Rate-limit window (five_hour / seven_day / seven_day_opus / seven_day_sonnet)
|
| 63 |
+
# ---------------------------------------------------------------------------
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
@dataclass
|
| 67 |
+
class RateLimitWindow:
|
| 68 |
+
"""A single rolling rate-limit window returned by the Anthropic usage API.
|
| 69 |
+
|
| 70 |
+
``used`` and ``limit`` are in Anthropic's internal token-equivalent units
|
| 71 |
+
(not raw tokens; Anthropic weights tokens differently per model family).
|
| 72 |
+
``utilization_pct`` is the authoritative 0–100 % figure from the API.
|
| 73 |
+
"""
|
| 74 |
+
|
| 75 |
+
used: int = 0
|
| 76 |
+
limit: int = 0
|
| 77 |
+
utilization_pct: float = 0.0
|
| 78 |
+
resets_at: datetime | None = None
|
| 79 |
+
|
| 80 |
+
@classmethod
|
| 81 |
+
def from_api_dict(cls, data: dict[str, Any]) -> RateLimitWindow:
|
| 82 |
+
return cls(
|
| 83 |
+
used=int(data.get("used") or 0),
|
| 84 |
+
limit=int(data.get("limit") or 0),
|
| 85 |
+
utilization_pct=float(data.get("utilization") or 0.0),
|
| 86 |
+
resets_at=_parse_timestamp(data.get("resets_at")),
|
| 87 |
+
)
|
| 88 |
+
|
| 89 |
+
def seconds_to_reset(self, *, now: datetime | None = None) -> float | None:
|
| 90 |
+
if self.resets_at is None:
|
| 91 |
+
return None
|
| 92 |
+
return max((self.resets_at - (now or _utc_now())).total_seconds(), 0.0)
|
| 93 |
+
|
| 94 |
+
def to_dict(self) -> dict[str, Any]:
|
| 95 |
+
return {
|
| 96 |
+
"used": self.used,
|
| 97 |
+
"limit": self.limit,
|
| 98 |
+
"utilization_pct": round(self.utilization_pct, 2),
|
| 99 |
+
"resets_at": _to_utc_iso(self.resets_at) if self.resets_at else None,
|
| 100 |
+
"seconds_to_reset": self.seconds_to_reset(),
|
| 101 |
+
}
|
| 102 |
+
|
| 103 |
+
|
| 104 |
+
# ---------------------------------------------------------------------------
|
| 105 |
+
# Extra-usage / overage block
|
| 106 |
+
# ---------------------------------------------------------------------------
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
@dataclass
|
| 110 |
+
class ExtraUsage:
|
| 111 |
+
"""Overage / extra-usage block from the Anthropic usage API.
|
| 112 |
+
|
| 113 |
+
``monthly_limit_cents`` and ``used_credits_cents`` are in US cents as
|
| 114 |
+
returned by the API (divide by 100 for USD).
|
| 115 |
+
"""
|
| 116 |
+
|
| 117 |
+
is_enabled: bool = False
|
| 118 |
+
monthly_limit_cents: int | None = None
|
| 119 |
+
used_credits_cents: int | None = None
|
| 120 |
+
utilization_pct: float | None = None
|
| 121 |
+
|
| 122 |
+
@classmethod
|
| 123 |
+
def from_api_dict(cls, data: dict[str, Any]) -> ExtraUsage:
|
| 124 |
+
return cls(
|
| 125 |
+
is_enabled=bool(data.get("is_enabled", False)),
|
| 126 |
+
monthly_limit_cents=_safe_int(data.get("monthly_limit")),
|
| 127 |
+
used_credits_cents=_safe_int(data.get("used_credits")),
|
| 128 |
+
utilization_pct=_safe_float(data.get("utilization")),
|
| 129 |
+
)
|
| 130 |
+
|
| 131 |
+
@property
|
| 132 |
+
def monthly_limit_usd(self) -> float | None:
|
| 133 |
+
if self.monthly_limit_cents is None:
|
| 134 |
+
return None
|
| 135 |
+
return self.monthly_limit_cents / 100.0
|
| 136 |
+
|
| 137 |
+
@property
|
| 138 |
+
def used_credits_usd(self) -> float | None:
|
| 139 |
+
if self.used_credits_cents is None:
|
| 140 |
+
return None
|
| 141 |
+
return self.used_credits_cents / 100.0
|
| 142 |
+
|
| 143 |
+
def to_dict(self) -> dict[str, Any]:
|
| 144 |
+
return {
|
| 145 |
+
"is_enabled": self.is_enabled,
|
| 146 |
+
"monthly_limit_usd": round(self.monthly_limit_usd, 2)
|
| 147 |
+
if self.monthly_limit_usd is not None
|
| 148 |
+
else None,
|
| 149 |
+
"used_credits_usd": round(self.used_credits_usd, 4)
|
| 150 |
+
if self.used_credits_usd is not None
|
| 151 |
+
else None,
|
| 152 |
+
"utilization_pct": round(self.utilization_pct, 2)
|
| 153 |
+
if self.utilization_pct is not None
|
| 154 |
+
else None,
|
| 155 |
+
}
|
| 156 |
+
|
| 157 |
+
|
| 158 |
+
# ---------------------------------------------------------------------------
|
| 159 |
+
# Full snapshot from one API poll
|
| 160 |
+
# ---------------------------------------------------------------------------
|
| 161 |
+
|
| 162 |
+
|
| 163 |
+
@dataclass
|
| 164 |
+
class SubscriptionSnapshot:
|
| 165 |
+
"""One complete poll of GET /api/oauth/usage."""
|
| 166 |
+
|
| 167 |
+
five_hour: RateLimitWindow = field(default_factory=RateLimitWindow)
|
| 168 |
+
seven_day: RateLimitWindow = field(default_factory=RateLimitWindow)
|
| 169 |
+
seven_day_opus: RateLimitWindow | None = None
|
| 170 |
+
seven_day_sonnet: RateLimitWindow | None = None
|
| 171 |
+
extra_usage: ExtraUsage = field(default_factory=ExtraUsage)
|
| 172 |
+
polled_at: datetime = field(default_factory=_utc_now)
|
| 173 |
+
token_prefix: str = ""
|
| 174 |
+
"""First 8 chars of the OAuth token (for multi-account detection)."""
|
| 175 |
+
|
| 176 |
+
@classmethod
|
| 177 |
+
def from_api_response(cls, data: dict[str, Any], *, token: str = "") -> SubscriptionSnapshot:
|
| 178 |
+
snap = cls(token_prefix=token[:8] if token else "")
|
| 179 |
+
if "five_hour" in data and data["five_hour"]:
|
| 180 |
+
snap.five_hour = RateLimitWindow.from_api_dict(data["five_hour"])
|
| 181 |
+
if "seven_day" in data and data["seven_day"]:
|
| 182 |
+
snap.seven_day = RateLimitWindow.from_api_dict(data["seven_day"])
|
| 183 |
+
if "seven_day_opus" in data and data["seven_day_opus"]:
|
| 184 |
+
snap.seven_day_opus = RateLimitWindow.from_api_dict(data["seven_day_opus"])
|
| 185 |
+
if "seven_day_sonnet" in data and data["seven_day_sonnet"]:
|
| 186 |
+
snap.seven_day_sonnet = RateLimitWindow.from_api_dict(data["seven_day_sonnet"])
|
| 187 |
+
if "extra_usage" in data and data["extra_usage"]:
|
| 188 |
+
snap.extra_usage = ExtraUsage.from_api_dict(data["extra_usage"])
|
| 189 |
+
return snap
|
| 190 |
+
|
| 191 |
+
def to_dict(self) -> dict[str, Any]:
|
| 192 |
+
d: dict[str, Any] = {
|
| 193 |
+
"five_hour": self.five_hour.to_dict(),
|
| 194 |
+
"seven_day": self.seven_day.to_dict(),
|
| 195 |
+
"extra_usage": self.extra_usage.to_dict(),
|
| 196 |
+
"polled_at": _to_utc_iso(self.polled_at),
|
| 197 |
+
"token_prefix": self.token_prefix,
|
| 198 |
+
}
|
| 199 |
+
if self.seven_day_opus:
|
| 200 |
+
d["seven_day_opus"] = self.seven_day_opus.to_dict()
|
| 201 |
+
if self.seven_day_sonnet:
|
| 202 |
+
d["seven_day_sonnet"] = self.seven_day_sonnet.to_dict()
|
| 203 |
+
return d
|
| 204 |
+
|
| 205 |
+
|
| 206 |
+
# ---------------------------------------------------------------------------
|
| 207 |
+
# Transcript-based window token breakdown
|
| 208 |
+
# ---------------------------------------------------------------------------
|
| 209 |
+
|
| 210 |
+
|
| 211 |
+
@dataclass
|
| 212 |
+
class WindowTokens:
|
| 213 |
+
"""Token breakdown from Claude transcript JSONL files for one time window."""
|
| 214 |
+
|
| 215 |
+
input: int = 0
|
| 216 |
+
output: int = 0
|
| 217 |
+
cache_reads: int = 0
|
| 218 |
+
cache_writes_5m: int = 0
|
| 219 |
+
cache_writes_1h: int = 0
|
| 220 |
+
cache_writes_total: int = 0
|
| 221 |
+
by_model: dict[str, dict[str, int]] = field(default_factory=dict)
|
| 222 |
+
weighted_token_equivalent: float = 0.0
|
| 223 |
+
"""Sonnet-normalised weighted total (opus×2, sonnet×1, haiku×0.5)."""
|
| 224 |
+
|
| 225 |
+
def total_raw(self) -> int:
|
| 226 |
+
return self.input + self.output + self.cache_reads + self.cache_writes_total
|
| 227 |
+
|
| 228 |
+
def to_dict(self) -> dict[str, Any]:
|
| 229 |
+
return {
|
| 230 |
+
"input": self.input,
|
| 231 |
+
"output": self.output,
|
| 232 |
+
"cache_reads": self.cache_reads,
|
| 233 |
+
"cache_writes_5m": self.cache_writes_5m,
|
| 234 |
+
"cache_writes_1h": self.cache_writes_1h,
|
| 235 |
+
"cache_writes_total": self.cache_writes_total,
|
| 236 |
+
"total_raw": self.total_raw(),
|
| 237 |
+
"weighted_token_equivalent": round(self.weighted_token_equivalent, 1),
|
| 238 |
+
"by_model": self.by_model,
|
| 239 |
+
}
|
| 240 |
+
|
| 241 |
+
|
| 242 |
+
# ---------------------------------------------------------------------------
|
| 243 |
+
# Headroom contribution estimate
|
| 244 |
+
# ---------------------------------------------------------------------------
|
| 245 |
+
|
| 246 |
+
|
| 247 |
+
@dataclass
|
| 248 |
+
class HeadroomContribution:
|
| 249 |
+
"""Tokens conserved within the current 5h window by Headroom's layers.
|
| 250 |
+
|
| 251 |
+
These are cumulative counters reset when the 5h window rolls over.
|
| 252 |
+
"""
|
| 253 |
+
|
| 254 |
+
tokens_submitted: int = 0
|
| 255 |
+
"""Raw input tokens actually forwarded to Anthropic by the proxy."""
|
| 256 |
+
|
| 257 |
+
tokens_saved_compression: int = 0
|
| 258 |
+
"""Input tokens removed by proxy compression."""
|
| 259 |
+
|
| 260 |
+
tokens_saved_rtk: int = 0
|
| 261 |
+
"""Tokens avoided by CLI filtering (rtk) before reaching context."""
|
| 262 |
+
|
| 263 |
+
tokens_saved_cache_reads: int = 0
|
| 264 |
+
"""Input tokens served from Anthropic prefix-cache (discounted reads)."""
|
| 265 |
+
|
| 266 |
+
compression_savings_usd: float = 0.0
|
| 267 |
+
cache_savings_usd: float = 0.0
|
| 268 |
+
|
| 269 |
+
def total_saved(self) -> int:
|
| 270 |
+
return self.tokens_saved_compression + self.tokens_saved_rtk + self.tokens_saved_cache_reads
|
| 271 |
+
|
| 272 |
+
def total_savings_usd(self) -> float:
|
| 273 |
+
return self.compression_savings_usd + self.cache_savings_usd
|
| 274 |
+
|
| 275 |
+
def raw_without_headroom(self) -> int:
|
| 276 |
+
return self.tokens_submitted + self.tokens_saved_compression + self.tokens_saved_rtk
|
| 277 |
+
|
| 278 |
+
def efficiency_pct(self) -> float:
|
| 279 |
+
raw = self.raw_without_headroom()
|
| 280 |
+
if raw == 0:
|
| 281 |
+
return 0.0
|
| 282 |
+
return round(self.total_saved() / raw * 100, 1)
|
| 283 |
+
|
| 284 |
+
def to_dict(self) -> dict[str, Any]:
|
| 285 |
+
return {
|
| 286 |
+
"tokens_submitted": self.tokens_submitted,
|
| 287 |
+
"tokens_saved": {
|
| 288 |
+
"compression": self.tokens_saved_compression,
|
| 289 |
+
"rtk": self.tokens_saved_rtk,
|
| 290 |
+
"cache_reads": self.tokens_saved_cache_reads,
|
| 291 |
+
"total": self.total_saved(),
|
| 292 |
+
},
|
| 293 |
+
"raw_without_headroom": self.raw_without_headroom(),
|
| 294 |
+
"efficiency_pct": self.efficiency_pct(),
|
| 295 |
+
"savings_usd": {
|
| 296 |
+
"compression": round(self.compression_savings_usd, 4),
|
| 297 |
+
"cache": round(self.cache_savings_usd, 4),
|
| 298 |
+
"total": round(self.total_savings_usd(), 4),
|
| 299 |
+
},
|
| 300 |
+
}
|
| 301 |
+
|
| 302 |
+
|
| 303 |
+
# ---------------------------------------------------------------------------
|
| 304 |
+
# Anomaly / discrepancy record
|
| 305 |
+
# ---------------------------------------------------------------------------
|
| 306 |
+
|
| 307 |
+
|
| 308 |
+
@dataclass
|
| 309 |
+
class WindowDiscrepancy:
|
| 310 |
+
"""Detected anomaly between expected and API-reported utilization."""
|
| 311 |
+
|
| 312 |
+
kind: str
|
| 313 |
+
"""'surge_pricing' | 'cache_miss' | 'none'"""
|
| 314 |
+
|
| 315 |
+
description: str = ""
|
| 316 |
+
severity: str = "info"
|
| 317 |
+
"""'info' | 'warning' | 'alert'"""
|
| 318 |
+
|
| 319 |
+
expected_utilization_pct: float | None = None
|
| 320 |
+
actual_utilization_pct: float | None = None
|
| 321 |
+
delta_pct: float | None = None
|
| 322 |
+
|
| 323 |
+
def to_dict(self) -> dict[str, Any]:
|
| 324 |
+
return {
|
| 325 |
+
"kind": self.kind,
|
| 326 |
+
"description": self.description,
|
| 327 |
+
"severity": self.severity,
|
| 328 |
+
"expected_utilization_pct": self.expected_utilization_pct,
|
| 329 |
+
"actual_utilization_pct": self.actual_utilization_pct,
|
| 330 |
+
"delta_pct": self.delta_pct,
|
| 331 |
+
}
|
| 332 |
+
|
| 333 |
+
|
| 334 |
+
# ---------------------------------------------------------------------------
|
| 335 |
+
# Full tracker state
|
| 336 |
+
# ---------------------------------------------------------------------------
|
| 337 |
+
|
| 338 |
+
|
| 339 |
+
@dataclass
|
| 340 |
+
class SubscriptionState:
|
| 341 |
+
"""Persistent state for the subscription tracker."""
|
| 342 |
+
|
| 343 |
+
latest: SubscriptionSnapshot | None = None
|
| 344 |
+
window_tokens: WindowTokens | None = None
|
| 345 |
+
"""Transcript-derived token breakdown for the current 5h window."""
|
| 346 |
+
|
| 347 |
+
contribution: HeadroomContribution = field(default_factory=HeadroomContribution)
|
| 348 |
+
discrepancies: list[WindowDiscrepancy] = field(default_factory=list)
|
| 349 |
+
history: list[SubscriptionSnapshot] = field(default_factory=list)
|
| 350 |
+
|
| 351 |
+
poll_count: int = 0
|
| 352 |
+
poll_errors: int = 0
|
| 353 |
+
last_error: str | None = None
|
| 354 |
+
last_active_at: datetime | None = None
|
| 355 |
+
|
| 356 |
+
_MAX_HISTORY: int = field(default=100, init=False, repr=False)
|
| 357 |
+
_MAX_DISCREPANCIES: int = field(default=20, init=False, repr=False)
|
| 358 |
+
|
| 359 |
+
def add_snapshot(self, snapshot: SubscriptionSnapshot) -> None:
|
| 360 |
+
self.latest = snapshot
|
| 361 |
+
self.history.append(snapshot)
|
| 362 |
+
if len(self.history) > self._MAX_HISTORY:
|
| 363 |
+
self.history = self.history[-self._MAX_HISTORY :]
|
| 364 |
+
self.poll_count += 1
|
| 365 |
+
|
| 366 |
+
def mark_error(self, msg: str) -> None:
|
| 367 |
+
self.poll_errors += 1
|
| 368 |
+
self.last_error = msg
|
| 369 |
+
|
| 370 |
+
def add_discrepancy(self, d: WindowDiscrepancy) -> None:
|
| 371 |
+
self.discrepancies.append(d)
|
| 372 |
+
if len(self.discrepancies) > self._MAX_DISCREPANCIES:
|
| 373 |
+
self.discrepancies = self.discrepancies[-self._MAX_DISCREPANCIES :]
|
| 374 |
+
|
| 375 |
+
def is_active(self, *, active_window_s: float = 60.0) -> bool:
|
| 376 |
+
if self.last_active_at is None:
|
| 377 |
+
return False
|
| 378 |
+
return (_utc_now() - self.last_active_at).total_seconds() <= active_window_s
|
| 379 |
+
|
| 380 |
+
def to_dict(self) -> dict[str, Any]:
|
| 381 |
+
return {
|
| 382 |
+
"latest": self.latest.to_dict() if self.latest else None,
|
| 383 |
+
"window_tokens": self.window_tokens.to_dict() if self.window_tokens else None,
|
| 384 |
+
"contribution": self.contribution.to_dict(),
|
| 385 |
+
"discrepancies": [d.to_dict() for d in self.discrepancies[-5:]],
|
| 386 |
+
"poll_count": self.poll_count,
|
| 387 |
+
"poll_errors": self.poll_errors,
|
| 388 |
+
"last_error": self.last_error,
|
| 389 |
+
"last_active_at": _to_utc_iso(self.last_active_at) if self.last_active_at else None,
|
| 390 |
+
}
|
| 391 |
+
|
| 392 |
+
def to_persist_dict(self) -> dict[str, Any]:
|
| 393 |
+
d = self.to_dict()
|
| 394 |
+
d["history"] = [s.to_dict() for s in self.history[-20:]]
|
| 395 |
+
return d
|
|
@@ -1,189 +1,189 @@
|
|
| 1 |
-
"""Parse Claude Code transcript JSONL files for per-window token breakdowns.
|
| 2 |
-
|
| 3 |
-
Mirrors the approach in the ClaudeCacheTTLStatusLine TypeScript reference
|
| 4 |
-
implementation (session-tracking.ts). Reads ~/.claude/projects/**/*.jsonl
|
| 5 |
-
and aggregates token usage for entries whose timestamp falls within a window.
|
| 6 |
-
|
| 7 |
-
Model weights (Sonnet-normalised, empirical estimates):
|
| 8 |
-
opus: 2.0× (higher rate-limit cost)
|
| 9 |
-
sonnet: 1.0× (baseline)
|
| 10 |
-
haiku: 0.5× (cheaper, lower rate-limit cost)
|
| 11 |
-
|
| 12 |
-
The weighted_token_equivalent lets callers detect surge pricing by comparing
|
| 13 |
-
it against the API-reported utilisation × window_limit.
|
| 14 |
-
"""
|
| 15 |
-
|
| 16 |
-
from __future__ import annotations
|
| 17 |
-
|
| 18 |
-
import json
|
| 19 |
-
import logging
|
| 20 |
-
import os
|
| 21 |
-
from pathlib import Path
|
| 22 |
-
from typing import Any
|
| 23 |
-
|
| 24 |
-
from headroom.subscription.models import WindowTokens
|
| 25 |
-
|
| 26 |
-
logger = logging.getLogger(__name__)
|
| 27 |
-
|
| 28 |
-
# Maximum bytes to read per transcript file (10 MB cap — generous, typical files <1 MB)
|
| 29 |
-
_MAX_FILE_BYTES = 10 * 1024 * 1024
|
| 30 |
-
|
| 31 |
-
# Sonnet-normalised model family weights
|
| 32 |
-
MODEL_FAMILY_WEIGHTS: dict[str, float] = {
|
| 33 |
-
"opus": 2.0,
|
| 34 |
-
"sonnet": 1.0,
|
| 35 |
-
"haiku": 0.5,
|
| 36 |
-
}
|
| 37 |
-
DEFAULT_MODEL_WEIGHT: float = 1.0
|
| 38 |
-
|
| 39 |
-
|
| 40 |
-
def _claude_config_dir() -> Path:
|
| 41 |
-
base = os.environ.get("CLAUDE_CONFIG_DIR") or os.path.expanduser("~/.claude")
|
| 42 |
-
return Path(base)
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
def get_model_weight(model_id: str) -> float:
|
| 46 |
-
"""Return the Sonnet-normalised weight for a model ID.
|
| 47 |
-
|
| 48 |
-
Matches against known family names using a word-boundary check.
|
| 49 |
-
Falls back to DEFAULT_MODEL_WEIGHT for unrecognised models.
|
| 50 |
-
"""
|
| 51 |
-
lower = model_id.lower()
|
| 52 |
-
import re
|
| 53 |
-
|
| 54 |
-
for family, weight in MODEL_FAMILY_WEIGHTS.items():
|
| 55 |
-
if re.search(rf"(?<![a-z]){family}(?![a-z])", lower):
|
| 56 |
-
return weight
|
| 57 |
-
return DEFAULT_MODEL_WEIGHT
|
| 58 |
-
|
| 59 |
-
|
| 60 |
-
def find_transcript_files() -> list[Path]:
|
| 61 |
-
"""Return all .jsonl files under ~/.claude/projects."""
|
| 62 |
-
projects = _claude_config_dir() / "projects"
|
| 63 |
-
results: list[Path] = []
|
| 64 |
-
_walk_jsonl(projects, results)
|
| 65 |
-
return results
|
| 66 |
-
|
| 67 |
-
|
| 68 |
-
def _walk_jsonl(directory: Path, results: list[Path]) -> None:
|
| 69 |
-
try:
|
| 70 |
-
entries = list(directory.iterdir())
|
| 71 |
-
except (OSError, PermissionError):
|
| 72 |
-
return
|
| 73 |
-
for entry in entries:
|
| 74 |
-
try:
|
| 75 |
-
if entry.is_dir():
|
| 76 |
-
_walk_jsonl(entry, results)
|
| 77 |
-
elif entry.suffix == ".jsonl":
|
| 78 |
-
results.append(entry)
|
| 79 |
-
except OSError:
|
| 80 |
-
continue
|
| 81 |
-
|
| 82 |
-
|
| 83 |
-
def _read_transcript_lines(path: Path) -> list[str]:
|
| 84 |
-
try:
|
| 85 |
-
size = path.stat().st_size
|
| 86 |
-
read_size = min(size, _MAX_FILE_BYTES)
|
| 87 |
-
with path.open("rb") as fh:
|
| 88 |
-
raw = fh.read(read_size)
|
| 89 |
-
return [line for line in raw.decode("utf-8", errors="replace").splitlines() if line.strip()]
|
| 90 |
-
except Exception:
|
| 91 |
-
return []
|
| 92 |
-
|
| 93 |
-
|
| 94 |
-
def _add_usage_to_tokens(dest: WindowTokens, usage: dict[str, Any]) -> None:
|
| 95 |
-
dest.input += int(usage.get("input_tokens") or 0)
|
| 96 |
-
dest.output += int(usage.get("output_tokens") or 0)
|
| 97 |
-
dest.cache_reads += int(usage.get("cache_read_input_tokens") or 0)
|
| 98 |
-
|
| 99 |
-
cache_creation = usage.get("cache_creation") or {}
|
| 100 |
-
w5m = int(cache_creation.get("ephemeral_5m_input_tokens") or 0)
|
| 101 |
-
w1h = int(cache_creation.get("ephemeral_1h_input_tokens") or 0)
|
| 102 |
-
total_writes = int(usage.get("cache_creation_input_tokens") or (w5m + w1h))
|
| 103 |
-
|
| 104 |
-
dest.cache_writes_5m += w5m
|
| 105 |
-
dest.cache_writes_1h += w1h
|
| 106 |
-
dest.cache_writes_total += total_writes
|
| 107 |
-
|
| 108 |
-
|
| 109 |
-
def compute_window_tokens(start_ts: float, end_ts: float) -> WindowTokens:
|
| 110 |
-
"""Sum transcript token usage for entries in [start_ts, end_ts).
|
| 111 |
-
|
| 112 |
-
Args:
|
| 113 |
-
start_ts: Window start as a Unix timestamp (seconds).
|
| 114 |
-
end_ts: Window end as a Unix timestamp (seconds).
|
| 115 |
-
|
| 116 |
-
Returns:
|
| 117 |
-
:class:`WindowTokens` with aggregate + per-model breakdown and
|
| 118 |
-
``weighted_token_equivalent`` (Sonnet-normalised).
|
| 119 |
-
"""
|
| 120 |
-
totals = WindowTokens()
|
| 121 |
-
by_model: dict[str, WindowTokens] = {}
|
| 122 |
-
unattributed = WindowTokens()
|
| 123 |
-
|
| 124 |
-
for path in find_transcript_files():
|
| 125 |
-
for line in _read_transcript_lines(path):
|
| 126 |
-
try:
|
| 127 |
-
entry: dict[str, Any] = json.loads(line)
|
| 128 |
-
except (json.JSONDecodeError, ValueError):
|
| 129 |
-
continue
|
| 130 |
-
|
| 131 |
-
ts_str = entry.get("timestamp")
|
| 132 |
-
if not ts_str:
|
| 133 |
-
continue
|
| 134 |
-
try:
|
| 135 |
-
from datetime import datetime
|
| 136 |
-
|
| 137 |
-
dt = datetime.fromisoformat(str(ts_str).replace("Z", "+00:00"))
|
| 138 |
-
ts = dt.timestamp()
|
| 139 |
-
except (ValueError, TypeError):
|
| 140 |
-
continue
|
| 141 |
-
|
| 142 |
-
if ts < start_ts or ts >= end_ts:
|
| 143 |
-
continue
|
| 144 |
-
|
| 145 |
-
msg = entry.get("message") or {}
|
| 146 |
-
usage = msg.get("usage")
|
| 147 |
-
if not usage:
|
| 148 |
-
continue
|
| 149 |
-
|
| 150 |
-
_add_usage_to_tokens(totals, usage)
|
| 151 |
-
|
| 152 |
-
model_id: str | None = msg.get("model")
|
| 153 |
-
if model_id:
|
| 154 |
-
if model_id not in by_model:
|
| 155 |
-
by_model[model_id] = WindowTokens()
|
| 156 |
-
_add_usage_to_tokens(by_model[model_id], usage)
|
| 157 |
-
else:
|
| 158 |
-
_add_usage_to_tokens(unattributed, usage)
|
| 159 |
-
|
| 160 |
-
# Compute Sonnet-normalised weighted equivalent
|
| 161 |
-
model_weights: dict[str, float] = {}
|
| 162 |
-
weighted = 0.0
|
| 163 |
-
|
| 164 |
-
for model_id, model_tokens in by_model.items():
|
| 165 |
-
w = get_model_weight(model_id)
|
| 166 |
-
model_weights[model_id] = w
|
| 167 |
-
weighted += _total_token_count(model_tokens) * w
|
| 168 |
-
|
| 169 |
-
weighted += _total_token_count(unattributed) * DEFAULT_MODEL_WEIGHT
|
| 170 |
-
|
| 171 |
-
totals.weighted_token_equivalent = weighted
|
| 172 |
-
totals.by_model = {mid: _window_tokens_to_dict(mt) for mid, mt in by_model.items()}
|
| 173 |
-
|
| 174 |
-
return totals
|
| 175 |
-
|
| 176 |
-
|
| 177 |
-
def _total_token_count(t: WindowTokens) -> int:
|
| 178 |
-
return t.input + t.output + t.cache_reads + t.cache_writes_total
|
| 179 |
-
|
| 180 |
-
|
| 181 |
-
def _window_tokens_to_dict(t: WindowTokens) -> dict[str, int]:
|
| 182 |
-
return {
|
| 183 |
-
"input": t.input,
|
| 184 |
-
"output": t.output,
|
| 185 |
-
"cache_reads": t.cache_reads,
|
| 186 |
-
"cache_writes_5m": t.cache_writes_5m,
|
| 187 |
-
"cache_writes_1h": t.cache_writes_1h,
|
| 188 |
-
"cache_writes_total": t.cache_writes_total,
|
| 189 |
-
}
|
|
|
|
| 1 |
+
"""Parse Claude Code transcript JSONL files for per-window token breakdowns.
|
| 2 |
+
|
| 3 |
+
Mirrors the approach in the ClaudeCacheTTLStatusLine TypeScript reference
|
| 4 |
+
implementation (session-tracking.ts). Reads ~/.claude/projects/**/*.jsonl
|
| 5 |
+
and aggregates token usage for entries whose timestamp falls within a window.
|
| 6 |
+
|
| 7 |
+
Model weights (Sonnet-normalised, empirical estimates):
|
| 8 |
+
opus: 2.0× (higher rate-limit cost)
|
| 9 |
+
sonnet: 1.0× (baseline)
|
| 10 |
+
haiku: 0.5× (cheaper, lower rate-limit cost)
|
| 11 |
+
|
| 12 |
+
The weighted_token_equivalent lets callers detect surge pricing by comparing
|
| 13 |
+
it against the API-reported utilisation × window_limit.
|
| 14 |
+
"""
|
| 15 |
+
|
| 16 |
+
from __future__ import annotations
|
| 17 |
+
|
| 18 |
+
import json
|
| 19 |
+
import logging
|
| 20 |
+
import os
|
| 21 |
+
from pathlib import Path
|
| 22 |
+
from typing import Any
|
| 23 |
+
|
| 24 |
+
from headroom.subscription.models import WindowTokens
|
| 25 |
+
|
| 26 |
+
logger = logging.getLogger(__name__)
|
| 27 |
+
|
| 28 |
+
# Maximum bytes to read per transcript file (10 MB cap — generous, typical files <1 MB)
|
| 29 |
+
_MAX_FILE_BYTES = 10 * 1024 * 1024
|
| 30 |
+
|
| 31 |
+
# Sonnet-normalised model family weights
|
| 32 |
+
MODEL_FAMILY_WEIGHTS: dict[str, float] = {
|
| 33 |
+
"opus": 2.0,
|
| 34 |
+
"sonnet": 1.0,
|
| 35 |
+
"haiku": 0.5,
|
| 36 |
+
}
|
| 37 |
+
DEFAULT_MODEL_WEIGHT: float = 1.0
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
def _claude_config_dir() -> Path:
|
| 41 |
+
base = os.environ.get("CLAUDE_CONFIG_DIR") or os.path.expanduser("~/.claude")
|
| 42 |
+
return Path(base)
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def get_model_weight(model_id: str) -> float:
|
| 46 |
+
"""Return the Sonnet-normalised weight for a model ID.
|
| 47 |
+
|
| 48 |
+
Matches against known family names using a word-boundary check.
|
| 49 |
+
Falls back to DEFAULT_MODEL_WEIGHT for unrecognised models.
|
| 50 |
+
"""
|
| 51 |
+
lower = model_id.lower()
|
| 52 |
+
import re
|
| 53 |
+
|
| 54 |
+
for family, weight in MODEL_FAMILY_WEIGHTS.items():
|
| 55 |
+
if re.search(rf"(?<![a-z]){family}(?![a-z])", lower):
|
| 56 |
+
return weight
|
| 57 |
+
return DEFAULT_MODEL_WEIGHT
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
def find_transcript_files() -> list[Path]:
|
| 61 |
+
"""Return all .jsonl files under ~/.claude/projects."""
|
| 62 |
+
projects = _claude_config_dir() / "projects"
|
| 63 |
+
results: list[Path] = []
|
| 64 |
+
_walk_jsonl(projects, results)
|
| 65 |
+
return results
|
| 66 |
+
|
| 67 |
+
|
| 68 |
+
def _walk_jsonl(directory: Path, results: list[Path]) -> None:
|
| 69 |
+
try:
|
| 70 |
+
entries = list(directory.iterdir())
|
| 71 |
+
except (OSError, PermissionError):
|
| 72 |
+
return
|
| 73 |
+
for entry in entries:
|
| 74 |
+
try:
|
| 75 |
+
if entry.is_dir():
|
| 76 |
+
_walk_jsonl(entry, results)
|
| 77 |
+
elif entry.suffix == ".jsonl":
|
| 78 |
+
results.append(entry)
|
| 79 |
+
except OSError:
|
| 80 |
+
continue
|
| 81 |
+
|
| 82 |
+
|
| 83 |
+
def _read_transcript_lines(path: Path) -> list[str]:
|
| 84 |
+
try:
|
| 85 |
+
size = path.stat().st_size
|
| 86 |
+
read_size = min(size, _MAX_FILE_BYTES)
|
| 87 |
+
with path.open("rb") as fh:
|
| 88 |
+
raw = fh.read(read_size)
|
| 89 |
+
return [line for line in raw.decode("utf-8", errors="replace").splitlines() if line.strip()]
|
| 90 |
+
except Exception:
|
| 91 |
+
return []
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def _add_usage_to_tokens(dest: WindowTokens, usage: dict[str, Any]) -> None:
|
| 95 |
+
dest.input += int(usage.get("input_tokens") or 0)
|
| 96 |
+
dest.output += int(usage.get("output_tokens") or 0)
|
| 97 |
+
dest.cache_reads += int(usage.get("cache_read_input_tokens") or 0)
|
| 98 |
+
|
| 99 |
+
cache_creation = usage.get("cache_creation") or {}
|
| 100 |
+
w5m = int(cache_creation.get("ephemeral_5m_input_tokens") or 0)
|
| 101 |
+
w1h = int(cache_creation.get("ephemeral_1h_input_tokens") or 0)
|
| 102 |
+
total_writes = int(usage.get("cache_creation_input_tokens") or (w5m + w1h))
|
| 103 |
+
|
| 104 |
+
dest.cache_writes_5m += w5m
|
| 105 |
+
dest.cache_writes_1h += w1h
|
| 106 |
+
dest.cache_writes_total += total_writes
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
def compute_window_tokens(start_ts: float, end_ts: float) -> WindowTokens:
|
| 110 |
+
"""Sum transcript token usage for entries in [start_ts, end_ts).
|
| 111 |
+
|
| 112 |
+
Args:
|
| 113 |
+
start_ts: Window start as a Unix timestamp (seconds).
|
| 114 |
+
end_ts: Window end as a Unix timestamp (seconds).
|
| 115 |
+
|
| 116 |
+
Returns:
|
| 117 |
+
:class:`WindowTokens` with aggregate + per-model breakdown and
|
| 118 |
+
``weighted_token_equivalent`` (Sonnet-normalised).
|
| 119 |
+
"""
|
| 120 |
+
totals = WindowTokens()
|
| 121 |
+
by_model: dict[str, WindowTokens] = {}
|
| 122 |
+
unattributed = WindowTokens()
|
| 123 |
+
|
| 124 |
+
for path in find_transcript_files():
|
| 125 |
+
for line in _read_transcript_lines(path):
|
| 126 |
+
try:
|
| 127 |
+
entry: dict[str, Any] = json.loads(line)
|
| 128 |
+
except (json.JSONDecodeError, ValueError):
|
| 129 |
+
continue
|
| 130 |
+
|
| 131 |
+
ts_str = entry.get("timestamp")
|
| 132 |
+
if not ts_str:
|
| 133 |
+
continue
|
| 134 |
+
try:
|
| 135 |
+
from datetime import datetime
|
| 136 |
+
|
| 137 |
+
dt = datetime.fromisoformat(str(ts_str).replace("Z", "+00:00"))
|
| 138 |
+
ts = dt.timestamp()
|
| 139 |
+
except (ValueError, TypeError):
|
| 140 |
+
continue
|
| 141 |
+
|
| 142 |
+
if ts < start_ts or ts >= end_ts:
|
| 143 |
+
continue
|
| 144 |
+
|
| 145 |
+
msg = entry.get("message") or {}
|
| 146 |
+
usage = msg.get("usage")
|
| 147 |
+
if not usage:
|
| 148 |
+
continue
|
| 149 |
+
|
| 150 |
+
_add_usage_to_tokens(totals, usage)
|
| 151 |
+
|
| 152 |
+
model_id: str | None = msg.get("model")
|
| 153 |
+
if model_id:
|
| 154 |
+
if model_id not in by_model:
|
| 155 |
+
by_model[model_id] = WindowTokens()
|
| 156 |
+
_add_usage_to_tokens(by_model[model_id], usage)
|
| 157 |
+
else:
|
| 158 |
+
_add_usage_to_tokens(unattributed, usage)
|
| 159 |
+
|
| 160 |
+
# Compute Sonnet-normalised weighted equivalent
|
| 161 |
+
model_weights: dict[str, float] = {}
|
| 162 |
+
weighted = 0.0
|
| 163 |
+
|
| 164 |
+
for model_id, model_tokens in by_model.items():
|
| 165 |
+
w = get_model_weight(model_id)
|
| 166 |
+
model_weights[model_id] = w
|
| 167 |
+
weighted += _total_token_count(model_tokens) * w
|
| 168 |
+
|
| 169 |
+
weighted += _total_token_count(unattributed) * DEFAULT_MODEL_WEIGHT
|
| 170 |
+
|
| 171 |
+
totals.weighted_token_equivalent = weighted
|
| 172 |
+
totals.by_model = {mid: _window_tokens_to_dict(mt) for mid, mt in by_model.items()}
|
| 173 |
+
|
| 174 |
+
return totals
|
| 175 |
+
|
| 176 |
+
|
| 177 |
+
def _total_token_count(t: WindowTokens) -> int:
|
| 178 |
+
return t.input + t.output + t.cache_reads + t.cache_writes_total
|
| 179 |
+
|
| 180 |
+
|
| 181 |
+
def _window_tokens_to_dict(t: WindowTokens) -> dict[str, int]:
|
| 182 |
+
return {
|
| 183 |
+
"input": t.input,
|
| 184 |
+
"output": t.output,
|
| 185 |
+
"cache_reads": t.cache_reads,
|
| 186 |
+
"cache_writes_5m": t.cache_writes_5m,
|
| 187 |
+
"cache_writes_1h": t.cache_writes_1h,
|
| 188 |
+
"cache_writes_total": t.cache_writes_total,
|
| 189 |
+
}
|
|
@@ -1,464 +1,464 @@
|
|
| 1 |
-
"""Background subscription window tracker for Anthropic OAuth accounts.
|
| 2 |
-
|
| 3 |
-
Polls GET https://api.anthropic.com/api/oauth/usage on a configurable interval
|
| 4 |
-
while there has been at least one active OAuth session within the last minute.
|
| 5 |
-
Falls back to a stored token from ~/.claude/.credentials.json when no live
|
| 6 |
-
request has come through the proxy recently.
|
| 7 |
-
|
| 8 |
-
Architecture:
|
| 9 |
-
- Single asyncio.Task polling loop (started in start(), stopped via asyncio.Event)
|
| 10 |
-
- Thread-safe state updates via threading.Lock (consistent with headroom patterns)
|
| 11 |
-
- Atomic JSON persistence via tempfile + os.replace()
|
| 12 |
-
- Module-level singleton via get_subscription_tracker() / configure_subscription_tracker()
|
| 13 |
-
|
| 14 |
-
Also reads Claude transcript JSONL files (via session_tracking module) to provide
|
| 15 |
-
token breakdowns per window that enable:
|
| 16 |
-
- Headroom efficiency metrics (tokens saved = raw - what proxy sent)
|
| 17 |
-
- Surge pricing detection (API utilization vs expected from weighted tokens)
|
| 18 |
-
- Cache miss detection (low cache_reads despite high input tokens)
|
| 19 |
-
"""
|
| 20 |
-
|
| 21 |
-
from __future__ import annotations
|
| 22 |
-
|
| 23 |
-
import asyncio
|
| 24 |
-
import json
|
| 25 |
-
import logging
|
| 26 |
-
import os
|
| 27 |
-
import tempfile
|
| 28 |
-
import threading
|
| 29 |
-
from pathlib import Path
|
| 30 |
-
from typing import Any
|
| 31 |
-
|
| 32 |
-
from headroom import paths as _paths
|
| 33 |
-
from headroom.subscription.base import QuotaTracker
|
| 34 |
-
from headroom.subscription.client import SubscriptionClient
|
| 35 |
-
from headroom.subscription.models import (
|
| 36 |
-
HeadroomContribution,
|
| 37 |
-
SubscriptionSnapshot,
|
| 38 |
-
SubscriptionState,
|
| 39 |
-
WindowDiscrepancy,
|
| 40 |
-
WindowTokens,
|
| 41 |
-
_utc_now,
|
| 42 |
-
)
|
| 43 |
-
|
| 44 |
-
logger = logging.getLogger(__name__)
|
| 45 |
-
|
| 46 |
-
_DEFAULT_POLL_INTERVAL_S = 300
|
| 47 |
-
_DEFAULT_ACTIVE_WINDOW_S = 60
|
| 48 |
-
_PERSIST_FILE_ENV = _paths.HEADROOM_SUBSCRIPTION_STATE_PATH_ENV
|
| 49 |
-
_DEFAULT_PERSIST_DIR = ".headroom"
|
| 50 |
-
_DEFAULT_PERSIST_FILE = "subscription_state.json"
|
| 51 |
-
|
| 52 |
-
# Surge pricing threshold: if actual utilization is >N% higher than expected,
|
| 53 |
-
# flag it as a potential surge pricing event.
|
| 54 |
-
_SURGE_THRESHOLD_PCT = 15.0
|
| 55 |
-
|
| 56 |
-
# Cache miss threshold: if cache_reads < N% of total input when we expect
|
| 57 |
-
# heavy caching (>50k input tokens in window), flag it.
|
| 58 |
-
_CACHE_MISS_RATIO_THRESHOLD = 0.10
|
| 59 |
-
|
| 60 |
-
|
| 61 |
-
def _get_persist_path() -> Path:
|
| 62 |
-
return _paths.subscription_state_path()
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
class SubscriptionTracker(QuotaTracker):
|
| 66 |
-
"""Background tracker for Anthropic Claude Code subscription windows.
|
| 67 |
-
|
| 68 |
-
Implements :class:`~headroom.subscription.base.QuotaTracker` so it can
|
| 69 |
-
be registered with :func:`~headroom.subscription.base.get_quota_registry`
|
| 70 |
-
alongside the Codex and Copilot trackers.
|
| 71 |
-
|
| 72 |
-
Args:
|
| 73 |
-
poll_interval_s: Seconds between polls while active (1–3600, default 300).
|
| 74 |
-
active_window_s: Seconds since last notify_active call that keeps
|
| 75 |
-
polling alive (default 60 = 1 minute).
|
| 76 |
-
enabled: Set to ``False`` to disable tracking (mirrors
|
| 77 |
-
``ProxyConfig.subscription_tracking_enabled``).
|
| 78 |
-
persist_path: Where to persist state across restarts.
|
| 79 |
-
client: Injected client (for testing); defaults to SubscriptionClient().
|
| 80 |
-
"""
|
| 81 |
-
|
| 82 |
-
# QuotaTracker identity
|
| 83 |
-
key = "subscription_window"
|
| 84 |
-
label = "Anthropic Claude Code"
|
| 85 |
-
|
| 86 |
-
def __init__(
|
| 87 |
-
self,
|
| 88 |
-
poll_interval_s: int = _DEFAULT_POLL_INTERVAL_S,
|
| 89 |
-
active_window_s: float = _DEFAULT_ACTIVE_WINDOW_S,
|
| 90 |
-
enabled: bool = True,
|
| 91 |
-
persist_path: Path | None = None,
|
| 92 |
-
client: SubscriptionClient | None = None,
|
| 93 |
-
) -> None:
|
| 94 |
-
self._enabled = enabled
|
| 95 |
-
self._poll_interval_s = max(1, min(poll_interval_s, 3600))
|
| 96 |
-
self._active_window_s = max(5.0, active_window_s)
|
| 97 |
-
self._persist_path = persist_path or _get_persist_path()
|
| 98 |
-
self._client = client or SubscriptionClient()
|
| 99 |
-
|
| 100 |
-
self._lock = threading.Lock()
|
| 101 |
-
self._state = SubscriptionState()
|
| 102 |
-
self._current_token: str | None = None
|
| 103 |
-
self._full_tokens: dict[str, int] = {} # token_prefix -> count of requests
|
| 104 |
-
|
| 105 |
-
self._stop_event: asyncio.Event | None = None
|
| 106 |
-
self._poll_task: asyncio.Task[None] | None = None
|
| 107 |
-
|
| 108 |
-
self._load_persisted_state()
|
| 109 |
-
|
| 110 |
-
# ------------------------------------------------------------------
|
| 111 |
-
# QuotaTracker interface
|
| 112 |
-
# ------------------------------------------------------------------
|
| 113 |
-
|
| 114 |
-
def is_available(self) -> bool:
|
| 115 |
-
"""Returns ``True`` when subscription tracking is enabled in config."""
|
| 116 |
-
return self._enabled
|
| 117 |
-
|
| 118 |
-
def get_stats(self) -> dict[str, Any] | None:
|
| 119 |
-
"""Return current tracker state dict for ``/stats``."""
|
| 120 |
-
return self.state
|
| 121 |
-
|
| 122 |
-
# ------------------------------------------------------------------
|
| 123 |
-
# Lifecycle
|
| 124 |
-
# ------------------------------------------------------------------
|
| 125 |
-
|
| 126 |
-
async def start(self) -> None:
|
| 127 |
-
"""Start the background polling loop."""
|
| 128 |
-
if self._poll_task and not self._poll_task.done():
|
| 129 |
-
return
|
| 130 |
-
self._stop_event = asyncio.Event()
|
| 131 |
-
self._poll_task = asyncio.create_task(self._poll_loop(), name="subscription-tracker")
|
| 132 |
-
logger.info("Subscription tracker started (poll_interval=%ds)", self._poll_interval_s)
|
| 133 |
-
|
| 134 |
-
async def stop(self) -> None:
|
| 135 |
-
"""Stop the background polling loop and persist current state."""
|
| 136 |
-
if self._stop_event:
|
| 137 |
-
self._stop_event.set()
|
| 138 |
-
if self._poll_task:
|
| 139 |
-
try:
|
| 140 |
-
await asyncio.wait_for(self._poll_task, timeout=5.0)
|
| 141 |
-
except (asyncio.TimeoutError, asyncio.CancelledError):
|
| 142 |
-
self._poll_task.cancel()
|
| 143 |
-
self._persist_state()
|
| 144 |
-
logger.info("Subscription tracker stopped")
|
| 145 |
-
|
| 146 |
-
# ------------------------------------------------------------------
|
| 147 |
-
# Proxy integration hooks
|
| 148 |
-
# ------------------------------------------------------------------
|
| 149 |
-
|
| 150 |
-
def notify_active(self, token: str) -> None:
|
| 151 |
-
"""Called by the proxy handler when an OAuth request comes through.
|
| 152 |
-
|
| 153 |
-
Stores the token for polling and marks the tracker as recently active.
|
| 154 |
-
Only processes Bearer tokens that look like OAuth (not API keys).
|
| 155 |
-
"""
|
| 156 |
-
if not token or not token.startswith("Bearer "):
|
| 157 |
-
return
|
| 158 |
-
raw = token[len("Bearer ") :]
|
| 159 |
-
# Skip raw API keys (not OAuth tokens)
|
| 160 |
-
if raw.startswith("sk-ant-api"):
|
| 161 |
-
return
|
| 162 |
-
with self._lock:
|
| 163 |
-
self._current_token = raw
|
| 164 |
-
self._state.last_active_at = _utc_now()
|
| 165 |
-
prefix = raw[:8]
|
| 166 |
-
self._full_tokens[prefix] = self._full_tokens.get(prefix, 0) + 1
|
| 167 |
-
|
| 168 |
-
def update_contribution(
|
| 169 |
-
self,
|
| 170 |
-
*,
|
| 171 |
-
tokens_submitted: int = 0,
|
| 172 |
-
tokens_saved_compression: int = 0,
|
| 173 |
-
tokens_saved_rtk: int = 0,
|
| 174 |
-
tokens_saved_cache_reads: int = 0,
|
| 175 |
-
compression_savings_usd: float = 0.0,
|
| 176 |
-
cache_savings_usd: float = 0.0,
|
| 177 |
-
) -> None:
|
| 178 |
-
"""Update headroom contribution counters for the current session window.
|
| 179 |
-
|
| 180 |
-
Called after each proxy request completes with the actual token deltas.
|
| 181 |
-
"""
|
| 182 |
-
with self._lock:
|
| 183 |
-
c = self._state.contribution
|
| 184 |
-
c.tokens_submitted += max(tokens_submitted, 0)
|
| 185 |
-
c.tokens_saved_compression += max(tokens_saved_compression, 0)
|
| 186 |
-
c.tokens_saved_rtk += max(tokens_saved_rtk, 0)
|
| 187 |
-
c.tokens_saved_cache_reads += max(tokens_saved_cache_reads, 0)
|
| 188 |
-
c.compression_savings_usd += max(compression_savings_usd, 0.0)
|
| 189 |
-
c.cache_savings_usd += max(cache_savings_usd, 0.0)
|
| 190 |
-
|
| 191 |
-
# ------------------------------------------------------------------
|
| 192 |
-
# State access
|
| 193 |
-
# ------------------------------------------------------------------
|
| 194 |
-
|
| 195 |
-
@property
|
| 196 |
-
def state(self) -> dict[str, Any]:
|
| 197 |
-
"""Return current tracker state as a serialisable dict."""
|
| 198 |
-
with self._lock:
|
| 199 |
-
return self._state.to_dict()
|
| 200 |
-
|
| 201 |
-
@property
|
| 202 |
-
def latest_snapshot(self) -> SubscriptionSnapshot | None:
|
| 203 |
-
with self._lock:
|
| 204 |
-
return self._state.latest
|
| 205 |
-
|
| 206 |
-
def is_active(self) -> bool:
|
| 207 |
-
with self._lock:
|
| 208 |
-
return self._state.is_active(active_window_s=self._active_window_s)
|
| 209 |
-
|
| 210 |
-
# ------------------------------------------------------------------
|
| 211 |
-
# Poll loop
|
| 212 |
-
# ------------------------------------------------------------------
|
| 213 |
-
|
| 214 |
-
async def _poll_loop(self) -> None:
|
| 215 |
-
assert self._stop_event is not None
|
| 216 |
-
while not self._stop_event.is_set():
|
| 217 |
-
try:
|
| 218 |
-
await self._maybe_poll()
|
| 219 |
-
except Exception as exc:
|
| 220 |
-
logger.warning("Subscription tracker poll error: %s", exc)
|
| 221 |
-
try:
|
| 222 |
-
# NOTE: do NOT wrap in asyncio.shield() — shield prevents the
|
| 223 |
-
# inner Event.wait() from being cancelled when wait_for times
|
| 224 |
-
# out, leaking one Task per poll interval. Over hours the
|
| 225 |
-
# accumulated idle waiters bog down the event loop scheduler
|
| 226 |
-
# (observed as the "aged proxy degradation" in 2026-04-17).
|
| 227 |
-
await asyncio.wait_for(
|
| 228 |
-
self._stop_event.wait(),
|
| 229 |
-
timeout=self._poll_interval_s,
|
| 230 |
-
)
|
| 231 |
-
break # stop event was set
|
| 232 |
-
except asyncio.TimeoutError:
|
| 233 |
-
pass # normal: poll interval elapsed
|
| 234 |
-
|
| 235 |
-
async def _maybe_poll(self) -> None:
|
| 236 |
-
with self._lock:
|
| 237 |
-
is_active = self._state.is_active(active_window_s=self._active_window_s)
|
| 238 |
-
token = self._current_token
|
| 239 |
-
|
| 240 |
-
if not is_active:
|
| 241 |
-
# Try background poll using credentials file token
|
| 242 |
-
from headroom.subscription.client import read_cached_oauth_token
|
| 243 |
-
|
| 244 |
-
bg_token = read_cached_oauth_token()
|
| 245 |
-
if not bg_token:
|
| 246 |
-
return
|
| 247 |
-
token = token or bg_token
|
| 248 |
-
|
| 249 |
-
snapshot = await self._client.fetch(token)
|
| 250 |
-
if snapshot is None:
|
| 251 |
-
with self._lock:
|
| 252 |
-
self._state.mark_error("fetch returned None")
|
| 253 |
-
return
|
| 254 |
-
|
| 255 |
-
# Read transcript-based window tokens
|
| 256 |
-
window_tokens = _compute_window_tokens_for_snapshot(snapshot)
|
| 257 |
-
|
| 258 |
-
# Detect anomalies
|
| 259 |
-
discrepancies = _detect_discrepancies(snapshot, window_tokens)
|
| 260 |
-
|
| 261 |
-
with self._lock:
|
| 262 |
-
self._state.add_snapshot(snapshot)
|
| 263 |
-
self._state.window_tokens = window_tokens
|
| 264 |
-
for d in discrepancies:
|
| 265 |
-
self._state.add_discrepancy(d)
|
| 266 |
-
self._state.last_error = None
|
| 267 |
-
# Reset contribution when 5h window rolls over
|
| 268 |
-
self._maybe_reset_contribution(snapshot)
|
| 269 |
-
|
| 270 |
-
self._persist_state()
|
| 271 |
-
logger.debug(
|
| 272 |
-
"Subscription poll: 5h=%.1f%% 7d=%.1f%%",
|
| 273 |
-
snapshot.five_hour.utilization_pct,
|
| 274 |
-
snapshot.seven_day.utilization_pct,
|
| 275 |
-
)
|
| 276 |
-
|
| 277 |
-
# Update OTEL metrics if configured
|
| 278 |
-
try:
|
| 279 |
-
from headroom.observability.metrics import get_otel_metrics
|
| 280 |
-
|
| 281 |
-
get_otel_metrics().record_subscription_window(self._state.to_dict())
|
| 282 |
-
except Exception:
|
| 283 |
-
pass
|
| 284 |
-
|
| 285 |
-
def _maybe_reset_contribution(self, snapshot: SubscriptionSnapshot) -> None:
|
| 286 |
-
"""Reset contribution counters when the 5h window rolls over."""
|
| 287 |
-
prev = self._state.history[-2] if len(self._state.history) >= 2 else None
|
| 288 |
-
if prev is None:
|
| 289 |
-
return
|
| 290 |
-
prev_resets_at = prev.five_hour.resets_at
|
| 291 |
-
curr_resets_at = snapshot.five_hour.resets_at
|
| 292 |
-
if (
|
| 293 |
-
prev_resets_at is not None
|
| 294 |
-
and curr_resets_at is not None
|
| 295 |
-
and curr_resets_at != prev_resets_at
|
| 296 |
-
):
|
| 297 |
-
logger.info("5h window rolled over; resetting headroom contribution counters")
|
| 298 |
-
self._state.contribution = HeadroomContribution()
|
| 299 |
-
|
| 300 |
-
# ------------------------------------------------------------------
|
| 301 |
-
# Persistence
|
| 302 |
-
# ------------------------------------------------------------------
|
| 303 |
-
|
| 304 |
-
def _persist_state(self) -> None:
|
| 305 |
-
try:
|
| 306 |
-
self._persist_path.parent.mkdir(parents=True, exist_ok=True)
|
| 307 |
-
with self._lock:
|
| 308 |
-
data = self._state.to_persist_dict()
|
| 309 |
-
with tempfile.NamedTemporaryFile(
|
| 310 |
-
mode="w",
|
| 311 |
-
dir=self._persist_path.parent,
|
| 312 |
-
delete=False,
|
| 313 |
-
suffix=".tmp",
|
| 314 |
-
encoding="utf-8",
|
| 315 |
-
) as fh:
|
| 316 |
-
json.dump(data, fh, indent=2)
|
| 317 |
-
tmp_path = fh.name
|
| 318 |
-
os.replace(tmp_path, self._persist_path)
|
| 319 |
-
except Exception as exc:
|
| 320 |
-
logger.debug("Failed to persist subscription state: %s", exc)
|
| 321 |
-
|
| 322 |
-
def _load_persisted_state(self) -> None:
|
| 323 |
-
try:
|
| 324 |
-
with open(self._persist_path, encoding="utf-8") as fh:
|
| 325 |
-
raw = json.load(fh)
|
| 326 |
-
# Restore only the contribution counters and poll counts for now;
|
| 327 |
-
# snapshot data is re-fetched on first active poll.
|
| 328 |
-
contrib = raw.get("contribution", {})
|
| 329 |
-
c = self._state.contribution
|
| 330 |
-
c.tokens_submitted = int(contrib.get("tokens_submitted", 0))
|
| 331 |
-
saved = contrib.get("tokens_saved", {})
|
| 332 |
-
c.tokens_saved_compression = int(saved.get("compression", 0))
|
| 333 |
-
c.tokens_saved_rtk = int(saved.get("rtk", 0))
|
| 334 |
-
c.tokens_saved_cache_reads = int(saved.get("cache_reads", 0))
|
| 335 |
-
savings_usd = contrib.get("savings_usd", {})
|
| 336 |
-
c.compression_savings_usd = float(savings_usd.get("compression", 0.0))
|
| 337 |
-
c.cache_savings_usd = float(savings_usd.get("cache", 0.0))
|
| 338 |
-
self._state.poll_count = int(raw.get("poll_count", 0))
|
| 339 |
-
logger.debug("Loaded persisted subscription state from %s", self._persist_path)
|
| 340 |
-
except FileNotFoundError:
|
| 341 |
-
pass
|
| 342 |
-
except Exception as exc:
|
| 343 |
-
logger.debug("Could not load persisted subscription state: %s", exc)
|
| 344 |
-
|
| 345 |
-
|
| 346 |
-
# ---------------------------------------------------------------------------
|
| 347 |
-
# Transcript-based window token computation
|
| 348 |
-
# ---------------------------------------------------------------------------
|
| 349 |
-
|
| 350 |
-
|
| 351 |
-
def _compute_window_tokens_for_snapshot(snapshot: SubscriptionSnapshot) -> WindowTokens:
|
| 352 |
-
"""Read Claude transcript files and sum tokens for the current 5h window."""
|
| 353 |
-
try:
|
| 354 |
-
from headroom.subscription import session_tracking
|
| 355 |
-
|
| 356 |
-
resets_at = snapshot.five_hour.resets_at
|
| 357 |
-
if resets_at is None:
|
| 358 |
-
return WindowTokens()
|
| 359 |
-
window_duration_s = 5 * 3600 # 5-hour window
|
| 360 |
-
start_ts = resets_at.timestamp() - window_duration_s
|
| 361 |
-
end_ts = resets_at.timestamp()
|
| 362 |
-
return session_tracking.compute_window_tokens(start_ts, end_ts)
|
| 363 |
-
except Exception as exc:
|
| 364 |
-
logger.debug("Could not compute window tokens from transcripts: %s", exc)
|
| 365 |
-
return WindowTokens()
|
| 366 |
-
|
| 367 |
-
|
| 368 |
-
# ---------------------------------------------------------------------------
|
| 369 |
-
# Anomaly detection
|
| 370 |
-
# ---------------------------------------------------------------------------
|
| 371 |
-
|
| 372 |
-
|
| 373 |
-
def _detect_discrepancies(
|
| 374 |
-
snapshot: SubscriptionSnapshot,
|
| 375 |
-
window_tokens: WindowTokens,
|
| 376 |
-
) -> list[WindowDiscrepancy]:
|
| 377 |
-
"""Detect surge pricing or cache miss anomalies in the snapshot."""
|
| 378 |
-
discrepancies: list[WindowDiscrepancy] = []
|
| 379 |
-
|
| 380 |
-
if snapshot.five_hour.limit > 0 and window_tokens.weighted_token_equivalent > 0:
|
| 381 |
-
expected_pct = window_tokens.weighted_token_equivalent / snapshot.five_hour.limit * 100.0
|
| 382 |
-
actual_pct = snapshot.five_hour.utilization_pct
|
| 383 |
-
delta = actual_pct - expected_pct
|
| 384 |
-
|
| 385 |
-
if delta > _SURGE_THRESHOLD_PCT:
|
| 386 |
-
discrepancies.append(
|
| 387 |
-
WindowDiscrepancy(
|
| 388 |
-
kind="surge_pricing",
|
| 389 |
-
description=(
|
| 390 |
-
f"API 5h utilization ({actual_pct:.1f}%) is "
|
| 391 |
-
f"{delta:.1f}% higher than transcript-implied "
|
| 392 |
-
f"({expected_pct:.1f}%); possible surge weighting."
|
| 393 |
-
),
|
| 394 |
-
severity="warning" if delta < 30 else "alert",
|
| 395 |
-
expected_utilization_pct=round(expected_pct, 2),
|
| 396 |
-
actual_utilization_pct=round(actual_pct, 2),
|
| 397 |
-
delta_pct=round(delta, 2),
|
| 398 |
-
)
|
| 399 |
-
)
|
| 400 |
-
|
| 401 |
-
total_input = window_tokens.input
|
| 402 |
-
total_cache_reads = window_tokens.cache_reads
|
| 403 |
-
if total_input > 50_000 and total_cache_reads < total_input * _CACHE_MISS_RATIO_THRESHOLD:
|
| 404 |
-
cache_ratio = total_cache_reads / total_input if total_input else 0
|
| 405 |
-
discrepancies.append(
|
| 406 |
-
WindowDiscrepancy(
|
| 407 |
-
kind="cache_miss",
|
| 408 |
-
description=(
|
| 409 |
-
f"Cache-read ratio is {cache_ratio:.1%} (threshold "
|
| 410 |
-
f"{_CACHE_MISS_RATIO_THRESHOLD:.0%}); system may not be "
|
| 411 |
-
"using prefix cache effectively."
|
| 412 |
-
),
|
| 413 |
-
severity="warning",
|
| 414 |
-
expected_utilization_pct=None,
|
| 415 |
-
actual_utilization_pct=None,
|
| 416 |
-
delta_pct=None,
|
| 417 |
-
)
|
| 418 |
-
)
|
| 419 |
-
|
| 420 |
-
return discrepancies
|
| 421 |
-
|
| 422 |
-
|
| 423 |
-
# ---------------------------------------------------------------------------
|
| 424 |
-
# Module-level singleton
|
| 425 |
-
# ---------------------------------------------------------------------------
|
| 426 |
-
|
| 427 |
-
_tracker_lock = threading.Lock()
|
| 428 |
-
_tracker_instance: SubscriptionTracker | None = None
|
| 429 |
-
|
| 430 |
-
|
| 431 |
-
def get_subscription_tracker() -> SubscriptionTracker | None:
|
| 432 |
-
"""Return the global singleton tracker, or None if not configured."""
|
| 433 |
-
return _tracker_instance
|
| 434 |
-
|
| 435 |
-
|
| 436 |
-
def configure_subscription_tracker(
|
| 437 |
-
poll_interval_s: int = _DEFAULT_POLL_INTERVAL_S,
|
| 438 |
-
active_window_s: float = _DEFAULT_ACTIVE_WINDOW_S,
|
| 439 |
-
enabled: bool = True,
|
| 440 |
-
persist_path: Path | None = None,
|
| 441 |
-
client: SubscriptionClient | None = None,
|
| 442 |
-
) -> SubscriptionTracker:
|
| 443 |
-
"""Create (or return existing) global tracker singleton."""
|
| 444 |
-
global _tracker_instance
|
| 445 |
-
with _tracker_lock:
|
| 446 |
-
if _tracker_instance is None:
|
| 447 |
-
_tracker_instance = SubscriptionTracker(
|
| 448 |
-
poll_interval_s=poll_interval_s,
|
| 449 |
-
active_window_s=active_window_s,
|
| 450 |
-
enabled=enabled,
|
| 451 |
-
persist_path=persist_path,
|
| 452 |
-
client=client,
|
| 453 |
-
)
|
| 454 |
-
return _tracker_instance
|
| 455 |
-
|
| 456 |
-
|
| 457 |
-
async def shutdown_subscription_tracker() -> None:
|
| 458 |
-
"""Stop and clean up the global tracker."""
|
| 459 |
-
global _tracker_instance
|
| 460 |
-
with _tracker_lock:
|
| 461 |
-
tracker = _tracker_instance
|
| 462 |
-
_tracker_instance = None
|
| 463 |
-
if tracker:
|
| 464 |
-
await tracker.stop()
|
|
|
|
| 1 |
+
"""Background subscription window tracker for Anthropic OAuth accounts.
|
| 2 |
+
|
| 3 |
+
Polls GET https://api.anthropic.com/api/oauth/usage on a configurable interval
|
| 4 |
+
while there has been at least one active OAuth session within the last minute.
|
| 5 |
+
Falls back to a stored token from ~/.claude/.credentials.json when no live
|
| 6 |
+
request has come through the proxy recently.
|
| 7 |
+
|
| 8 |
+
Architecture:
|
| 9 |
+
- Single asyncio.Task polling loop (started in start(), stopped via asyncio.Event)
|
| 10 |
+
- Thread-safe state updates via threading.Lock (consistent with headroom patterns)
|
| 11 |
+
- Atomic JSON persistence via tempfile + os.replace()
|
| 12 |
+
- Module-level singleton via get_subscription_tracker() / configure_subscription_tracker()
|
| 13 |
+
|
| 14 |
+
Also reads Claude transcript JSONL files (via session_tracking module) to provide
|
| 15 |
+
token breakdowns per window that enable:
|
| 16 |
+
- Headroom efficiency metrics (tokens saved = raw - what proxy sent)
|
| 17 |
+
- Surge pricing detection (API utilization vs expected from weighted tokens)
|
| 18 |
+
- Cache miss detection (low cache_reads despite high input tokens)
|
| 19 |
+
"""
|
| 20 |
+
|
| 21 |
+
from __future__ import annotations
|
| 22 |
+
|
| 23 |
+
import asyncio
|
| 24 |
+
import json
|
| 25 |
+
import logging
|
| 26 |
+
import os
|
| 27 |
+
import tempfile
|
| 28 |
+
import threading
|
| 29 |
+
from pathlib import Path
|
| 30 |
+
from typing import Any
|
| 31 |
+
|
| 32 |
+
from headroom import paths as _paths
|
| 33 |
+
from headroom.subscription.base import QuotaTracker
|
| 34 |
+
from headroom.subscription.client import SubscriptionClient
|
| 35 |
+
from headroom.subscription.models import (
|
| 36 |
+
HeadroomContribution,
|
| 37 |
+
SubscriptionSnapshot,
|
| 38 |
+
SubscriptionState,
|
| 39 |
+
WindowDiscrepancy,
|
| 40 |
+
WindowTokens,
|
| 41 |
+
_utc_now,
|
| 42 |
+
)
|
| 43 |
+
|
| 44 |
+
logger = logging.getLogger(__name__)
|
| 45 |
+
|
| 46 |
+
_DEFAULT_POLL_INTERVAL_S = 300
|
| 47 |
+
_DEFAULT_ACTIVE_WINDOW_S = 60
|
| 48 |
+
_PERSIST_FILE_ENV = _paths.HEADROOM_SUBSCRIPTION_STATE_PATH_ENV
|
| 49 |
+
_DEFAULT_PERSIST_DIR = ".headroom"
|
| 50 |
+
_DEFAULT_PERSIST_FILE = "subscription_state.json"
|
| 51 |
+
|
| 52 |
+
# Surge pricing threshold: if actual utilization is >N% higher than expected,
|
| 53 |
+
# flag it as a potential surge pricing event.
|
| 54 |
+
_SURGE_THRESHOLD_PCT = 15.0
|
| 55 |
+
|
| 56 |
+
# Cache miss threshold: if cache_reads < N% of total input when we expect
|
| 57 |
+
# heavy caching (>50k input tokens in window), flag it.
|
| 58 |
+
_CACHE_MISS_RATIO_THRESHOLD = 0.10
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
def _get_persist_path() -> Path:
|
| 62 |
+
return _paths.subscription_state_path()
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
class SubscriptionTracker(QuotaTracker):
|
| 66 |
+
"""Background tracker for Anthropic Claude Code subscription windows.
|
| 67 |
+
|
| 68 |
+
Implements :class:`~headroom.subscription.base.QuotaTracker` so it can
|
| 69 |
+
be registered with :func:`~headroom.subscription.base.get_quota_registry`
|
| 70 |
+
alongside the Codex and Copilot trackers.
|
| 71 |
+
|
| 72 |
+
Args:
|
| 73 |
+
poll_interval_s: Seconds between polls while active (1–3600, default 300).
|
| 74 |
+
active_window_s: Seconds since last notify_active call that keeps
|
| 75 |
+
polling alive (default 60 = 1 minute).
|
| 76 |
+
enabled: Set to ``False`` to disable tracking (mirrors
|
| 77 |
+
``ProxyConfig.subscription_tracking_enabled``).
|
| 78 |
+
persist_path: Where to persist state across restarts.
|
| 79 |
+
client: Injected client (for testing); defaults to SubscriptionClient().
|
| 80 |
+
"""
|
| 81 |
+
|
| 82 |
+
# QuotaTracker identity
|
| 83 |
+
key = "subscription_window"
|
| 84 |
+
label = "Anthropic Claude Code"
|
| 85 |
+
|
| 86 |
+
def __init__(
|
| 87 |
+
self,
|
| 88 |
+
poll_interval_s: int = _DEFAULT_POLL_INTERVAL_S,
|
| 89 |
+
active_window_s: float = _DEFAULT_ACTIVE_WINDOW_S,
|
| 90 |
+
enabled: bool = True,
|
| 91 |
+
persist_path: Path | None = None,
|
| 92 |
+
client: SubscriptionClient | None = None,
|
| 93 |
+
) -> None:
|
| 94 |
+
self._enabled = enabled
|
| 95 |
+
self._poll_interval_s = max(1, min(poll_interval_s, 3600))
|
| 96 |
+
self._active_window_s = max(5.0, active_window_s)
|
| 97 |
+
self._persist_path = persist_path or _get_persist_path()
|
| 98 |
+
self._client = client or SubscriptionClient()
|
| 99 |
+
|
| 100 |
+
self._lock = threading.Lock()
|
| 101 |
+
self._state = SubscriptionState()
|
| 102 |
+
self._current_token: str | None = None
|
| 103 |
+
self._full_tokens: dict[str, int] = {} # token_prefix -> count of requests
|
| 104 |
+
|
| 105 |
+
self._stop_event: asyncio.Event | None = None
|
| 106 |
+
self._poll_task: asyncio.Task[None] | None = None
|
| 107 |
+
|
| 108 |
+
self._load_persisted_state()
|
| 109 |
+
|
| 110 |
+
# ------------------------------------------------------------------
|
| 111 |
+
# QuotaTracker interface
|
| 112 |
+
# ------------------------------------------------------------------
|
| 113 |
+
|
| 114 |
+
def is_available(self) -> bool:
|
| 115 |
+
"""Returns ``True`` when subscription tracking is enabled in config."""
|
| 116 |
+
return self._enabled
|
| 117 |
+
|
| 118 |
+
def get_stats(self) -> dict[str, Any] | None:
|
| 119 |
+
"""Return current tracker state dict for ``/stats``."""
|
| 120 |
+
return self.state
|
| 121 |
+
|
| 122 |
+
# ------------------------------------------------------------------
|
| 123 |
+
# Lifecycle
|
| 124 |
+
# ------------------------------------------------------------------
|
| 125 |
+
|
| 126 |
+
async def start(self) -> None:
|
| 127 |
+
"""Start the background polling loop."""
|
| 128 |
+
if self._poll_task and not self._poll_task.done():
|
| 129 |
+
return
|
| 130 |
+
self._stop_event = asyncio.Event()
|
| 131 |
+
self._poll_task = asyncio.create_task(self._poll_loop(), name="subscription-tracker")
|
| 132 |
+
logger.info("Subscription tracker started (poll_interval=%ds)", self._poll_interval_s)
|
| 133 |
+
|
| 134 |
+
async def stop(self) -> None:
|
| 135 |
+
"""Stop the background polling loop and persist current state."""
|
| 136 |
+
if self._stop_event:
|
| 137 |
+
self._stop_event.set()
|
| 138 |
+
if self._poll_task:
|
| 139 |
+
try:
|
| 140 |
+
await asyncio.wait_for(self._poll_task, timeout=5.0)
|
| 141 |
+
except (asyncio.TimeoutError, asyncio.CancelledError):
|
| 142 |
+
self._poll_task.cancel()
|
| 143 |
+
self._persist_state()
|
| 144 |
+
logger.info("Subscription tracker stopped")
|
| 145 |
+
|
| 146 |
+
# ------------------------------------------------------------------
|
| 147 |
+
# Proxy integration hooks
|
| 148 |
+
# ------------------------------------------------------------------
|
| 149 |
+
|
| 150 |
+
def notify_active(self, token: str) -> None:
|
| 151 |
+
"""Called by the proxy handler when an OAuth request comes through.
|
| 152 |
+
|
| 153 |
+
Stores the token for polling and marks the tracker as recently active.
|
| 154 |
+
Only processes Bearer tokens that look like OAuth (not API keys).
|
| 155 |
+
"""
|
| 156 |
+
if not token or not token.startswith("Bearer "):
|
| 157 |
+
return
|
| 158 |
+
raw = token[len("Bearer ") :]
|
| 159 |
+
# Skip raw API keys (not OAuth tokens)
|
| 160 |
+
if raw.startswith("sk-ant-api"):
|
| 161 |
+
return
|
| 162 |
+
with self._lock:
|
| 163 |
+
self._current_token = raw
|
| 164 |
+
self._state.last_active_at = _utc_now()
|
| 165 |
+
prefix = raw[:8]
|
| 166 |
+
self._full_tokens[prefix] = self._full_tokens.get(prefix, 0) + 1
|
| 167 |
+
|
| 168 |
+
def update_contribution(
|
| 169 |
+
self,
|
| 170 |
+
*,
|
| 171 |
+
tokens_submitted: int = 0,
|
| 172 |
+
tokens_saved_compression: int = 0,
|
| 173 |
+
tokens_saved_rtk: int = 0,
|
| 174 |
+
tokens_saved_cache_reads: int = 0,
|
| 175 |
+
compression_savings_usd: float = 0.0,
|
| 176 |
+
cache_savings_usd: float = 0.0,
|
| 177 |
+
) -> None:
|
| 178 |
+
"""Update headroom contribution counters for the current session window.
|
| 179 |
+
|
| 180 |
+
Called after each proxy request completes with the actual token deltas.
|
| 181 |
+
"""
|
| 182 |
+
with self._lock:
|
| 183 |
+
c = self._state.contribution
|
| 184 |
+
c.tokens_submitted += max(tokens_submitted, 0)
|
| 185 |
+
c.tokens_saved_compression += max(tokens_saved_compression, 0)
|
| 186 |
+
c.tokens_saved_rtk += max(tokens_saved_rtk, 0)
|
| 187 |
+
c.tokens_saved_cache_reads += max(tokens_saved_cache_reads, 0)
|
| 188 |
+
c.compression_savings_usd += max(compression_savings_usd, 0.0)
|
| 189 |
+
c.cache_savings_usd += max(cache_savings_usd, 0.0)
|
| 190 |
+
|
| 191 |
+
# ------------------------------------------------------------------
|
| 192 |
+
# State access
|
| 193 |
+
# ------------------------------------------------------------------
|
| 194 |
+
|
| 195 |
+
@property
|
| 196 |
+
def state(self) -> dict[str, Any]:
|
| 197 |
+
"""Return current tracker state as a serialisable dict."""
|
| 198 |
+
with self._lock:
|
| 199 |
+
return self._state.to_dict()
|
| 200 |
+
|
| 201 |
+
@property
|
| 202 |
+
def latest_snapshot(self) -> SubscriptionSnapshot | None:
|
| 203 |
+
with self._lock:
|
| 204 |
+
return self._state.latest
|
| 205 |
+
|
| 206 |
+
def is_active(self) -> bool:
|
| 207 |
+
with self._lock:
|
| 208 |
+
return self._state.is_active(active_window_s=self._active_window_s)
|
| 209 |
+
|
| 210 |
+
# ------------------------------------------------------------------
|
| 211 |
+
# Poll loop
|
| 212 |
+
# ------------------------------------------------------------------
|
| 213 |
+
|
| 214 |
+
async def _poll_loop(self) -> None:
|
| 215 |
+
assert self._stop_event is not None
|
| 216 |
+
while not self._stop_event.is_set():
|
| 217 |
+
try:
|
| 218 |
+
await self._maybe_poll()
|
| 219 |
+
except Exception as exc:
|
| 220 |
+
logger.warning("Subscription tracker poll error: %s", exc)
|
| 221 |
+
try:
|
| 222 |
+
# NOTE: do NOT wrap in asyncio.shield() — shield prevents the
|
| 223 |
+
# inner Event.wait() from being cancelled when wait_for times
|
| 224 |
+
# out, leaking one Task per poll interval. Over hours the
|
| 225 |
+
# accumulated idle waiters bog down the event loop scheduler
|
| 226 |
+
# (observed as the "aged proxy degradation" in 2026-04-17).
|
| 227 |
+
await asyncio.wait_for(
|
| 228 |
+
self._stop_event.wait(),
|
| 229 |
+
timeout=self._poll_interval_s,
|
| 230 |
+
)
|
| 231 |
+
break # stop event was set
|
| 232 |
+
except asyncio.TimeoutError:
|
| 233 |
+
pass # normal: poll interval elapsed
|
| 234 |
+
|
| 235 |
+
async def _maybe_poll(self) -> None:
|
| 236 |
+
with self._lock:
|
| 237 |
+
is_active = self._state.is_active(active_window_s=self._active_window_s)
|
| 238 |
+
token = self._current_token
|
| 239 |
+
|
| 240 |
+
if not is_active:
|
| 241 |
+
# Try background poll using credentials file token
|
| 242 |
+
from headroom.subscription.client import read_cached_oauth_token
|
| 243 |
+
|
| 244 |
+
bg_token = read_cached_oauth_token()
|
| 245 |
+
if not bg_token:
|
| 246 |
+
return
|
| 247 |
+
token = token or bg_token
|
| 248 |
+
|
| 249 |
+
snapshot = await self._client.fetch(token)
|
| 250 |
+
if snapshot is None:
|
| 251 |
+
with self._lock:
|
| 252 |
+
self._state.mark_error("fetch returned None")
|
| 253 |
+
return
|
| 254 |
+
|
| 255 |
+
# Read transcript-based window tokens
|
| 256 |
+
window_tokens = _compute_window_tokens_for_snapshot(snapshot)
|
| 257 |
+
|
| 258 |
+
# Detect anomalies
|
| 259 |
+
discrepancies = _detect_discrepancies(snapshot, window_tokens)
|
| 260 |
+
|
| 261 |
+
with self._lock:
|
| 262 |
+
self._state.add_snapshot(snapshot)
|
| 263 |
+
self._state.window_tokens = window_tokens
|
| 264 |
+
for d in discrepancies:
|
| 265 |
+
self._state.add_discrepancy(d)
|
| 266 |
+
self._state.last_error = None
|
| 267 |
+
# Reset contribution when 5h window rolls over
|
| 268 |
+
self._maybe_reset_contribution(snapshot)
|
| 269 |
+
|
| 270 |
+
self._persist_state()
|
| 271 |
+
logger.debug(
|
| 272 |
+
"Subscription poll: 5h=%.1f%% 7d=%.1f%%",
|
| 273 |
+
snapshot.five_hour.utilization_pct,
|
| 274 |
+
snapshot.seven_day.utilization_pct,
|
| 275 |
+
)
|
| 276 |
+
|
| 277 |
+
# Update OTEL metrics if configured
|
| 278 |
+
try:
|
| 279 |
+
from headroom.observability.metrics import get_otel_metrics
|
| 280 |
+
|
| 281 |
+
get_otel_metrics().record_subscription_window(self._state.to_dict())
|
| 282 |
+
except Exception:
|
| 283 |
+
pass
|
| 284 |
+
|
| 285 |
+
def _maybe_reset_contribution(self, snapshot: SubscriptionSnapshot) -> None:
|
| 286 |
+
"""Reset contribution counters when the 5h window rolls over."""
|
| 287 |
+
prev = self._state.history[-2] if len(self._state.history) >= 2 else None
|
| 288 |
+
if prev is None:
|
| 289 |
+
return
|
| 290 |
+
prev_resets_at = prev.five_hour.resets_at
|
| 291 |
+
curr_resets_at = snapshot.five_hour.resets_at
|
| 292 |
+
if (
|
| 293 |
+
prev_resets_at is not None
|
| 294 |
+
and curr_resets_at is not None
|
| 295 |
+
and curr_resets_at != prev_resets_at
|
| 296 |
+
):
|
| 297 |
+
logger.info("5h window rolled over; resetting headroom contribution counters")
|
| 298 |
+
self._state.contribution = HeadroomContribution()
|
| 299 |
+
|
| 300 |
+
# ------------------------------------------------------------------
|
| 301 |
+
# Persistence
|
| 302 |
+
# ------------------------------------------------------------------
|
| 303 |
+
|
| 304 |
+
def _persist_state(self) -> None:
|
| 305 |
+
try:
|
| 306 |
+
self._persist_path.parent.mkdir(parents=True, exist_ok=True)
|
| 307 |
+
with self._lock:
|
| 308 |
+
data = self._state.to_persist_dict()
|
| 309 |
+
with tempfile.NamedTemporaryFile(
|
| 310 |
+
mode="w",
|
| 311 |
+
dir=self._persist_path.parent,
|
| 312 |
+
delete=False,
|
| 313 |
+
suffix=".tmp",
|
| 314 |
+
encoding="utf-8",
|
| 315 |
+
) as fh:
|
| 316 |
+
json.dump(data, fh, indent=2)
|
| 317 |
+
tmp_path = fh.name
|
| 318 |
+
os.replace(tmp_path, self._persist_path)
|
| 319 |
+
except Exception as exc:
|
| 320 |
+
logger.debug("Failed to persist subscription state: %s", exc)
|
| 321 |
+
|
| 322 |
+
def _load_persisted_state(self) -> None:
|
| 323 |
+
try:
|
| 324 |
+
with open(self._persist_path, encoding="utf-8") as fh:
|
| 325 |
+
raw = json.load(fh)
|
| 326 |
+
# Restore only the contribution counters and poll counts for now;
|
| 327 |
+
# snapshot data is re-fetched on first active poll.
|
| 328 |
+
contrib = raw.get("contribution", {})
|
| 329 |
+
c = self._state.contribution
|
| 330 |
+
c.tokens_submitted = int(contrib.get("tokens_submitted", 0))
|
| 331 |
+
saved = contrib.get("tokens_saved", {})
|
| 332 |
+
c.tokens_saved_compression = int(saved.get("compression", 0))
|
| 333 |
+
c.tokens_saved_rtk = int(saved.get("rtk", 0))
|
| 334 |
+
c.tokens_saved_cache_reads = int(saved.get("cache_reads", 0))
|
| 335 |
+
savings_usd = contrib.get("savings_usd", {})
|
| 336 |
+
c.compression_savings_usd = float(savings_usd.get("compression", 0.0))
|
| 337 |
+
c.cache_savings_usd = float(savings_usd.get("cache", 0.0))
|
| 338 |
+
self._state.poll_count = int(raw.get("poll_count", 0))
|
| 339 |
+
logger.debug("Loaded persisted subscription state from %s", self._persist_path)
|
| 340 |
+
except FileNotFoundError:
|
| 341 |
+
pass
|
| 342 |
+
except Exception as exc:
|
| 343 |
+
logger.debug("Could not load persisted subscription state: %s", exc)
|
| 344 |
+
|
| 345 |
+
|
| 346 |
+
# ---------------------------------------------------------------------------
|
| 347 |
+
# Transcript-based window token computation
|
| 348 |
+
# ---------------------------------------------------------------------------
|
| 349 |
+
|
| 350 |
+
|
| 351 |
+
def _compute_window_tokens_for_snapshot(snapshot: SubscriptionSnapshot) -> WindowTokens:
|
| 352 |
+
"""Read Claude transcript files and sum tokens for the current 5h window."""
|
| 353 |
+
try:
|
| 354 |
+
from headroom.subscription import session_tracking
|
| 355 |
+
|
| 356 |
+
resets_at = snapshot.five_hour.resets_at
|
| 357 |
+
if resets_at is None:
|
| 358 |
+
return WindowTokens()
|
| 359 |
+
window_duration_s = 5 * 3600 # 5-hour window
|
| 360 |
+
start_ts = resets_at.timestamp() - window_duration_s
|
| 361 |
+
end_ts = resets_at.timestamp()
|
| 362 |
+
return session_tracking.compute_window_tokens(start_ts, end_ts)
|
| 363 |
+
except Exception as exc:
|
| 364 |
+
logger.debug("Could not compute window tokens from transcripts: %s", exc)
|
| 365 |
+
return WindowTokens()
|
| 366 |
+
|
| 367 |
+
|
| 368 |
+
# ---------------------------------------------------------------------------
|
| 369 |
+
# Anomaly detection
|
| 370 |
+
# ---------------------------------------------------------------------------
|
| 371 |
+
|
| 372 |
+
|
| 373 |
+
def _detect_discrepancies(
|
| 374 |
+
snapshot: SubscriptionSnapshot,
|
| 375 |
+
window_tokens: WindowTokens,
|
| 376 |
+
) -> list[WindowDiscrepancy]:
|
| 377 |
+
"""Detect surge pricing or cache miss anomalies in the snapshot."""
|
| 378 |
+
discrepancies: list[WindowDiscrepancy] = []
|
| 379 |
+
|
| 380 |
+
if snapshot.five_hour.limit > 0 and window_tokens.weighted_token_equivalent > 0:
|
| 381 |
+
expected_pct = window_tokens.weighted_token_equivalent / snapshot.five_hour.limit * 100.0
|
| 382 |
+
actual_pct = snapshot.five_hour.utilization_pct
|
| 383 |
+
delta = actual_pct - expected_pct
|
| 384 |
+
|
| 385 |
+
if delta > _SURGE_THRESHOLD_PCT:
|
| 386 |
+
discrepancies.append(
|
| 387 |
+
WindowDiscrepancy(
|
| 388 |
+
kind="surge_pricing",
|
| 389 |
+
description=(
|
| 390 |
+
f"API 5h utilization ({actual_pct:.1f}%) is "
|
| 391 |
+
f"{delta:.1f}% higher than transcript-implied "
|
| 392 |
+
f"({expected_pct:.1f}%); possible surge weighting."
|
| 393 |
+
),
|
| 394 |
+
severity="warning" if delta < 30 else "alert",
|
| 395 |
+
expected_utilization_pct=round(expected_pct, 2),
|
| 396 |
+
actual_utilization_pct=round(actual_pct, 2),
|
| 397 |
+
delta_pct=round(delta, 2),
|
| 398 |
+
)
|
| 399 |
+
)
|
| 400 |
+
|
| 401 |
+
total_input = window_tokens.input
|
| 402 |
+
total_cache_reads = window_tokens.cache_reads
|
| 403 |
+
if total_input > 50_000 and total_cache_reads < total_input * _CACHE_MISS_RATIO_THRESHOLD:
|
| 404 |
+
cache_ratio = total_cache_reads / total_input if total_input else 0
|
| 405 |
+
discrepancies.append(
|
| 406 |
+
WindowDiscrepancy(
|
| 407 |
+
kind="cache_miss",
|
| 408 |
+
description=(
|
| 409 |
+
f"Cache-read ratio is {cache_ratio:.1%} (threshold "
|
| 410 |
+
f"{_CACHE_MISS_RATIO_THRESHOLD:.0%}); system may not be "
|
| 411 |
+
"using prefix cache effectively."
|
| 412 |
+
),
|
| 413 |
+
severity="warning",
|
| 414 |
+
expected_utilization_pct=None,
|
| 415 |
+
actual_utilization_pct=None,
|
| 416 |
+
delta_pct=None,
|
| 417 |
+
)
|
| 418 |
+
)
|
| 419 |
+
|
| 420 |
+
return discrepancies
|
| 421 |
+
|
| 422 |
+
|
| 423 |
+
# ---------------------------------------------------------------------------
|
| 424 |
+
# Module-level singleton
|
| 425 |
+
# ---------------------------------------------------------------------------
|
| 426 |
+
|
| 427 |
+
_tracker_lock = threading.Lock()
|
| 428 |
+
_tracker_instance: SubscriptionTracker | None = None
|
| 429 |
+
|
| 430 |
+
|
| 431 |
+
def get_subscription_tracker() -> SubscriptionTracker | None:
|
| 432 |
+
"""Return the global singleton tracker, or None if not configured."""
|
| 433 |
+
return _tracker_instance
|
| 434 |
+
|
| 435 |
+
|
| 436 |
+
def configure_subscription_tracker(
|
| 437 |
+
poll_interval_s: int = _DEFAULT_POLL_INTERVAL_S,
|
| 438 |
+
active_window_s: float = _DEFAULT_ACTIVE_WINDOW_S,
|
| 439 |
+
enabled: bool = True,
|
| 440 |
+
persist_path: Path | None = None,
|
| 441 |
+
client: SubscriptionClient | None = None,
|
| 442 |
+
) -> SubscriptionTracker:
|
| 443 |
+
"""Create (or return existing) global tracker singleton."""
|
| 444 |
+
global _tracker_instance
|
| 445 |
+
with _tracker_lock:
|
| 446 |
+
if _tracker_instance is None:
|
| 447 |
+
_tracker_instance = SubscriptionTracker(
|
| 448 |
+
poll_interval_s=poll_interval_s,
|
| 449 |
+
active_window_s=active_window_s,
|
| 450 |
+
enabled=enabled,
|
| 451 |
+
persist_path=persist_path,
|
| 452 |
+
client=client,
|
| 453 |
+
)
|
| 454 |
+
return _tracker_instance
|
| 455 |
+
|
| 456 |
+
|
| 457 |
+
async def shutdown_subscription_tracker() -> None:
|
| 458 |
+
"""Stop and clean up the global tracker."""
|
| 459 |
+
global _tracker_instance
|
| 460 |
+
with _tracker_lock:
|
| 461 |
+
tracker = _tracker_instance
|
| 462 |
+
_tracker_instance = None
|
| 463 |
+
if tracker:
|
| 464 |
+
await tracker.stop()
|
|
The diff for this file is too large to render.
See raw diff
|
|
|
|
@@ -1,203 +1,203 @@
|
|
| 1 |
-
#!/usr/bin/env python3
|
| 2 |
-
"""Generate changelog from conventional commits."""
|
| 3 |
-
|
| 4 |
-
from __future__ import annotations
|
| 5 |
-
|
| 6 |
-
import argparse
|
| 7 |
-
import re
|
| 8 |
-
import subprocess
|
| 9 |
-
from datetime import date
|
| 10 |
-
from pathlib import Path
|
| 11 |
-
from typing import NamedTuple
|
| 12 |
-
|
| 13 |
-
ROOT = Path(__file__).parent.parent
|
| 14 |
-
|
| 15 |
-
COMMIT_PATTERN = re.compile(
|
| 16 |
-
r"^(feat|fix|ci|chore|perf|refactor|docs|style|test)(\(.+\))?(!)?:\s*(.+)$"
|
| 17 |
-
)
|
| 18 |
-
BREAKING_CHANGE_PATTERN = re.compile(r"^BREAKING CHANGE:\s*(.+)$", re.MULTILINE)
|
| 19 |
-
COMMIT_ENTRY_PATTERN = re.compile(r"^(.+?)(?:\n(.+))?\|(\w+)$", re.MULTILINE)
|
| 20 |
-
FIELD_SEP = "\x1f"
|
| 21 |
-
RECORD_SEP = "\x1e"
|
| 22 |
-
GIT_LOG_FORMAT = "%s%x1f%b%x1f%h%x1e"
|
| 23 |
-
|
| 24 |
-
TYPE_LABELS: dict[str, str] = {
|
| 25 |
-
"feat": "Features",
|
| 26 |
-
"fix": "Bug Fixes",
|
| 27 |
-
"ci": "CI/CD",
|
| 28 |
-
"chore": "Chores",
|
| 29 |
-
"perf": "Performance",
|
| 30 |
-
"refactor": "Refactors",
|
| 31 |
-
"docs": "Documentation",
|
| 32 |
-
"style": "Styles",
|
| 33 |
-
"test": "Tests",
|
| 34 |
-
"other": "Other Changes",
|
| 35 |
-
}
|
| 36 |
-
|
| 37 |
-
|
| 38 |
-
class ParsedCommit(NamedTuple):
|
| 39 |
-
type: str
|
| 40 |
-
scope: str | None
|
| 41 |
-
breaking: bool
|
| 42 |
-
message: str
|
| 43 |
-
hash: str
|
| 44 |
-
|
| 45 |
-
|
| 46 |
-
def iter_commit_entries(log_output: str) -> list[tuple[str, str, str]]:
|
| 47 |
-
"""Split raw git log output into (subject, body, hash) tuples."""
|
| 48 |
-
|
| 49 |
-
if not log_output.strip():
|
| 50 |
-
return []
|
| 51 |
-
|
| 52 |
-
if RECORD_SEP in log_output and FIELD_SEP in log_output:
|
| 53 |
-
entries: list[tuple[str, str, str]] = []
|
| 54 |
-
for raw_entry in log_output.split(RECORD_SEP):
|
| 55 |
-
if not raw_entry:
|
| 56 |
-
continue
|
| 57 |
-
if FIELD_SEP not in raw_entry:
|
| 58 |
-
continue
|
| 59 |
-
subject, body_and_hash = raw_entry.split(FIELD_SEP, 1)
|
| 60 |
-
if FIELD_SEP not in body_and_hash:
|
| 61 |
-
continue
|
| 62 |
-
body, commit_hash = body_and_hash.rsplit(FIELD_SEP, 1)
|
| 63 |
-
entries.append((subject.strip(), body.strip(), commit_hash.strip()))
|
| 64 |
-
return entries
|
| 65 |
-
|
| 66 |
-
return [
|
| 67 |
-
(
|
| 68 |
-
match.group(1).strip(),
|
| 69 |
-
(match.group(2) or "").strip(),
|
| 70 |
-
match.group(3).strip(),
|
| 71 |
-
)
|
| 72 |
-
for match in COMMIT_ENTRY_PATTERN.finditer(log_output)
|
| 73 |
-
]
|
| 74 |
-
|
| 75 |
-
|
| 76 |
-
def get_merge_summary(subject: str, body: str) -> str:
|
| 77 |
-
"""Return the first meaningful summary line for a merge commit."""
|
| 78 |
-
|
| 79 |
-
if not subject.startswith("Merge "):
|
| 80 |
-
return ""
|
| 81 |
-
|
| 82 |
-
for line in body.splitlines():
|
| 83 |
-
stripped = line.strip()
|
| 84 |
-
if stripped:
|
| 85 |
-
return stripped
|
| 86 |
-
return ""
|
| 87 |
-
|
| 88 |
-
|
| 89 |
-
def parse_commits(log_output: str) -> list[ParsedCommit]:
|
| 90 |
-
"""Parse git log output into structured commits."""
|
| 91 |
-
|
| 92 |
-
commits: list[ParsedCommit] = []
|
| 93 |
-
|
| 94 |
-
for subject, body, commit_hash in iter_commit_entries(log_output):
|
| 95 |
-
is_breaking = bool(BREAKING_CHANGE_PATTERN.search(body))
|
| 96 |
-
merge_summary = get_merge_summary(subject, body)
|
| 97 |
-
candidates = [subject]
|
| 98 |
-
if merge_summary:
|
| 99 |
-
candidates.insert(0, merge_summary)
|
| 100 |
-
|
| 101 |
-
for candidate in candidates:
|
| 102 |
-
commit_match = COMMIT_PATTERN.match(candidate)
|
| 103 |
-
if not commit_match:
|
| 104 |
-
continue
|
| 105 |
-
|
| 106 |
-
scope = commit_match.group(2)
|
| 107 |
-
if scope:
|
| 108 |
-
scope = scope[1:-1]
|
| 109 |
-
commits.append(
|
| 110 |
-
ParsedCommit(
|
| 111 |
-
type=commit_match.group(1),
|
| 112 |
-
scope=scope,
|
| 113 |
-
breaking=is_breaking or bool(commit_match.group(3)),
|
| 114 |
-
message=commit_match.group(4),
|
| 115 |
-
hash=commit_hash,
|
| 116 |
-
)
|
| 117 |
-
)
|
| 118 |
-
break
|
| 119 |
-
else:
|
| 120 |
-
fallback_message = merge_summary or subject
|
| 121 |
-
if not fallback_message or fallback_message.startswith("Merge "):
|
| 122 |
-
continue
|
| 123 |
-
commits.append(
|
| 124 |
-
ParsedCommit(
|
| 125 |
-
type="other",
|
| 126 |
-
scope=None,
|
| 127 |
-
breaking=is_breaking,
|
| 128 |
-
message=fallback_message,
|
| 129 |
-
hash=commit_hash,
|
| 130 |
-
)
|
| 131 |
-
)
|
| 132 |
-
|
| 133 |
-
return commits
|
| 134 |
-
|
| 135 |
-
|
| 136 |
-
def generate_changelog(version: str, commits: list[ParsedCommit]) -> str:
|
| 137 |
-
"""Generate markdown changelog from parsed commits."""
|
| 138 |
-
today = date.today().isoformat()
|
| 139 |
-
lines = [f"## [{version}] - {today}", ""]
|
| 140 |
-
|
| 141 |
-
# Collect breaking changes
|
| 142 |
-
breaking_commits = [c for c in commits if c.breaking]
|
| 143 |
-
if breaking_commits:
|
| 144 |
-
lines.append("### Breaking Changes")
|
| 145 |
-
for commit in breaking_commits:
|
| 146 |
-
if commit.scope:
|
| 147 |
-
lines.append(f"- **{commit.scope}**: {commit.message} ({commit.hash})")
|
| 148 |
-
else:
|
| 149 |
-
lines.append(f"- {commit.message} ({commit.hash})")
|
| 150 |
-
lines.append("")
|
| 151 |
-
|
| 152 |
-
# Group by type
|
| 153 |
-
by_type: dict[str, list[ParsedCommit]] = {}
|
| 154 |
-
for commit in commits:
|
| 155 |
-
by_type.setdefault(commit.type, []).append(commit)
|
| 156 |
-
|
| 157 |
-
for commit_type, label in TYPE_LABELS.items():
|
| 158 |
-
type_commits = by_type.get(commit_type, [])
|
| 159 |
-
if not type_commits:
|
| 160 |
-
continue
|
| 161 |
-
lines.append(f"### {label}")
|
| 162 |
-
for commit in type_commits:
|
| 163 |
-
if commit.scope:
|
| 164 |
-
lines.append(f"- **{commit.scope}**: {commit.message} ({commit.hash})")
|
| 165 |
-
else:
|
| 166 |
-
lines.append(f"- {commit.message} ({commit.hash})")
|
| 167 |
-
lines.append("")
|
| 168 |
-
|
| 169 |
-
return "\n".join(lines) + "\n"
|
| 170 |
-
|
| 171 |
-
|
| 172 |
-
def run_git_log(since: str | None, cwd: Path) -> str:
|
| 173 |
-
"""Run git log command and return output."""
|
| 174 |
-
cmd = ["git", "log", "--first-parent", f"--pretty=format:{GIT_LOG_FORMAT}"]
|
| 175 |
-
if since:
|
| 176 |
-
cmd.append(f"{since}..HEAD")
|
| 177 |
-
else:
|
| 178 |
-
cmd.append("HEAD")
|
| 179 |
-
result = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)
|
| 180 |
-
return result.stdout
|
| 181 |
-
|
| 182 |
-
|
| 183 |
-
def main() -> None:
|
| 184 |
-
parser = argparse.ArgumentParser(description="Generate changelog from conventional commits")
|
| 185 |
-
parser.add_argument("--version", required=True, help="Version number (e.g., 0.6.0)")
|
| 186 |
-
parser.add_argument("--since", help="Starting tag (exclusive)")
|
| 187 |
-
parser.add_argument("--dry-run", action="store_true", help="Print to stdout instead of writing")
|
| 188 |
-
args = parser.parse_args()
|
| 189 |
-
|
| 190 |
-
log_output = run_git_log(args.since, ROOT)
|
| 191 |
-
commits = parse_commits(log_output)
|
| 192 |
-
changelog = generate_changelog(args.version, commits)
|
| 193 |
-
|
| 194 |
-
if args.dry_run:
|
| 195 |
-
print(changelog)
|
| 196 |
-
else:
|
| 197 |
-
output_path = ROOT / ".changelog.md"
|
| 198 |
-
output_path.write_text(changelog, encoding="utf-8")
|
| 199 |
-
print(f"Changelog written to {output_path}")
|
| 200 |
-
|
| 201 |
-
|
| 202 |
-
if __name__ == "__main__":
|
| 203 |
-
main()
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""Generate changelog from conventional commits."""
|
| 3 |
+
|
| 4 |
+
from __future__ import annotations
|
| 5 |
+
|
| 6 |
+
import argparse
|
| 7 |
+
import re
|
| 8 |
+
import subprocess
|
| 9 |
+
from datetime import date
|
| 10 |
+
from pathlib import Path
|
| 11 |
+
from typing import NamedTuple
|
| 12 |
+
|
| 13 |
+
ROOT = Path(__file__).parent.parent
|
| 14 |
+
|
| 15 |
+
COMMIT_PATTERN = re.compile(
|
| 16 |
+
r"^(feat|fix|ci|chore|perf|refactor|docs|style|test)(\(.+\))?(!)?:\s*(.+)$"
|
| 17 |
+
)
|
| 18 |
+
BREAKING_CHANGE_PATTERN = re.compile(r"^BREAKING CHANGE:\s*(.+)$", re.MULTILINE)
|
| 19 |
+
COMMIT_ENTRY_PATTERN = re.compile(r"^(.+?)(?:\n(.+))?\|(\w+)$", re.MULTILINE)
|
| 20 |
+
FIELD_SEP = "\x1f"
|
| 21 |
+
RECORD_SEP = "\x1e"
|
| 22 |
+
GIT_LOG_FORMAT = "%s%x1f%b%x1f%h%x1e"
|
| 23 |
+
|
| 24 |
+
TYPE_LABELS: dict[str, str] = {
|
| 25 |
+
"feat": "Features",
|
| 26 |
+
"fix": "Bug Fixes",
|
| 27 |
+
"ci": "CI/CD",
|
| 28 |
+
"chore": "Chores",
|
| 29 |
+
"perf": "Performance",
|
| 30 |
+
"refactor": "Refactors",
|
| 31 |
+
"docs": "Documentation",
|
| 32 |
+
"style": "Styles",
|
| 33 |
+
"test": "Tests",
|
| 34 |
+
"other": "Other Changes",
|
| 35 |
+
}
|
| 36 |
+
|
| 37 |
+
|
| 38 |
+
class ParsedCommit(NamedTuple):
|
| 39 |
+
type: str
|
| 40 |
+
scope: str | None
|
| 41 |
+
breaking: bool
|
| 42 |
+
message: str
|
| 43 |
+
hash: str
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def iter_commit_entries(log_output: str) -> list[tuple[str, str, str]]:
|
| 47 |
+
"""Split raw git log output into (subject, body, hash) tuples."""
|
| 48 |
+
|
| 49 |
+
if not log_output.strip():
|
| 50 |
+
return []
|
| 51 |
+
|
| 52 |
+
if RECORD_SEP in log_output and FIELD_SEP in log_output:
|
| 53 |
+
entries: list[tuple[str, str, str]] = []
|
| 54 |
+
for raw_entry in log_output.split(RECORD_SEP):
|
| 55 |
+
if not raw_entry:
|
| 56 |
+
continue
|
| 57 |
+
if FIELD_SEP not in raw_entry:
|
| 58 |
+
continue
|
| 59 |
+
subject, body_and_hash = raw_entry.split(FIELD_SEP, 1)
|
| 60 |
+
if FIELD_SEP not in body_and_hash:
|
| 61 |
+
continue
|
| 62 |
+
body, commit_hash = body_and_hash.rsplit(FIELD_SEP, 1)
|
| 63 |
+
entries.append((subject.strip(), body.strip(), commit_hash.strip()))
|
| 64 |
+
return entries
|
| 65 |
+
|
| 66 |
+
return [
|
| 67 |
+
(
|
| 68 |
+
match.group(1).strip(),
|
| 69 |
+
(match.group(2) or "").strip(),
|
| 70 |
+
match.group(3).strip(),
|
| 71 |
+
)
|
| 72 |
+
for match in COMMIT_ENTRY_PATTERN.finditer(log_output)
|
| 73 |
+
]
|
| 74 |
+
|
| 75 |
+
|
| 76 |
+
def get_merge_summary(subject: str, body: str) -> str:
|
| 77 |
+
"""Return the first meaningful summary line for a merge commit."""
|
| 78 |
+
|
| 79 |
+
if not subject.startswith("Merge "):
|
| 80 |
+
return ""
|
| 81 |
+
|
| 82 |
+
for line in body.splitlines():
|
| 83 |
+
stripped = line.strip()
|
| 84 |
+
if stripped:
|
| 85 |
+
return stripped
|
| 86 |
+
return ""
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def parse_commits(log_output: str) -> list[ParsedCommit]:
|
| 90 |
+
"""Parse git log output into structured commits."""
|
| 91 |
+
|
| 92 |
+
commits: list[ParsedCommit] = []
|
| 93 |
+
|
| 94 |
+
for subject, body, commit_hash in iter_commit_entries(log_output):
|
| 95 |
+
is_breaking = bool(BREAKING_CHANGE_PATTERN.search(body))
|
| 96 |
+
merge_summary = get_merge_summary(subject, body)
|
| 97 |
+
candidates = [subject]
|
| 98 |
+
if merge_summary:
|
| 99 |
+
candidates.insert(0, merge_summary)
|
| 100 |
+
|
| 101 |
+
for candidate in candidates:
|
| 102 |
+
commit_match = COMMIT_PATTERN.match(candidate)
|
| 103 |
+
if not commit_match:
|
| 104 |
+
continue
|
| 105 |
+
|
| 106 |
+
scope = commit_match.group(2)
|
| 107 |
+
if scope:
|
| 108 |
+
scope = scope[1:-1]
|
| 109 |
+
commits.append(
|
| 110 |
+
ParsedCommit(
|
| 111 |
+
type=commit_match.group(1),
|
| 112 |
+
scope=scope,
|
| 113 |
+
breaking=is_breaking or bool(commit_match.group(3)),
|
| 114 |
+
message=commit_match.group(4),
|
| 115 |
+
hash=commit_hash,
|
| 116 |
+
)
|
| 117 |
+
)
|
| 118 |
+
break
|
| 119 |
+
else:
|
| 120 |
+
fallback_message = merge_summary or subject
|
| 121 |
+
if not fallback_message or fallback_message.startswith("Merge "):
|
| 122 |
+
continue
|
| 123 |
+
commits.append(
|
| 124 |
+
ParsedCommit(
|
| 125 |
+
type="other",
|
| 126 |
+
scope=None,
|
| 127 |
+
breaking=is_breaking,
|
| 128 |
+
message=fallback_message,
|
| 129 |
+
hash=commit_hash,
|
| 130 |
+
)
|
| 131 |
+
)
|
| 132 |
+
|
| 133 |
+
return commits
|
| 134 |
+
|
| 135 |
+
|
| 136 |
+
def generate_changelog(version: str, commits: list[ParsedCommit]) -> str:
|
| 137 |
+
"""Generate markdown changelog from parsed commits."""
|
| 138 |
+
today = date.today().isoformat()
|
| 139 |
+
lines = [f"## [{version}] - {today}", ""]
|
| 140 |
+
|
| 141 |
+
# Collect breaking changes
|
| 142 |
+
breaking_commits = [c for c in commits if c.breaking]
|
| 143 |
+
if breaking_commits:
|
| 144 |
+
lines.append("### Breaking Changes")
|
| 145 |
+
for commit in breaking_commits:
|
| 146 |
+
if commit.scope:
|
| 147 |
+
lines.append(f"- **{commit.scope}**: {commit.message} ({commit.hash})")
|
| 148 |
+
else:
|
| 149 |
+
lines.append(f"- {commit.message} ({commit.hash})")
|
| 150 |
+
lines.append("")
|
| 151 |
+
|
| 152 |
+
# Group by type
|
| 153 |
+
by_type: dict[str, list[ParsedCommit]] = {}
|
| 154 |
+
for commit in commits:
|
| 155 |
+
by_type.setdefault(commit.type, []).append(commit)
|
| 156 |
+
|
| 157 |
+
for commit_type, label in TYPE_LABELS.items():
|
| 158 |
+
type_commits = by_type.get(commit_type, [])
|
| 159 |
+
if not type_commits:
|
| 160 |
+
continue
|
| 161 |
+
lines.append(f"### {label}")
|
| 162 |
+
for commit in type_commits:
|
| 163 |
+
if commit.scope:
|
| 164 |
+
lines.append(f"- **{commit.scope}**: {commit.message} ({commit.hash})")
|
| 165 |
+
else:
|
| 166 |
+
lines.append(f"- {commit.message} ({commit.hash})")
|
| 167 |
+
lines.append("")
|
| 168 |
+
|
| 169 |
+
return "\n".join(lines) + "\n"
|
| 170 |
+
|
| 171 |
+
|
| 172 |
+
def run_git_log(since: str | None, cwd: Path) -> str:
|
| 173 |
+
"""Run git log command and return output."""
|
| 174 |
+
cmd = ["git", "log", "--first-parent", f"--pretty=format:{GIT_LOG_FORMAT}"]
|
| 175 |
+
if since:
|
| 176 |
+
cmd.append(f"{since}..HEAD")
|
| 177 |
+
else:
|
| 178 |
+
cmd.append("HEAD")
|
| 179 |
+
result = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)
|
| 180 |
+
return result.stdout
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
def main() -> None:
|
| 184 |
+
parser = argparse.ArgumentParser(description="Generate changelog from conventional commits")
|
| 185 |
+
parser.add_argument("--version", required=True, help="Version number (e.g., 0.6.0)")
|
| 186 |
+
parser.add_argument("--since", help="Starting tag (exclusive)")
|
| 187 |
+
parser.add_argument("--dry-run", action="store_true", help="Print to stdout instead of writing")
|
| 188 |
+
args = parser.parse_args()
|
| 189 |
+
|
| 190 |
+
log_output = run_git_log(args.since, ROOT)
|
| 191 |
+
commits = parse_commits(log_output)
|
| 192 |
+
changelog = generate_changelog(args.version, commits)
|
| 193 |
+
|
| 194 |
+
if args.dry_run:
|
| 195 |
+
print(changelog)
|
| 196 |
+
else:
|
| 197 |
+
output_path = ROOT / ".changelog.md"
|
| 198 |
+
output_path.write_text(changelog, encoding="utf-8")
|
| 199 |
+
print(f"Changelog written to {output_path}")
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
if __name__ == "__main__":
|
| 203 |
+
main()
|
|
@@ -1,55 +1,55 @@
|
|
| 1 |
-
"""Sync plugin manifest versions to the repo's computed release semver."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import subprocess
|
| 6 |
-
import sys
|
| 7 |
-
from pathlib import Path
|
| 8 |
-
|
| 9 |
-
ROOT = Path(__file__).resolve().parent.parent
|
| 10 |
-
if str(ROOT) not in sys.path:
|
| 11 |
-
sys.path.insert(0, str(ROOT))
|
| 12 |
-
|
| 13 |
-
from headroom.release_version import ( # noqa: E402
|
| 14 |
-
compute_release_version,
|
| 15 |
-
determine_bump_level,
|
| 16 |
-
find_latest_release_tag,
|
| 17 |
-
get_canonical_version,
|
| 18 |
-
list_release_commits,
|
| 19 |
-
list_release_tags,
|
| 20 |
-
)
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
def compute_repo_semver(root: Path) -> str:
|
| 24 |
-
"""Return the npm-style semver for the repo's next release."""
|
| 25 |
-
tags = list_release_tags(root)
|
| 26 |
-
previous_tag = find_latest_release_tag(tags) or ""
|
| 27 |
-
level = determine_bump_level(list_release_commits(root, previous_tag))
|
| 28 |
-
info = compute_release_version(
|
| 29 |
-
canonical_version=get_canonical_version(root),
|
| 30 |
-
level=level,
|
| 31 |
-
tags=tags,
|
| 32 |
-
)
|
| 33 |
-
return info.npm_version
|
| 34 |
-
|
| 35 |
-
|
| 36 |
-
def main() -> None:
|
| 37 |
-
root = ROOT
|
| 38 |
-
version = compute_repo_semver(root)
|
| 39 |
-
subprocess.run(
|
| 40 |
-
[
|
| 41 |
-
sys.executable,
|
| 42 |
-
str(root / "scripts" / "version-sync.py"),
|
| 43 |
-
"--root",
|
| 44 |
-
str(root),
|
| 45 |
-
"--version",
|
| 46 |
-
version,
|
| 47 |
-
"--plugin-manifests-only",
|
| 48 |
-
],
|
| 49 |
-
cwd=root,
|
| 50 |
-
check=True,
|
| 51 |
-
)
|
| 52 |
-
|
| 53 |
-
|
| 54 |
-
if __name__ == "__main__":
|
| 55 |
-
main()
|
|
|
|
| 1 |
+
"""Sync plugin manifest versions to the repo's computed release semver."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import subprocess
|
| 6 |
+
import sys
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
|
| 9 |
+
ROOT = Path(__file__).resolve().parent.parent
|
| 10 |
+
if str(ROOT) not in sys.path:
|
| 11 |
+
sys.path.insert(0, str(ROOT))
|
| 12 |
+
|
| 13 |
+
from headroom.release_version import ( # noqa: E402
|
| 14 |
+
compute_release_version,
|
| 15 |
+
determine_bump_level,
|
| 16 |
+
find_latest_release_tag,
|
| 17 |
+
get_canonical_version,
|
| 18 |
+
list_release_commits,
|
| 19 |
+
list_release_tags,
|
| 20 |
+
)
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
def compute_repo_semver(root: Path) -> str:
|
| 24 |
+
"""Return the npm-style semver for the repo's next release."""
|
| 25 |
+
tags = list_release_tags(root)
|
| 26 |
+
previous_tag = find_latest_release_tag(tags) or ""
|
| 27 |
+
level = determine_bump_level(list_release_commits(root, previous_tag))
|
| 28 |
+
info = compute_release_version(
|
| 29 |
+
canonical_version=get_canonical_version(root),
|
| 30 |
+
level=level,
|
| 31 |
+
tags=tags,
|
| 32 |
+
)
|
| 33 |
+
return info.npm_version
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
def main() -> None:
|
| 37 |
+
root = ROOT
|
| 38 |
+
version = compute_repo_semver(root)
|
| 39 |
+
subprocess.run(
|
| 40 |
+
[
|
| 41 |
+
sys.executable,
|
| 42 |
+
str(root / "scripts" / "version-sync.py"),
|
| 43 |
+
"--root",
|
| 44 |
+
str(root),
|
| 45 |
+
"--version",
|
| 46 |
+
version,
|
| 47 |
+
"--plugin-manifests-only",
|
| 48 |
+
],
|
| 49 |
+
cwd=root,
|
| 50 |
+
check=True,
|
| 51 |
+
)
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
if __name__ == "__main__":
|
| 55 |
+
main()
|
|
@@ -1,302 +1,302 @@
|
|
| 1 |
-
"""Tests for changelog-gen.py."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import importlib.util
|
| 6 |
-
from pathlib import Path
|
| 7 |
-
|
| 8 |
-
import pytest
|
| 9 |
-
|
| 10 |
-
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 11 |
-
|
| 12 |
-
# Load changelog_gen module from scripts directory (filename has hyphen)
|
| 13 |
-
_spec = importlib.util.spec_from_file_location(
|
| 14 |
-
"changelog_gen", ROOT / "scripts" / "changelog-gen.py"
|
| 15 |
-
)
|
| 16 |
-
if _spec is None:
|
| 17 |
-
raise ImportError("Could not load changelog_gen module")
|
| 18 |
-
_changelog_gen = importlib.util.module_from_spec(_spec)
|
| 19 |
-
if _spec.loader is None:
|
| 20 |
-
raise ImportError("Could not load changelog_gen module")
|
| 21 |
-
_spec.loader.exec_module(_changelog_gen)
|
| 22 |
-
|
| 23 |
-
COMMIT_PATTERN = _changelog_gen.COMMIT_PATTERN
|
| 24 |
-
BREAKING_CHANGE_PATTERN = _changelog_gen.BREAKING_CHANGE_PATTERN
|
| 25 |
-
FIELD_SEP = _changelog_gen.FIELD_SEP
|
| 26 |
-
RECORD_SEP = _changelog_gen.RECORD_SEP
|
| 27 |
-
ParsedCommit = _changelog_gen.ParsedCommit
|
| 28 |
-
generate_changelog = _changelog_gen.generate_changelog
|
| 29 |
-
iter_commit_entries = _changelog_gen.iter_commit_entries
|
| 30 |
-
parse_commits = _changelog_gen.parse_commits
|
| 31 |
-
|
| 32 |
-
|
| 33 |
-
def make_log_entry(subject: str, commit_hash: str, body: str = "") -> str:
|
| 34 |
-
return f"{subject}{FIELD_SEP}{body}{FIELD_SEP}{commit_hash}{RECORD_SEP}"
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
class TestParseCommits:
|
| 38 |
-
"""Tests for parse_commits function."""
|
| 39 |
-
|
| 40 |
-
def test_parses_feat_commit(self) -> None:
|
| 41 |
-
log_output = make_log_entry("feat(core): add feature", "abc1234")
|
| 42 |
-
commits = parse_commits(log_output)
|
| 43 |
-
assert len(commits) == 1
|
| 44 |
-
assert commits[0].type == "feat"
|
| 45 |
-
assert commits[0].scope == "core"
|
| 46 |
-
assert commits[0].message == "add feature"
|
| 47 |
-
assert commits[0].hash == "abc1234"
|
| 48 |
-
assert commits[0].breaking is False
|
| 49 |
-
|
| 50 |
-
def test_parses_fix_commit(self) -> None:
|
| 51 |
-
log_output = make_log_entry("fix(ui): fix bug", "def5678")
|
| 52 |
-
commits = parse_commits(log_output)
|
| 53 |
-
assert len(commits) == 1
|
| 54 |
-
assert commits[0].type == "fix"
|
| 55 |
-
assert commits[0].scope == "ui"
|
| 56 |
-
assert commits[0].message == "fix bug"
|
| 57 |
-
assert commits[0].hash == "def5678"
|
| 58 |
-
|
| 59 |
-
def test_parses_ci_commit(self) -> None:
|
| 60 |
-
log_output = make_log_entry("ci: update github actions", "xyz789")
|
| 61 |
-
commits = parse_commits(log_output)
|
| 62 |
-
assert len(commits) == 1
|
| 63 |
-
assert commits[0].type == "ci"
|
| 64 |
-
assert commits[0].scope is None
|
| 65 |
-
assert commits[0].message == "update github actions"
|
| 66 |
-
assert commits[0].hash == "xyz789"
|
| 67 |
-
|
| 68 |
-
def test_parses_chore_commit(self) -> None:
|
| 69 |
-
log_output = make_log_entry("chore: cleanup", "xyz999")
|
| 70 |
-
commits = parse_commits(log_output)
|
| 71 |
-
assert len(commits) == 1
|
| 72 |
-
assert commits[0].type == "chore"
|
| 73 |
-
assert commits[0].scope is None
|
| 74 |
-
|
| 75 |
-
def test_parses_perf_commit(self) -> None:
|
| 76 |
-
log_output = make_log_entry("perf(dashboard): improve performance", "abc111")
|
| 77 |
-
commits = parse_commits(log_output)
|
| 78 |
-
assert len(commits) == 1
|
| 79 |
-
assert commits[0].type == "perf"
|
| 80 |
-
assert commits[0].scope == "dashboard"
|
| 81 |
-
|
| 82 |
-
def test_parses_refactor_commit(self) -> None:
|
| 83 |
-
log_output = make_log_entry("refactor(api): refactor endpoint", "abc222")
|
| 84 |
-
commits = parse_commits(log_output)
|
| 85 |
-
assert len(commits) == 1
|
| 86 |
-
assert commits[0].type == "refactor"
|
| 87 |
-
assert commits[0].scope == "api"
|
| 88 |
-
|
| 89 |
-
def test_parses_docs_commit(self) -> None:
|
| 90 |
-
log_output = make_log_entry("docs: update readme", "abc333")
|
| 91 |
-
commits = parse_commits(log_output)
|
| 92 |
-
assert len(commits) == 1
|
| 93 |
-
assert commits[0].type == "docs"
|
| 94 |
-
|
| 95 |
-
def test_parses_style_commit(self) -> None:
|
| 96 |
-
log_output = make_log_entry("style: format code", "abc444")
|
| 97 |
-
commits = parse_commits(log_output)
|
| 98 |
-
assert len(commits) == 1
|
| 99 |
-
assert commits[0].type == "style"
|
| 100 |
-
|
| 101 |
-
def test_parses_test_commit(self) -> None:
|
| 102 |
-
log_output = make_log_entry("test: add tests for feature", "abc555")
|
| 103 |
-
commits = parse_commits(log_output)
|
| 104 |
-
assert len(commits) == 1
|
| 105 |
-
assert commits[0].type == "test"
|
| 106 |
-
|
| 107 |
-
def test_detects_breaking_change_in_body(self) -> None:
|
| 108 |
-
log_output = make_log_entry(
|
| 109 |
-
"feat(core): add feature",
|
| 110 |
-
"abc666",
|
| 111 |
-
"BREAKING CHANGE: api changed",
|
| 112 |
-
)
|
| 113 |
-
commits = parse_commits(log_output)
|
| 114 |
-
assert len(commits) == 1
|
| 115 |
-
assert commits[0].breaking is True
|
| 116 |
-
|
| 117 |
-
def test_detects_breaking_change_exclamation(self) -> None:
|
| 118 |
-
log_output = make_log_entry("feat(core)!: api changed", "abc777")
|
| 119 |
-
commits = parse_commits(log_output)
|
| 120 |
-
assert len(commits) == 1
|
| 121 |
-
assert commits[0].breaking is True
|
| 122 |
-
|
| 123 |
-
def test_no_scope_no_problem(self) -> None:
|
| 124 |
-
log_output = make_log_entry("feat: simple feature", "abc888")
|
| 125 |
-
commits = parse_commits(log_output)
|
| 126 |
-
assert len(commits) == 1
|
| 127 |
-
assert commits[0].scope is None
|
| 128 |
-
assert commits[0].message == "simple feature"
|
| 129 |
-
|
| 130 |
-
def test_uses_pr_title_from_merge_commit_body(self) -> None:
|
| 131 |
-
log_output = make_log_entry(
|
| 132 |
-
"Merge pull request #173 from JerrettDavis/fix/pipeline-permissions-and-docs",
|
| 133 |
-
"73f6673",
|
| 134 |
-
"fix: repair release and docs pipelines",
|
| 135 |
-
)
|
| 136 |
-
commits = parse_commits(log_output)
|
| 137 |
-
assert len(commits) == 1
|
| 138 |
-
assert commits[0].type == "fix"
|
| 139 |
-
assert commits[0].message == "repair release and docs pipelines"
|
| 140 |
-
|
| 141 |
-
def test_falls_back_to_other_changes_for_non_conventional_merge(self) -> None:
|
| 142 |
-
log_output = make_log_entry(
|
| 143 |
-
"Merge pull request #186 from skorokithakis/patch-1",
|
| 144 |
-
"1e80ee3",
|
| 145 |
-
"Add support for custom Anthropic API URL",
|
| 146 |
-
)
|
| 147 |
-
commits = parse_commits(log_output)
|
| 148 |
-
assert len(commits) == 1
|
| 149 |
-
assert commits[0].type == "other"
|
| 150 |
-
assert commits[0].message == "Add support for custom Anthropic API URL"
|
| 151 |
-
|
| 152 |
-
def test_iter_commit_entries_parses_real_git_log_delimiters(self) -> None:
|
| 153 |
-
log_output = make_log_entry(
|
| 154 |
-
"fix: patch release flow", "abc1234", "BREAKING CHANGE: no"
|
| 155 |
-
) + make_log_entry("docs: update readme", "def5678")
|
| 156 |
-
assert iter_commit_entries(log_output) == [
|
| 157 |
-
("fix: patch release flow", "BREAKING CHANGE: no", "abc1234"),
|
| 158 |
-
("docs: update readme", "", "def5678"),
|
| 159 |
-
]
|
| 160 |
-
|
| 161 |
-
def test_iter_commit_entries_keeps_field_separator_inside_body(self) -> None:
|
| 162 |
-
body = f"line one{FIELD_SEP}line two"
|
| 163 |
-
log_output = make_log_entry("fix: patch release flow", "abc1234", body)
|
| 164 |
-
assert iter_commit_entries(log_output) == [
|
| 165 |
-
("fix: patch release flow", body, "abc1234"),
|
| 166 |
-
]
|
| 167 |
-
|
| 168 |
-
|
| 169 |
-
class TestGenerateChangelog:
|
| 170 |
-
"""Tests for generate_changelog function."""
|
| 171 |
-
|
| 172 |
-
def test_generates_version_header(self) -> None:
|
| 173 |
-
commits = [
|
| 174 |
-
ParsedCommit(type="feat", scope=None, breaking=False, message="test", hash="abc123")
|
| 175 |
-
]
|
| 176 |
-
result = generate_changelog("0.6.0", commits)
|
| 177 |
-
assert "## [0.6.0]" in result
|
| 178 |
-
|
| 179 |
-
def test_includes_date(self) -> None:
|
| 180 |
-
commits = []
|
| 181 |
-
result = generate_changelog("0.6.0", commits)
|
| 182 |
-
import re
|
| 183 |
-
|
| 184 |
-
date_match = re.search(r"\d{4}-\d{2}-\d{2}", result)
|
| 185 |
-
assert date_match is not None
|
| 186 |
-
|
| 187 |
-
def test_groups_by_type(self) -> None:
|
| 188 |
-
commits = [
|
| 189 |
-
ParsedCommit(
|
| 190 |
-
type="feat", scope=None, breaking=False, message="add feature", hash="abc123"
|
| 191 |
-
),
|
| 192 |
-
ParsedCommit(type="fix", scope=None, breaking=False, message="fix bug", hash="def456"),
|
| 193 |
-
]
|
| 194 |
-
result = generate_changelog("0.6.0", commits)
|
| 195 |
-
assert "### Features" in result
|
| 196 |
-
assert "### Bug Fixes" in result
|
| 197 |
-
assert "- add feature (abc123)" in result
|
| 198 |
-
assert "- fix bug (def456)" in result
|
| 199 |
-
|
| 200 |
-
def test_includes_scope_in_bullet(self) -> None:
|
| 201 |
-
commits = [
|
| 202 |
-
ParsedCommit(
|
| 203 |
-
type="feat", scope="core", breaking=False, message="add feature", hash="abc123"
|
| 204 |
-
)
|
| 205 |
-
]
|
| 206 |
-
result = generate_changelog("0.6.0", commits)
|
| 207 |
-
assert "- **core**: add feature (abc123)" in result
|
| 208 |
-
|
| 209 |
-
def test_breaking_change_section_when_present(self) -> None:
|
| 210 |
-
commits = [
|
| 211 |
-
ParsedCommit(
|
| 212 |
-
type="feat", scope="core", breaking=True, message="api changed", hash="abc123"
|
| 213 |
-
)
|
| 214 |
-
]
|
| 215 |
-
result = generate_changelog("0.6.0", commits)
|
| 216 |
-
assert "### Breaking Changes" in result
|
| 217 |
-
assert "**core**" in result
|
| 218 |
-
|
| 219 |
-
def test_no_breaking_change_section_when_none(self) -> None:
|
| 220 |
-
commits = [
|
| 221 |
-
ParsedCommit(
|
| 222 |
-
type="feat", scope=None, breaking=False, message="add feature", hash="abc123"
|
| 223 |
-
)
|
| 224 |
-
]
|
| 225 |
-
result = generate_changelog("0.6.0", commits)
|
| 226 |
-
assert "Breaking Changes" not in result
|
| 227 |
-
|
| 228 |
-
def test_includes_other_changes_section(self) -> None:
|
| 229 |
-
commits = [
|
| 230 |
-
ParsedCommit(
|
| 231 |
-
type="other",
|
| 232 |
-
scope=None,
|
| 233 |
-
breaking=False,
|
| 234 |
-
message="Add support for custom Anthropic API URL",
|
| 235 |
-
hash="abc123",
|
| 236 |
-
)
|
| 237 |
-
]
|
| 238 |
-
result = generate_changelog("0.6.0", commits)
|
| 239 |
-
assert "### Other Changes" in result
|
| 240 |
-
assert "- Add support for custom Anthropic API URL (abc123)" in result
|
| 241 |
-
|
| 242 |
-
|
| 243 |
-
class TestIntegrationWithMock:
|
| 244 |
-
"""Integration tests with mocked subprocess.run."""
|
| 245 |
-
|
| 246 |
-
def test_full_flow_with_mocked_git(self) -> None:
|
| 247 |
-
log_output = (
|
| 248 |
-
make_log_entry("feat(core): add new feature", "abc1234")
|
| 249 |
-
+ make_log_entry("fix(ui): fix bug", "def5678")
|
| 250 |
-
+ make_log_entry("ci: update github actions", "xyz789")
|
| 251 |
-
+ make_log_entry(
|
| 252 |
-
"feat(outer): breaking change", "bbb111", "BREAKING CHANGE: this is breaking"
|
| 253 |
-
)
|
| 254 |
-
+ make_log_entry("chore: cleanup", "yyy999")
|
| 255 |
-
)
|
| 256 |
-
commits = parse_commits(log_output)
|
| 257 |
-
|
| 258 |
-
assert len(commits) == 5
|
| 259 |
-
assert any(c.type == "feat" and c.scope == "core" for c in commits)
|
| 260 |
-
assert any(c.type == "fix" and c.scope == "ui" for c in commits)
|
| 261 |
-
assert any(c.type == "ci" for c in commits)
|
| 262 |
-
assert any(c.type == "feat" and c.breaking for c in commits)
|
| 263 |
-
assert any(c.type == "chore" for c in commits)
|
| 264 |
-
|
| 265 |
-
changelog = generate_changelog("0.7.0", commits)
|
| 266 |
-
assert "## [0.7.0]" in changelog
|
| 267 |
-
assert "### Features" in changelog
|
| 268 |
-
assert "### Bug Fixes" in changelog
|
| 269 |
-
assert "### CI/CD" in changelog
|
| 270 |
-
assert "### Breaking Changes" in changelog
|
| 271 |
-
|
| 272 |
-
|
| 273 |
-
class TestCommitPattern:
|
| 274 |
-
"""Tests for the commit regex pattern."""
|
| 275 |
-
|
| 276 |
-
def test_feat_with_scope(self) -> None:
|
| 277 |
-
match = COMMIT_PATTERN.match("feat(core): add feature")
|
| 278 |
-
assert match is not None
|
| 279 |
-
assert match.group(1) == "feat"
|
| 280 |
-
assert match.group(2) == "(core)"
|
| 281 |
-
assert match.group(4) == "add feature"
|
| 282 |
-
|
| 283 |
-
def test_fix_without_scope(self) -> None:
|
| 284 |
-
match = COMMIT_PATTERN.match("fix: fix bug")
|
| 285 |
-
assert match is not None
|
| 286 |
-
assert match.group(1) == "fix"
|
| 287 |
-
assert match.group(2) is None
|
| 288 |
-
assert match.group(4) == "fix bug"
|
| 289 |
-
|
| 290 |
-
def test_with_exclamation(self) -> None:
|
| 291 |
-
match = COMMIT_PATTERN.match("feat(core)!: api changed")
|
| 292 |
-
assert match is not None
|
| 293 |
-
assert match.group(3) == "!"
|
| 294 |
-
|
| 295 |
-
def test_without_exclamation(self) -> None:
|
| 296 |
-
match = COMMIT_PATTERN.match("feat(core): add feature")
|
| 297 |
-
assert match is not None
|
| 298 |
-
assert match.group(3) is None
|
| 299 |
-
|
| 300 |
-
|
| 301 |
-
if __name__ == "__main__":
|
| 302 |
-
pytest.main([__file__, "-v"])
|
|
|
|
| 1 |
+
"""Tests for changelog-gen.py."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import importlib.util
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
import pytest
|
| 9 |
+
|
| 10 |
+
ROOT = Path(__file__).resolve().parent.parent.parent
|
| 11 |
+
|
| 12 |
+
# Load changelog_gen module from scripts directory (filename has hyphen)
|
| 13 |
+
_spec = importlib.util.spec_from_file_location(
|
| 14 |
+
"changelog_gen", ROOT / "scripts" / "changelog-gen.py"
|
| 15 |
+
)
|
| 16 |
+
if _spec is None:
|
| 17 |
+
raise ImportError("Could not load changelog_gen module")
|
| 18 |
+
_changelog_gen = importlib.util.module_from_spec(_spec)
|
| 19 |
+
if _spec.loader is None:
|
| 20 |
+
raise ImportError("Could not load changelog_gen module")
|
| 21 |
+
_spec.loader.exec_module(_changelog_gen)
|
| 22 |
+
|
| 23 |
+
COMMIT_PATTERN = _changelog_gen.COMMIT_PATTERN
|
| 24 |
+
BREAKING_CHANGE_PATTERN = _changelog_gen.BREAKING_CHANGE_PATTERN
|
| 25 |
+
FIELD_SEP = _changelog_gen.FIELD_SEP
|
| 26 |
+
RECORD_SEP = _changelog_gen.RECORD_SEP
|
| 27 |
+
ParsedCommit = _changelog_gen.ParsedCommit
|
| 28 |
+
generate_changelog = _changelog_gen.generate_changelog
|
| 29 |
+
iter_commit_entries = _changelog_gen.iter_commit_entries
|
| 30 |
+
parse_commits = _changelog_gen.parse_commits
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
def make_log_entry(subject: str, commit_hash: str, body: str = "") -> str:
|
| 34 |
+
return f"{subject}{FIELD_SEP}{body}{FIELD_SEP}{commit_hash}{RECORD_SEP}"
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
class TestParseCommits:
|
| 38 |
+
"""Tests for parse_commits function."""
|
| 39 |
+
|
| 40 |
+
def test_parses_feat_commit(self) -> None:
|
| 41 |
+
log_output = make_log_entry("feat(core): add feature", "abc1234")
|
| 42 |
+
commits = parse_commits(log_output)
|
| 43 |
+
assert len(commits) == 1
|
| 44 |
+
assert commits[0].type == "feat"
|
| 45 |
+
assert commits[0].scope == "core"
|
| 46 |
+
assert commits[0].message == "add feature"
|
| 47 |
+
assert commits[0].hash == "abc1234"
|
| 48 |
+
assert commits[0].breaking is False
|
| 49 |
+
|
| 50 |
+
def test_parses_fix_commit(self) -> None:
|
| 51 |
+
log_output = make_log_entry("fix(ui): fix bug", "def5678")
|
| 52 |
+
commits = parse_commits(log_output)
|
| 53 |
+
assert len(commits) == 1
|
| 54 |
+
assert commits[0].type == "fix"
|
| 55 |
+
assert commits[0].scope == "ui"
|
| 56 |
+
assert commits[0].message == "fix bug"
|
| 57 |
+
assert commits[0].hash == "def5678"
|
| 58 |
+
|
| 59 |
+
def test_parses_ci_commit(self) -> None:
|
| 60 |
+
log_output = make_log_entry("ci: update github actions", "xyz789")
|
| 61 |
+
commits = parse_commits(log_output)
|
| 62 |
+
assert len(commits) == 1
|
| 63 |
+
assert commits[0].type == "ci"
|
| 64 |
+
assert commits[0].scope is None
|
| 65 |
+
assert commits[0].message == "update github actions"
|
| 66 |
+
assert commits[0].hash == "xyz789"
|
| 67 |
+
|
| 68 |
+
def test_parses_chore_commit(self) -> None:
|
| 69 |
+
log_output = make_log_entry("chore: cleanup", "xyz999")
|
| 70 |
+
commits = parse_commits(log_output)
|
| 71 |
+
assert len(commits) == 1
|
| 72 |
+
assert commits[0].type == "chore"
|
| 73 |
+
assert commits[0].scope is None
|
| 74 |
+
|
| 75 |
+
def test_parses_perf_commit(self) -> None:
|
| 76 |
+
log_output = make_log_entry("perf(dashboard): improve performance", "abc111")
|
| 77 |
+
commits = parse_commits(log_output)
|
| 78 |
+
assert len(commits) == 1
|
| 79 |
+
assert commits[0].type == "perf"
|
| 80 |
+
assert commits[0].scope == "dashboard"
|
| 81 |
+
|
| 82 |
+
def test_parses_refactor_commit(self) -> None:
|
| 83 |
+
log_output = make_log_entry("refactor(api): refactor endpoint", "abc222")
|
| 84 |
+
commits = parse_commits(log_output)
|
| 85 |
+
assert len(commits) == 1
|
| 86 |
+
assert commits[0].type == "refactor"
|
| 87 |
+
assert commits[0].scope == "api"
|
| 88 |
+
|
| 89 |
+
def test_parses_docs_commit(self) -> None:
|
| 90 |
+
log_output = make_log_entry("docs: update readme", "abc333")
|
| 91 |
+
commits = parse_commits(log_output)
|
| 92 |
+
assert len(commits) == 1
|
| 93 |
+
assert commits[0].type == "docs"
|
| 94 |
+
|
| 95 |
+
def test_parses_style_commit(self) -> None:
|
| 96 |
+
log_output = make_log_entry("style: format code", "abc444")
|
| 97 |
+
commits = parse_commits(log_output)
|
| 98 |
+
assert len(commits) == 1
|
| 99 |
+
assert commits[0].type == "style"
|
| 100 |
+
|
| 101 |
+
def test_parses_test_commit(self) -> None:
|
| 102 |
+
log_output = make_log_entry("test: add tests for feature", "abc555")
|
| 103 |
+
commits = parse_commits(log_output)
|
| 104 |
+
assert len(commits) == 1
|
| 105 |
+
assert commits[0].type == "test"
|
| 106 |
+
|
| 107 |
+
def test_detects_breaking_change_in_body(self) -> None:
|
| 108 |
+
log_output = make_log_entry(
|
| 109 |
+
"feat(core): add feature",
|
| 110 |
+
"abc666",
|
| 111 |
+
"BREAKING CHANGE: api changed",
|
| 112 |
+
)
|
| 113 |
+
commits = parse_commits(log_output)
|
| 114 |
+
assert len(commits) == 1
|
| 115 |
+
assert commits[0].breaking is True
|
| 116 |
+
|
| 117 |
+
def test_detects_breaking_change_exclamation(self) -> None:
|
| 118 |
+
log_output = make_log_entry("feat(core)!: api changed", "abc777")
|
| 119 |
+
commits = parse_commits(log_output)
|
| 120 |
+
assert len(commits) == 1
|
| 121 |
+
assert commits[0].breaking is True
|
| 122 |
+
|
| 123 |
+
def test_no_scope_no_problem(self) -> None:
|
| 124 |
+
log_output = make_log_entry("feat: simple feature", "abc888")
|
| 125 |
+
commits = parse_commits(log_output)
|
| 126 |
+
assert len(commits) == 1
|
| 127 |
+
assert commits[0].scope is None
|
| 128 |
+
assert commits[0].message == "simple feature"
|
| 129 |
+
|
| 130 |
+
def test_uses_pr_title_from_merge_commit_body(self) -> None:
|
| 131 |
+
log_output = make_log_entry(
|
| 132 |
+
"Merge pull request #173 from JerrettDavis/fix/pipeline-permissions-and-docs",
|
| 133 |
+
"73f6673",
|
| 134 |
+
"fix: repair release and docs pipelines",
|
| 135 |
+
)
|
| 136 |
+
commits = parse_commits(log_output)
|
| 137 |
+
assert len(commits) == 1
|
| 138 |
+
assert commits[0].type == "fix"
|
| 139 |
+
assert commits[0].message == "repair release and docs pipelines"
|
| 140 |
+
|
| 141 |
+
def test_falls_back_to_other_changes_for_non_conventional_merge(self) -> None:
|
| 142 |
+
log_output = make_log_entry(
|
| 143 |
+
"Merge pull request #186 from skorokithakis/patch-1",
|
| 144 |
+
"1e80ee3",
|
| 145 |
+
"Add support for custom Anthropic API URL",
|
| 146 |
+
)
|
| 147 |
+
commits = parse_commits(log_output)
|
| 148 |
+
assert len(commits) == 1
|
| 149 |
+
assert commits[0].type == "other"
|
| 150 |
+
assert commits[0].message == "Add support for custom Anthropic API URL"
|
| 151 |
+
|
| 152 |
+
def test_iter_commit_entries_parses_real_git_log_delimiters(self) -> None:
|
| 153 |
+
log_output = make_log_entry(
|
| 154 |
+
"fix: patch release flow", "abc1234", "BREAKING CHANGE: no"
|
| 155 |
+
) + make_log_entry("docs: update readme", "def5678")
|
| 156 |
+
assert iter_commit_entries(log_output) == [
|
| 157 |
+
("fix: patch release flow", "BREAKING CHANGE: no", "abc1234"),
|
| 158 |
+
("docs: update readme", "", "def5678"),
|
| 159 |
+
]
|
| 160 |
+
|
| 161 |
+
def test_iter_commit_entries_keeps_field_separator_inside_body(self) -> None:
|
| 162 |
+
body = f"line one{FIELD_SEP}line two"
|
| 163 |
+
log_output = make_log_entry("fix: patch release flow", "abc1234", body)
|
| 164 |
+
assert iter_commit_entries(log_output) == [
|
| 165 |
+
("fix: patch release flow", body, "abc1234"),
|
| 166 |
+
]
|
| 167 |
+
|
| 168 |
+
|
| 169 |
+
class TestGenerateChangelog:
|
| 170 |
+
"""Tests for generate_changelog function."""
|
| 171 |
+
|
| 172 |
+
def test_generates_version_header(self) -> None:
|
| 173 |
+
commits = [
|
| 174 |
+
ParsedCommit(type="feat", scope=None, breaking=False, message="test", hash="abc123")
|
| 175 |
+
]
|
| 176 |
+
result = generate_changelog("0.6.0", commits)
|
| 177 |
+
assert "## [0.6.0]" in result
|
| 178 |
+
|
| 179 |
+
def test_includes_date(self) -> None:
|
| 180 |
+
commits = []
|
| 181 |
+
result = generate_changelog("0.6.0", commits)
|
| 182 |
+
import re
|
| 183 |
+
|
| 184 |
+
date_match = re.search(r"\d{4}-\d{2}-\d{2}", result)
|
| 185 |
+
assert date_match is not None
|
| 186 |
+
|
| 187 |
+
def test_groups_by_type(self) -> None:
|
| 188 |
+
commits = [
|
| 189 |
+
ParsedCommit(
|
| 190 |
+
type="feat", scope=None, breaking=False, message="add feature", hash="abc123"
|
| 191 |
+
),
|
| 192 |
+
ParsedCommit(type="fix", scope=None, breaking=False, message="fix bug", hash="def456"),
|
| 193 |
+
]
|
| 194 |
+
result = generate_changelog("0.6.0", commits)
|
| 195 |
+
assert "### Features" in result
|
| 196 |
+
assert "### Bug Fixes" in result
|
| 197 |
+
assert "- add feature (abc123)" in result
|
| 198 |
+
assert "- fix bug (def456)" in result
|
| 199 |
+
|
| 200 |
+
def test_includes_scope_in_bullet(self) -> None:
|
| 201 |
+
commits = [
|
| 202 |
+
ParsedCommit(
|
| 203 |
+
type="feat", scope="core", breaking=False, message="add feature", hash="abc123"
|
| 204 |
+
)
|
| 205 |
+
]
|
| 206 |
+
result = generate_changelog("0.6.0", commits)
|
| 207 |
+
assert "- **core**: add feature (abc123)" in result
|
| 208 |
+
|
| 209 |
+
def test_breaking_change_section_when_present(self) -> None:
|
| 210 |
+
commits = [
|
| 211 |
+
ParsedCommit(
|
| 212 |
+
type="feat", scope="core", breaking=True, message="api changed", hash="abc123"
|
| 213 |
+
)
|
| 214 |
+
]
|
| 215 |
+
result = generate_changelog("0.6.0", commits)
|
| 216 |
+
assert "### Breaking Changes" in result
|
| 217 |
+
assert "**core**" in result
|
| 218 |
+
|
| 219 |
+
def test_no_breaking_change_section_when_none(self) -> None:
|
| 220 |
+
commits = [
|
| 221 |
+
ParsedCommit(
|
| 222 |
+
type="feat", scope=None, breaking=False, message="add feature", hash="abc123"
|
| 223 |
+
)
|
| 224 |
+
]
|
| 225 |
+
result = generate_changelog("0.6.0", commits)
|
| 226 |
+
assert "Breaking Changes" not in result
|
| 227 |
+
|
| 228 |
+
def test_includes_other_changes_section(self) -> None:
|
| 229 |
+
commits = [
|
| 230 |
+
ParsedCommit(
|
| 231 |
+
type="other",
|
| 232 |
+
scope=None,
|
| 233 |
+
breaking=False,
|
| 234 |
+
message="Add support for custom Anthropic API URL",
|
| 235 |
+
hash="abc123",
|
| 236 |
+
)
|
| 237 |
+
]
|
| 238 |
+
result = generate_changelog("0.6.0", commits)
|
| 239 |
+
assert "### Other Changes" in result
|
| 240 |
+
assert "- Add support for custom Anthropic API URL (abc123)" in result
|
| 241 |
+
|
| 242 |
+
|
| 243 |
+
class TestIntegrationWithMock:
|
| 244 |
+
"""Integration tests with mocked subprocess.run."""
|
| 245 |
+
|
| 246 |
+
def test_full_flow_with_mocked_git(self) -> None:
|
| 247 |
+
log_output = (
|
| 248 |
+
make_log_entry("feat(core): add new feature", "abc1234")
|
| 249 |
+
+ make_log_entry("fix(ui): fix bug", "def5678")
|
| 250 |
+
+ make_log_entry("ci: update github actions", "xyz789")
|
| 251 |
+
+ make_log_entry(
|
| 252 |
+
"feat(outer): breaking change", "bbb111", "BREAKING CHANGE: this is breaking"
|
| 253 |
+
)
|
| 254 |
+
+ make_log_entry("chore: cleanup", "yyy999")
|
| 255 |
+
)
|
| 256 |
+
commits = parse_commits(log_output)
|
| 257 |
+
|
| 258 |
+
assert len(commits) == 5
|
| 259 |
+
assert any(c.type == "feat" and c.scope == "core" for c in commits)
|
| 260 |
+
assert any(c.type == "fix" and c.scope == "ui" for c in commits)
|
| 261 |
+
assert any(c.type == "ci" for c in commits)
|
| 262 |
+
assert any(c.type == "feat" and c.breaking for c in commits)
|
| 263 |
+
assert any(c.type == "chore" for c in commits)
|
| 264 |
+
|
| 265 |
+
changelog = generate_changelog("0.7.0", commits)
|
| 266 |
+
assert "## [0.7.0]" in changelog
|
| 267 |
+
assert "### Features" in changelog
|
| 268 |
+
assert "### Bug Fixes" in changelog
|
| 269 |
+
assert "### CI/CD" in changelog
|
| 270 |
+
assert "### Breaking Changes" in changelog
|
| 271 |
+
|
| 272 |
+
|
| 273 |
+
class TestCommitPattern:
|
| 274 |
+
"""Tests for the commit regex pattern."""
|
| 275 |
+
|
| 276 |
+
def test_feat_with_scope(self) -> None:
|
| 277 |
+
match = COMMIT_PATTERN.match("feat(core): add feature")
|
| 278 |
+
assert match is not None
|
| 279 |
+
assert match.group(1) == "feat"
|
| 280 |
+
assert match.group(2) == "(core)"
|
| 281 |
+
assert match.group(4) == "add feature"
|
| 282 |
+
|
| 283 |
+
def test_fix_without_scope(self) -> None:
|
| 284 |
+
match = COMMIT_PATTERN.match("fix: fix bug")
|
| 285 |
+
assert match is not None
|
| 286 |
+
assert match.group(1) == "fix"
|
| 287 |
+
assert match.group(2) is None
|
| 288 |
+
assert match.group(4) == "fix bug"
|
| 289 |
+
|
| 290 |
+
def test_with_exclamation(self) -> None:
|
| 291 |
+
match = COMMIT_PATTERN.match("feat(core)!: api changed")
|
| 292 |
+
assert match is not None
|
| 293 |
+
assert match.group(3) == "!"
|
| 294 |
+
|
| 295 |
+
def test_without_exclamation(self) -> None:
|
| 296 |
+
match = COMMIT_PATTERN.match("feat(core): add feature")
|
| 297 |
+
assert match is not None
|
| 298 |
+
assert match.group(3) is None
|
| 299 |
+
|
| 300 |
+
|
| 301 |
+
if __name__ == "__main__":
|
| 302 |
+
pytest.main([__file__, "-v"])
|
|
@@ -1,68 +1,68 @@
|
|
| 1 |
-
"""Tests for sync-plugin-versions.py."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import importlib.util
|
| 6 |
-
from pathlib import Path
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
def _load_module():
|
| 10 |
-
script = Path(__file__).parent.parent / "sync-plugin-versions.py"
|
| 11 |
-
spec = importlib.util.spec_from_file_location("sync_plugin_versions", script)
|
| 12 |
-
assert spec is not None
|
| 13 |
-
assert spec.loader is not None
|
| 14 |
-
module = importlib.util.module_from_spec(spec)
|
| 15 |
-
spec.loader.exec_module(module)
|
| 16 |
-
return module
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
def test_compute_repo_semver_uses_release_helpers(monkeypatch) -> None:
|
| 20 |
-
module = _load_module()
|
| 21 |
-
calls: dict[str, object] = {}
|
| 22 |
-
|
| 23 |
-
monkeypatch.setattr(module, "list_release_tags", lambda root: ["v0.9.0"])
|
| 24 |
-
monkeypatch.setattr(module, "find_latest_release_tag", lambda tags: "v0.9.0")
|
| 25 |
-
monkeypatch.setattr(module, "list_release_commits", lambda root, tag: ["feat: add init"])
|
| 26 |
-
monkeypatch.setattr(module, "determine_bump_level", lambda commits: "minor")
|
| 27 |
-
monkeypatch.setattr(module, "get_canonical_version", lambda root: "0.5.25")
|
| 28 |
-
|
| 29 |
-
def fake_compute_release_version(*, canonical_version: str, level: str, tags: list[str]):
|
| 30 |
-
calls["canonical_version"] = canonical_version
|
| 31 |
-
calls["level"] = level
|
| 32 |
-
calls["tags"] = tags
|
| 33 |
-
return type("Info", (), {"npm_version": "0.10.0"})()
|
| 34 |
-
|
| 35 |
-
monkeypatch.setattr(module, "compute_release_version", fake_compute_release_version)
|
| 36 |
-
|
| 37 |
-
assert module.compute_repo_semver(Path("repo")) == "0.10.0"
|
| 38 |
-
assert calls == {
|
| 39 |
-
"canonical_version": "0.5.25",
|
| 40 |
-
"level": "minor",
|
| 41 |
-
"tags": ["v0.9.0"],
|
| 42 |
-
}
|
| 43 |
-
|
| 44 |
-
|
| 45 |
-
def test_main_runs_plugin_only_version_sync(monkeypatch) -> None:
|
| 46 |
-
module = _load_module()
|
| 47 |
-
commands: list[list[str]] = []
|
| 48 |
-
|
| 49 |
-
monkeypatch.setattr(module, "compute_repo_semver", lambda root: "0.10.0")
|
| 50 |
-
monkeypatch.setattr(
|
| 51 |
-
module.subprocess,
|
| 52 |
-
"run",
|
| 53 |
-
lambda command, cwd, check: commands.append(command),
|
| 54 |
-
)
|
| 55 |
-
|
| 56 |
-
module.main()
|
| 57 |
-
|
| 58 |
-
assert commands == [
|
| 59 |
-
[
|
| 60 |
-
module.sys.executable,
|
| 61 |
-
str(module.ROOT / "scripts" / "version-sync.py"),
|
| 62 |
-
"--root",
|
| 63 |
-
str(module.ROOT),
|
| 64 |
-
"--version",
|
| 65 |
-
"0.10.0",
|
| 66 |
-
"--plugin-manifests-only",
|
| 67 |
-
]
|
| 68 |
-
]
|
|
|
|
| 1 |
+
"""Tests for sync-plugin-versions.py."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import importlib.util
|
| 6 |
+
from pathlib import Path
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
def _load_module():
|
| 10 |
+
script = Path(__file__).parent.parent / "sync-plugin-versions.py"
|
| 11 |
+
spec = importlib.util.spec_from_file_location("sync_plugin_versions", script)
|
| 12 |
+
assert spec is not None
|
| 13 |
+
assert spec.loader is not None
|
| 14 |
+
module = importlib.util.module_from_spec(spec)
|
| 15 |
+
spec.loader.exec_module(module)
|
| 16 |
+
return module
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
def test_compute_repo_semver_uses_release_helpers(monkeypatch) -> None:
|
| 20 |
+
module = _load_module()
|
| 21 |
+
calls: dict[str, object] = {}
|
| 22 |
+
|
| 23 |
+
monkeypatch.setattr(module, "list_release_tags", lambda root: ["v0.9.0"])
|
| 24 |
+
monkeypatch.setattr(module, "find_latest_release_tag", lambda tags: "v0.9.0")
|
| 25 |
+
monkeypatch.setattr(module, "list_release_commits", lambda root, tag: ["feat: add init"])
|
| 26 |
+
monkeypatch.setattr(module, "determine_bump_level", lambda commits: "minor")
|
| 27 |
+
monkeypatch.setattr(module, "get_canonical_version", lambda root: "0.5.25")
|
| 28 |
+
|
| 29 |
+
def fake_compute_release_version(*, canonical_version: str, level: str, tags: list[str]):
|
| 30 |
+
calls["canonical_version"] = canonical_version
|
| 31 |
+
calls["level"] = level
|
| 32 |
+
calls["tags"] = tags
|
| 33 |
+
return type("Info", (), {"npm_version": "0.10.0"})()
|
| 34 |
+
|
| 35 |
+
monkeypatch.setattr(module, "compute_release_version", fake_compute_release_version)
|
| 36 |
+
|
| 37 |
+
assert module.compute_repo_semver(Path("repo")) == "0.10.0"
|
| 38 |
+
assert calls == {
|
| 39 |
+
"canonical_version": "0.5.25",
|
| 40 |
+
"level": "minor",
|
| 41 |
+
"tags": ["v0.9.0"],
|
| 42 |
+
}
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
def test_main_runs_plugin_only_version_sync(monkeypatch) -> None:
|
| 46 |
+
module = _load_module()
|
| 47 |
+
commands: list[list[str]] = []
|
| 48 |
+
|
| 49 |
+
monkeypatch.setattr(module, "compute_repo_semver", lambda root: "0.10.0")
|
| 50 |
+
monkeypatch.setattr(
|
| 51 |
+
module.subprocess,
|
| 52 |
+
"run",
|
| 53 |
+
lambda command, cwd, check: commands.append(command),
|
| 54 |
+
)
|
| 55 |
+
|
| 56 |
+
module.main()
|
| 57 |
+
|
| 58 |
+
assert commands == [
|
| 59 |
+
[
|
| 60 |
+
module.sys.executable,
|
| 61 |
+
str(module.ROOT / "scripts" / "version-sync.py"),
|
| 62 |
+
"--root",
|
| 63 |
+
str(module.ROOT),
|
| 64 |
+
"--version",
|
| 65 |
+
"0.10.0",
|
| 66 |
+
"--plugin-manifests-only",
|
| 67 |
+
]
|
| 68 |
+
]
|
|
@@ -1,384 +1,384 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
from types import SimpleNamespace
|
| 4 |
-
|
| 5 |
-
import pytest
|
| 6 |
-
|
| 7 |
-
from headroom.backends import anyllm
|
| 8 |
-
from headroom.backends.base import BackendResponse, StreamEvent
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
class FakeAsyncStream:
|
| 12 |
-
def __init__(self, items) -> None: # noqa: ANN001
|
| 13 |
-
self._items = list(items)
|
| 14 |
-
|
| 15 |
-
def __aiter__(self):
|
| 16 |
-
self._iter = iter(self._items)
|
| 17 |
-
return self
|
| 18 |
-
|
| 19 |
-
async def __anext__(self):
|
| 20 |
-
try:
|
| 21 |
-
return next(self._iter)
|
| 22 |
-
except StopIteration as exc:
|
| 23 |
-
raise StopAsyncIteration from exc
|
| 24 |
-
|
| 25 |
-
|
| 26 |
-
class FakeAnyLLMInstance:
|
| 27 |
-
def __init__(self) -> None:
|
| 28 |
-
self.calls: list[dict[str, object]] = []
|
| 29 |
-
self.response = None
|
| 30 |
-
self.raise_error: Exception | None = None
|
| 31 |
-
|
| 32 |
-
async def acompletion(self, **kwargs): # noqa: ANN003
|
| 33 |
-
self.calls.append(kwargs)
|
| 34 |
-
if self.raise_error is not None:
|
| 35 |
-
raise self.raise_error
|
| 36 |
-
return self.response
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
def make_backend(
|
| 40 |
-
monkeypatch: pytest.MonkeyPatch, provider: str = "groq"
|
| 41 |
-
) -> tuple[anyllm.AnyLLMBackend, FakeAnyLLMInstance]:
|
| 42 |
-
fake_instance = FakeAnyLLMInstance()
|
| 43 |
-
|
| 44 |
-
class FakeAnyLLM:
|
| 45 |
-
@staticmethod
|
| 46 |
-
def create(requested_provider: str):
|
| 47 |
-
assert requested_provider == provider
|
| 48 |
-
return fake_instance
|
| 49 |
-
|
| 50 |
-
monkeypatch.setattr(anyllm, "ANYLLM_AVAILABLE", True)
|
| 51 |
-
monkeypatch.setattr(anyllm, "AnyLLM", FakeAnyLLM)
|
| 52 |
-
return anyllm.AnyLLMBackend(provider=provider.upper()), fake_instance
|
| 53 |
-
|
| 54 |
-
|
| 55 |
-
def make_choice(
|
| 56 |
-
content: str = "hello", finish_reason: str = "stop", tool_calls=None, index: int = 0
|
| 57 |
-
):
|
| 58 |
-
return SimpleNamespace(
|
| 59 |
-
index=index,
|
| 60 |
-
finish_reason=finish_reason,
|
| 61 |
-
message=SimpleNamespace(role="assistant", content=content, tool_calls=tool_calls),
|
| 62 |
-
)
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
def make_response(*choices, usage=None):
|
| 66 |
-
return SimpleNamespace(
|
| 67 |
-
id="resp_123",
|
| 68 |
-
created=123456,
|
| 69 |
-
choices=list(choices),
|
| 70 |
-
usage=usage,
|
| 71 |
-
)
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
def make_tool_call(tool_id: str, name: str, arguments):
|
| 75 |
-
return SimpleNamespace(id=tool_id, function=SimpleNamespace(name=name, arguments=arguments))
|
| 76 |
-
|
| 77 |
-
|
| 78 |
-
def test_init_raises_without_anyllm() -> None:
|
| 79 |
-
original_available = anyllm.ANYLLM_AVAILABLE
|
| 80 |
-
try:
|
| 81 |
-
anyllm.ANYLLM_AVAILABLE = False
|
| 82 |
-
with pytest.raises(ImportError):
|
| 83 |
-
anyllm.AnyLLMBackend()
|
| 84 |
-
finally:
|
| 85 |
-
anyllm.ANYLLM_AVAILABLE = original_available
|
| 86 |
-
|
| 87 |
-
|
| 88 |
-
def test_init_name_and_basic_methods(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 89 |
-
backend, instance = make_backend(monkeypatch, provider="groq")
|
| 90 |
-
|
| 91 |
-
assert backend.provider == "groq"
|
| 92 |
-
assert backend.name == "anyllm-groq"
|
| 93 |
-
assert backend.map_model_id("claude-3-5") == "claude-3-5"
|
| 94 |
-
assert backend.supports_model("anything") is True
|
| 95 |
-
assert backend.llm is instance
|
| 96 |
-
|
| 97 |
-
|
| 98 |
-
def test_convert_content_blocks_and_messages(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 99 |
-
backend, _instance = make_backend(monkeypatch)
|
| 100 |
-
|
| 101 |
-
assert backend._convert_content_blocks([{"type": "text", "text": "hello"}]) == "hello"
|
| 102 |
-
assert backend._convert_content_blocks(
|
| 103 |
-
[
|
| 104 |
-
{"type": "text", "text": "caption"},
|
| 105 |
-
{
|
| 106 |
-
"type": "image",
|
| 107 |
-
"source": {"type": "base64", "media_type": "image/jpeg", "data": "abc"},
|
| 108 |
-
},
|
| 109 |
-
{"type": "image", "source": {"type": "url", "url": "https://example.com/img.png"}},
|
| 110 |
-
]
|
| 111 |
-
) == [
|
| 112 |
-
{"type": "text", "text": "caption"},
|
| 113 |
-
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,abc"}},
|
| 114 |
-
{"type": "image_url", "image_url": {"url": "https://example.com/img.png"}},
|
| 115 |
-
]
|
| 116 |
-
assert backend._convert_content_blocks([{"type": "tool_use", "id": "ignored"}]) == ""
|
| 117 |
-
|
| 118 |
-
converted = backend._convert_messages(
|
| 119 |
-
[
|
| 120 |
-
{"role": "user", "content": "plain text"},
|
| 121 |
-
{
|
| 122 |
-
"role": "assistant",
|
| 123 |
-
"content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}],
|
| 124 |
-
},
|
| 125 |
-
{
|
| 126 |
-
"role": "user",
|
| 127 |
-
"content": [
|
| 128 |
-
{"type": "text", "text": "look"},
|
| 129 |
-
{"type": "image", "source": {"type": "url", "url": "https://example.com"}},
|
| 130 |
-
],
|
| 131 |
-
},
|
| 132 |
-
{"role": "user", "content": 123},
|
| 133 |
-
]
|
| 134 |
-
)
|
| 135 |
-
|
| 136 |
-
assert converted == [
|
| 137 |
-
{"role": "user", "content": "plain text"},
|
| 138 |
-
{"role": "assistant", "content": "a\nb"},
|
| 139 |
-
{
|
| 140 |
-
"role": "user",
|
| 141 |
-
"content": [
|
| 142 |
-
{"type": "text", "text": "look"},
|
| 143 |
-
{"type": "image_url", "image_url": {"url": "https://example.com"}},
|
| 144 |
-
],
|
| 145 |
-
},
|
| 146 |
-
]
|
| 147 |
-
|
| 148 |
-
|
| 149 |
-
def test_to_anthropic_response_maps_tool_calls_and_usage(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 150 |
-
backend, _instance = make_backend(monkeypatch)
|
| 151 |
-
response = make_response(
|
| 152 |
-
make_choice(
|
| 153 |
-
content="hello",
|
| 154 |
-
finish_reason="tool_calls",
|
| 155 |
-
tool_calls=[
|
| 156 |
-
make_tool_call("tc1", "memory_save", '{"content":"python"}'),
|
| 157 |
-
make_tool_call("tc2", "memory_search", {"query": "python"}),
|
| 158 |
-
],
|
| 159 |
-
),
|
| 160 |
-
usage=SimpleNamespace(prompt_tokens=12, completion_tokens=7),
|
| 161 |
-
)
|
| 162 |
-
|
| 163 |
-
converted = backend._to_anthropic_response(response, "claude-sonnet")
|
| 164 |
-
|
| 165 |
-
assert converted["type"] == "message"
|
| 166 |
-
assert converted["role"] == "assistant"
|
| 167 |
-
assert converted["model"] == "claude-sonnet"
|
| 168 |
-
assert converted["stop_reason"] == "tool_use"
|
| 169 |
-
assert converted["usage"] == {"input_tokens": 12, "output_tokens": 7}
|
| 170 |
-
assert converted["content"][0] == {"type": "text", "text": "hello"}
|
| 171 |
-
assert converted["content"][1]["input"] == {"content": "python"}
|
| 172 |
-
assert converted["content"][2]["input"] == {"query": "python"}
|
| 173 |
-
|
| 174 |
-
|
| 175 |
-
@pytest.mark.asyncio
|
| 176 |
-
async def test_send_message_builds_anthropic_response(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 177 |
-
backend, instance = make_backend(monkeypatch)
|
| 178 |
-
instance.response = make_response(
|
| 179 |
-
make_choice("done", "stop"),
|
| 180 |
-
usage=SimpleNamespace(prompt_tokens=4, completion_tokens=6),
|
| 181 |
-
)
|
| 182 |
-
|
| 183 |
-
result = await backend.send_message(
|
| 184 |
-
{
|
| 185 |
-
"model": "claude-3-7-sonnet",
|
| 186 |
-
"messages": [{"role": "user", "content": [{"type": "text", "text": "hello"}]}],
|
| 187 |
-
"system": [{"text": "system rule"}, "extra"],
|
| 188 |
-
"max_tokens": 200,
|
| 189 |
-
"temperature": 0.3,
|
| 190 |
-
"top_p": 0.8,
|
| 191 |
-
"stop_sequences": ["END"],
|
| 192 |
-
"tools": [{"name": "t"}],
|
| 193 |
-
"tool_choice": {"type": "auto"},
|
| 194 |
-
},
|
| 195 |
-
{},
|
| 196 |
-
)
|
| 197 |
-
|
| 198 |
-
assert isinstance(result, BackendResponse)
|
| 199 |
-
assert result.status_code == 200
|
| 200 |
-
assert result.headers == {"content-type": "application/json"}
|
| 201 |
-
assert result.body["content"][0]["text"] == "done"
|
| 202 |
-
assert instance.calls[0]["messages"][0] == {"role": "system", "content": "system rule extra"}
|
| 203 |
-
assert instance.calls[0]["stop"] == ["END"]
|
| 204 |
-
|
| 205 |
-
|
| 206 |
-
@pytest.mark.asyncio
|
| 207 |
-
async def test_send_message_returns_error_response(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 208 |
-
backend, instance = make_backend(monkeypatch)
|
| 209 |
-
instance.raise_error = RuntimeError("authentication api_key missing")
|
| 210 |
-
|
| 211 |
-
result = await backend.send_message({"messages": []}, {})
|
| 212 |
-
|
| 213 |
-
assert result.status_code == 401
|
| 214 |
-
assert result.body["error"]["type"] == "authentication_error"
|
| 215 |
-
assert result.error == "authentication api_key missing"
|
| 216 |
-
|
| 217 |
-
|
| 218 |
-
@pytest.mark.asyncio
|
| 219 |
-
async def test_stream_message_yields_events_and_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 220 |
-
backend, instance = make_backend(monkeypatch)
|
| 221 |
-
instance.response = FakeAsyncStream(
|
| 222 |
-
[
|
| 223 |
-
SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content="hel"))]),
|
| 224 |
-
SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content="lo"))]),
|
| 225 |
-
SimpleNamespace(choices=[]),
|
| 226 |
-
]
|
| 227 |
-
)
|
| 228 |
-
|
| 229 |
-
events = [
|
| 230 |
-
event
|
| 231 |
-
async for event in backend.stream_message(
|
| 232 |
-
{"model": "claude", "messages": [], "system": "sys"}, {}
|
| 233 |
-
)
|
| 234 |
-
]
|
| 235 |
-
|
| 236 |
-
assert [event.event_type for event in events] == [
|
| 237 |
-
"message_start",
|
| 238 |
-
"content_block_start",
|
| 239 |
-
"content_block_delta",
|
| 240 |
-
"content_block_delta",
|
| 241 |
-
"content_block_stop",
|
| 242 |
-
"message_delta",
|
| 243 |
-
"message_stop",
|
| 244 |
-
]
|
| 245 |
-
assert events[0].data["message"]["model"] == "claude"
|
| 246 |
-
assert events[5].data["usage"] == {"output_tokens": 2}
|
| 247 |
-
assert instance.calls[0]["stream"] is True
|
| 248 |
-
assert instance.calls[0]["messages"][0] == {"role": "system", "content": "sys"}
|
| 249 |
-
|
| 250 |
-
backend_error, instance_error = make_backend(monkeypatch, provider="openai")
|
| 251 |
-
instance_error.raise_error = RuntimeError("stream broke")
|
| 252 |
-
error_events = [event async for event in backend_error.stream_message({"messages": []}, {})]
|
| 253 |
-
assert error_events[-1].event_type == "error"
|
| 254 |
-
assert error_events[-1].data["error"]["message"] == "stream broke"
|
| 255 |
-
|
| 256 |
-
|
| 257 |
-
@pytest.mark.asyncio
|
| 258 |
-
async def test_send_openai_message_maps_choices_and_tool_calls(
|
| 259 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 260 |
-
) -> None:
|
| 261 |
-
backend, instance = make_backend(monkeypatch)
|
| 262 |
-
instance.response = make_response(
|
| 263 |
-
make_choice(
|
| 264 |
-
content="answer",
|
| 265 |
-
finish_reason="stop",
|
| 266 |
-
tool_calls=[
|
| 267 |
-
make_tool_call("tc1", "memory_search", '{"query":"python"}'),
|
| 268 |
-
SimpleNamespace(id="tc2", function=None),
|
| 269 |
-
],
|
| 270 |
-
index=0,
|
| 271 |
-
),
|
| 272 |
-
usage=SimpleNamespace(prompt_tokens=2, completion_tokens=3, total_tokens=5),
|
| 273 |
-
)
|
| 274 |
-
|
| 275 |
-
result = await backend.send_openai_message(
|
| 276 |
-
{
|
| 277 |
-
"model": "gpt-4o",
|
| 278 |
-
"messages": [{"role": "user", "content": "hi"}],
|
| 279 |
-
"max_tokens": 50,
|
| 280 |
-
"temperature": 0.2,
|
| 281 |
-
"top_p": 0.9,
|
| 282 |
-
"stop": ["END"],
|
| 283 |
-
"tools": [{"name": "memory"}],
|
| 284 |
-
"tool_choice": "auto",
|
| 285 |
-
"response_format": {"type": "json_object"},
|
| 286 |
-
"seed": 1,
|
| 287 |
-
"n": 2,
|
| 288 |
-
},
|
| 289 |
-
{},
|
| 290 |
-
)
|
| 291 |
-
|
| 292 |
-
assert result.status_code == 200
|
| 293 |
-
assert result.body["object"] == "chat.completion"
|
| 294 |
-
assert (
|
| 295 |
-
result.body["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "memory_search"
|
| 296 |
-
)
|
| 297 |
-
assert result.body["choices"][0]["message"]["tool_calls"][1] == {
|
| 298 |
-
"id": "tc2",
|
| 299 |
-
"type": "function",
|
| 300 |
-
}
|
| 301 |
-
assert result.body["usage"] == {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}
|
| 302 |
-
|
| 303 |
-
|
| 304 |
-
@pytest.mark.asyncio
|
| 305 |
-
async def test_send_openai_message_returns_error_response(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 306 |
-
backend, instance = make_backend(monkeypatch)
|
| 307 |
-
instance.raise_error = RuntimeError("model not found")
|
| 308 |
-
|
| 309 |
-
result = await backend.send_openai_message({"messages": []}, {})
|
| 310 |
-
|
| 311 |
-
assert result.status_code == 404
|
| 312 |
-
assert result.body["error"]["type"] == "model_not_found"
|
| 313 |
-
|
| 314 |
-
|
| 315 |
-
@pytest.mark.asyncio
|
| 316 |
-
async def test_stream_openai_message_yields_sse_chunks_and_done(
|
| 317 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 318 |
-
) -> None:
|
| 319 |
-
backend, instance = make_backend(monkeypatch)
|
| 320 |
-
instance.response = FakeAsyncStream(
|
| 321 |
-
[
|
| 322 |
-
SimpleNamespace(
|
| 323 |
-
model_dump=lambda **kwargs: {
|
| 324 |
-
"id": "chunk1",
|
| 325 |
-
"choices": [{"delta": {"content": "a"}}],
|
| 326 |
-
}
|
| 327 |
-
),
|
| 328 |
-
SimpleNamespace(
|
| 329 |
-
model_dump=lambda **kwargs: {
|
| 330 |
-
"id": "chunk2",
|
| 331 |
-
"choices": [{"delta": {"content": "b"}}],
|
| 332 |
-
}
|
| 333 |
-
),
|
| 334 |
-
]
|
| 335 |
-
)
|
| 336 |
-
|
| 337 |
-
chunks = [
|
| 338 |
-
chunk
|
| 339 |
-
async for chunk in backend.stream_openai_message(
|
| 340 |
-
{
|
| 341 |
-
"messages": [{"role": "user", "content": "hi"}],
|
| 342 |
-
"stream_options": {"include_usage": True},
|
| 343 |
-
},
|
| 344 |
-
{},
|
| 345 |
-
)
|
| 346 |
-
]
|
| 347 |
-
|
| 348 |
-
assert chunks[0].startswith("data: {")
|
| 349 |
-
assert chunks[-1] == "data: [DONE]\n\n"
|
| 350 |
-
assert instance.calls[0]["stream"] is True
|
| 351 |
-
assert instance.calls[0]["stream_options"] == {"include_usage": True}
|
| 352 |
-
|
| 353 |
-
backend_error, instance_error = make_backend(monkeypatch, provider="anthropic")
|
| 354 |
-
instance_error.raise_error = RuntimeError("rate limit hit")
|
| 355 |
-
error_chunks = [
|
| 356 |
-
chunk async for chunk in backend_error.stream_openai_message({"messages": []}, {})
|
| 357 |
-
]
|
| 358 |
-
assert '"backend_error"' in error_chunks[0]
|
| 359 |
-
assert error_chunks[-1] == "data: [DONE]\n\n"
|
| 360 |
-
|
| 361 |
-
|
| 362 |
-
def test_error_response_classifies_common_failures(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 363 |
-
backend, _instance = make_backend(monkeypatch)
|
| 364 |
-
|
| 365 |
-
auth = backend._error_response(RuntimeError("authentication api key missing"))
|
| 366 |
-
rate = backend._error_response(RuntimeError("rate limit exceeded"), openai_format=True)
|
| 367 |
-
model = backend._error_response(RuntimeError("model not found"), openai_format=True)
|
| 368 |
-
generic = backend._error_response(RuntimeError("other error"))
|
| 369 |
-
|
| 370 |
-
assert auth.status_code == 401
|
| 371 |
-
assert auth.body["error"]["type"] == "authentication_error"
|
| 372 |
-
assert rate.status_code == 429
|
| 373 |
-
assert rate.body["error"]["type"] == "rate_limit_exceeded"
|
| 374 |
-
assert model.status_code == 404
|
| 375 |
-
assert model.body["error"]["type"] == "model_not_found"
|
| 376 |
-
assert generic.status_code == 500
|
| 377 |
-
assert generic.body["error"]["type"] == "api_error"
|
| 378 |
-
|
| 379 |
-
|
| 380 |
-
@pytest.mark.asyncio
|
| 381 |
-
async def test_close_is_noop(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 382 |
-
backend, _instance = make_backend(monkeypatch)
|
| 383 |
-
assert await backend.close() is None
|
| 384 |
-
assert isinstance(StreamEvent(event_type="message_start", data={}), StreamEvent)
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
from types import SimpleNamespace
|
| 4 |
+
|
| 5 |
+
import pytest
|
| 6 |
+
|
| 7 |
+
from headroom.backends import anyllm
|
| 8 |
+
from headroom.backends.base import BackendResponse, StreamEvent
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
class FakeAsyncStream:
|
| 12 |
+
def __init__(self, items) -> None: # noqa: ANN001
|
| 13 |
+
self._items = list(items)
|
| 14 |
+
|
| 15 |
+
def __aiter__(self):
|
| 16 |
+
self._iter = iter(self._items)
|
| 17 |
+
return self
|
| 18 |
+
|
| 19 |
+
async def __anext__(self):
|
| 20 |
+
try:
|
| 21 |
+
return next(self._iter)
|
| 22 |
+
except StopIteration as exc:
|
| 23 |
+
raise StopAsyncIteration from exc
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
class FakeAnyLLMInstance:
|
| 27 |
+
def __init__(self) -> None:
|
| 28 |
+
self.calls: list[dict[str, object]] = []
|
| 29 |
+
self.response = None
|
| 30 |
+
self.raise_error: Exception | None = None
|
| 31 |
+
|
| 32 |
+
async def acompletion(self, **kwargs): # noqa: ANN003
|
| 33 |
+
self.calls.append(kwargs)
|
| 34 |
+
if self.raise_error is not None:
|
| 35 |
+
raise self.raise_error
|
| 36 |
+
return self.response
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
def make_backend(
|
| 40 |
+
monkeypatch: pytest.MonkeyPatch, provider: str = "groq"
|
| 41 |
+
) -> tuple[anyllm.AnyLLMBackend, FakeAnyLLMInstance]:
|
| 42 |
+
fake_instance = FakeAnyLLMInstance()
|
| 43 |
+
|
| 44 |
+
class FakeAnyLLM:
|
| 45 |
+
@staticmethod
|
| 46 |
+
def create(requested_provider: str):
|
| 47 |
+
assert requested_provider == provider
|
| 48 |
+
return fake_instance
|
| 49 |
+
|
| 50 |
+
monkeypatch.setattr(anyllm, "ANYLLM_AVAILABLE", True)
|
| 51 |
+
monkeypatch.setattr(anyllm, "AnyLLM", FakeAnyLLM)
|
| 52 |
+
return anyllm.AnyLLMBackend(provider=provider.upper()), fake_instance
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
def make_choice(
|
| 56 |
+
content: str = "hello", finish_reason: str = "stop", tool_calls=None, index: int = 0
|
| 57 |
+
):
|
| 58 |
+
return SimpleNamespace(
|
| 59 |
+
index=index,
|
| 60 |
+
finish_reason=finish_reason,
|
| 61 |
+
message=SimpleNamespace(role="assistant", content=content, tool_calls=tool_calls),
|
| 62 |
+
)
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def make_response(*choices, usage=None):
|
| 66 |
+
return SimpleNamespace(
|
| 67 |
+
id="resp_123",
|
| 68 |
+
created=123456,
|
| 69 |
+
choices=list(choices),
|
| 70 |
+
usage=usage,
|
| 71 |
+
)
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def make_tool_call(tool_id: str, name: str, arguments):
|
| 75 |
+
return SimpleNamespace(id=tool_id, function=SimpleNamespace(name=name, arguments=arguments))
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def test_init_raises_without_anyllm() -> None:
|
| 79 |
+
original_available = anyllm.ANYLLM_AVAILABLE
|
| 80 |
+
try:
|
| 81 |
+
anyllm.ANYLLM_AVAILABLE = False
|
| 82 |
+
with pytest.raises(ImportError):
|
| 83 |
+
anyllm.AnyLLMBackend()
|
| 84 |
+
finally:
|
| 85 |
+
anyllm.ANYLLM_AVAILABLE = original_available
|
| 86 |
+
|
| 87 |
+
|
| 88 |
+
def test_init_name_and_basic_methods(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 89 |
+
backend, instance = make_backend(monkeypatch, provider="groq")
|
| 90 |
+
|
| 91 |
+
assert backend.provider == "groq"
|
| 92 |
+
assert backend.name == "anyllm-groq"
|
| 93 |
+
assert backend.map_model_id("claude-3-5") == "claude-3-5"
|
| 94 |
+
assert backend.supports_model("anything") is True
|
| 95 |
+
assert backend.llm is instance
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
def test_convert_content_blocks_and_messages(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 99 |
+
backend, _instance = make_backend(monkeypatch)
|
| 100 |
+
|
| 101 |
+
assert backend._convert_content_blocks([{"type": "text", "text": "hello"}]) == "hello"
|
| 102 |
+
assert backend._convert_content_blocks(
|
| 103 |
+
[
|
| 104 |
+
{"type": "text", "text": "caption"},
|
| 105 |
+
{
|
| 106 |
+
"type": "image",
|
| 107 |
+
"source": {"type": "base64", "media_type": "image/jpeg", "data": "abc"},
|
| 108 |
+
},
|
| 109 |
+
{"type": "image", "source": {"type": "url", "url": "https://example.com/img.png"}},
|
| 110 |
+
]
|
| 111 |
+
) == [
|
| 112 |
+
{"type": "text", "text": "caption"},
|
| 113 |
+
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,abc"}},
|
| 114 |
+
{"type": "image_url", "image_url": {"url": "https://example.com/img.png"}},
|
| 115 |
+
]
|
| 116 |
+
assert backend._convert_content_blocks([{"type": "tool_use", "id": "ignored"}]) == ""
|
| 117 |
+
|
| 118 |
+
converted = backend._convert_messages(
|
| 119 |
+
[
|
| 120 |
+
{"role": "user", "content": "plain text"},
|
| 121 |
+
{
|
| 122 |
+
"role": "assistant",
|
| 123 |
+
"content": [{"type": "text", "text": "a"}, {"type": "text", "text": "b"}],
|
| 124 |
+
},
|
| 125 |
+
{
|
| 126 |
+
"role": "user",
|
| 127 |
+
"content": [
|
| 128 |
+
{"type": "text", "text": "look"},
|
| 129 |
+
{"type": "image", "source": {"type": "url", "url": "https://example.com"}},
|
| 130 |
+
],
|
| 131 |
+
},
|
| 132 |
+
{"role": "user", "content": 123},
|
| 133 |
+
]
|
| 134 |
+
)
|
| 135 |
+
|
| 136 |
+
assert converted == [
|
| 137 |
+
{"role": "user", "content": "plain text"},
|
| 138 |
+
{"role": "assistant", "content": "a\nb"},
|
| 139 |
+
{
|
| 140 |
+
"role": "user",
|
| 141 |
+
"content": [
|
| 142 |
+
{"type": "text", "text": "look"},
|
| 143 |
+
{"type": "image_url", "image_url": {"url": "https://example.com"}},
|
| 144 |
+
],
|
| 145 |
+
},
|
| 146 |
+
]
|
| 147 |
+
|
| 148 |
+
|
| 149 |
+
def test_to_anthropic_response_maps_tool_calls_and_usage(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 150 |
+
backend, _instance = make_backend(monkeypatch)
|
| 151 |
+
response = make_response(
|
| 152 |
+
make_choice(
|
| 153 |
+
content="hello",
|
| 154 |
+
finish_reason="tool_calls",
|
| 155 |
+
tool_calls=[
|
| 156 |
+
make_tool_call("tc1", "memory_save", '{"content":"python"}'),
|
| 157 |
+
make_tool_call("tc2", "memory_search", {"query": "python"}),
|
| 158 |
+
],
|
| 159 |
+
),
|
| 160 |
+
usage=SimpleNamespace(prompt_tokens=12, completion_tokens=7),
|
| 161 |
+
)
|
| 162 |
+
|
| 163 |
+
converted = backend._to_anthropic_response(response, "claude-sonnet")
|
| 164 |
+
|
| 165 |
+
assert converted["type"] == "message"
|
| 166 |
+
assert converted["role"] == "assistant"
|
| 167 |
+
assert converted["model"] == "claude-sonnet"
|
| 168 |
+
assert converted["stop_reason"] == "tool_use"
|
| 169 |
+
assert converted["usage"] == {"input_tokens": 12, "output_tokens": 7}
|
| 170 |
+
assert converted["content"][0] == {"type": "text", "text": "hello"}
|
| 171 |
+
assert converted["content"][1]["input"] == {"content": "python"}
|
| 172 |
+
assert converted["content"][2]["input"] == {"query": "python"}
|
| 173 |
+
|
| 174 |
+
|
| 175 |
+
@pytest.mark.asyncio
|
| 176 |
+
async def test_send_message_builds_anthropic_response(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 177 |
+
backend, instance = make_backend(monkeypatch)
|
| 178 |
+
instance.response = make_response(
|
| 179 |
+
make_choice("done", "stop"),
|
| 180 |
+
usage=SimpleNamespace(prompt_tokens=4, completion_tokens=6),
|
| 181 |
+
)
|
| 182 |
+
|
| 183 |
+
result = await backend.send_message(
|
| 184 |
+
{
|
| 185 |
+
"model": "claude-3-7-sonnet",
|
| 186 |
+
"messages": [{"role": "user", "content": [{"type": "text", "text": "hello"}]}],
|
| 187 |
+
"system": [{"text": "system rule"}, "extra"],
|
| 188 |
+
"max_tokens": 200,
|
| 189 |
+
"temperature": 0.3,
|
| 190 |
+
"top_p": 0.8,
|
| 191 |
+
"stop_sequences": ["END"],
|
| 192 |
+
"tools": [{"name": "t"}],
|
| 193 |
+
"tool_choice": {"type": "auto"},
|
| 194 |
+
},
|
| 195 |
+
{},
|
| 196 |
+
)
|
| 197 |
+
|
| 198 |
+
assert isinstance(result, BackendResponse)
|
| 199 |
+
assert result.status_code == 200
|
| 200 |
+
assert result.headers == {"content-type": "application/json"}
|
| 201 |
+
assert result.body["content"][0]["text"] == "done"
|
| 202 |
+
assert instance.calls[0]["messages"][0] == {"role": "system", "content": "system rule extra"}
|
| 203 |
+
assert instance.calls[0]["stop"] == ["END"]
|
| 204 |
+
|
| 205 |
+
|
| 206 |
+
@pytest.mark.asyncio
|
| 207 |
+
async def test_send_message_returns_error_response(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 208 |
+
backend, instance = make_backend(monkeypatch)
|
| 209 |
+
instance.raise_error = RuntimeError("authentication api_key missing")
|
| 210 |
+
|
| 211 |
+
result = await backend.send_message({"messages": []}, {})
|
| 212 |
+
|
| 213 |
+
assert result.status_code == 401
|
| 214 |
+
assert result.body["error"]["type"] == "authentication_error"
|
| 215 |
+
assert result.error == "authentication api_key missing"
|
| 216 |
+
|
| 217 |
+
|
| 218 |
+
@pytest.mark.asyncio
|
| 219 |
+
async def test_stream_message_yields_events_and_error(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 220 |
+
backend, instance = make_backend(monkeypatch)
|
| 221 |
+
instance.response = FakeAsyncStream(
|
| 222 |
+
[
|
| 223 |
+
SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content="hel"))]),
|
| 224 |
+
SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content="lo"))]),
|
| 225 |
+
SimpleNamespace(choices=[]),
|
| 226 |
+
]
|
| 227 |
+
)
|
| 228 |
+
|
| 229 |
+
events = [
|
| 230 |
+
event
|
| 231 |
+
async for event in backend.stream_message(
|
| 232 |
+
{"model": "claude", "messages": [], "system": "sys"}, {}
|
| 233 |
+
)
|
| 234 |
+
]
|
| 235 |
+
|
| 236 |
+
assert [event.event_type for event in events] == [
|
| 237 |
+
"message_start",
|
| 238 |
+
"content_block_start",
|
| 239 |
+
"content_block_delta",
|
| 240 |
+
"content_block_delta",
|
| 241 |
+
"content_block_stop",
|
| 242 |
+
"message_delta",
|
| 243 |
+
"message_stop",
|
| 244 |
+
]
|
| 245 |
+
assert events[0].data["message"]["model"] == "claude"
|
| 246 |
+
assert events[5].data["usage"] == {"output_tokens": 2}
|
| 247 |
+
assert instance.calls[0]["stream"] is True
|
| 248 |
+
assert instance.calls[0]["messages"][0] == {"role": "system", "content": "sys"}
|
| 249 |
+
|
| 250 |
+
backend_error, instance_error = make_backend(monkeypatch, provider="openai")
|
| 251 |
+
instance_error.raise_error = RuntimeError("stream broke")
|
| 252 |
+
error_events = [event async for event in backend_error.stream_message({"messages": []}, {})]
|
| 253 |
+
assert error_events[-1].event_type == "error"
|
| 254 |
+
assert error_events[-1].data["error"]["message"] == "stream broke"
|
| 255 |
+
|
| 256 |
+
|
| 257 |
+
@pytest.mark.asyncio
|
| 258 |
+
async def test_send_openai_message_maps_choices_and_tool_calls(
|
| 259 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 260 |
+
) -> None:
|
| 261 |
+
backend, instance = make_backend(monkeypatch)
|
| 262 |
+
instance.response = make_response(
|
| 263 |
+
make_choice(
|
| 264 |
+
content="answer",
|
| 265 |
+
finish_reason="stop",
|
| 266 |
+
tool_calls=[
|
| 267 |
+
make_tool_call("tc1", "memory_search", '{"query":"python"}'),
|
| 268 |
+
SimpleNamespace(id="tc2", function=None),
|
| 269 |
+
],
|
| 270 |
+
index=0,
|
| 271 |
+
),
|
| 272 |
+
usage=SimpleNamespace(prompt_tokens=2, completion_tokens=3, total_tokens=5),
|
| 273 |
+
)
|
| 274 |
+
|
| 275 |
+
result = await backend.send_openai_message(
|
| 276 |
+
{
|
| 277 |
+
"model": "gpt-4o",
|
| 278 |
+
"messages": [{"role": "user", "content": "hi"}],
|
| 279 |
+
"max_tokens": 50,
|
| 280 |
+
"temperature": 0.2,
|
| 281 |
+
"top_p": 0.9,
|
| 282 |
+
"stop": ["END"],
|
| 283 |
+
"tools": [{"name": "memory"}],
|
| 284 |
+
"tool_choice": "auto",
|
| 285 |
+
"response_format": {"type": "json_object"},
|
| 286 |
+
"seed": 1,
|
| 287 |
+
"n": 2,
|
| 288 |
+
},
|
| 289 |
+
{},
|
| 290 |
+
)
|
| 291 |
+
|
| 292 |
+
assert result.status_code == 200
|
| 293 |
+
assert result.body["object"] == "chat.completion"
|
| 294 |
+
assert (
|
| 295 |
+
result.body["choices"][0]["message"]["tool_calls"][0]["function"]["name"] == "memory_search"
|
| 296 |
+
)
|
| 297 |
+
assert result.body["choices"][0]["message"]["tool_calls"][1] == {
|
| 298 |
+
"id": "tc2",
|
| 299 |
+
"type": "function",
|
| 300 |
+
}
|
| 301 |
+
assert result.body["usage"] == {"prompt_tokens": 2, "completion_tokens": 3, "total_tokens": 5}
|
| 302 |
+
|
| 303 |
+
|
| 304 |
+
@pytest.mark.asyncio
|
| 305 |
+
async def test_send_openai_message_returns_error_response(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 306 |
+
backend, instance = make_backend(monkeypatch)
|
| 307 |
+
instance.raise_error = RuntimeError("model not found")
|
| 308 |
+
|
| 309 |
+
result = await backend.send_openai_message({"messages": []}, {})
|
| 310 |
+
|
| 311 |
+
assert result.status_code == 404
|
| 312 |
+
assert result.body["error"]["type"] == "model_not_found"
|
| 313 |
+
|
| 314 |
+
|
| 315 |
+
@pytest.mark.asyncio
|
| 316 |
+
async def test_stream_openai_message_yields_sse_chunks_and_done(
|
| 317 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 318 |
+
) -> None:
|
| 319 |
+
backend, instance = make_backend(monkeypatch)
|
| 320 |
+
instance.response = FakeAsyncStream(
|
| 321 |
+
[
|
| 322 |
+
SimpleNamespace(
|
| 323 |
+
model_dump=lambda **kwargs: {
|
| 324 |
+
"id": "chunk1",
|
| 325 |
+
"choices": [{"delta": {"content": "a"}}],
|
| 326 |
+
}
|
| 327 |
+
),
|
| 328 |
+
SimpleNamespace(
|
| 329 |
+
model_dump=lambda **kwargs: {
|
| 330 |
+
"id": "chunk2",
|
| 331 |
+
"choices": [{"delta": {"content": "b"}}],
|
| 332 |
+
}
|
| 333 |
+
),
|
| 334 |
+
]
|
| 335 |
+
)
|
| 336 |
+
|
| 337 |
+
chunks = [
|
| 338 |
+
chunk
|
| 339 |
+
async for chunk in backend.stream_openai_message(
|
| 340 |
+
{
|
| 341 |
+
"messages": [{"role": "user", "content": "hi"}],
|
| 342 |
+
"stream_options": {"include_usage": True},
|
| 343 |
+
},
|
| 344 |
+
{},
|
| 345 |
+
)
|
| 346 |
+
]
|
| 347 |
+
|
| 348 |
+
assert chunks[0].startswith("data: {")
|
| 349 |
+
assert chunks[-1] == "data: [DONE]\n\n"
|
| 350 |
+
assert instance.calls[0]["stream"] is True
|
| 351 |
+
assert instance.calls[0]["stream_options"] == {"include_usage": True}
|
| 352 |
+
|
| 353 |
+
backend_error, instance_error = make_backend(monkeypatch, provider="anthropic")
|
| 354 |
+
instance_error.raise_error = RuntimeError("rate limit hit")
|
| 355 |
+
error_chunks = [
|
| 356 |
+
chunk async for chunk in backend_error.stream_openai_message({"messages": []}, {})
|
| 357 |
+
]
|
| 358 |
+
assert '"backend_error"' in error_chunks[0]
|
| 359 |
+
assert error_chunks[-1] == "data: [DONE]\n\n"
|
| 360 |
+
|
| 361 |
+
|
| 362 |
+
def test_error_response_classifies_common_failures(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 363 |
+
backend, _instance = make_backend(monkeypatch)
|
| 364 |
+
|
| 365 |
+
auth = backend._error_response(RuntimeError("authentication api key missing"))
|
| 366 |
+
rate = backend._error_response(RuntimeError("rate limit exceeded"), openai_format=True)
|
| 367 |
+
model = backend._error_response(RuntimeError("model not found"), openai_format=True)
|
| 368 |
+
generic = backend._error_response(RuntimeError("other error"))
|
| 369 |
+
|
| 370 |
+
assert auth.status_code == 401
|
| 371 |
+
assert auth.body["error"]["type"] == "authentication_error"
|
| 372 |
+
assert rate.status_code == 429
|
| 373 |
+
assert rate.body["error"]["type"] == "rate_limit_exceeded"
|
| 374 |
+
assert model.status_code == 404
|
| 375 |
+
assert model.body["error"]["type"] == "model_not_found"
|
| 376 |
+
assert generic.status_code == 500
|
| 377 |
+
assert generic.body["error"]["type"] == "api_error"
|
| 378 |
+
|
| 379 |
+
|
| 380 |
+
@pytest.mark.asyncio
|
| 381 |
+
async def test_close_is_noop(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 382 |
+
backend, _instance = make_backend(monkeypatch)
|
| 383 |
+
assert await backend.close() is None
|
| 384 |
+
assert isinstance(StreamEvent(event_type="message_start", data={}), StreamEvent)
|
|
@@ -1,126 +1,126 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
import sys
|
| 4 |
-
from dataclasses import dataclass
|
| 5 |
-
|
| 6 |
-
from headroom.ccr.batch_store import (
|
| 7 |
-
BatchContext,
|
| 8 |
-
BatchContextStore,
|
| 9 |
-
BatchRequestContext,
|
| 10 |
-
get_batch_context_store,
|
| 11 |
-
reset_batch_context_store,
|
| 12 |
-
)
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
def test_batch_context_defaults_and_expiry(monkeypatch) -> None:
|
| 16 |
-
monkeypatch.setattr("headroom.ccr.batch_store.time.time", lambda: 100.0)
|
| 17 |
-
context = BatchContext(batch_id="batch-1", provider="anthropic", created_at=100.0)
|
| 18 |
-
assert context.expires_at == 100.0 + 86400
|
| 19 |
-
assert context.is_expired is False
|
| 20 |
-
|
| 21 |
-
request = BatchRequestContext(
|
| 22 |
-
custom_id="req-1",
|
| 23 |
-
messages=[{"role": "user", "content": "hello"}],
|
| 24 |
-
tools=[{"name": "tool"}],
|
| 25 |
-
model="gpt-4o",
|
| 26 |
-
system_instruction="system",
|
| 27 |
-
extras={"x": 1},
|
| 28 |
-
)
|
| 29 |
-
context.add_request(request)
|
| 30 |
-
assert context.get_request("req-1") is request
|
| 31 |
-
assert context.get_request("missing") is None
|
| 32 |
-
|
| 33 |
-
monkeypatch.setattr("headroom.ccr.batch_store.time.time", lambda: context.expires_at + 1)
|
| 34 |
-
assert context.is_expired is True
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
async def test_batch_context_store_core_operations(monkeypatch) -> None:
|
| 38 |
-
now = {"value": 100.0}
|
| 39 |
-
monkeypatch.setattr("headroom.ccr.batch_store.time.time", lambda: now["value"])
|
| 40 |
-
|
| 41 |
-
store = BatchContextStore(ttl=10, max_contexts=2)
|
| 42 |
-
first = BatchContext(batch_id="b1", provider="anthropic")
|
| 43 |
-
second = BatchContext(batch_id="b2", provider="google")
|
| 44 |
-
third = BatchContext(batch_id="b3", provider="openai")
|
| 45 |
-
|
| 46 |
-
await store.store(first)
|
| 47 |
-
now["value"] = 101.0
|
| 48 |
-
await store.store(second)
|
| 49 |
-
assert (await store.get("b1")) is first
|
| 50 |
-
|
| 51 |
-
now["value"] = 102.0
|
| 52 |
-
await store.store(third)
|
| 53 |
-
assert await store.get("b1") is None
|
| 54 |
-
assert await store.get("b2") is second
|
| 55 |
-
assert await store.get("b3") is third
|
| 56 |
-
|
| 57 |
-
assert await store.remove("b2") is True
|
| 58 |
-
assert await store.remove("b2") is False
|
| 59 |
-
|
| 60 |
-
|
| 61 |
-
async def test_batch_context_store_cleanup_stats_and_memory_stats(monkeypatch) -> None:
|
| 62 |
-
now = {"value": 200.0}
|
| 63 |
-
monkeypatch.setattr("headroom.ccr.batch_store.time.time", lambda: now["value"])
|
| 64 |
-
store = BatchContextStore(ttl=5, max_contexts=10)
|
| 65 |
-
|
| 66 |
-
first = BatchContext(batch_id="b1", provider="anthropic")
|
| 67 |
-
first.add_request(
|
| 68 |
-
BatchRequestContext(custom_id="r1", messages=[{"content": "alpha"}], tools=[])
|
| 69 |
-
)
|
| 70 |
-
second = BatchContext(batch_id="b2", provider="google")
|
| 71 |
-
second.add_request(
|
| 72 |
-
BatchRequestContext(custom_id="r2", messages=[{"content": ["nested"]}], tools=[{}])
|
| 73 |
-
)
|
| 74 |
-
|
| 75 |
-
await store.store(first)
|
| 76 |
-
await store.store(second)
|
| 77 |
-
assert await store.stats() == {
|
| 78 |
-
"total_contexts": 2,
|
| 79 |
-
"max_contexts": 10,
|
| 80 |
-
"ttl_seconds": 5,
|
| 81 |
-
"providers": {"anthropic": 1, "google": 1},
|
| 82 |
-
}
|
| 83 |
-
|
| 84 |
-
now["value"] = 210.0
|
| 85 |
-
assert await store.cleanup_expired() == 2
|
| 86 |
-
assert (await store.stats())["total_contexts"] == 0
|
| 87 |
-
|
| 88 |
-
@dataclass
|
| 89 |
-
class FakeComponentStats:
|
| 90 |
-
name: str
|
| 91 |
-
entry_count: int
|
| 92 |
-
size_bytes: int
|
| 93 |
-
budget_bytes: int | None
|
| 94 |
-
hits: int
|
| 95 |
-
misses: int
|
| 96 |
-
evictions: int
|
| 97 |
-
|
| 98 |
-
monkeypatch.setitem(
|
| 99 |
-
sys.modules,
|
| 100 |
-
"headroom.memory.tracker",
|
| 101 |
-
type("TrackerModule", (), {"ComponentStats": FakeComponentStats}),
|
| 102 |
-
)
|
| 103 |
-
|
| 104 |
-
now["value"] = 220.0
|
| 105 |
-
third = BatchContext(batch_id="b3", provider="openai")
|
| 106 |
-
third.add_request(
|
| 107 |
-
BatchRequestContext(
|
| 108 |
-
custom_id="r3", messages=[{"content": "payload"}], tools=[{"name": "t"}]
|
| 109 |
-
)
|
| 110 |
-
)
|
| 111 |
-
await store.store(third)
|
| 112 |
-
stats = store.get_memory_stats()
|
| 113 |
-
assert stats.name == "batch_context_store"
|
| 114 |
-
assert stats.entry_count == 1
|
| 115 |
-
assert stats.size_bytes > 0
|
| 116 |
-
|
| 117 |
-
|
| 118 |
-
def test_global_batch_context_store_reset() -> None:
|
| 119 |
-
reset_batch_context_store()
|
| 120 |
-
store_one = get_batch_context_store()
|
| 121 |
-
store_two = get_batch_context_store()
|
| 122 |
-
assert store_one is store_two
|
| 123 |
-
|
| 124 |
-
reset_batch_context_store()
|
| 125 |
-
store_three = get_batch_context_store()
|
| 126 |
-
assert store_three is not store_one
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import sys
|
| 4 |
+
from dataclasses import dataclass
|
| 5 |
+
|
| 6 |
+
from headroom.ccr.batch_store import (
|
| 7 |
+
BatchContext,
|
| 8 |
+
BatchContextStore,
|
| 9 |
+
BatchRequestContext,
|
| 10 |
+
get_batch_context_store,
|
| 11 |
+
reset_batch_context_store,
|
| 12 |
+
)
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def test_batch_context_defaults_and_expiry(monkeypatch) -> None:
|
| 16 |
+
monkeypatch.setattr("headroom.ccr.batch_store.time.time", lambda: 100.0)
|
| 17 |
+
context = BatchContext(batch_id="batch-1", provider="anthropic", created_at=100.0)
|
| 18 |
+
assert context.expires_at == 100.0 + 86400
|
| 19 |
+
assert context.is_expired is False
|
| 20 |
+
|
| 21 |
+
request = BatchRequestContext(
|
| 22 |
+
custom_id="req-1",
|
| 23 |
+
messages=[{"role": "user", "content": "hello"}],
|
| 24 |
+
tools=[{"name": "tool"}],
|
| 25 |
+
model="gpt-4o",
|
| 26 |
+
system_instruction="system",
|
| 27 |
+
extras={"x": 1},
|
| 28 |
+
)
|
| 29 |
+
context.add_request(request)
|
| 30 |
+
assert context.get_request("req-1") is request
|
| 31 |
+
assert context.get_request("missing") is None
|
| 32 |
+
|
| 33 |
+
monkeypatch.setattr("headroom.ccr.batch_store.time.time", lambda: context.expires_at + 1)
|
| 34 |
+
assert context.is_expired is True
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
async def test_batch_context_store_core_operations(monkeypatch) -> None:
|
| 38 |
+
now = {"value": 100.0}
|
| 39 |
+
monkeypatch.setattr("headroom.ccr.batch_store.time.time", lambda: now["value"])
|
| 40 |
+
|
| 41 |
+
store = BatchContextStore(ttl=10, max_contexts=2)
|
| 42 |
+
first = BatchContext(batch_id="b1", provider="anthropic")
|
| 43 |
+
second = BatchContext(batch_id="b2", provider="google")
|
| 44 |
+
third = BatchContext(batch_id="b3", provider="openai")
|
| 45 |
+
|
| 46 |
+
await store.store(first)
|
| 47 |
+
now["value"] = 101.0
|
| 48 |
+
await store.store(second)
|
| 49 |
+
assert (await store.get("b1")) is first
|
| 50 |
+
|
| 51 |
+
now["value"] = 102.0
|
| 52 |
+
await store.store(third)
|
| 53 |
+
assert await store.get("b1") is None
|
| 54 |
+
assert await store.get("b2") is second
|
| 55 |
+
assert await store.get("b3") is third
|
| 56 |
+
|
| 57 |
+
assert await store.remove("b2") is True
|
| 58 |
+
assert await store.remove("b2") is False
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
async def test_batch_context_store_cleanup_stats_and_memory_stats(monkeypatch) -> None:
|
| 62 |
+
now = {"value": 200.0}
|
| 63 |
+
monkeypatch.setattr("headroom.ccr.batch_store.time.time", lambda: now["value"])
|
| 64 |
+
store = BatchContextStore(ttl=5, max_contexts=10)
|
| 65 |
+
|
| 66 |
+
first = BatchContext(batch_id="b1", provider="anthropic")
|
| 67 |
+
first.add_request(
|
| 68 |
+
BatchRequestContext(custom_id="r1", messages=[{"content": "alpha"}], tools=[])
|
| 69 |
+
)
|
| 70 |
+
second = BatchContext(batch_id="b2", provider="google")
|
| 71 |
+
second.add_request(
|
| 72 |
+
BatchRequestContext(custom_id="r2", messages=[{"content": ["nested"]}], tools=[{}])
|
| 73 |
+
)
|
| 74 |
+
|
| 75 |
+
await store.store(first)
|
| 76 |
+
await store.store(second)
|
| 77 |
+
assert await store.stats() == {
|
| 78 |
+
"total_contexts": 2,
|
| 79 |
+
"max_contexts": 10,
|
| 80 |
+
"ttl_seconds": 5,
|
| 81 |
+
"providers": {"anthropic": 1, "google": 1},
|
| 82 |
+
}
|
| 83 |
+
|
| 84 |
+
now["value"] = 210.0
|
| 85 |
+
assert await store.cleanup_expired() == 2
|
| 86 |
+
assert (await store.stats())["total_contexts"] == 0
|
| 87 |
+
|
| 88 |
+
@dataclass
|
| 89 |
+
class FakeComponentStats:
|
| 90 |
+
name: str
|
| 91 |
+
entry_count: int
|
| 92 |
+
size_bytes: int
|
| 93 |
+
budget_bytes: int | None
|
| 94 |
+
hits: int
|
| 95 |
+
misses: int
|
| 96 |
+
evictions: int
|
| 97 |
+
|
| 98 |
+
monkeypatch.setitem(
|
| 99 |
+
sys.modules,
|
| 100 |
+
"headroom.memory.tracker",
|
| 101 |
+
type("TrackerModule", (), {"ComponentStats": FakeComponentStats}),
|
| 102 |
+
)
|
| 103 |
+
|
| 104 |
+
now["value"] = 220.0
|
| 105 |
+
third = BatchContext(batch_id="b3", provider="openai")
|
| 106 |
+
third.add_request(
|
| 107 |
+
BatchRequestContext(
|
| 108 |
+
custom_id="r3", messages=[{"content": "payload"}], tools=[{"name": "t"}]
|
| 109 |
+
)
|
| 110 |
+
)
|
| 111 |
+
await store.store(third)
|
| 112 |
+
stats = store.get_memory_stats()
|
| 113 |
+
assert stats.name == "batch_context_store"
|
| 114 |
+
assert stats.entry_count == 1
|
| 115 |
+
assert stats.size_bytes > 0
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
def test_global_batch_context_store_reset() -> None:
|
| 119 |
+
reset_batch_context_store()
|
| 120 |
+
store_one = get_batch_context_store()
|
| 121 |
+
store_two = get_batch_context_store()
|
| 122 |
+
assert store_one is store_two
|
| 123 |
+
|
| 124 |
+
reset_batch_context_store()
|
| 125 |
+
store_three = get_batch_context_store()
|
| 126 |
+
assert store_three is not store_one
|
|
@@ -1,372 +1,372 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
import json
|
| 4 |
-
from typing import Any
|
| 5 |
-
|
| 6 |
-
import pytest
|
| 7 |
-
|
| 8 |
-
from headroom.ccr.response_handler import (
|
| 9 |
-
CCRResponseHandler,
|
| 10 |
-
CCRToolCall,
|
| 11 |
-
CCRToolResult,
|
| 12 |
-
StreamingCCRBuffer,
|
| 13 |
-
StreamingCCRHandler,
|
| 14 |
-
)
|
| 15 |
-
from headroom.ccr.tool_injection import CCR_TOOL_NAME
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
class FakeStore:
|
| 19 |
-
def __init__(
|
| 20 |
-
self, *, search_error: Exception | None = None, retrieve_error: Exception | None = None
|
| 21 |
-
) -> None:
|
| 22 |
-
self.search_error = search_error
|
| 23 |
-
self.retrieve_error = retrieve_error
|
| 24 |
-
|
| 25 |
-
def search(self, hash_key: str, query: str) -> list[dict[str, str]]:
|
| 26 |
-
if self.search_error:
|
| 27 |
-
raise self.search_error
|
| 28 |
-
return [{"id": "1", "text": query}]
|
| 29 |
-
|
| 30 |
-
def retrieve(self, hash_key: str):
|
| 31 |
-
if self.retrieve_error:
|
| 32 |
-
raise self.retrieve_error
|
| 33 |
-
return {"unexpected": True}
|
| 34 |
-
|
| 35 |
-
|
| 36 |
-
async def _async_iter(items: list[bytes]):
|
| 37 |
-
for item in items:
|
| 38 |
-
yield item
|
| 39 |
-
|
| 40 |
-
|
| 41 |
-
def test_extract_tool_calls_google_and_invalid_shapes() -> None:
|
| 42 |
-
handler = CCRResponseHandler()
|
| 43 |
-
google_response = {
|
| 44 |
-
"candidates": [
|
| 45 |
-
{
|
| 46 |
-
"content": {
|
| 47 |
-
"parts": [
|
| 48 |
-
{"text": "hello"},
|
| 49 |
-
{"functionCall": {"name": CCR_TOOL_NAME, "args": {"hash": "abc"}}},
|
| 50 |
-
]
|
| 51 |
-
}
|
| 52 |
-
}
|
| 53 |
-
]
|
| 54 |
-
}
|
| 55 |
-
assert handler._extract_tool_calls(google_response, "google") == [
|
| 56 |
-
{"functionCall": {"name": CCR_TOOL_NAME, "args": {"hash": "abc"}}}
|
| 57 |
-
]
|
| 58 |
-
assert handler._extract_tool_calls({"content": "bad"}, "anthropic") == []
|
| 59 |
-
with pytest.raises(IndexError):
|
| 60 |
-
handler._extract_tool_calls({"choices": []}, "openai")
|
| 61 |
-
assert handler._extract_tool_calls({"candidates": []}, "google") == []
|
| 62 |
-
|
| 63 |
-
|
| 64 |
-
def test_parse_ccr_tool_calls_google_and_other_calls() -> None:
|
| 65 |
-
handler = CCRResponseHandler()
|
| 66 |
-
response = {
|
| 67 |
-
"candidates": [
|
| 68 |
-
{
|
| 69 |
-
"content": {
|
| 70 |
-
"parts": [
|
| 71 |
-
{
|
| 72 |
-
"functionCall": {
|
| 73 |
-
"name": CCR_TOOL_NAME,
|
| 74 |
-
"args": {
|
| 75 |
-
"hash": "aaaaaaaaaaaaaaaaaaaaaaaa",
|
| 76 |
-
"query": "pizza",
|
| 77 |
-
},
|
| 78 |
-
}
|
| 79 |
-
},
|
| 80 |
-
{"functionCall": {"name": "other_tool", "args": {}}},
|
| 81 |
-
]
|
| 82 |
-
}
|
| 83 |
-
}
|
| 84 |
-
]
|
| 85 |
-
}
|
| 86 |
-
ccr_calls, other_calls = handler._parse_ccr_tool_calls(response, "google")
|
| 87 |
-
assert ccr_calls == [
|
| 88 |
-
CCRToolCall(
|
| 89 |
-
tool_call_id=CCR_TOOL_NAME,
|
| 90 |
-
hash_key="aaaaaaaaaaaaaaaaaaaaaaaa",
|
| 91 |
-
query="pizza",
|
| 92 |
-
)
|
| 93 |
-
]
|
| 94 |
-
assert other_calls == [{"functionCall": {"name": "other_tool", "args": {}}}]
|
| 95 |
-
|
| 96 |
-
|
| 97 |
-
def test_execute_retrieval_error_paths(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 98 |
-
handler = CCRResponseHandler()
|
| 99 |
-
monkeypatch.setattr(
|
| 100 |
-
"headroom.ccr.response_handler.get_compression_store",
|
| 101 |
-
lambda: FakeStore(search_error=RuntimeError("search boom")),
|
| 102 |
-
)
|
| 103 |
-
search_result = handler._execute_retrieval(
|
| 104 |
-
CCRToolCall(tool_call_id="t1", hash_key="abc", query="find")
|
| 105 |
-
)
|
| 106 |
-
assert search_result.success is False
|
| 107 |
-
assert "Retrieval failed: search boom" in search_result.content
|
| 108 |
-
|
| 109 |
-
monkeypatch.setattr(
|
| 110 |
-
"headroom.ccr.response_handler.get_compression_store",
|
| 111 |
-
lambda: FakeStore(retrieve_error=RuntimeError("retrieve boom")),
|
| 112 |
-
)
|
| 113 |
-
retrieve_result = handler._execute_retrieval(CCRToolCall(tool_call_id="t2", hash_key="abc"))
|
| 114 |
-
assert retrieve_result.success is False
|
| 115 |
-
assert "Retrieval failed: retrieve boom" in retrieve_result.content
|
| 116 |
-
|
| 117 |
-
|
| 118 |
-
def test_create_tool_result_message_google_and_generic_formats() -> None:
|
| 119 |
-
handler = CCRResponseHandler()
|
| 120 |
-
results = [
|
| 121 |
-
CCRToolResult(tool_call_id="headroom_retrieve", content='{"count": 1}', success=True)
|
| 122 |
-
]
|
| 123 |
-
google_message = handler._create_tool_result_message(results, "google")
|
| 124 |
-
assert google_message == {
|
| 125 |
-
"role": "user",
|
| 126 |
-
"parts": [{"functionResponse": {"name": "headroom_retrieve", "response": {"count": 1}}}],
|
| 127 |
-
}
|
| 128 |
-
|
| 129 |
-
generic_message = handler._create_tool_result_message(
|
| 130 |
-
[CCRToolResult(tool_call_id="tool-1", content="not-json", success=False)],
|
| 131 |
-
"other",
|
| 132 |
-
)
|
| 133 |
-
assert generic_message["role"] == "tool"
|
| 134 |
-
assert json.loads(generic_message["content"]) == [
|
| 135 |
-
{"tool_call_id": "tool-1", "result": "not-json"}
|
| 136 |
-
]
|
| 137 |
-
|
| 138 |
-
invalid_google = handler._create_tool_result_message(
|
| 139 |
-
[CCRToolResult(tool_call_id="headroom_retrieve", content="not-json", success=True)],
|
| 140 |
-
"google",
|
| 141 |
-
)
|
| 142 |
-
assert invalid_google["parts"][0]["functionResponse"]["response"] == {"content": "not-json"}
|
| 143 |
-
|
| 144 |
-
|
| 145 |
-
def test_extract_assistant_message_google_and_generic() -> None:
|
| 146 |
-
handler = CCRResponseHandler()
|
| 147 |
-
google_message = handler._extract_assistant_message(
|
| 148 |
-
{"candidates": [{"content": {"parts": [{"text": "hello"}]}}]},
|
| 149 |
-
"google",
|
| 150 |
-
)
|
| 151 |
-
assert google_message == {"role": "model", "parts": [{"text": "hello"}]}
|
| 152 |
-
|
| 153 |
-
assert handler._extract_assistant_message({}, "google") == {"role": "model", "parts": []}
|
| 154 |
-
assert handler._extract_assistant_message({"content": "plain"}, "other") == {
|
| 155 |
-
"role": "assistant",
|
| 156 |
-
"content": "plain",
|
| 157 |
-
}
|
| 158 |
-
|
| 159 |
-
|
| 160 |
-
@pytest.mark.asyncio
|
| 161 |
-
async def test_handle_response_openai_success_and_failure(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 162 |
-
handler = CCRResponseHandler()
|
| 163 |
-
initial_response = {
|
| 164 |
-
"choices": [
|
| 165 |
-
{
|
| 166 |
-
"message": {
|
| 167 |
-
"role": "assistant",
|
| 168 |
-
"content": None,
|
| 169 |
-
"tool_calls": [
|
| 170 |
-
{
|
| 171 |
-
"id": "call_1",
|
| 172 |
-
"type": "function",
|
| 173 |
-
"function": {
|
| 174 |
-
"name": CCR_TOOL_NAME,
|
| 175 |
-
"arguments": '{"hash":"aaaaaaaaaaaaaaaaaaaaaaaa"}',
|
| 176 |
-
},
|
| 177 |
-
}
|
| 178 |
-
],
|
| 179 |
-
}
|
| 180 |
-
}
|
| 181 |
-
]
|
| 182 |
-
}
|
| 183 |
-
monkeypatch.setattr(
|
| 184 |
-
handler,
|
| 185 |
-
"_execute_retrieval",
|
| 186 |
-
lambda call: CCRToolResult(
|
| 187 |
-
tool_call_id=call.tool_call_id,
|
| 188 |
-
content='{"hash":"aaaaaaaaaaaaaaaaaaaaaaaa"}',
|
| 189 |
-
success=True,
|
| 190 |
-
),
|
| 191 |
-
)
|
| 192 |
-
|
| 193 |
-
captured_messages: list[list[dict[str, Any]]] = []
|
| 194 |
-
|
| 195 |
-
async def success_api_call(messages, tools):
|
| 196 |
-
captured_messages.append(messages)
|
| 197 |
-
return {"choices": [{"message": {"role": "assistant", "content": "done"}}]}
|
| 198 |
-
|
| 199 |
-
result = await handler.handle_response(
|
| 200 |
-
initial_response, [{"role": "user", "content": "hi"}], [], success_api_call, "openai"
|
| 201 |
-
)
|
| 202 |
-
assert result == {"choices": [{"message": {"role": "assistant", "content": "done"}}]}
|
| 203 |
-
assert captured_messages[0][1]["role"] == "assistant"
|
| 204 |
-
assert captured_messages[0][2]["role"] == "tool"
|
| 205 |
-
assert handler.get_stats()["total_retrievals"] == 1
|
| 206 |
-
|
| 207 |
-
async def failing_api_call(messages, tools):
|
| 208 |
-
raise RuntimeError("continuation failed")
|
| 209 |
-
|
| 210 |
-
failed = await handler.handle_response(initial_response, [], [], failing_api_call, "openai")
|
| 211 |
-
assert failed == initial_response
|
| 212 |
-
|
| 213 |
-
|
| 214 |
-
def test_streaming_buffer_and_parse_sse_helpers() -> None:
|
| 215 |
-
buffer = StreamingCCRBuffer()
|
| 216 |
-
assert buffer.add_chunk(b"plain") is False
|
| 217 |
-
assert buffer.get_accumulated() == b"plain"
|
| 218 |
-
|
| 219 |
-
handler = StreamingCCRHandler(CCRResponseHandler(), provider="anthropic")
|
| 220 |
-
anthropic_data = b"\n".join(
|
| 221 |
-
[
|
| 222 |
-
b'data: {"type":"content_block_start","content_block":{"type":"text","text":"Hel"}}',
|
| 223 |
-
b'data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"lo"}}',
|
| 224 |
-
b'data: {"type":"content_block_stop"}',
|
| 225 |
-
b'data: {"type":"content_block_start","content_block":{"type":"tool_use","id":"tool_1","name":"headroom_retrieve"}}',
|
| 226 |
-
b'data: {"type":"content_block_delta","delta":{"type":"input_json_delta","partial_json":"{\\"hash\\":\\"abc\\"}"}}',
|
| 227 |
-
b'data: {"type":"content_block_stop"}',
|
| 228 |
-
b'data: {"type":"message_delta","delta":{"stop_reason":"tool_use"}}',
|
| 229 |
-
b"data: [DONE]",
|
| 230 |
-
]
|
| 231 |
-
)
|
| 232 |
-
parsed = handler._parse_sse_stream(anthropic_data)
|
| 233 |
-
assert parsed["content"][0] == {"type": "text", "text": "Hello"}
|
| 234 |
-
assert parsed["content"][1]["name"] == "headroom_retrieve"
|
| 235 |
-
assert parsed["content"][1]["input"] == {"hash": "abc"}
|
| 236 |
-
assert parsed["stop_reason"] == "tool_use"
|
| 237 |
-
|
| 238 |
-
openai_handler = StreamingCCRHandler(CCRResponseHandler(), provider="openai")
|
| 239 |
-
parsed_openai = openai_handler._reconstruct_openai_response(
|
| 240 |
-
[
|
| 241 |
-
{"choices": [{"delta": {"content": "Hi"}}]},
|
| 242 |
-
{
|
| 243 |
-
"choices": [
|
| 244 |
-
{
|
| 245 |
-
"delta": {
|
| 246 |
-
"tool_calls": [
|
| 247 |
-
{
|
| 248 |
-
"index": 0,
|
| 249 |
-
"id": "call_1",
|
| 250 |
-
"function": {
|
| 251 |
-
"name": "headroom_retrieve",
|
| 252 |
-
"arguments": '{"hash":"aaaaaaaaaaaa',
|
| 253 |
-
},
|
| 254 |
-
}
|
| 255 |
-
]
|
| 256 |
-
}
|
| 257 |
-
}
|
| 258 |
-
]
|
| 259 |
-
},
|
| 260 |
-
{
|
| 261 |
-
"choices": [
|
| 262 |
-
{
|
| 263 |
-
"delta": {
|
| 264 |
-
"tool_calls": [
|
| 265 |
-
{
|
| 266 |
-
"index": 0,
|
| 267 |
-
"function": {"arguments": 'aaaaaaaaaaaa"}'},
|
| 268 |
-
}
|
| 269 |
-
]
|
| 270 |
-
}
|
| 271 |
-
}
|
| 272 |
-
]
|
| 273 |
-
},
|
| 274 |
-
]
|
| 275 |
-
)
|
| 276 |
-
message = parsed_openai["choices"][0]["message"]
|
| 277 |
-
assert message["content"] == "Hi"
|
| 278 |
-
assert message["tool_calls"][0]["id"] == "call_1"
|
| 279 |
-
assert message["tool_calls"][0]["function"]["arguments"] == (
|
| 280 |
-
'{"hash":"aaaaaaaaaaaaaaaaaaaaaaaa"}'
|
| 281 |
-
)
|
| 282 |
-
|
| 283 |
-
|
| 284 |
-
@pytest.mark.asyncio
|
| 285 |
-
async def test_streaming_handler_process_stream_pass_through_and_ccr(
|
| 286 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 287 |
-
) -> None:
|
| 288 |
-
response_handler = CCRResponseHandler()
|
| 289 |
-
handler = StreamingCCRHandler(response_handler, provider="anthropic")
|
| 290 |
-
|
| 291 |
-
passthrough_chunks = [
|
| 292 |
-
b'data: {"type":"content_block_delta","delta":{"text":"hello"}}',
|
| 293 |
-
b'data: {"stop_reason":"end_turn"}',
|
| 294 |
-
]
|
| 295 |
-
yielded = [
|
| 296 |
-
chunk
|
| 297 |
-
async for chunk in handler.process_stream(
|
| 298 |
-
_async_iter(passthrough_chunks), [], None, lambda m, t: None
|
| 299 |
-
)
|
| 300 |
-
]
|
| 301 |
-
assert yielded == passthrough_chunks
|
| 302 |
-
|
| 303 |
-
ccr_handler = StreamingCCRHandler(response_handler, provider="anthropic")
|
| 304 |
-
monkeypatch.setattr(
|
| 305 |
-
ccr_handler,
|
| 306 |
-
"_parse_sse_stream",
|
| 307 |
-
lambda data: {
|
| 308 |
-
"content": [
|
| 309 |
-
{
|
| 310 |
-
"type": "tool_use",
|
| 311 |
-
"id": "tool_1",
|
| 312 |
-
"name": CCR_TOOL_NAME,
|
| 313 |
-
"input": {"hash": "abc"},
|
| 314 |
-
}
|
| 315 |
-
]
|
| 316 |
-
},
|
| 317 |
-
)
|
| 318 |
-
|
| 319 |
-
async def fake_handle_response(response, messages, tools, api_call_fn, provider): # noqa: ANN001
|
| 320 |
-
return {"content": [{"type": "text", "text": "done"}]}
|
| 321 |
-
|
| 322 |
-
async def fake_response_to_sse(response): # noqa: ANN001
|
| 323 |
-
yield b"event: message_start\n"
|
| 324 |
-
yield b"event: message_stop\n"
|
| 325 |
-
|
| 326 |
-
monkeypatch.setattr(response_handler, "handle_response", fake_handle_response)
|
| 327 |
-
monkeypatch.setattr(ccr_handler, "_response_to_sse", fake_response_to_sse)
|
| 328 |
-
|
| 329 |
-
ccr_chunks = [
|
| 330 |
-
b'{"type":"tool_use","name":"headroom_retrieve"',
|
| 331 |
-
b',"stop_reason":"tool_use"}',
|
| 332 |
-
b"tail",
|
| 333 |
-
]
|
| 334 |
-
streamed = [
|
| 335 |
-
chunk
|
| 336 |
-
async for chunk in ccr_handler.process_stream(
|
| 337 |
-
_async_iter(ccr_chunks), [], None, lambda m, t: None
|
| 338 |
-
)
|
| 339 |
-
]
|
| 340 |
-
assert streamed == [b"event: message_start\n", b"event: message_stop\n"]
|
| 341 |
-
|
| 342 |
-
|
| 343 |
-
@pytest.mark.asyncio
|
| 344 |
-
async def test_streaming_handler_falls_back_to_buffer_on_processing_error(
|
| 345 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 346 |
-
) -> None:
|
| 347 |
-
response_handler = CCRResponseHandler()
|
| 348 |
-
handler = StreamingCCRHandler(response_handler, provider="openai")
|
| 349 |
-
monkeypatch.setattr(
|
| 350 |
-
handler,
|
| 351 |
-
"_parse_sse_stream",
|
| 352 |
-
lambda data: (_ for _ in ()).throw(RuntimeError("parse failed")),
|
| 353 |
-
)
|
| 354 |
-
|
| 355 |
-
chunks = [b'{"type":"tool_use","name":"headroom_retrieve"', b',"stop_reason":"tool_use"}']
|
| 356 |
-
streamed = [
|
| 357 |
-
chunk
|
| 358 |
-
async for chunk in handler.process_stream(_async_iter(chunks), [], None, lambda m, t: None)
|
| 359 |
-
]
|
| 360 |
-
assert streamed == [b"".join(chunks)]
|
| 361 |
-
|
| 362 |
-
|
| 363 |
-
@pytest.mark.asyncio
|
| 364 |
-
async def test_response_to_sse_formats() -> None:
|
| 365 |
-
anthropic = StreamingCCRHandler(CCRResponseHandler(), provider="anthropic")
|
| 366 |
-
anthropic_chunks = [chunk async for chunk in anthropic._response_to_sse({"content": []})]
|
| 367 |
-
assert anthropic_chunks[0] == b"event: message_start\n"
|
| 368 |
-
assert anthropic_chunks[-1] == b'data: {"type": "message_stop"}\n\n'
|
| 369 |
-
|
| 370 |
-
openai = StreamingCCRHandler(CCRResponseHandler(), provider="openai")
|
| 371 |
-
openai_chunks = [chunk async for chunk in openai._response_to_sse({"choices": []})]
|
| 372 |
-
assert openai_chunks == [b'data: {"choices": []}\n\n', b"data: [DONE]\n\n"]
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
from typing import Any
|
| 5 |
+
|
| 6 |
+
import pytest
|
| 7 |
+
|
| 8 |
+
from headroom.ccr.response_handler import (
|
| 9 |
+
CCRResponseHandler,
|
| 10 |
+
CCRToolCall,
|
| 11 |
+
CCRToolResult,
|
| 12 |
+
StreamingCCRBuffer,
|
| 13 |
+
StreamingCCRHandler,
|
| 14 |
+
)
|
| 15 |
+
from headroom.ccr.tool_injection import CCR_TOOL_NAME
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class FakeStore:
|
| 19 |
+
def __init__(
|
| 20 |
+
self, *, search_error: Exception | None = None, retrieve_error: Exception | None = None
|
| 21 |
+
) -> None:
|
| 22 |
+
self.search_error = search_error
|
| 23 |
+
self.retrieve_error = retrieve_error
|
| 24 |
+
|
| 25 |
+
def search(self, hash_key: str, query: str) -> list[dict[str, str]]:
|
| 26 |
+
if self.search_error:
|
| 27 |
+
raise self.search_error
|
| 28 |
+
return [{"id": "1", "text": query}]
|
| 29 |
+
|
| 30 |
+
def retrieve(self, hash_key: str):
|
| 31 |
+
if self.retrieve_error:
|
| 32 |
+
raise self.retrieve_error
|
| 33 |
+
return {"unexpected": True}
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
async def _async_iter(items: list[bytes]):
|
| 37 |
+
for item in items:
|
| 38 |
+
yield item
|
| 39 |
+
|
| 40 |
+
|
| 41 |
+
def test_extract_tool_calls_google_and_invalid_shapes() -> None:
|
| 42 |
+
handler = CCRResponseHandler()
|
| 43 |
+
google_response = {
|
| 44 |
+
"candidates": [
|
| 45 |
+
{
|
| 46 |
+
"content": {
|
| 47 |
+
"parts": [
|
| 48 |
+
{"text": "hello"},
|
| 49 |
+
{"functionCall": {"name": CCR_TOOL_NAME, "args": {"hash": "abc"}}},
|
| 50 |
+
]
|
| 51 |
+
}
|
| 52 |
+
}
|
| 53 |
+
]
|
| 54 |
+
}
|
| 55 |
+
assert handler._extract_tool_calls(google_response, "google") == [
|
| 56 |
+
{"functionCall": {"name": CCR_TOOL_NAME, "args": {"hash": "abc"}}}
|
| 57 |
+
]
|
| 58 |
+
assert handler._extract_tool_calls({"content": "bad"}, "anthropic") == []
|
| 59 |
+
with pytest.raises(IndexError):
|
| 60 |
+
handler._extract_tool_calls({"choices": []}, "openai")
|
| 61 |
+
assert handler._extract_tool_calls({"candidates": []}, "google") == []
|
| 62 |
+
|
| 63 |
+
|
| 64 |
+
def test_parse_ccr_tool_calls_google_and_other_calls() -> None:
|
| 65 |
+
handler = CCRResponseHandler()
|
| 66 |
+
response = {
|
| 67 |
+
"candidates": [
|
| 68 |
+
{
|
| 69 |
+
"content": {
|
| 70 |
+
"parts": [
|
| 71 |
+
{
|
| 72 |
+
"functionCall": {
|
| 73 |
+
"name": CCR_TOOL_NAME,
|
| 74 |
+
"args": {
|
| 75 |
+
"hash": "aaaaaaaaaaaaaaaaaaaaaaaa",
|
| 76 |
+
"query": "pizza",
|
| 77 |
+
},
|
| 78 |
+
}
|
| 79 |
+
},
|
| 80 |
+
{"functionCall": {"name": "other_tool", "args": {}}},
|
| 81 |
+
]
|
| 82 |
+
}
|
| 83 |
+
}
|
| 84 |
+
]
|
| 85 |
+
}
|
| 86 |
+
ccr_calls, other_calls = handler._parse_ccr_tool_calls(response, "google")
|
| 87 |
+
assert ccr_calls == [
|
| 88 |
+
CCRToolCall(
|
| 89 |
+
tool_call_id=CCR_TOOL_NAME,
|
| 90 |
+
hash_key="aaaaaaaaaaaaaaaaaaaaaaaa",
|
| 91 |
+
query="pizza",
|
| 92 |
+
)
|
| 93 |
+
]
|
| 94 |
+
assert other_calls == [{"functionCall": {"name": "other_tool", "args": {}}}]
|
| 95 |
+
|
| 96 |
+
|
| 97 |
+
def test_execute_retrieval_error_paths(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 98 |
+
handler = CCRResponseHandler()
|
| 99 |
+
monkeypatch.setattr(
|
| 100 |
+
"headroom.ccr.response_handler.get_compression_store",
|
| 101 |
+
lambda: FakeStore(search_error=RuntimeError("search boom")),
|
| 102 |
+
)
|
| 103 |
+
search_result = handler._execute_retrieval(
|
| 104 |
+
CCRToolCall(tool_call_id="t1", hash_key="abc", query="find")
|
| 105 |
+
)
|
| 106 |
+
assert search_result.success is False
|
| 107 |
+
assert "Retrieval failed: search boom" in search_result.content
|
| 108 |
+
|
| 109 |
+
monkeypatch.setattr(
|
| 110 |
+
"headroom.ccr.response_handler.get_compression_store",
|
| 111 |
+
lambda: FakeStore(retrieve_error=RuntimeError("retrieve boom")),
|
| 112 |
+
)
|
| 113 |
+
retrieve_result = handler._execute_retrieval(CCRToolCall(tool_call_id="t2", hash_key="abc"))
|
| 114 |
+
assert retrieve_result.success is False
|
| 115 |
+
assert "Retrieval failed: retrieve boom" in retrieve_result.content
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
def test_create_tool_result_message_google_and_generic_formats() -> None:
|
| 119 |
+
handler = CCRResponseHandler()
|
| 120 |
+
results = [
|
| 121 |
+
CCRToolResult(tool_call_id="headroom_retrieve", content='{"count": 1}', success=True)
|
| 122 |
+
]
|
| 123 |
+
google_message = handler._create_tool_result_message(results, "google")
|
| 124 |
+
assert google_message == {
|
| 125 |
+
"role": "user",
|
| 126 |
+
"parts": [{"functionResponse": {"name": "headroom_retrieve", "response": {"count": 1}}}],
|
| 127 |
+
}
|
| 128 |
+
|
| 129 |
+
generic_message = handler._create_tool_result_message(
|
| 130 |
+
[CCRToolResult(tool_call_id="tool-1", content="not-json", success=False)],
|
| 131 |
+
"other",
|
| 132 |
+
)
|
| 133 |
+
assert generic_message["role"] == "tool"
|
| 134 |
+
assert json.loads(generic_message["content"]) == [
|
| 135 |
+
{"tool_call_id": "tool-1", "result": "not-json"}
|
| 136 |
+
]
|
| 137 |
+
|
| 138 |
+
invalid_google = handler._create_tool_result_message(
|
| 139 |
+
[CCRToolResult(tool_call_id="headroom_retrieve", content="not-json", success=True)],
|
| 140 |
+
"google",
|
| 141 |
+
)
|
| 142 |
+
assert invalid_google["parts"][0]["functionResponse"]["response"] == {"content": "not-json"}
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def test_extract_assistant_message_google_and_generic() -> None:
|
| 146 |
+
handler = CCRResponseHandler()
|
| 147 |
+
google_message = handler._extract_assistant_message(
|
| 148 |
+
{"candidates": [{"content": {"parts": [{"text": "hello"}]}}]},
|
| 149 |
+
"google",
|
| 150 |
+
)
|
| 151 |
+
assert google_message == {"role": "model", "parts": [{"text": "hello"}]}
|
| 152 |
+
|
| 153 |
+
assert handler._extract_assistant_message({}, "google") == {"role": "model", "parts": []}
|
| 154 |
+
assert handler._extract_assistant_message({"content": "plain"}, "other") == {
|
| 155 |
+
"role": "assistant",
|
| 156 |
+
"content": "plain",
|
| 157 |
+
}
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
@pytest.mark.asyncio
|
| 161 |
+
async def test_handle_response_openai_success_and_failure(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 162 |
+
handler = CCRResponseHandler()
|
| 163 |
+
initial_response = {
|
| 164 |
+
"choices": [
|
| 165 |
+
{
|
| 166 |
+
"message": {
|
| 167 |
+
"role": "assistant",
|
| 168 |
+
"content": None,
|
| 169 |
+
"tool_calls": [
|
| 170 |
+
{
|
| 171 |
+
"id": "call_1",
|
| 172 |
+
"type": "function",
|
| 173 |
+
"function": {
|
| 174 |
+
"name": CCR_TOOL_NAME,
|
| 175 |
+
"arguments": '{"hash":"aaaaaaaaaaaaaaaaaaaaaaaa"}',
|
| 176 |
+
},
|
| 177 |
+
}
|
| 178 |
+
],
|
| 179 |
+
}
|
| 180 |
+
}
|
| 181 |
+
]
|
| 182 |
+
}
|
| 183 |
+
monkeypatch.setattr(
|
| 184 |
+
handler,
|
| 185 |
+
"_execute_retrieval",
|
| 186 |
+
lambda call: CCRToolResult(
|
| 187 |
+
tool_call_id=call.tool_call_id,
|
| 188 |
+
content='{"hash":"aaaaaaaaaaaaaaaaaaaaaaaa"}',
|
| 189 |
+
success=True,
|
| 190 |
+
),
|
| 191 |
+
)
|
| 192 |
+
|
| 193 |
+
captured_messages: list[list[dict[str, Any]]] = []
|
| 194 |
+
|
| 195 |
+
async def success_api_call(messages, tools):
|
| 196 |
+
captured_messages.append(messages)
|
| 197 |
+
return {"choices": [{"message": {"role": "assistant", "content": "done"}}]}
|
| 198 |
+
|
| 199 |
+
result = await handler.handle_response(
|
| 200 |
+
initial_response, [{"role": "user", "content": "hi"}], [], success_api_call, "openai"
|
| 201 |
+
)
|
| 202 |
+
assert result == {"choices": [{"message": {"role": "assistant", "content": "done"}}]}
|
| 203 |
+
assert captured_messages[0][1]["role"] == "assistant"
|
| 204 |
+
assert captured_messages[0][2]["role"] == "tool"
|
| 205 |
+
assert handler.get_stats()["total_retrievals"] == 1
|
| 206 |
+
|
| 207 |
+
async def failing_api_call(messages, tools):
|
| 208 |
+
raise RuntimeError("continuation failed")
|
| 209 |
+
|
| 210 |
+
failed = await handler.handle_response(initial_response, [], [], failing_api_call, "openai")
|
| 211 |
+
assert failed == initial_response
|
| 212 |
+
|
| 213 |
+
|
| 214 |
+
def test_streaming_buffer_and_parse_sse_helpers() -> None:
|
| 215 |
+
buffer = StreamingCCRBuffer()
|
| 216 |
+
assert buffer.add_chunk(b"plain") is False
|
| 217 |
+
assert buffer.get_accumulated() == b"plain"
|
| 218 |
+
|
| 219 |
+
handler = StreamingCCRHandler(CCRResponseHandler(), provider="anthropic")
|
| 220 |
+
anthropic_data = b"\n".join(
|
| 221 |
+
[
|
| 222 |
+
b'data: {"type":"content_block_start","content_block":{"type":"text","text":"Hel"}}',
|
| 223 |
+
b'data: {"type":"content_block_delta","delta":{"type":"text_delta","text":"lo"}}',
|
| 224 |
+
b'data: {"type":"content_block_stop"}',
|
| 225 |
+
b'data: {"type":"content_block_start","content_block":{"type":"tool_use","id":"tool_1","name":"headroom_retrieve"}}',
|
| 226 |
+
b'data: {"type":"content_block_delta","delta":{"type":"input_json_delta","partial_json":"{\\"hash\\":\\"abc\\"}"}}',
|
| 227 |
+
b'data: {"type":"content_block_stop"}',
|
| 228 |
+
b'data: {"type":"message_delta","delta":{"stop_reason":"tool_use"}}',
|
| 229 |
+
b"data: [DONE]",
|
| 230 |
+
]
|
| 231 |
+
)
|
| 232 |
+
parsed = handler._parse_sse_stream(anthropic_data)
|
| 233 |
+
assert parsed["content"][0] == {"type": "text", "text": "Hello"}
|
| 234 |
+
assert parsed["content"][1]["name"] == "headroom_retrieve"
|
| 235 |
+
assert parsed["content"][1]["input"] == {"hash": "abc"}
|
| 236 |
+
assert parsed["stop_reason"] == "tool_use"
|
| 237 |
+
|
| 238 |
+
openai_handler = StreamingCCRHandler(CCRResponseHandler(), provider="openai")
|
| 239 |
+
parsed_openai = openai_handler._reconstruct_openai_response(
|
| 240 |
+
[
|
| 241 |
+
{"choices": [{"delta": {"content": "Hi"}}]},
|
| 242 |
+
{
|
| 243 |
+
"choices": [
|
| 244 |
+
{
|
| 245 |
+
"delta": {
|
| 246 |
+
"tool_calls": [
|
| 247 |
+
{
|
| 248 |
+
"index": 0,
|
| 249 |
+
"id": "call_1",
|
| 250 |
+
"function": {
|
| 251 |
+
"name": "headroom_retrieve",
|
| 252 |
+
"arguments": '{"hash":"aaaaaaaaaaaa',
|
| 253 |
+
},
|
| 254 |
+
}
|
| 255 |
+
]
|
| 256 |
+
}
|
| 257 |
+
}
|
| 258 |
+
]
|
| 259 |
+
},
|
| 260 |
+
{
|
| 261 |
+
"choices": [
|
| 262 |
+
{
|
| 263 |
+
"delta": {
|
| 264 |
+
"tool_calls": [
|
| 265 |
+
{
|
| 266 |
+
"index": 0,
|
| 267 |
+
"function": {"arguments": 'aaaaaaaaaaaa"}'},
|
| 268 |
+
}
|
| 269 |
+
]
|
| 270 |
+
}
|
| 271 |
+
}
|
| 272 |
+
]
|
| 273 |
+
},
|
| 274 |
+
]
|
| 275 |
+
)
|
| 276 |
+
message = parsed_openai["choices"][0]["message"]
|
| 277 |
+
assert message["content"] == "Hi"
|
| 278 |
+
assert message["tool_calls"][0]["id"] == "call_1"
|
| 279 |
+
assert message["tool_calls"][0]["function"]["arguments"] == (
|
| 280 |
+
'{"hash":"aaaaaaaaaaaaaaaaaaaaaaaa"}'
|
| 281 |
+
)
|
| 282 |
+
|
| 283 |
+
|
| 284 |
+
@pytest.mark.asyncio
|
| 285 |
+
async def test_streaming_handler_process_stream_pass_through_and_ccr(
|
| 286 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 287 |
+
) -> None:
|
| 288 |
+
response_handler = CCRResponseHandler()
|
| 289 |
+
handler = StreamingCCRHandler(response_handler, provider="anthropic")
|
| 290 |
+
|
| 291 |
+
passthrough_chunks = [
|
| 292 |
+
b'data: {"type":"content_block_delta","delta":{"text":"hello"}}',
|
| 293 |
+
b'data: {"stop_reason":"end_turn"}',
|
| 294 |
+
]
|
| 295 |
+
yielded = [
|
| 296 |
+
chunk
|
| 297 |
+
async for chunk in handler.process_stream(
|
| 298 |
+
_async_iter(passthrough_chunks), [], None, lambda m, t: None
|
| 299 |
+
)
|
| 300 |
+
]
|
| 301 |
+
assert yielded == passthrough_chunks
|
| 302 |
+
|
| 303 |
+
ccr_handler = StreamingCCRHandler(response_handler, provider="anthropic")
|
| 304 |
+
monkeypatch.setattr(
|
| 305 |
+
ccr_handler,
|
| 306 |
+
"_parse_sse_stream",
|
| 307 |
+
lambda data: {
|
| 308 |
+
"content": [
|
| 309 |
+
{
|
| 310 |
+
"type": "tool_use",
|
| 311 |
+
"id": "tool_1",
|
| 312 |
+
"name": CCR_TOOL_NAME,
|
| 313 |
+
"input": {"hash": "abc"},
|
| 314 |
+
}
|
| 315 |
+
]
|
| 316 |
+
},
|
| 317 |
+
)
|
| 318 |
+
|
| 319 |
+
async def fake_handle_response(response, messages, tools, api_call_fn, provider): # noqa: ANN001
|
| 320 |
+
return {"content": [{"type": "text", "text": "done"}]}
|
| 321 |
+
|
| 322 |
+
async def fake_response_to_sse(response): # noqa: ANN001
|
| 323 |
+
yield b"event: message_start\n"
|
| 324 |
+
yield b"event: message_stop\n"
|
| 325 |
+
|
| 326 |
+
monkeypatch.setattr(response_handler, "handle_response", fake_handle_response)
|
| 327 |
+
monkeypatch.setattr(ccr_handler, "_response_to_sse", fake_response_to_sse)
|
| 328 |
+
|
| 329 |
+
ccr_chunks = [
|
| 330 |
+
b'{"type":"tool_use","name":"headroom_retrieve"',
|
| 331 |
+
b',"stop_reason":"tool_use"}',
|
| 332 |
+
b"tail",
|
| 333 |
+
]
|
| 334 |
+
streamed = [
|
| 335 |
+
chunk
|
| 336 |
+
async for chunk in ccr_handler.process_stream(
|
| 337 |
+
_async_iter(ccr_chunks), [], None, lambda m, t: None
|
| 338 |
+
)
|
| 339 |
+
]
|
| 340 |
+
assert streamed == [b"event: message_start\n", b"event: message_stop\n"]
|
| 341 |
+
|
| 342 |
+
|
| 343 |
+
@pytest.mark.asyncio
|
| 344 |
+
async def test_streaming_handler_falls_back_to_buffer_on_processing_error(
|
| 345 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 346 |
+
) -> None:
|
| 347 |
+
response_handler = CCRResponseHandler()
|
| 348 |
+
handler = StreamingCCRHandler(response_handler, provider="openai")
|
| 349 |
+
monkeypatch.setattr(
|
| 350 |
+
handler,
|
| 351 |
+
"_parse_sse_stream",
|
| 352 |
+
lambda data: (_ for _ in ()).throw(RuntimeError("parse failed")),
|
| 353 |
+
)
|
| 354 |
+
|
| 355 |
+
chunks = [b'{"type":"tool_use","name":"headroom_retrieve"', b',"stop_reason":"tool_use"}']
|
| 356 |
+
streamed = [
|
| 357 |
+
chunk
|
| 358 |
+
async for chunk in handler.process_stream(_async_iter(chunks), [], None, lambda m, t: None)
|
| 359 |
+
]
|
| 360 |
+
assert streamed == [b"".join(chunks)]
|
| 361 |
+
|
| 362 |
+
|
| 363 |
+
@pytest.mark.asyncio
|
| 364 |
+
async def test_response_to_sse_formats() -> None:
|
| 365 |
+
anthropic = StreamingCCRHandler(CCRResponseHandler(), provider="anthropic")
|
| 366 |
+
anthropic_chunks = [chunk async for chunk in anthropic._response_to_sse({"content": []})]
|
| 367 |
+
assert anthropic_chunks[0] == b"event: message_start\n"
|
| 368 |
+
assert anthropic_chunks[-1] == b'data: {"type": "message_stop"}\n\n'
|
| 369 |
+
|
| 370 |
+
openai = StreamingCCRHandler(CCRResponseHandler(), provider="openai")
|
| 371 |
+
openai_chunks = [chunk async for chunk in openai._response_to_sse({"choices": []})]
|
| 372 |
+
assert openai_chunks == [b'data: {"choices": []}\n\n', b"data: [DONE]\n\n"]
|
|
@@ -1,335 +1,335 @@
|
|
| 1 |
-
"""Tests for `headroom wrap copilot` command."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import importlib
|
| 6 |
-
import sys
|
| 7 |
-
import types
|
| 8 |
-
from pathlib import Path
|
| 9 |
-
from unittest.mock import patch
|
| 10 |
-
|
| 11 |
-
import click
|
| 12 |
-
import pytest
|
| 13 |
-
from click.testing import CliRunner
|
| 14 |
-
|
| 15 |
-
from headroom.copilot_auth import DEFAULT_API_URL
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
@pytest.fixture
|
| 19 |
-
def runner() -> CliRunner:
|
| 20 |
-
return CliRunner()
|
| 21 |
-
|
| 22 |
-
|
| 23 |
-
@pytest.fixture
|
| 24 |
-
def wrap_modules(monkeypatch: pytest.MonkeyPatch) -> tuple[types.ModuleType, click.Group]:
|
| 25 |
-
headroom_pkg = sys.modules.get("headroom")
|
| 26 |
-
saved_headroom_cli_attr = (
|
| 27 |
-
headroom_pkg.cli if headroom_pkg is not None and hasattr(headroom_pkg, "cli") else None
|
| 28 |
-
)
|
| 29 |
-
saved_modules = {
|
| 30 |
-
name: sys.modules.get(name)
|
| 31 |
-
for name in ("headroom.cli", "headroom.cli.main", "headroom.cli.wrap")
|
| 32 |
-
}
|
| 33 |
-
|
| 34 |
-
fake_main_module = types.ModuleType("headroom.cli.main")
|
| 35 |
-
fake_main_module.main = click.Group()
|
| 36 |
-
sys.modules["headroom.cli.main"] = fake_main_module
|
| 37 |
-
sys.modules.pop("headroom.cli", None)
|
| 38 |
-
sys.modules.pop("headroom.cli.wrap", None)
|
| 39 |
-
|
| 40 |
-
wrap_cli = importlib.import_module("headroom.cli.wrap")
|
| 41 |
-
monkeypatch.setattr(wrap_cli, "_check_proxy", lambda _port: False)
|
| 42 |
-
|
| 43 |
-
try:
|
| 44 |
-
yield wrap_cli, fake_main_module.main
|
| 45 |
-
finally:
|
| 46 |
-
for name in ("headroom.cli.wrap", "headroom.cli.main", "headroom.cli"):
|
| 47 |
-
sys.modules.pop(name, None)
|
| 48 |
-
for name, module in saved_modules.items():
|
| 49 |
-
if module is not None:
|
| 50 |
-
sys.modules[name] = module
|
| 51 |
-
if saved_modules["headroom.cli"] is not None:
|
| 52 |
-
cli_pkg = saved_modules["headroom.cli"]
|
| 53 |
-
if saved_modules["headroom.cli.main"] is not None:
|
| 54 |
-
cli_pkg.main = saved_modules["headroom.cli.main"]
|
| 55 |
-
if saved_modules["headroom.cli.wrap"] is not None:
|
| 56 |
-
cli_pkg.wrap = saved_modules["headroom.cli.wrap"]
|
| 57 |
-
if headroom_pkg is not None:
|
| 58 |
-
if saved_headroom_cli_attr is None:
|
| 59 |
-
if hasattr(headroom_pkg, "cli"):
|
| 60 |
-
delattr(headroom_pkg, "cli")
|
| 61 |
-
else:
|
| 62 |
-
headroom_pkg.cli = saved_headroom_cli_attr
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
def test_wrap_copilot_auto_anthropic_injects_instructions(
|
| 66 |
-
runner: CliRunner,
|
| 67 |
-
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 68 |
-
tmp_path: Path,
|
| 69 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 70 |
-
) -> None:
|
| 71 |
-
wrap_cli, main = wrap_modules
|
| 72 |
-
monkeypatch.chdir(tmp_path)
|
| 73 |
-
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test-dummy")
|
| 74 |
-
captured: dict[str, object] = {}
|
| 75 |
-
|
| 76 |
-
def fake_launch_tool(**kwargs): # noqa: ANN003
|
| 77 |
-
captured.update(kwargs)
|
| 78 |
-
|
| 79 |
-
with (
|
| 80 |
-
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
| 81 |
-
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
| 82 |
-
patch("headroom.cli.wrap._ensure_rtk_binary", return_value=Path("/tmp/rtk")),
|
| 83 |
-
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
| 84 |
-
):
|
| 85 |
-
result = runner.invoke(
|
| 86 |
-
main,
|
| 87 |
-
["wrap", "copilot", "--", "--model", "claude-sonnet-4-20250514"],
|
| 88 |
-
)
|
| 89 |
-
|
| 90 |
-
assert result.exit_code == 0, result.output
|
| 91 |
-
instructions = tmp_path / ".github" / "copilot-instructions.md"
|
| 92 |
-
assert instructions.exists()
|
| 93 |
-
content = instructions.read_text()
|
| 94 |
-
assert wrap_cli._RTK_MARKER in content
|
| 95 |
-
assert "RTK (Rust Token Killer)" in content
|
| 96 |
-
|
| 97 |
-
env = captured["env"]
|
| 98 |
-
assert isinstance(env, dict)
|
| 99 |
-
assert env["COPILOT_PROVIDER_TYPE"] == "anthropic"
|
| 100 |
-
assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787"
|
| 101 |
-
assert "COPILOT_PROVIDER_WIRE_API" not in env
|
| 102 |
-
assert captured["agent_type"] == "copilot"
|
| 103 |
-
assert captured["tool_label"] == "COPILOT"
|
| 104 |
-
assert captured["args"] == ("--model", "claude-sonnet-4-20250514")
|
| 105 |
-
|
| 106 |
-
|
| 107 |
-
def test_wrap_copilot_openai_backend_sets_completions_env(
|
| 108 |
-
runner: CliRunner,
|
| 109 |
-
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 110 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 111 |
-
) -> None:
|
| 112 |
-
_wrap_cli, main = wrap_modules
|
| 113 |
-
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-dummy")
|
| 114 |
-
captured: dict[str, object] = {}
|
| 115 |
-
|
| 116 |
-
def fake_launch_tool(**kwargs): # noqa: ANN003
|
| 117 |
-
captured.update(kwargs)
|
| 118 |
-
|
| 119 |
-
with (
|
| 120 |
-
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
| 121 |
-
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
| 122 |
-
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
| 123 |
-
):
|
| 124 |
-
result = runner.invoke(
|
| 125 |
-
main,
|
| 126 |
-
[
|
| 127 |
-
"wrap",
|
| 128 |
-
"copilot",
|
| 129 |
-
"--no-rtk",
|
| 130 |
-
"--backend",
|
| 131 |
-
"anyllm",
|
| 132 |
-
"--anyllm-provider",
|
| 133 |
-
"groq",
|
| 134 |
-
"--region",
|
| 135 |
-
"us-central1",
|
| 136 |
-
"--",
|
| 137 |
-
"--model",
|
| 138 |
-
"gpt-4o",
|
| 139 |
-
],
|
| 140 |
-
)
|
| 141 |
-
|
| 142 |
-
assert result.exit_code == 0, result.output
|
| 143 |
-
|
| 144 |
-
env = captured["env"]
|
| 145 |
-
assert isinstance(env, dict)
|
| 146 |
-
assert env["COPILOT_PROVIDER_TYPE"] == "openai"
|
| 147 |
-
assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1"
|
| 148 |
-
assert env["COPILOT_PROVIDER_WIRE_API"] == "completions"
|
| 149 |
-
assert captured["backend"] == "anyllm"
|
| 150 |
-
assert captured["anyllm_provider"] == "groq"
|
| 151 |
-
assert captured["region"] == "us-central1"
|
| 152 |
-
assert captured["args"] == ("--model", "gpt-4o")
|
| 153 |
-
|
| 154 |
-
|
| 155 |
-
def test_wrap_copilot_auto_detects_running_proxy_backend(
|
| 156 |
-
runner: CliRunner,
|
| 157 |
-
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 158 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 159 |
-
) -> None:
|
| 160 |
-
_wrap_cli, main = wrap_modules
|
| 161 |
-
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-dummy")
|
| 162 |
-
captured: dict[str, object] = {}
|
| 163 |
-
|
| 164 |
-
def fake_launch_tool(**kwargs): # noqa: ANN003
|
| 165 |
-
captured.update(kwargs)
|
| 166 |
-
|
| 167 |
-
with (
|
| 168 |
-
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
| 169 |
-
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
| 170 |
-
patch("headroom.cli.wrap._check_proxy", return_value=True),
|
| 171 |
-
patch("headroom.cli.wrap._detect_running_proxy_backend", return_value="anyllm"),
|
| 172 |
-
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
| 173 |
-
):
|
| 174 |
-
result = runner.invoke(
|
| 175 |
-
main,
|
| 176 |
-
["wrap", "copilot", "--no-rtk", "--", "--model", "gpt-4o"],
|
| 177 |
-
)
|
| 178 |
-
|
| 179 |
-
assert result.exit_code == 0, result.output
|
| 180 |
-
env = captured["env"]
|
| 181 |
-
assert isinstance(env, dict)
|
| 182 |
-
assert env["COPILOT_PROVIDER_TYPE"] == "openai"
|
| 183 |
-
assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1"
|
| 184 |
-
assert env["COPILOT_PROVIDER_WIRE_API"] == "completions"
|
| 185 |
-
|
| 186 |
-
|
| 187 |
-
def test_wrap_copilot_prefers_existing_oauth_session(
|
| 188 |
-
runner: CliRunner,
|
| 189 |
-
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 190 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 191 |
-
) -> None:
|
| 192 |
-
_wrap_cli, main = wrap_modules
|
| 193 |
-
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test-dummy")
|
| 194 |
-
captured: dict[str, object] = {}
|
| 195 |
-
|
| 196 |
-
def fake_launch_tool(**kwargs): # noqa: ANN003
|
| 197 |
-
captured.update(kwargs)
|
| 198 |
-
|
| 199 |
-
with patch("headroom.cli.wrap.shutil.which", return_value="copilot"):
|
| 200 |
-
with patch("headroom.cli.wrap.resolve_client_bearer_token", return_value="gho-existing"):
|
| 201 |
-
with patch("headroom.cli.wrap.has_oauth_auth", return_value=True):
|
| 202 |
-
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool):
|
| 203 |
-
result = runner.invoke(
|
| 204 |
-
main,
|
| 205 |
-
["wrap", "copilot", "--no-rtk", "--", "--model", "claude-sonnet-4.6"],
|
| 206 |
-
)
|
| 207 |
-
|
| 208 |
-
assert result.exit_code == 0, result.output
|
| 209 |
-
env = captured["env"]
|
| 210 |
-
assert isinstance(env, dict)
|
| 211 |
-
assert env["COPILOT_PROVIDER_TYPE"] == "openai"
|
| 212 |
-
assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1"
|
| 213 |
-
assert env["COPILOT_PROVIDER_WIRE_API"] == "completions"
|
| 214 |
-
assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "gho-existing"
|
| 215 |
-
assert "COPILOT_PROVIDER_API_KEY" not in env
|
| 216 |
-
assert captured["openai_api_url"] == DEFAULT_API_URL
|
| 217 |
-
|
| 218 |
-
|
| 219 |
-
def test_wrap_copilot_translated_backend_still_requires_byok(
|
| 220 |
-
runner: CliRunner,
|
| 221 |
-
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 222 |
-
) -> None:
|
| 223 |
-
_wrap_cli, main = wrap_modules
|
| 224 |
-
with patch("headroom.cli.wrap.shutil.which", return_value="copilot"):
|
| 225 |
-
with patch("headroom.cli.wrap.has_oauth_auth", return_value=True):
|
| 226 |
-
result = runner.invoke(
|
| 227 |
-
main,
|
| 228 |
-
[
|
| 229 |
-
"wrap",
|
| 230 |
-
"copilot",
|
| 231 |
-
"--no-rtk",
|
| 232 |
-
"--backend",
|
| 233 |
-
"anyllm",
|
| 234 |
-
"--",
|
| 235 |
-
"--model",
|
| 236 |
-
"gpt-4o",
|
| 237 |
-
],
|
| 238 |
-
)
|
| 239 |
-
|
| 240 |
-
assert result.exit_code == 1
|
| 241 |
-
assert "Copilot BYOK mode requires a provider API key" in result.output
|
| 242 |
-
|
| 243 |
-
|
| 244 |
-
def test_wrap_copilot_rejects_wire_api_for_anthropic_provider(
|
| 245 |
-
runner: CliRunner,
|
| 246 |
-
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 247 |
-
) -> None:
|
| 248 |
-
_wrap_cli, main = wrap_modules
|
| 249 |
-
with patch("headroom.cli.wrap.shutil.which", return_value="copilot"):
|
| 250 |
-
result = runner.invoke(
|
| 251 |
-
main,
|
| 252 |
-
[
|
| 253 |
-
"wrap",
|
| 254 |
-
"copilot",
|
| 255 |
-
"--wire-api",
|
| 256 |
-
"responses",
|
| 257 |
-
"--",
|
| 258 |
-
"--model",
|
| 259 |
-
"claude-sonnet-4-20250514",
|
| 260 |
-
],
|
| 261 |
-
)
|
| 262 |
-
|
| 263 |
-
assert result.exit_code != 0
|
| 264 |
-
assert "--wire-api is only valid" in result.output
|
| 265 |
-
|
| 266 |
-
|
| 267 |
-
def test_wrap_copilot_rejects_responses_for_translated_backends(
|
| 268 |
-
runner: CliRunner,
|
| 269 |
-
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 270 |
-
) -> None:
|
| 271 |
-
_wrap_cli, main = wrap_modules
|
| 272 |
-
with patch("headroom.cli.wrap.shutil.which", return_value="copilot"):
|
| 273 |
-
result = runner.invoke(
|
| 274 |
-
main,
|
| 275 |
-
[
|
| 276 |
-
"wrap",
|
| 277 |
-
"copilot",
|
| 278 |
-
"--backend",
|
| 279 |
-
"anyllm",
|
| 280 |
-
"--wire-api",
|
| 281 |
-
"responses",
|
| 282 |
-
"--",
|
| 283 |
-
"--model",
|
| 284 |
-
"gpt-4o",
|
| 285 |
-
],
|
| 286 |
-
)
|
| 287 |
-
|
| 288 |
-
assert result.exit_code != 0
|
| 289 |
-
assert "not supported with translated backends" in result.output
|
| 290 |
-
|
| 291 |
-
|
| 292 |
-
def test_wrap_copilot_clears_stale_wire_api_in_anthropic_mode(
|
| 293 |
-
runner: CliRunner,
|
| 294 |
-
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 295 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 296 |
-
) -> None:
|
| 297 |
-
_wrap_cli, main = wrap_modules
|
| 298 |
-
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test-dummy")
|
| 299 |
-
captured: dict[str, object] = {}
|
| 300 |
-
|
| 301 |
-
def fake_launch_tool(**kwargs): # noqa: ANN003
|
| 302 |
-
captured.update(kwargs)
|
| 303 |
-
|
| 304 |
-
with (
|
| 305 |
-
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
| 306 |
-
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
| 307 |
-
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
| 308 |
-
):
|
| 309 |
-
result = runner.invoke(
|
| 310 |
-
main,
|
| 311 |
-
["wrap", "copilot", "--no-rtk", "--", "--model", "claude-sonnet-4-20250514"],
|
| 312 |
-
env={
|
| 313 |
-
"COPILOT_PROVIDER_WIRE_API": "responses",
|
| 314 |
-
"ANTHROPIC_API_KEY": "sk-test-dummy",
|
| 315 |
-
},
|
| 316 |
-
)
|
| 317 |
-
|
| 318 |
-
assert result.exit_code == 0, result.output
|
| 319 |
-
env = captured["env"]
|
| 320 |
-
assert isinstance(env, dict)
|
| 321 |
-
assert env["COPILOT_PROVIDER_TYPE"] == "anthropic"
|
| 322 |
-
assert "COPILOT_PROVIDER_WIRE_API" not in env
|
| 323 |
-
|
| 324 |
-
|
| 325 |
-
def test_wrap_copilot_fails_when_binary_missing(
|
| 326 |
-
runner: CliRunner,
|
| 327 |
-
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 328 |
-
) -> None:
|
| 329 |
-
_wrap_cli, main = wrap_modules
|
| 330 |
-
with patch("headroom.cli.wrap.shutil.which", return_value=None):
|
| 331 |
-
result = runner.invoke(main, ["wrap", "copilot", "--", "--model", "gpt-4o"])
|
| 332 |
-
|
| 333 |
-
assert result.exit_code == 1
|
| 334 |
-
assert "'copilot' not found in PATH" in result.output
|
| 335 |
-
assert "Install GitHub Copilot CLI" in result.output
|
|
|
|
| 1 |
+
"""Tests for `headroom wrap copilot` command."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import importlib
|
| 6 |
+
import sys
|
| 7 |
+
import types
|
| 8 |
+
from pathlib import Path
|
| 9 |
+
from unittest.mock import patch
|
| 10 |
+
|
| 11 |
+
import click
|
| 12 |
+
import pytest
|
| 13 |
+
from click.testing import CliRunner
|
| 14 |
+
|
| 15 |
+
from headroom.copilot_auth import DEFAULT_API_URL
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
@pytest.fixture
|
| 19 |
+
def runner() -> CliRunner:
|
| 20 |
+
return CliRunner()
|
| 21 |
+
|
| 22 |
+
|
| 23 |
+
@pytest.fixture
|
| 24 |
+
def wrap_modules(monkeypatch: pytest.MonkeyPatch) -> tuple[types.ModuleType, click.Group]:
|
| 25 |
+
headroom_pkg = sys.modules.get("headroom")
|
| 26 |
+
saved_headroom_cli_attr = (
|
| 27 |
+
headroom_pkg.cli if headroom_pkg is not None and hasattr(headroom_pkg, "cli") else None
|
| 28 |
+
)
|
| 29 |
+
saved_modules = {
|
| 30 |
+
name: sys.modules.get(name)
|
| 31 |
+
for name in ("headroom.cli", "headroom.cli.main", "headroom.cli.wrap")
|
| 32 |
+
}
|
| 33 |
+
|
| 34 |
+
fake_main_module = types.ModuleType("headroom.cli.main")
|
| 35 |
+
fake_main_module.main = click.Group()
|
| 36 |
+
sys.modules["headroom.cli.main"] = fake_main_module
|
| 37 |
+
sys.modules.pop("headroom.cli", None)
|
| 38 |
+
sys.modules.pop("headroom.cli.wrap", None)
|
| 39 |
+
|
| 40 |
+
wrap_cli = importlib.import_module("headroom.cli.wrap")
|
| 41 |
+
monkeypatch.setattr(wrap_cli, "_check_proxy", lambda _port: False)
|
| 42 |
+
|
| 43 |
+
try:
|
| 44 |
+
yield wrap_cli, fake_main_module.main
|
| 45 |
+
finally:
|
| 46 |
+
for name in ("headroom.cli.wrap", "headroom.cli.main", "headroom.cli"):
|
| 47 |
+
sys.modules.pop(name, None)
|
| 48 |
+
for name, module in saved_modules.items():
|
| 49 |
+
if module is not None:
|
| 50 |
+
sys.modules[name] = module
|
| 51 |
+
if saved_modules["headroom.cli"] is not None:
|
| 52 |
+
cli_pkg = saved_modules["headroom.cli"]
|
| 53 |
+
if saved_modules["headroom.cli.main"] is not None:
|
| 54 |
+
cli_pkg.main = saved_modules["headroom.cli.main"]
|
| 55 |
+
if saved_modules["headroom.cli.wrap"] is not None:
|
| 56 |
+
cli_pkg.wrap = saved_modules["headroom.cli.wrap"]
|
| 57 |
+
if headroom_pkg is not None:
|
| 58 |
+
if saved_headroom_cli_attr is None:
|
| 59 |
+
if hasattr(headroom_pkg, "cli"):
|
| 60 |
+
delattr(headroom_pkg, "cli")
|
| 61 |
+
else:
|
| 62 |
+
headroom_pkg.cli = saved_headroom_cli_attr
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
def test_wrap_copilot_auto_anthropic_injects_instructions(
|
| 66 |
+
runner: CliRunner,
|
| 67 |
+
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 68 |
+
tmp_path: Path,
|
| 69 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 70 |
+
) -> None:
|
| 71 |
+
wrap_cli, main = wrap_modules
|
| 72 |
+
monkeypatch.chdir(tmp_path)
|
| 73 |
+
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test-dummy")
|
| 74 |
+
captured: dict[str, object] = {}
|
| 75 |
+
|
| 76 |
+
def fake_launch_tool(**kwargs): # noqa: ANN003
|
| 77 |
+
captured.update(kwargs)
|
| 78 |
+
|
| 79 |
+
with (
|
| 80 |
+
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
| 81 |
+
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
| 82 |
+
patch("headroom.cli.wrap._ensure_rtk_binary", return_value=Path("/tmp/rtk")),
|
| 83 |
+
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
| 84 |
+
):
|
| 85 |
+
result = runner.invoke(
|
| 86 |
+
main,
|
| 87 |
+
["wrap", "copilot", "--", "--model", "claude-sonnet-4-20250514"],
|
| 88 |
+
)
|
| 89 |
+
|
| 90 |
+
assert result.exit_code == 0, result.output
|
| 91 |
+
instructions = tmp_path / ".github" / "copilot-instructions.md"
|
| 92 |
+
assert instructions.exists()
|
| 93 |
+
content = instructions.read_text()
|
| 94 |
+
assert wrap_cli._RTK_MARKER in content
|
| 95 |
+
assert "RTK (Rust Token Killer)" in content
|
| 96 |
+
|
| 97 |
+
env = captured["env"]
|
| 98 |
+
assert isinstance(env, dict)
|
| 99 |
+
assert env["COPILOT_PROVIDER_TYPE"] == "anthropic"
|
| 100 |
+
assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787"
|
| 101 |
+
assert "COPILOT_PROVIDER_WIRE_API" not in env
|
| 102 |
+
assert captured["agent_type"] == "copilot"
|
| 103 |
+
assert captured["tool_label"] == "COPILOT"
|
| 104 |
+
assert captured["args"] == ("--model", "claude-sonnet-4-20250514")
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
def test_wrap_copilot_openai_backend_sets_completions_env(
|
| 108 |
+
runner: CliRunner,
|
| 109 |
+
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 110 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 111 |
+
) -> None:
|
| 112 |
+
_wrap_cli, main = wrap_modules
|
| 113 |
+
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-dummy")
|
| 114 |
+
captured: dict[str, object] = {}
|
| 115 |
+
|
| 116 |
+
def fake_launch_tool(**kwargs): # noqa: ANN003
|
| 117 |
+
captured.update(kwargs)
|
| 118 |
+
|
| 119 |
+
with (
|
| 120 |
+
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
| 121 |
+
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
| 122 |
+
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
| 123 |
+
):
|
| 124 |
+
result = runner.invoke(
|
| 125 |
+
main,
|
| 126 |
+
[
|
| 127 |
+
"wrap",
|
| 128 |
+
"copilot",
|
| 129 |
+
"--no-rtk",
|
| 130 |
+
"--backend",
|
| 131 |
+
"anyllm",
|
| 132 |
+
"--anyllm-provider",
|
| 133 |
+
"groq",
|
| 134 |
+
"--region",
|
| 135 |
+
"us-central1",
|
| 136 |
+
"--",
|
| 137 |
+
"--model",
|
| 138 |
+
"gpt-4o",
|
| 139 |
+
],
|
| 140 |
+
)
|
| 141 |
+
|
| 142 |
+
assert result.exit_code == 0, result.output
|
| 143 |
+
|
| 144 |
+
env = captured["env"]
|
| 145 |
+
assert isinstance(env, dict)
|
| 146 |
+
assert env["COPILOT_PROVIDER_TYPE"] == "openai"
|
| 147 |
+
assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1"
|
| 148 |
+
assert env["COPILOT_PROVIDER_WIRE_API"] == "completions"
|
| 149 |
+
assert captured["backend"] == "anyllm"
|
| 150 |
+
assert captured["anyllm_provider"] == "groq"
|
| 151 |
+
assert captured["region"] == "us-central1"
|
| 152 |
+
assert captured["args"] == ("--model", "gpt-4o")
|
| 153 |
+
|
| 154 |
+
|
| 155 |
+
def test_wrap_copilot_auto_detects_running_proxy_backend(
|
| 156 |
+
runner: CliRunner,
|
| 157 |
+
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 158 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 159 |
+
) -> None:
|
| 160 |
+
_wrap_cli, main = wrap_modules
|
| 161 |
+
monkeypatch.setenv("OPENAI_API_KEY", "sk-test-dummy")
|
| 162 |
+
captured: dict[str, object] = {}
|
| 163 |
+
|
| 164 |
+
def fake_launch_tool(**kwargs): # noqa: ANN003
|
| 165 |
+
captured.update(kwargs)
|
| 166 |
+
|
| 167 |
+
with (
|
| 168 |
+
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
| 169 |
+
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
| 170 |
+
patch("headroom.cli.wrap._check_proxy", return_value=True),
|
| 171 |
+
patch("headroom.cli.wrap._detect_running_proxy_backend", return_value="anyllm"),
|
| 172 |
+
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
| 173 |
+
):
|
| 174 |
+
result = runner.invoke(
|
| 175 |
+
main,
|
| 176 |
+
["wrap", "copilot", "--no-rtk", "--", "--model", "gpt-4o"],
|
| 177 |
+
)
|
| 178 |
+
|
| 179 |
+
assert result.exit_code == 0, result.output
|
| 180 |
+
env = captured["env"]
|
| 181 |
+
assert isinstance(env, dict)
|
| 182 |
+
assert env["COPILOT_PROVIDER_TYPE"] == "openai"
|
| 183 |
+
assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1"
|
| 184 |
+
assert env["COPILOT_PROVIDER_WIRE_API"] == "completions"
|
| 185 |
+
|
| 186 |
+
|
| 187 |
+
def test_wrap_copilot_prefers_existing_oauth_session(
|
| 188 |
+
runner: CliRunner,
|
| 189 |
+
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 190 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 191 |
+
) -> None:
|
| 192 |
+
_wrap_cli, main = wrap_modules
|
| 193 |
+
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test-dummy")
|
| 194 |
+
captured: dict[str, object] = {}
|
| 195 |
+
|
| 196 |
+
def fake_launch_tool(**kwargs): # noqa: ANN003
|
| 197 |
+
captured.update(kwargs)
|
| 198 |
+
|
| 199 |
+
with patch("headroom.cli.wrap.shutil.which", return_value="copilot"):
|
| 200 |
+
with patch("headroom.cli.wrap.resolve_client_bearer_token", return_value="gho-existing"):
|
| 201 |
+
with patch("headroom.cli.wrap.has_oauth_auth", return_value=True):
|
| 202 |
+
with patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool):
|
| 203 |
+
result = runner.invoke(
|
| 204 |
+
main,
|
| 205 |
+
["wrap", "copilot", "--no-rtk", "--", "--model", "claude-sonnet-4.6"],
|
| 206 |
+
)
|
| 207 |
+
|
| 208 |
+
assert result.exit_code == 0, result.output
|
| 209 |
+
env = captured["env"]
|
| 210 |
+
assert isinstance(env, dict)
|
| 211 |
+
assert env["COPILOT_PROVIDER_TYPE"] == "openai"
|
| 212 |
+
assert env["COPILOT_PROVIDER_BASE_URL"] == "http://127.0.0.1:8787/v1"
|
| 213 |
+
assert env["COPILOT_PROVIDER_WIRE_API"] == "completions"
|
| 214 |
+
assert env["COPILOT_PROVIDER_BEARER_TOKEN"] == "gho-existing"
|
| 215 |
+
assert "COPILOT_PROVIDER_API_KEY" not in env
|
| 216 |
+
assert captured["openai_api_url"] == DEFAULT_API_URL
|
| 217 |
+
|
| 218 |
+
|
| 219 |
+
def test_wrap_copilot_translated_backend_still_requires_byok(
|
| 220 |
+
runner: CliRunner,
|
| 221 |
+
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 222 |
+
) -> None:
|
| 223 |
+
_wrap_cli, main = wrap_modules
|
| 224 |
+
with patch("headroom.cli.wrap.shutil.which", return_value="copilot"):
|
| 225 |
+
with patch("headroom.cli.wrap.has_oauth_auth", return_value=True):
|
| 226 |
+
result = runner.invoke(
|
| 227 |
+
main,
|
| 228 |
+
[
|
| 229 |
+
"wrap",
|
| 230 |
+
"copilot",
|
| 231 |
+
"--no-rtk",
|
| 232 |
+
"--backend",
|
| 233 |
+
"anyllm",
|
| 234 |
+
"--",
|
| 235 |
+
"--model",
|
| 236 |
+
"gpt-4o",
|
| 237 |
+
],
|
| 238 |
+
)
|
| 239 |
+
|
| 240 |
+
assert result.exit_code == 1
|
| 241 |
+
assert "Copilot BYOK mode requires a provider API key" in result.output
|
| 242 |
+
|
| 243 |
+
|
| 244 |
+
def test_wrap_copilot_rejects_wire_api_for_anthropic_provider(
|
| 245 |
+
runner: CliRunner,
|
| 246 |
+
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 247 |
+
) -> None:
|
| 248 |
+
_wrap_cli, main = wrap_modules
|
| 249 |
+
with patch("headroom.cli.wrap.shutil.which", return_value="copilot"):
|
| 250 |
+
result = runner.invoke(
|
| 251 |
+
main,
|
| 252 |
+
[
|
| 253 |
+
"wrap",
|
| 254 |
+
"copilot",
|
| 255 |
+
"--wire-api",
|
| 256 |
+
"responses",
|
| 257 |
+
"--",
|
| 258 |
+
"--model",
|
| 259 |
+
"claude-sonnet-4-20250514",
|
| 260 |
+
],
|
| 261 |
+
)
|
| 262 |
+
|
| 263 |
+
assert result.exit_code != 0
|
| 264 |
+
assert "--wire-api is only valid" in result.output
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
def test_wrap_copilot_rejects_responses_for_translated_backends(
|
| 268 |
+
runner: CliRunner,
|
| 269 |
+
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 270 |
+
) -> None:
|
| 271 |
+
_wrap_cli, main = wrap_modules
|
| 272 |
+
with patch("headroom.cli.wrap.shutil.which", return_value="copilot"):
|
| 273 |
+
result = runner.invoke(
|
| 274 |
+
main,
|
| 275 |
+
[
|
| 276 |
+
"wrap",
|
| 277 |
+
"copilot",
|
| 278 |
+
"--backend",
|
| 279 |
+
"anyllm",
|
| 280 |
+
"--wire-api",
|
| 281 |
+
"responses",
|
| 282 |
+
"--",
|
| 283 |
+
"--model",
|
| 284 |
+
"gpt-4o",
|
| 285 |
+
],
|
| 286 |
+
)
|
| 287 |
+
|
| 288 |
+
assert result.exit_code != 0
|
| 289 |
+
assert "not supported with translated backends" in result.output
|
| 290 |
+
|
| 291 |
+
|
| 292 |
+
def test_wrap_copilot_clears_stale_wire_api_in_anthropic_mode(
|
| 293 |
+
runner: CliRunner,
|
| 294 |
+
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 295 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 296 |
+
) -> None:
|
| 297 |
+
_wrap_cli, main = wrap_modules
|
| 298 |
+
monkeypatch.setenv("ANTHROPIC_API_KEY", "sk-test-dummy")
|
| 299 |
+
captured: dict[str, object] = {}
|
| 300 |
+
|
| 301 |
+
def fake_launch_tool(**kwargs): # noqa: ANN003
|
| 302 |
+
captured.update(kwargs)
|
| 303 |
+
|
| 304 |
+
with (
|
| 305 |
+
patch("headroom.cli.wrap.shutil.which", return_value="copilot"),
|
| 306 |
+
patch("headroom.cli.wrap.has_oauth_auth", return_value=False),
|
| 307 |
+
patch("headroom.cli.wrap._launch_tool", side_effect=fake_launch_tool),
|
| 308 |
+
):
|
| 309 |
+
result = runner.invoke(
|
| 310 |
+
main,
|
| 311 |
+
["wrap", "copilot", "--no-rtk", "--", "--model", "claude-sonnet-4-20250514"],
|
| 312 |
+
env={
|
| 313 |
+
"COPILOT_PROVIDER_WIRE_API": "responses",
|
| 314 |
+
"ANTHROPIC_API_KEY": "sk-test-dummy",
|
| 315 |
+
},
|
| 316 |
+
)
|
| 317 |
+
|
| 318 |
+
assert result.exit_code == 0, result.output
|
| 319 |
+
env = captured["env"]
|
| 320 |
+
assert isinstance(env, dict)
|
| 321 |
+
assert env["COPILOT_PROVIDER_TYPE"] == "anthropic"
|
| 322 |
+
assert "COPILOT_PROVIDER_WIRE_API" not in env
|
| 323 |
+
|
| 324 |
+
|
| 325 |
+
def test_wrap_copilot_fails_when_binary_missing(
|
| 326 |
+
runner: CliRunner,
|
| 327 |
+
wrap_modules: tuple[types.ModuleType, click.Group],
|
| 328 |
+
) -> None:
|
| 329 |
+
_wrap_cli, main = wrap_modules
|
| 330 |
+
with patch("headroom.cli.wrap.shutil.which", return_value=None):
|
| 331 |
+
result = runner.invoke(main, ["wrap", "copilot", "--", "--model", "gpt-4o"])
|
| 332 |
+
|
| 333 |
+
assert result.exit_code == 1
|
| 334 |
+
assert "'copilot' not found in PATH" in result.output
|
| 335 |
+
assert "Install GitHub Copilot CLI" in result.output
|
|
@@ -1,266 +1,266 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
from pathlib import Path
|
| 4 |
-
from types import SimpleNamespace
|
| 5 |
-
|
| 6 |
-
import click
|
| 7 |
-
import click.shell_completion as click_shell_completion
|
| 8 |
-
import pytest
|
| 9 |
-
from click.testing import CliRunner
|
| 10 |
-
|
| 11 |
-
from headroom.cli.learn import _AgentChoice
|
| 12 |
-
from headroom.cli.main import main
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
@pytest.fixture
|
| 16 |
-
def runner() -> CliRunner:
|
| 17 |
-
return CliRunner()
|
| 18 |
-
|
| 19 |
-
|
| 20 |
-
class FakeWriter:
|
| 21 |
-
def __init__(self) -> None:
|
| 22 |
-
self.calls: list[tuple[list[object], object, bool]] = []
|
| 23 |
-
|
| 24 |
-
def write(self, recommendations, project, dry_run: bool): # noqa: ANN001, ANN201
|
| 25 |
-
self.calls.append((recommendations, project, dry_run))
|
| 26 |
-
return SimpleNamespace(
|
| 27 |
-
dry_run=dry_run,
|
| 28 |
-
content_by_file={
|
| 29 |
-
Path(project.project_path) / "AGENTS.md": "<!-- headroom -->\nRule 1\nRule 2"
|
| 30 |
-
},
|
| 31 |
-
)
|
| 32 |
-
|
| 33 |
-
|
| 34 |
-
class FakePlugin:
|
| 35 |
-
def __init__(self, name: str, display_name: str, projects: list[object]) -> None:
|
| 36 |
-
self.name = name
|
| 37 |
-
self.display_name = display_name
|
| 38 |
-
self._projects = projects
|
| 39 |
-
self.writer = FakeWriter()
|
| 40 |
-
self.scan_calls: list[tuple[object, int]] = []
|
| 41 |
-
|
| 42 |
-
def detect(self) -> bool:
|
| 43 |
-
return True
|
| 44 |
-
|
| 45 |
-
def create_writer(self) -> FakeWriter:
|
| 46 |
-
return self.writer
|
| 47 |
-
|
| 48 |
-
def discover_projects(self) -> list[object]:
|
| 49 |
-
return self._projects
|
| 50 |
-
|
| 51 |
-
def scan_project(self, project, max_workers: int = 1): # noqa: ANN001, ANN201
|
| 52 |
-
self.scan_calls.append((project, max_workers))
|
| 53 |
-
return [SimpleNamespace(events=["event"], tool_calls=[], failure_count=0)]
|
| 54 |
-
|
| 55 |
-
|
| 56 |
-
class FakeAnalyzer:
|
| 57 |
-
def __init__(self, model: str | None = None) -> None:
|
| 58 |
-
self.model = model
|
| 59 |
-
self.calls: list[tuple[object, list[object]]] = []
|
| 60 |
-
|
| 61 |
-
def analyze(self, project, sessions): # noqa: ANN001, ANN201
|
| 62 |
-
self.calls.append((project, sessions))
|
| 63 |
-
return SimpleNamespace(
|
| 64 |
-
total_sessions=len(sessions),
|
| 65 |
-
total_calls=3,
|
| 66 |
-
total_failures=1,
|
| 67 |
-
failure_rate=1 / 3,
|
| 68 |
-
recommendations=[SimpleNamespace(section="Rules")],
|
| 69 |
-
)
|
| 70 |
-
|
| 71 |
-
|
| 72 |
-
def test_agent_choice_convert_and_shell_complete(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 73 |
-
choice = _AgentChoice()
|
| 74 |
-
monkeypatch.setattr(click, "shell_completion", click_shell_completion)
|
| 75 |
-
monkeypatch.setattr(
|
| 76 |
-
"headroom.learn.registry.get_registry",
|
| 77 |
-
lambda: {"codex": object(), "claude": object()},
|
| 78 |
-
)
|
| 79 |
-
monkeypatch.setattr(
|
| 80 |
-
"headroom.learn.registry.available_agent_names",
|
| 81 |
-
lambda: ["claude", "codex"],
|
| 82 |
-
)
|
| 83 |
-
|
| 84 |
-
assert choice.convert("auto", None, None) == "auto"
|
| 85 |
-
assert choice.convert("CODEX", None, None) == "codex"
|
| 86 |
-
with pytest.raises(Exception, match="Unknown agent: bad"):
|
| 87 |
-
choice.convert("bad", None, None)
|
| 88 |
-
|
| 89 |
-
completions = choice.shell_complete(None, None, "c") # type: ignore[arg-type]
|
| 90 |
-
assert [item.value for item in completions] == ["claude", "codex"]
|
| 91 |
-
assert choice.get_metavar(None) == "[auto|<agent>]" # type: ignore[arg-type]
|
| 92 |
-
|
| 93 |
-
|
| 94 |
-
def test_learn_exits_cleanly_when_model_detection_fails(
|
| 95 |
-
monkeypatch: pytest.MonkeyPatch, runner: CliRunner
|
| 96 |
-
) -> None:
|
| 97 |
-
monkeypatch.setattr(
|
| 98 |
-
"headroom.learn.analyzer._detect_default_model",
|
| 99 |
-
lambda: (_ for _ in ()).throw(RuntimeError("no model")),
|
| 100 |
-
)
|
| 101 |
-
|
| 102 |
-
result = runner.invoke(main, ["learn"], catch_exceptions=False)
|
| 103 |
-
|
| 104 |
-
assert result.exit_code == 1
|
| 105 |
-
assert "Error: no model" in result.output
|
| 106 |
-
|
| 107 |
-
|
| 108 |
-
def test_learn_auto_agent_reports_no_detected_plugins(
|
| 109 |
-
monkeypatch: pytest.MonkeyPatch, runner: CliRunner
|
| 110 |
-
) -> None:
|
| 111 |
-
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 112 |
-
monkeypatch.setattr("headroom.learn.registry.auto_detect_plugins", lambda: [])
|
| 113 |
-
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", FakeAnalyzer)
|
| 114 |
-
|
| 115 |
-
result = runner.invoke(main, ["learn"], catch_exceptions=False)
|
| 116 |
-
|
| 117 |
-
assert result.exit_code == 0
|
| 118 |
-
assert "No coding agent data found." in result.output
|
| 119 |
-
|
| 120 |
-
|
| 121 |
-
def test_learn_single_agent_shows_available_projects_when_cwd_missing(
|
| 122 |
-
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
|
| 123 |
-
) -> None:
|
| 124 |
-
project = SimpleNamespace(name="demo", project_path=tmp_path / "demo")
|
| 125 |
-
plugin = FakePlugin("codex", "Codex", [project])
|
| 126 |
-
|
| 127 |
-
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 128 |
-
monkeypatch.setattr("headroom.learn.registry.get_plugin", lambda name: plugin)
|
| 129 |
-
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", FakeAnalyzer)
|
| 130 |
-
|
| 131 |
-
with runner.isolated_filesystem(temp_dir=tmp_path):
|
| 132 |
-
result = runner.invoke(main, ["learn", "--agent", "codex"], catch_exceptions=False)
|
| 133 |
-
|
| 134 |
-
assert result.exit_code == 0
|
| 135 |
-
assert "No codex project data found for" in result.output
|
| 136 |
-
assert "Available codex projects:" in result.output
|
| 137 |
-
assert "demo" in result.output
|
| 138 |
-
|
| 139 |
-
|
| 140 |
-
def test_learn_project_lookup_and_apply_flow(
|
| 141 |
-
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
|
| 142 |
-
) -> None:
|
| 143 |
-
project_path = tmp_path / "project-a"
|
| 144 |
-
project_path.mkdir()
|
| 145 |
-
matched = SimpleNamespace(name="project-a", project_path=project_path)
|
| 146 |
-
unmatched = SimpleNamespace(name="project-b", project_path=tmp_path / "project-b")
|
| 147 |
-
plugin = FakePlugin("codex", "Codex", [matched, unmatched])
|
| 148 |
-
analyzer = FakeAnalyzer()
|
| 149 |
-
|
| 150 |
-
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 151 |
-
monkeypatch.setattr("headroom.learn.registry.get_plugin", lambda name: plugin)
|
| 152 |
-
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", lambda model=None: analyzer)
|
| 153 |
-
|
| 154 |
-
result = runner.invoke(
|
| 155 |
-
main,
|
| 156 |
-
["learn", "--agent", "codex", "--project", str(project_path), "--apply", "--workers", "4"],
|
| 157 |
-
catch_exceptions=False,
|
| 158 |
-
)
|
| 159 |
-
|
| 160 |
-
assert result.exit_code == 0, result.output
|
| 161 |
-
assert "Path: " in result.output
|
| 162 |
-
assert "Analyzing with gpt-4o..." in result.output
|
| 163 |
-
assert "Recommendations: 1" in result.output
|
| 164 |
-
assert "[WROTE]" in result.output
|
| 165 |
-
assert "Rule 1" in result.output
|
| 166 |
-
assert plugin.scan_calls == [(matched, 4)]
|
| 167 |
-
assert analyzer.calls[0][0] is matched
|
| 168 |
-
assert plugin.writer.calls[0][2] is False
|
| 169 |
-
|
| 170 |
-
|
| 171 |
-
def test_learn_reports_missing_requested_project_and_lists_discovered(
|
| 172 |
-
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
|
| 173 |
-
) -> None:
|
| 174 |
-
requested = tmp_path / "missing"
|
| 175 |
-
requested.mkdir()
|
| 176 |
-
discovered = SimpleNamespace(name="project-a", project_path=tmp_path / "project-a")
|
| 177 |
-
plugin = FakePlugin("claude", "Claude Code", [discovered])
|
| 178 |
-
|
| 179 |
-
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 180 |
-
monkeypatch.setattr("headroom.learn.registry.get_plugin", lambda name: plugin)
|
| 181 |
-
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", FakeAnalyzer)
|
| 182 |
-
|
| 183 |
-
result = runner.invoke(
|
| 184 |
-
main,
|
| 185 |
-
["learn", "--agent", "claude", "--project", str(requested)],
|
| 186 |
-
catch_exceptions=False,
|
| 187 |
-
)
|
| 188 |
-
|
| 189 |
-
assert result.exit_code == 0
|
| 190 |
-
assert f"No project data found for {requested.resolve()}" in result.output
|
| 191 |
-
assert "Available discovered projects:" in result.output
|
| 192 |
-
assert "[claude]" in result.output
|
| 193 |
-
|
| 194 |
-
|
| 195 |
-
def test_learn_analyze_all_uses_default_workers_and_prints_summary(
|
| 196 |
-
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
|
| 197 |
-
) -> None:
|
| 198 |
-
projects_a = [SimpleNamespace(name="a", project_path=tmp_path / "a")]
|
| 199 |
-
projects_b = [SimpleNamespace(name="b", project_path=tmp_path / "b")]
|
| 200 |
-
plugin_a = FakePlugin("codex", "Codex", projects_a)
|
| 201 |
-
plugin_b = FakePlugin("claude", "Claude Code", projects_b)
|
| 202 |
-
analyzer = FakeAnalyzer()
|
| 203 |
-
|
| 204 |
-
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 205 |
-
monkeypatch.setattr(
|
| 206 |
-
"headroom.learn.registry.auto_detect_plugins",
|
| 207 |
-
lambda: [plugin_a, plugin_b],
|
| 208 |
-
)
|
| 209 |
-
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", lambda model=None: analyzer)
|
| 210 |
-
monkeypatch.setattr("os.cpu_count", lambda: 12)
|
| 211 |
-
|
| 212 |
-
result = runner.invoke(main, ["learn", "--all"], catch_exceptions=False)
|
| 213 |
-
|
| 214 |
-
assert result.exit_code == 0, result.output
|
| 215 |
-
assert "Detected agents: Codex, Claude Code" in result.output
|
| 216 |
-
assert "Total: 2 projects, 2 failures, 2 recommendations" in result.output
|
| 217 |
-
assert plugin_a.scan_calls == [(projects_a[0], 8)]
|
| 218 |
-
assert plugin_b.scan_calls == [(projects_b[0], 8)]
|
| 219 |
-
|
| 220 |
-
|
| 221 |
-
def test_learn_handles_empty_sessions_and_no_pattern_outputs(
|
| 222 |
-
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
|
| 223 |
-
) -> None:
|
| 224 |
-
no_sessions = SimpleNamespace(name="empty", project_path=tmp_path / "empty")
|
| 225 |
-
no_failures = SimpleNamespace(name="clean", project_path=tmp_path / "clean")
|
| 226 |
-
no_actions = SimpleNamespace(name="no-actions", project_path=tmp_path / "no-actions")
|
| 227 |
-
|
| 228 |
-
class BranchingPlugin(FakePlugin):
|
| 229 |
-
def scan_project(self, project, max_workers: int = 1): # noqa: ANN001, ANN201
|
| 230 |
-
self.scan_calls.append((project, max_workers))
|
| 231 |
-
if project is no_sessions:
|
| 232 |
-
return []
|
| 233 |
-
return [SimpleNamespace(events=["event"], tool_calls=[], failure_count=0)]
|
| 234 |
-
|
| 235 |
-
class BranchingAnalyzer(FakeAnalyzer):
|
| 236 |
-
def analyze(self, project, sessions): # noqa: ANN001, ANN201
|
| 237 |
-
self.calls.append((project, sessions))
|
| 238 |
-
if project is no_failures:
|
| 239 |
-
return SimpleNamespace(
|
| 240 |
-
total_sessions=1,
|
| 241 |
-
total_calls=2,
|
| 242 |
-
total_failures=0,
|
| 243 |
-
failure_rate=0.0,
|
| 244 |
-
recommendations=[],
|
| 245 |
-
)
|
| 246 |
-
return SimpleNamespace(
|
| 247 |
-
total_sessions=1,
|
| 248 |
-
total_calls=2,
|
| 249 |
-
total_failures=1,
|
| 250 |
-
failure_rate=0.5,
|
| 251 |
-
recommendations=[],
|
| 252 |
-
)
|
| 253 |
-
|
| 254 |
-
plugin = BranchingPlugin("codex", "Codex", [no_sessions, no_failures, no_actions])
|
| 255 |
-
analyzer = BranchingAnalyzer()
|
| 256 |
-
|
| 257 |
-
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 258 |
-
monkeypatch.setattr("headroom.learn.registry.get_plugin", lambda name: plugin)
|
| 259 |
-
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", lambda model=None: analyzer)
|
| 260 |
-
|
| 261 |
-
result = runner.invoke(main, ["learn", "--agent", "codex", "--all"], catch_exceptions=False)
|
| 262 |
-
|
| 263 |
-
assert result.exit_code == 0, result.output
|
| 264 |
-
assert "No conversation data found." in result.output
|
| 265 |
-
assert "No failures or patterns found." in result.output
|
| 266 |
-
assert "No actionable patterns found." in result.output
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
from types import SimpleNamespace
|
| 5 |
+
|
| 6 |
+
import click
|
| 7 |
+
import click.shell_completion as click_shell_completion
|
| 8 |
+
import pytest
|
| 9 |
+
from click.testing import CliRunner
|
| 10 |
+
|
| 11 |
+
from headroom.cli.learn import _AgentChoice
|
| 12 |
+
from headroom.cli.main import main
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
@pytest.fixture
|
| 16 |
+
def runner() -> CliRunner:
|
| 17 |
+
return CliRunner()
|
| 18 |
+
|
| 19 |
+
|
| 20 |
+
class FakeWriter:
|
| 21 |
+
def __init__(self) -> None:
|
| 22 |
+
self.calls: list[tuple[list[object], object, bool]] = []
|
| 23 |
+
|
| 24 |
+
def write(self, recommendations, project, dry_run: bool): # noqa: ANN001, ANN201
|
| 25 |
+
self.calls.append((recommendations, project, dry_run))
|
| 26 |
+
return SimpleNamespace(
|
| 27 |
+
dry_run=dry_run,
|
| 28 |
+
content_by_file={
|
| 29 |
+
Path(project.project_path) / "AGENTS.md": "<!-- headroom -->\nRule 1\nRule 2"
|
| 30 |
+
},
|
| 31 |
+
)
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
class FakePlugin:
|
| 35 |
+
def __init__(self, name: str, display_name: str, projects: list[object]) -> None:
|
| 36 |
+
self.name = name
|
| 37 |
+
self.display_name = display_name
|
| 38 |
+
self._projects = projects
|
| 39 |
+
self.writer = FakeWriter()
|
| 40 |
+
self.scan_calls: list[tuple[object, int]] = []
|
| 41 |
+
|
| 42 |
+
def detect(self) -> bool:
|
| 43 |
+
return True
|
| 44 |
+
|
| 45 |
+
def create_writer(self) -> FakeWriter:
|
| 46 |
+
return self.writer
|
| 47 |
+
|
| 48 |
+
def discover_projects(self) -> list[object]:
|
| 49 |
+
return self._projects
|
| 50 |
+
|
| 51 |
+
def scan_project(self, project, max_workers: int = 1): # noqa: ANN001, ANN201
|
| 52 |
+
self.scan_calls.append((project, max_workers))
|
| 53 |
+
return [SimpleNamespace(events=["event"], tool_calls=[], failure_count=0)]
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
class FakeAnalyzer:
|
| 57 |
+
def __init__(self, model: str | None = None) -> None:
|
| 58 |
+
self.model = model
|
| 59 |
+
self.calls: list[tuple[object, list[object]]] = []
|
| 60 |
+
|
| 61 |
+
def analyze(self, project, sessions): # noqa: ANN001, ANN201
|
| 62 |
+
self.calls.append((project, sessions))
|
| 63 |
+
return SimpleNamespace(
|
| 64 |
+
total_sessions=len(sessions),
|
| 65 |
+
total_calls=3,
|
| 66 |
+
total_failures=1,
|
| 67 |
+
failure_rate=1 / 3,
|
| 68 |
+
recommendations=[SimpleNamespace(section="Rules")],
|
| 69 |
+
)
|
| 70 |
+
|
| 71 |
+
|
| 72 |
+
def test_agent_choice_convert_and_shell_complete(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 73 |
+
choice = _AgentChoice()
|
| 74 |
+
monkeypatch.setattr(click, "shell_completion", click_shell_completion)
|
| 75 |
+
monkeypatch.setattr(
|
| 76 |
+
"headroom.learn.registry.get_registry",
|
| 77 |
+
lambda: {"codex": object(), "claude": object()},
|
| 78 |
+
)
|
| 79 |
+
monkeypatch.setattr(
|
| 80 |
+
"headroom.learn.registry.available_agent_names",
|
| 81 |
+
lambda: ["claude", "codex"],
|
| 82 |
+
)
|
| 83 |
+
|
| 84 |
+
assert choice.convert("auto", None, None) == "auto"
|
| 85 |
+
assert choice.convert("CODEX", None, None) == "codex"
|
| 86 |
+
with pytest.raises(Exception, match="Unknown agent: bad"):
|
| 87 |
+
choice.convert("bad", None, None)
|
| 88 |
+
|
| 89 |
+
completions = choice.shell_complete(None, None, "c") # type: ignore[arg-type]
|
| 90 |
+
assert [item.value for item in completions] == ["claude", "codex"]
|
| 91 |
+
assert choice.get_metavar(None) == "[auto|<agent>]" # type: ignore[arg-type]
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def test_learn_exits_cleanly_when_model_detection_fails(
|
| 95 |
+
monkeypatch: pytest.MonkeyPatch, runner: CliRunner
|
| 96 |
+
) -> None:
|
| 97 |
+
monkeypatch.setattr(
|
| 98 |
+
"headroom.learn.analyzer._detect_default_model",
|
| 99 |
+
lambda: (_ for _ in ()).throw(RuntimeError("no model")),
|
| 100 |
+
)
|
| 101 |
+
|
| 102 |
+
result = runner.invoke(main, ["learn"], catch_exceptions=False)
|
| 103 |
+
|
| 104 |
+
assert result.exit_code == 1
|
| 105 |
+
assert "Error: no model" in result.output
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def test_learn_auto_agent_reports_no_detected_plugins(
|
| 109 |
+
monkeypatch: pytest.MonkeyPatch, runner: CliRunner
|
| 110 |
+
) -> None:
|
| 111 |
+
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 112 |
+
monkeypatch.setattr("headroom.learn.registry.auto_detect_plugins", lambda: [])
|
| 113 |
+
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", FakeAnalyzer)
|
| 114 |
+
|
| 115 |
+
result = runner.invoke(main, ["learn"], catch_exceptions=False)
|
| 116 |
+
|
| 117 |
+
assert result.exit_code == 0
|
| 118 |
+
assert "No coding agent data found." in result.output
|
| 119 |
+
|
| 120 |
+
|
| 121 |
+
def test_learn_single_agent_shows_available_projects_when_cwd_missing(
|
| 122 |
+
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
|
| 123 |
+
) -> None:
|
| 124 |
+
project = SimpleNamespace(name="demo", project_path=tmp_path / "demo")
|
| 125 |
+
plugin = FakePlugin("codex", "Codex", [project])
|
| 126 |
+
|
| 127 |
+
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 128 |
+
monkeypatch.setattr("headroom.learn.registry.get_plugin", lambda name: plugin)
|
| 129 |
+
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", FakeAnalyzer)
|
| 130 |
+
|
| 131 |
+
with runner.isolated_filesystem(temp_dir=tmp_path):
|
| 132 |
+
result = runner.invoke(main, ["learn", "--agent", "codex"], catch_exceptions=False)
|
| 133 |
+
|
| 134 |
+
assert result.exit_code == 0
|
| 135 |
+
assert "No codex project data found for" in result.output
|
| 136 |
+
assert "Available codex projects:" in result.output
|
| 137 |
+
assert "demo" in result.output
|
| 138 |
+
|
| 139 |
+
|
| 140 |
+
def test_learn_project_lookup_and_apply_flow(
|
| 141 |
+
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
|
| 142 |
+
) -> None:
|
| 143 |
+
project_path = tmp_path / "project-a"
|
| 144 |
+
project_path.mkdir()
|
| 145 |
+
matched = SimpleNamespace(name="project-a", project_path=project_path)
|
| 146 |
+
unmatched = SimpleNamespace(name="project-b", project_path=tmp_path / "project-b")
|
| 147 |
+
plugin = FakePlugin("codex", "Codex", [matched, unmatched])
|
| 148 |
+
analyzer = FakeAnalyzer()
|
| 149 |
+
|
| 150 |
+
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 151 |
+
monkeypatch.setattr("headroom.learn.registry.get_plugin", lambda name: plugin)
|
| 152 |
+
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", lambda model=None: analyzer)
|
| 153 |
+
|
| 154 |
+
result = runner.invoke(
|
| 155 |
+
main,
|
| 156 |
+
["learn", "--agent", "codex", "--project", str(project_path), "--apply", "--workers", "4"],
|
| 157 |
+
catch_exceptions=False,
|
| 158 |
+
)
|
| 159 |
+
|
| 160 |
+
assert result.exit_code == 0, result.output
|
| 161 |
+
assert "Path: " in result.output
|
| 162 |
+
assert "Analyzing with gpt-4o..." in result.output
|
| 163 |
+
assert "Recommendations: 1" in result.output
|
| 164 |
+
assert "[WROTE]" in result.output
|
| 165 |
+
assert "Rule 1" in result.output
|
| 166 |
+
assert plugin.scan_calls == [(matched, 4)]
|
| 167 |
+
assert analyzer.calls[0][0] is matched
|
| 168 |
+
assert plugin.writer.calls[0][2] is False
|
| 169 |
+
|
| 170 |
+
|
| 171 |
+
def test_learn_reports_missing_requested_project_and_lists_discovered(
|
| 172 |
+
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
|
| 173 |
+
) -> None:
|
| 174 |
+
requested = tmp_path / "missing"
|
| 175 |
+
requested.mkdir()
|
| 176 |
+
discovered = SimpleNamespace(name="project-a", project_path=tmp_path / "project-a")
|
| 177 |
+
plugin = FakePlugin("claude", "Claude Code", [discovered])
|
| 178 |
+
|
| 179 |
+
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 180 |
+
monkeypatch.setattr("headroom.learn.registry.get_plugin", lambda name: plugin)
|
| 181 |
+
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", FakeAnalyzer)
|
| 182 |
+
|
| 183 |
+
result = runner.invoke(
|
| 184 |
+
main,
|
| 185 |
+
["learn", "--agent", "claude", "--project", str(requested)],
|
| 186 |
+
catch_exceptions=False,
|
| 187 |
+
)
|
| 188 |
+
|
| 189 |
+
assert result.exit_code == 0
|
| 190 |
+
assert f"No project data found for {requested.resolve()}" in result.output
|
| 191 |
+
assert "Available discovered projects:" in result.output
|
| 192 |
+
assert "[claude]" in result.output
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
def test_learn_analyze_all_uses_default_workers_and_prints_summary(
|
| 196 |
+
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
|
| 197 |
+
) -> None:
|
| 198 |
+
projects_a = [SimpleNamespace(name="a", project_path=tmp_path / "a")]
|
| 199 |
+
projects_b = [SimpleNamespace(name="b", project_path=tmp_path / "b")]
|
| 200 |
+
plugin_a = FakePlugin("codex", "Codex", projects_a)
|
| 201 |
+
plugin_b = FakePlugin("claude", "Claude Code", projects_b)
|
| 202 |
+
analyzer = FakeAnalyzer()
|
| 203 |
+
|
| 204 |
+
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 205 |
+
monkeypatch.setattr(
|
| 206 |
+
"headroom.learn.registry.auto_detect_plugins",
|
| 207 |
+
lambda: [plugin_a, plugin_b],
|
| 208 |
+
)
|
| 209 |
+
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", lambda model=None: analyzer)
|
| 210 |
+
monkeypatch.setattr("os.cpu_count", lambda: 12)
|
| 211 |
+
|
| 212 |
+
result = runner.invoke(main, ["learn", "--all"], catch_exceptions=False)
|
| 213 |
+
|
| 214 |
+
assert result.exit_code == 0, result.output
|
| 215 |
+
assert "Detected agents: Codex, Claude Code" in result.output
|
| 216 |
+
assert "Total: 2 projects, 2 failures, 2 recommendations" in result.output
|
| 217 |
+
assert plugin_a.scan_calls == [(projects_a[0], 8)]
|
| 218 |
+
assert plugin_b.scan_calls == [(projects_b[0], 8)]
|
| 219 |
+
|
| 220 |
+
|
| 221 |
+
def test_learn_handles_empty_sessions_and_no_pattern_outputs(
|
| 222 |
+
monkeypatch: pytest.MonkeyPatch, runner: CliRunner, tmp_path: Path
|
| 223 |
+
) -> None:
|
| 224 |
+
no_sessions = SimpleNamespace(name="empty", project_path=tmp_path / "empty")
|
| 225 |
+
no_failures = SimpleNamespace(name="clean", project_path=tmp_path / "clean")
|
| 226 |
+
no_actions = SimpleNamespace(name="no-actions", project_path=tmp_path / "no-actions")
|
| 227 |
+
|
| 228 |
+
class BranchingPlugin(FakePlugin):
|
| 229 |
+
def scan_project(self, project, max_workers: int = 1): # noqa: ANN001, ANN201
|
| 230 |
+
self.scan_calls.append((project, max_workers))
|
| 231 |
+
if project is no_sessions:
|
| 232 |
+
return []
|
| 233 |
+
return [SimpleNamespace(events=["event"], tool_calls=[], failure_count=0)]
|
| 234 |
+
|
| 235 |
+
class BranchingAnalyzer(FakeAnalyzer):
|
| 236 |
+
def analyze(self, project, sessions): # noqa: ANN001, ANN201
|
| 237 |
+
self.calls.append((project, sessions))
|
| 238 |
+
if project is no_failures:
|
| 239 |
+
return SimpleNamespace(
|
| 240 |
+
total_sessions=1,
|
| 241 |
+
total_calls=2,
|
| 242 |
+
total_failures=0,
|
| 243 |
+
failure_rate=0.0,
|
| 244 |
+
recommendations=[],
|
| 245 |
+
)
|
| 246 |
+
return SimpleNamespace(
|
| 247 |
+
total_sessions=1,
|
| 248 |
+
total_calls=2,
|
| 249 |
+
total_failures=1,
|
| 250 |
+
failure_rate=0.5,
|
| 251 |
+
recommendations=[],
|
| 252 |
+
)
|
| 253 |
+
|
| 254 |
+
plugin = BranchingPlugin("codex", "Codex", [no_sessions, no_failures, no_actions])
|
| 255 |
+
analyzer = BranchingAnalyzer()
|
| 256 |
+
|
| 257 |
+
monkeypatch.setattr("headroom.learn.analyzer._detect_default_model", lambda: "gpt-4o")
|
| 258 |
+
monkeypatch.setattr("headroom.learn.registry.get_plugin", lambda name: plugin)
|
| 259 |
+
monkeypatch.setattr("headroom.learn.analyzer.SessionAnalyzer", lambda model=None: analyzer)
|
| 260 |
+
|
| 261 |
+
result = runner.invoke(main, ["learn", "--agent", "codex", "--all"], catch_exceptions=False)
|
| 262 |
+
|
| 263 |
+
assert result.exit_code == 0, result.output
|
| 264 |
+
assert "No conversation data found." in result.output
|
| 265 |
+
assert "No failures or patterns found." in result.output
|
| 266 |
+
assert "No actionable patterns found." in result.output
|
|
@@ -1,227 +1,227 @@
|
|
| 1 |
-
"""Unit tests for headroom.subscription.codex_rate_limits."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import time
|
| 6 |
-
|
| 7 |
-
from headroom.subscription.codex_rate_limits import (
|
| 8 |
-
CodexRateLimitState,
|
| 9 |
-
CodexRateLimitWindow,
|
| 10 |
-
parse_codex_rate_limits,
|
| 11 |
-
)
|
| 12 |
-
|
| 13 |
-
# ---------------------------------------------------------------------------
|
| 14 |
-
# CodexRateLimitWindow helpers
|
| 15 |
-
# ---------------------------------------------------------------------------
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
class TestCodexRateLimitWindow:
|
| 19 |
-
def test_window_label_minutes(self):
|
| 20 |
-
w = CodexRateLimitWindow(used_percent=10.0, window_minutes=45)
|
| 21 |
-
assert w.window_label == "45m"
|
| 22 |
-
|
| 23 |
-
def test_window_label_hours(self):
|
| 24 |
-
w = CodexRateLimitWindow(used_percent=10.0, window_minutes=60)
|
| 25 |
-
assert w.window_label == "1h"
|
| 26 |
-
|
| 27 |
-
def test_window_label_hours_with_minutes(self):
|
| 28 |
-
w = CodexRateLimitWindow(used_percent=10.0, window_minutes=90)
|
| 29 |
-
assert w.window_label == "1h30m"
|
| 30 |
-
|
| 31 |
-
def test_window_label_unknown(self):
|
| 32 |
-
w = CodexRateLimitWindow(used_percent=10.0, window_minutes=None)
|
| 33 |
-
assert w.window_label == "unknown"
|
| 34 |
-
|
| 35 |
-
def test_seconds_until_reset_future(self):
|
| 36 |
-
future = int(time.time()) + 3600
|
| 37 |
-
w = CodexRateLimitWindow(used_percent=10.0, resets_at=future)
|
| 38 |
-
secs = w.seconds_until_reset
|
| 39 |
-
assert secs is not None
|
| 40 |
-
assert 3590 <= secs <= 3600
|
| 41 |
-
|
| 42 |
-
def test_seconds_until_reset_past(self):
|
| 43 |
-
past = int(time.time()) - 100
|
| 44 |
-
w = CodexRateLimitWindow(used_percent=10.0, resets_at=past)
|
| 45 |
-
assert w.seconds_until_reset == 0
|
| 46 |
-
|
| 47 |
-
def test_seconds_until_reset_none(self):
|
| 48 |
-
w = CodexRateLimitWindow(used_percent=10.0, resets_at=None)
|
| 49 |
-
assert w.seconds_until_reset is None
|
| 50 |
-
|
| 51 |
-
def test_to_dict_keys(self):
|
| 52 |
-
w = CodexRateLimitWindow(used_percent=42.5, window_minutes=60, resets_at=9999999)
|
| 53 |
-
d = w.to_dict()
|
| 54 |
-
assert set(d.keys()) == {
|
| 55 |
-
"used_percent",
|
| 56 |
-
"window_minutes",
|
| 57 |
-
"window_label",
|
| 58 |
-
"resets_at",
|
| 59 |
-
"seconds_until_reset",
|
| 60 |
-
}
|
| 61 |
-
assert d["used_percent"] == 42.5
|
| 62 |
-
assert d["window_label"] == "1h"
|
| 63 |
-
|
| 64 |
-
|
| 65 |
-
# ---------------------------------------------------------------------------
|
| 66 |
-
# parse_codex_rate_limits
|
| 67 |
-
# ---------------------------------------------------------------------------
|
| 68 |
-
|
| 69 |
-
|
| 70 |
-
class TestParseCodexRateLimits:
|
| 71 |
-
def test_returns_none_for_empty_headers(self):
|
| 72 |
-
assert parse_codex_rate_limits({}) is None
|
| 73 |
-
|
| 74 |
-
def test_returns_none_for_non_codex_headers(self):
|
| 75 |
-
headers = {"content-type": "application/json", "x-request-id": "abc"}
|
| 76 |
-
assert parse_codex_rate_limits(headers) is None
|
| 77 |
-
|
| 78 |
-
def test_parses_primary_window(self):
|
| 79 |
-
headers = {
|
| 80 |
-
"x-codex-primary-used-percent": "35.5",
|
| 81 |
-
"x-codex-primary-window-minutes": "60",
|
| 82 |
-
"x-codex-primary-reset-at": "1704069000",
|
| 83 |
-
}
|
| 84 |
-
snap = parse_codex_rate_limits(headers)
|
| 85 |
-
assert snap is not None
|
| 86 |
-
assert snap.limit_id == "codex"
|
| 87 |
-
assert snap.primary is not None
|
| 88 |
-
assert snap.primary.used_percent == 35.5
|
| 89 |
-
assert snap.primary.window_minutes == 60
|
| 90 |
-
assert snap.primary.resets_at == 1704069000
|
| 91 |
-
assert snap.secondary is None
|
| 92 |
-
|
| 93 |
-
def test_parses_secondary_window(self):
|
| 94 |
-
headers = {
|
| 95 |
-
"x-codex-primary-used-percent": "10.0",
|
| 96 |
-
"x-codex-secondary-used-percent": "80.0",
|
| 97 |
-
"x-codex-secondary-window-minutes": "1440",
|
| 98 |
-
"x-codex-secondary-reset-at": "1704100000",
|
| 99 |
-
}
|
| 100 |
-
snap = parse_codex_rate_limits(headers)
|
| 101 |
-
assert snap is not None
|
| 102 |
-
assert snap.secondary is not None
|
| 103 |
-
assert snap.secondary.used_percent == 80.0
|
| 104 |
-
assert snap.secondary.window_minutes == 1440
|
| 105 |
-
|
| 106 |
-
def test_parses_credits(self):
|
| 107 |
-
headers = {
|
| 108 |
-
"x-codex-primary-used-percent": "5.0",
|
| 109 |
-
"x-codex-credits-has-credits": "true",
|
| 110 |
-
"x-codex-credits-unlimited": "false",
|
| 111 |
-
"x-codex-credits-balance": "$12.50",
|
| 112 |
-
}
|
| 113 |
-
snap = parse_codex_rate_limits(headers)
|
| 114 |
-
assert snap is not None
|
| 115 |
-
assert snap.credits is not None
|
| 116 |
-
assert snap.credits.has_credits is True
|
| 117 |
-
assert snap.credits.unlimited is False
|
| 118 |
-
assert snap.credits.balance == "$12.50"
|
| 119 |
-
|
| 120 |
-
def test_parses_unlimited_credits(self):
|
| 121 |
-
headers = {
|
| 122 |
-
"x-codex-primary-used-percent": "0.0",
|
| 123 |
-
"x-codex-credits-has-credits": "true",
|
| 124 |
-
"x-codex-credits-unlimited": "true",
|
| 125 |
-
}
|
| 126 |
-
snap = parse_codex_rate_limits(headers)
|
| 127 |
-
assert snap is not None
|
| 128 |
-
assert snap.credits is not None
|
| 129 |
-
assert snap.credits.unlimited is True
|
| 130 |
-
assert snap.credits.balance is None
|
| 131 |
-
|
| 132 |
-
def test_parses_limit_name(self):
|
| 133 |
-
headers = {
|
| 134 |
-
"x-codex-primary-used-percent": "20.0",
|
| 135 |
-
"x-codex-limit-name": "gpt-5.2-codex-sonic",
|
| 136 |
-
}
|
| 137 |
-
snap = parse_codex_rate_limits(headers)
|
| 138 |
-
assert snap is not None
|
| 139 |
-
assert snap.limit_name == "gpt-5.2-codex-sonic"
|
| 140 |
-
|
| 141 |
-
def test_parses_promo_message(self):
|
| 142 |
-
headers = {
|
| 143 |
-
"x-codex-primary-used-percent": "50.0",
|
| 144 |
-
"x-codex-promo-message": "Try our new model!",
|
| 145 |
-
}
|
| 146 |
-
snap = parse_codex_rate_limits(headers)
|
| 147 |
-
assert snap is not None
|
| 148 |
-
assert snap.promo_message == "Try our new model!"
|
| 149 |
-
|
| 150 |
-
def test_only_credits_header_triggers_parse(self):
|
| 151 |
-
headers = {
|
| 152 |
-
"x-codex-credits-has-credits": "true",
|
| 153 |
-
"x-codex-credits-unlimited": "false",
|
| 154 |
-
}
|
| 155 |
-
snap = parse_codex_rate_limits(headers)
|
| 156 |
-
assert snap is not None
|
| 157 |
-
assert snap.primary is None
|
| 158 |
-
assert snap.credits is not None
|
| 159 |
-
|
| 160 |
-
def test_invalid_float_ignored(self):
|
| 161 |
-
headers = {"x-codex-primary-used-percent": "not_a_number"}
|
| 162 |
-
assert parse_codex_rate_limits(headers) is None
|
| 163 |
-
|
| 164 |
-
def test_to_dict_structure(self):
|
| 165 |
-
headers = {
|
| 166 |
-
"x-codex-primary-used-percent": "42.0",
|
| 167 |
-
"x-codex-primary-window-minutes": "60",
|
| 168 |
-
}
|
| 169 |
-
snap = parse_codex_rate_limits(headers)
|
| 170 |
-
assert snap is not None
|
| 171 |
-
d = snap.to_dict()
|
| 172 |
-
assert "limit_id" in d
|
| 173 |
-
assert "primary" in d
|
| 174 |
-
assert "secondary" in d
|
| 175 |
-
assert "credits" in d
|
| 176 |
-
assert "captured_at" in d
|
| 177 |
-
|
| 178 |
-
|
| 179 |
-
# ---------------------------------------------------------------------------
|
| 180 |
-
# CodexRateLimitState
|
| 181 |
-
# ---------------------------------------------------------------------------
|
| 182 |
-
|
| 183 |
-
|
| 184 |
-
class TestCodexRateLimitState:
|
| 185 |
-
def test_initial_state_is_none(self):
|
| 186 |
-
state = CodexRateLimitState()
|
| 187 |
-
assert state.latest is None
|
| 188 |
-
assert state.get_stats() is None
|
| 189 |
-
|
| 190 |
-
def test_update_from_headers_stores_snapshot(self):
|
| 191 |
-
state = CodexRateLimitState()
|
| 192 |
-
headers = {
|
| 193 |
-
"x-codex-primary-used-percent": "55.0",
|
| 194 |
-
"x-codex-primary-window-minutes": "60",
|
| 195 |
-
}
|
| 196 |
-
state.update_from_headers(headers)
|
| 197 |
-
snap = state.latest
|
| 198 |
-
assert snap is not None
|
| 199 |
-
assert snap.primary is not None
|
| 200 |
-
assert snap.primary.used_percent == 55.0
|
| 201 |
-
|
| 202 |
-
def test_update_from_empty_headers_is_noop(self):
|
| 203 |
-
state = CodexRateLimitState()
|
| 204 |
-
state.update_from_headers({})
|
| 205 |
-
assert state.latest is None
|
| 206 |
-
|
| 207 |
-
def test_update_from_non_codex_headers_is_noop(self):
|
| 208 |
-
state = CodexRateLimitState()
|
| 209 |
-
state.update_from_headers({"content-type": "application/json"})
|
| 210 |
-
assert state.latest is None
|
| 211 |
-
|
| 212 |
-
def test_get_stats_returns_dict_when_data_present(self):
|
| 213 |
-
state = CodexRateLimitState()
|
| 214 |
-
state.update_from_headers({"x-codex-primary-used-percent": "10.0"})
|
| 215 |
-
stats = state.get_stats()
|
| 216 |
-
assert stats is not None
|
| 217 |
-
assert isinstance(stats, dict)
|
| 218 |
-
assert stats["limit_id"] == "codex"
|
| 219 |
-
|
| 220 |
-
def test_update_overwrites_previous_snapshot(self):
|
| 221 |
-
state = CodexRateLimitState()
|
| 222 |
-
state.update_from_headers({"x-codex-primary-used-percent": "10.0"})
|
| 223 |
-
state.update_from_headers({"x-codex-primary-used-percent": "90.0"})
|
| 224 |
-
snap = state.latest
|
| 225 |
-
assert snap is not None
|
| 226 |
-
assert snap.primary is not None
|
| 227 |
-
assert snap.primary.used_percent == 90.0
|
|
|
|
| 1 |
+
"""Unit tests for headroom.subscription.codex_rate_limits."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import time
|
| 6 |
+
|
| 7 |
+
from headroom.subscription.codex_rate_limits import (
|
| 8 |
+
CodexRateLimitState,
|
| 9 |
+
CodexRateLimitWindow,
|
| 10 |
+
parse_codex_rate_limits,
|
| 11 |
+
)
|
| 12 |
+
|
| 13 |
+
# ---------------------------------------------------------------------------
|
| 14 |
+
# CodexRateLimitWindow helpers
|
| 15 |
+
# ---------------------------------------------------------------------------
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class TestCodexRateLimitWindow:
|
| 19 |
+
def test_window_label_minutes(self):
|
| 20 |
+
w = CodexRateLimitWindow(used_percent=10.0, window_minutes=45)
|
| 21 |
+
assert w.window_label == "45m"
|
| 22 |
+
|
| 23 |
+
def test_window_label_hours(self):
|
| 24 |
+
w = CodexRateLimitWindow(used_percent=10.0, window_minutes=60)
|
| 25 |
+
assert w.window_label == "1h"
|
| 26 |
+
|
| 27 |
+
def test_window_label_hours_with_minutes(self):
|
| 28 |
+
w = CodexRateLimitWindow(used_percent=10.0, window_minutes=90)
|
| 29 |
+
assert w.window_label == "1h30m"
|
| 30 |
+
|
| 31 |
+
def test_window_label_unknown(self):
|
| 32 |
+
w = CodexRateLimitWindow(used_percent=10.0, window_minutes=None)
|
| 33 |
+
assert w.window_label == "unknown"
|
| 34 |
+
|
| 35 |
+
def test_seconds_until_reset_future(self):
|
| 36 |
+
future = int(time.time()) + 3600
|
| 37 |
+
w = CodexRateLimitWindow(used_percent=10.0, resets_at=future)
|
| 38 |
+
secs = w.seconds_until_reset
|
| 39 |
+
assert secs is not None
|
| 40 |
+
assert 3590 <= secs <= 3600
|
| 41 |
+
|
| 42 |
+
def test_seconds_until_reset_past(self):
|
| 43 |
+
past = int(time.time()) - 100
|
| 44 |
+
w = CodexRateLimitWindow(used_percent=10.0, resets_at=past)
|
| 45 |
+
assert w.seconds_until_reset == 0
|
| 46 |
+
|
| 47 |
+
def test_seconds_until_reset_none(self):
|
| 48 |
+
w = CodexRateLimitWindow(used_percent=10.0, resets_at=None)
|
| 49 |
+
assert w.seconds_until_reset is None
|
| 50 |
+
|
| 51 |
+
def test_to_dict_keys(self):
|
| 52 |
+
w = CodexRateLimitWindow(used_percent=42.5, window_minutes=60, resets_at=9999999)
|
| 53 |
+
d = w.to_dict()
|
| 54 |
+
assert set(d.keys()) == {
|
| 55 |
+
"used_percent",
|
| 56 |
+
"window_minutes",
|
| 57 |
+
"window_label",
|
| 58 |
+
"resets_at",
|
| 59 |
+
"seconds_until_reset",
|
| 60 |
+
}
|
| 61 |
+
assert d["used_percent"] == 42.5
|
| 62 |
+
assert d["window_label"] == "1h"
|
| 63 |
+
|
| 64 |
+
|
| 65 |
+
# ---------------------------------------------------------------------------
|
| 66 |
+
# parse_codex_rate_limits
|
| 67 |
+
# ---------------------------------------------------------------------------
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
class TestParseCodexRateLimits:
|
| 71 |
+
def test_returns_none_for_empty_headers(self):
|
| 72 |
+
assert parse_codex_rate_limits({}) is None
|
| 73 |
+
|
| 74 |
+
def test_returns_none_for_non_codex_headers(self):
|
| 75 |
+
headers = {"content-type": "application/json", "x-request-id": "abc"}
|
| 76 |
+
assert parse_codex_rate_limits(headers) is None
|
| 77 |
+
|
| 78 |
+
def test_parses_primary_window(self):
|
| 79 |
+
headers = {
|
| 80 |
+
"x-codex-primary-used-percent": "35.5",
|
| 81 |
+
"x-codex-primary-window-minutes": "60",
|
| 82 |
+
"x-codex-primary-reset-at": "1704069000",
|
| 83 |
+
}
|
| 84 |
+
snap = parse_codex_rate_limits(headers)
|
| 85 |
+
assert snap is not None
|
| 86 |
+
assert snap.limit_id == "codex"
|
| 87 |
+
assert snap.primary is not None
|
| 88 |
+
assert snap.primary.used_percent == 35.5
|
| 89 |
+
assert snap.primary.window_minutes == 60
|
| 90 |
+
assert snap.primary.resets_at == 1704069000
|
| 91 |
+
assert snap.secondary is None
|
| 92 |
+
|
| 93 |
+
def test_parses_secondary_window(self):
|
| 94 |
+
headers = {
|
| 95 |
+
"x-codex-primary-used-percent": "10.0",
|
| 96 |
+
"x-codex-secondary-used-percent": "80.0",
|
| 97 |
+
"x-codex-secondary-window-minutes": "1440",
|
| 98 |
+
"x-codex-secondary-reset-at": "1704100000",
|
| 99 |
+
}
|
| 100 |
+
snap = parse_codex_rate_limits(headers)
|
| 101 |
+
assert snap is not None
|
| 102 |
+
assert snap.secondary is not None
|
| 103 |
+
assert snap.secondary.used_percent == 80.0
|
| 104 |
+
assert snap.secondary.window_minutes == 1440
|
| 105 |
+
|
| 106 |
+
def test_parses_credits(self):
|
| 107 |
+
headers = {
|
| 108 |
+
"x-codex-primary-used-percent": "5.0",
|
| 109 |
+
"x-codex-credits-has-credits": "true",
|
| 110 |
+
"x-codex-credits-unlimited": "false",
|
| 111 |
+
"x-codex-credits-balance": "$12.50",
|
| 112 |
+
}
|
| 113 |
+
snap = parse_codex_rate_limits(headers)
|
| 114 |
+
assert snap is not None
|
| 115 |
+
assert snap.credits is not None
|
| 116 |
+
assert snap.credits.has_credits is True
|
| 117 |
+
assert snap.credits.unlimited is False
|
| 118 |
+
assert snap.credits.balance == "$12.50"
|
| 119 |
+
|
| 120 |
+
def test_parses_unlimited_credits(self):
|
| 121 |
+
headers = {
|
| 122 |
+
"x-codex-primary-used-percent": "0.0",
|
| 123 |
+
"x-codex-credits-has-credits": "true",
|
| 124 |
+
"x-codex-credits-unlimited": "true",
|
| 125 |
+
}
|
| 126 |
+
snap = parse_codex_rate_limits(headers)
|
| 127 |
+
assert snap is not None
|
| 128 |
+
assert snap.credits is not None
|
| 129 |
+
assert snap.credits.unlimited is True
|
| 130 |
+
assert snap.credits.balance is None
|
| 131 |
+
|
| 132 |
+
def test_parses_limit_name(self):
|
| 133 |
+
headers = {
|
| 134 |
+
"x-codex-primary-used-percent": "20.0",
|
| 135 |
+
"x-codex-limit-name": "gpt-5.2-codex-sonic",
|
| 136 |
+
}
|
| 137 |
+
snap = parse_codex_rate_limits(headers)
|
| 138 |
+
assert snap is not None
|
| 139 |
+
assert snap.limit_name == "gpt-5.2-codex-sonic"
|
| 140 |
+
|
| 141 |
+
def test_parses_promo_message(self):
|
| 142 |
+
headers = {
|
| 143 |
+
"x-codex-primary-used-percent": "50.0",
|
| 144 |
+
"x-codex-promo-message": "Try our new model!",
|
| 145 |
+
}
|
| 146 |
+
snap = parse_codex_rate_limits(headers)
|
| 147 |
+
assert snap is not None
|
| 148 |
+
assert snap.promo_message == "Try our new model!"
|
| 149 |
+
|
| 150 |
+
def test_only_credits_header_triggers_parse(self):
|
| 151 |
+
headers = {
|
| 152 |
+
"x-codex-credits-has-credits": "true",
|
| 153 |
+
"x-codex-credits-unlimited": "false",
|
| 154 |
+
}
|
| 155 |
+
snap = parse_codex_rate_limits(headers)
|
| 156 |
+
assert snap is not None
|
| 157 |
+
assert snap.primary is None
|
| 158 |
+
assert snap.credits is not None
|
| 159 |
+
|
| 160 |
+
def test_invalid_float_ignored(self):
|
| 161 |
+
headers = {"x-codex-primary-used-percent": "not_a_number"}
|
| 162 |
+
assert parse_codex_rate_limits(headers) is None
|
| 163 |
+
|
| 164 |
+
def test_to_dict_structure(self):
|
| 165 |
+
headers = {
|
| 166 |
+
"x-codex-primary-used-percent": "42.0",
|
| 167 |
+
"x-codex-primary-window-minutes": "60",
|
| 168 |
+
}
|
| 169 |
+
snap = parse_codex_rate_limits(headers)
|
| 170 |
+
assert snap is not None
|
| 171 |
+
d = snap.to_dict()
|
| 172 |
+
assert "limit_id" in d
|
| 173 |
+
assert "primary" in d
|
| 174 |
+
assert "secondary" in d
|
| 175 |
+
assert "credits" in d
|
| 176 |
+
assert "captured_at" in d
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
# ---------------------------------------------------------------------------
|
| 180 |
+
# CodexRateLimitState
|
| 181 |
+
# ---------------------------------------------------------------------------
|
| 182 |
+
|
| 183 |
+
|
| 184 |
+
class TestCodexRateLimitState:
|
| 185 |
+
def test_initial_state_is_none(self):
|
| 186 |
+
state = CodexRateLimitState()
|
| 187 |
+
assert state.latest is None
|
| 188 |
+
assert state.get_stats() is None
|
| 189 |
+
|
| 190 |
+
def test_update_from_headers_stores_snapshot(self):
|
| 191 |
+
state = CodexRateLimitState()
|
| 192 |
+
headers = {
|
| 193 |
+
"x-codex-primary-used-percent": "55.0",
|
| 194 |
+
"x-codex-primary-window-minutes": "60",
|
| 195 |
+
}
|
| 196 |
+
state.update_from_headers(headers)
|
| 197 |
+
snap = state.latest
|
| 198 |
+
assert snap is not None
|
| 199 |
+
assert snap.primary is not None
|
| 200 |
+
assert snap.primary.used_percent == 55.0
|
| 201 |
+
|
| 202 |
+
def test_update_from_empty_headers_is_noop(self):
|
| 203 |
+
state = CodexRateLimitState()
|
| 204 |
+
state.update_from_headers({})
|
| 205 |
+
assert state.latest is None
|
| 206 |
+
|
| 207 |
+
def test_update_from_non_codex_headers_is_noop(self):
|
| 208 |
+
state = CodexRateLimitState()
|
| 209 |
+
state.update_from_headers({"content-type": "application/json"})
|
| 210 |
+
assert state.latest is None
|
| 211 |
+
|
| 212 |
+
def test_get_stats_returns_dict_when_data_present(self):
|
| 213 |
+
state = CodexRateLimitState()
|
| 214 |
+
state.update_from_headers({"x-codex-primary-used-percent": "10.0"})
|
| 215 |
+
stats = state.get_stats()
|
| 216 |
+
assert stats is not None
|
| 217 |
+
assert isinstance(stats, dict)
|
| 218 |
+
assert stats["limit_id"] == "codex"
|
| 219 |
+
|
| 220 |
+
def test_update_overwrites_previous_snapshot(self):
|
| 221 |
+
state = CodexRateLimitState()
|
| 222 |
+
state.update_from_headers({"x-codex-primary-used-percent": "10.0"})
|
| 223 |
+
state.update_from_headers({"x-codex-primary-used-percent": "90.0"})
|
| 224 |
+
snap = state.latest
|
| 225 |
+
assert snap is not None
|
| 226 |
+
assert snap.primary is not None
|
| 227 |
+
assert snap.primary.used_percent == 90.0
|
|
@@ -1,274 +1,274 @@
|
|
| 1 |
-
"""Tests for the one-function compress() API and integrations."""
|
| 2 |
-
|
| 3 |
-
import json
|
| 4 |
-
|
| 5 |
-
import pytest
|
| 6 |
-
|
| 7 |
-
from headroom.compress import CompressResult, compress
|
| 8 |
-
from headroom.hooks import CompressionHooks
|
| 9 |
-
|
| 10 |
-
try:
|
| 11 |
-
from starlette.applications import Starlette
|
| 12 |
-
from starlette.requests import Request
|
| 13 |
-
from starlette.responses import JSONResponse
|
| 14 |
-
from starlette.routing import Route
|
| 15 |
-
from starlette.testclient import TestClient
|
| 16 |
-
|
| 17 |
-
from headroom.integrations.asgi import CompressionMiddleware
|
| 18 |
-
|
| 19 |
-
HAS_STARLETTE = True
|
| 20 |
-
except ImportError:
|
| 21 |
-
HAS_STARLETTE = False
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
# =============================================================================
|
| 25 |
-
# Tests: compress() function
|
| 26 |
-
# =============================================================================
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
class TestCompressFunction:
|
| 30 |
-
def test_empty_messages(self):
|
| 31 |
-
result = compress([], model="test")
|
| 32 |
-
assert result.messages == []
|
| 33 |
-
assert result.tokens_saved == 0
|
| 34 |
-
|
| 35 |
-
def test_small_messages_passthrough(self):
|
| 36 |
-
"""Small messages below compression threshold pass through unchanged."""
|
| 37 |
-
messages = [{"role": "user", "content": "hello"}]
|
| 38 |
-
result = compress(messages, model="gpt-4o")
|
| 39 |
-
assert result.messages[0]["content"] == "hello"
|
| 40 |
-
assert result.tokens_saved == 0
|
| 41 |
-
|
| 42 |
-
def test_returns_compress_result(self):
|
| 43 |
-
result = compress([{"role": "user", "content": "hi"}])
|
| 44 |
-
assert isinstance(result, CompressResult)
|
| 45 |
-
assert hasattr(result, "messages")
|
| 46 |
-
assert hasattr(result, "tokens_saved")
|
| 47 |
-
assert hasattr(result, "compression_ratio")
|
| 48 |
-
assert hasattr(result, "transforms_applied")
|
| 49 |
-
|
| 50 |
-
def test_large_tool_output_compressed(self):
|
| 51 |
-
"""Large JSON tool output should be compressed."""
|
| 52 |
-
big_data = json.dumps(
|
| 53 |
-
[
|
| 54 |
-
{"id": i, "status": "active", "name": f"item_{i}", "value": i * 17}
|
| 55 |
-
for i in range(200)
|
| 56 |
-
]
|
| 57 |
-
)
|
| 58 |
-
messages = [
|
| 59 |
-
{"role": "user", "content": "What are the top items?"},
|
| 60 |
-
{"role": "tool", "content": big_data, "tool_call_id": "call_1"},
|
| 61 |
-
]
|
| 62 |
-
result = compress(messages, model="gpt-4o")
|
| 63 |
-
assert result.tokens_after <= result.tokens_before
|
| 64 |
-
assert len(result.messages) == 2
|
| 65 |
-
|
| 66 |
-
def test_compact_json_counts_tokens_not_whitespace(self):
|
| 67 |
-
"""Compact JSON arrays should still compress under token thresholds."""
|
| 68 |
-
numbers = [42.0 + i * 0.1 for i in range(200)]
|
| 69 |
-
messages = [
|
| 70 |
-
{"role": "system", "content": "You are helpful."},
|
| 71 |
-
{"role": "user", "content": "Show metrics"},
|
| 72 |
-
{
|
| 73 |
-
"role": "assistant",
|
| 74 |
-
"content": None,
|
| 75 |
-
"tool_calls": [
|
| 76 |
-
{
|
| 77 |
-
"id": "call_1",
|
| 78 |
-
"type": "function",
|
| 79 |
-
"function": {"name": "get_metrics", "arguments": "{}"},
|
| 80 |
-
}
|
| 81 |
-
],
|
| 82 |
-
},
|
| 83 |
-
{"role": "tool", "tool_call_id": "call_1", "content": json.dumps(numbers)},
|
| 84 |
-
]
|
| 85 |
-
|
| 86 |
-
result = compress(messages, min_tokens_to_compress=250)
|
| 87 |
-
|
| 88 |
-
assert result.tokens_saved > 0
|
| 89 |
-
assert any(
|
| 90 |
-
transform.startswith("router:smart_crusher") for transform in result.transforms_applied
|
| 91 |
-
)
|
| 92 |
-
|
| 93 |
-
def test_optimize_false_passthrough(self):
|
| 94 |
-
"""optimize=False returns messages unchanged."""
|
| 95 |
-
messages = [{"role": "user", "content": "hello world " * 100}]
|
| 96 |
-
result = compress(messages, optimize=False)
|
| 97 |
-
assert result.messages is messages
|
| 98 |
-
assert result.tokens_saved == 0
|
| 99 |
-
|
| 100 |
-
def test_with_custom_hooks(self):
|
| 101 |
-
"""Hooks are called when provided."""
|
| 102 |
-
calls = []
|
| 103 |
-
|
| 104 |
-
class TrackingHooks(CompressionHooks):
|
| 105 |
-
def pre_compress(self, messages, ctx):
|
| 106 |
-
calls.append(("pre", len(messages)))
|
| 107 |
-
return messages
|
| 108 |
-
|
| 109 |
-
def compute_biases(self, messages, ctx):
|
| 110 |
-
calls.append(("biases", len(messages)))
|
| 111 |
-
return {}
|
| 112 |
-
|
| 113 |
-
def post_compress(self, event):
|
| 114 |
-
calls.append(("post", event.tokens_saved))
|
| 115 |
-
|
| 116 |
-
big_data = json.dumps([{"id": i, "status": "active"} for i in range(100)])
|
| 117 |
-
messages = [
|
| 118 |
-
{"role": "user", "content": "analyze"},
|
| 119 |
-
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
| 120 |
-
]
|
| 121 |
-
compress(messages, hooks=TrackingHooks())
|
| 122 |
-
|
| 123 |
-
assert any(c[0] == "pre" for c in calls)
|
| 124 |
-
assert any(c[0] == "biases" for c in calls)
|
| 125 |
-
|
| 126 |
-
|
| 127 |
-
class TestCompressResultFields:
|
| 128 |
-
def test_fields_populated(self):
|
| 129 |
-
big_data = json.dumps([{"id": i, "type": "log"} for i in range(100)])
|
| 130 |
-
messages = [
|
| 131 |
-
{"role": "user", "content": "summarize"},
|
| 132 |
-
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
| 133 |
-
]
|
| 134 |
-
result = compress(messages, model="claude-sonnet-4-5-20250929")
|
| 135 |
-
assert result.tokens_before > 0
|
| 136 |
-
assert result.tokens_after >= 0
|
| 137 |
-
assert result.tokens_saved >= 0
|
| 138 |
-
assert 0.0 <= result.compression_ratio <= 1.0
|
| 139 |
-
|
| 140 |
-
|
| 141 |
-
# =============================================================================
|
| 142 |
-
# Tests: ASGI CompressionMiddleware (requires starlette)
|
| 143 |
-
# =============================================================================
|
| 144 |
-
|
| 145 |
-
|
| 146 |
-
def _make_asgi_app(middleware_kwargs=None):
|
| 147 |
-
"""Create a test ASGI app with CompressionMiddleware."""
|
| 148 |
-
|
| 149 |
-
async def chat_endpoint(request: Request) -> JSONResponse:
|
| 150 |
-
body = await request.json()
|
| 151 |
-
return JSONResponse(
|
| 152 |
-
{
|
| 153 |
-
"model": "gpt-4o",
|
| 154 |
-
"choices": [{"message": {"content": "response"}}],
|
| 155 |
-
"usage": {"prompt_tokens": 10, "completion_tokens": 5},
|
| 156 |
-
"_message_count": len(body.get("messages", [])),
|
| 157 |
-
}
|
| 158 |
-
)
|
| 159 |
-
|
| 160 |
-
async def health(request: Request) -> JSONResponse:
|
| 161 |
-
return JSONResponse({"status": "ok"})
|
| 162 |
-
|
| 163 |
-
app = Starlette(
|
| 164 |
-
routes=[
|
| 165 |
-
Route("/health", health),
|
| 166 |
-
Route("/v1/chat/completions", chat_endpoint, methods=["POST"]),
|
| 167 |
-
Route("/v1/messages", chat_endpoint, methods=["POST"]),
|
| 168 |
-
]
|
| 169 |
-
)
|
| 170 |
-
app.add_middleware(CompressionMiddleware, **(middleware_kwargs or {}))
|
| 171 |
-
return app
|
| 172 |
-
|
| 173 |
-
|
| 174 |
-
@pytest.mark.skipif(not HAS_STARLETTE, reason="starlette not installed")
|
| 175 |
-
class TestASGIMiddleware:
|
| 176 |
-
def test_non_llm_paths_passthrough(self):
|
| 177 |
-
app = _make_asgi_app()
|
| 178 |
-
client = TestClient(app)
|
| 179 |
-
resp = client.get("/health")
|
| 180 |
-
assert resp.status_code == 200
|
| 181 |
-
assert resp.json()["status"] == "ok"
|
| 182 |
-
|
| 183 |
-
def test_small_messages_passthrough(self):
|
| 184 |
-
app = _make_asgi_app()
|
| 185 |
-
client = TestClient(app)
|
| 186 |
-
resp = client.post(
|
| 187 |
-
"/v1/chat/completions",
|
| 188 |
-
json={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]},
|
| 189 |
-
)
|
| 190 |
-
assert resp.status_code == 200
|
| 191 |
-
|
| 192 |
-
def test_large_messages_compressed(self):
|
| 193 |
-
"""Large tool output should be compressed by middleware."""
|
| 194 |
-
app = _make_asgi_app()
|
| 195 |
-
client = TestClient(app)
|
| 196 |
-
|
| 197 |
-
big_data = json.dumps([{"id": i, "status": "active"} for i in range(200)])
|
| 198 |
-
resp = client.post(
|
| 199 |
-
"/v1/chat/completions",
|
| 200 |
-
json={
|
| 201 |
-
"model": "gpt-4o",
|
| 202 |
-
"messages": [
|
| 203 |
-
{"role": "user", "content": "analyze"},
|
| 204 |
-
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
| 205 |
-
],
|
| 206 |
-
},
|
| 207 |
-
)
|
| 208 |
-
assert resp.status_code == 200
|
| 209 |
-
|
| 210 |
-
def test_anthropic_path(self):
|
| 211 |
-
"""Works with Anthropic /v1/messages path."""
|
| 212 |
-
app = _make_asgi_app()
|
| 213 |
-
client = TestClient(app)
|
| 214 |
-
resp = client.post(
|
| 215 |
-
"/v1/messages",
|
| 216 |
-
json={
|
| 217 |
-
"model": "claude-sonnet-4-5-20250929",
|
| 218 |
-
"messages": [{"role": "user", "content": "hello"}],
|
| 219 |
-
},
|
| 220 |
-
)
|
| 221 |
-
assert resp.status_code == 200
|
| 222 |
-
|
| 223 |
-
def test_get_requests_passthrough(self):
|
| 224 |
-
"""GET requests to LLM paths pass through."""
|
| 225 |
-
app = _make_asgi_app()
|
| 226 |
-
client = TestClient(app)
|
| 227 |
-
resp = client.get("/v1/chat/completions")
|
| 228 |
-
assert resp.status_code in (200, 405)
|
| 229 |
-
|
| 230 |
-
|
| 231 |
-
# =============================================================================
|
| 232 |
-
# Tests: LiteLLM Callback
|
| 233 |
-
# =============================================================================
|
| 234 |
-
|
| 235 |
-
|
| 236 |
-
class TestLiteLLMCallback:
|
| 237 |
-
def test_callback_imports(self):
|
| 238 |
-
"""Verify the callback can be imported."""
|
| 239 |
-
from headroom.integrations.litellm_callback import HeadroomCallback
|
| 240 |
-
|
| 241 |
-
callback = HeadroomCallback()
|
| 242 |
-
assert callback.total_tokens_saved == 0
|
| 243 |
-
|
| 244 |
-
def test_callback_compresses_messages(self):
|
| 245 |
-
"""Callback compresses messages in pre_call_hook."""
|
| 246 |
-
import asyncio
|
| 247 |
-
|
| 248 |
-
from headroom.integrations.litellm_callback import HeadroomCallback
|
| 249 |
-
|
| 250 |
-
callback = HeadroomCallback()
|
| 251 |
-
|
| 252 |
-
big_data = json.dumps([{"id": i, "status": "active"} for i in range(200)])
|
| 253 |
-
data = {
|
| 254 |
-
"model": "gpt-4o",
|
| 255 |
-
"messages": [
|
| 256 |
-
{"role": "user", "content": "analyze"},
|
| 257 |
-
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
| 258 |
-
],
|
| 259 |
-
}
|
| 260 |
-
|
| 261 |
-
result = asyncio.run(callback.async_pre_call_hook("key", data, "completion"))
|
| 262 |
-
assert result is data
|
| 263 |
-
|
| 264 |
-
def test_callback_ignores_non_completion(self):
|
| 265 |
-
"""Non-completion calls are passed through."""
|
| 266 |
-
import asyncio
|
| 267 |
-
|
| 268 |
-
from headroom.integrations.litellm_callback import HeadroomCallback
|
| 269 |
-
|
| 270 |
-
callback = HeadroomCallback()
|
| 271 |
-
data = {"messages": [{"role": "user", "content": "hi"}]}
|
| 272 |
-
|
| 273 |
-
result = asyncio.run(callback.async_pre_call_hook("key", data, "embedding"))
|
| 274 |
-
assert result is data
|
|
|
|
| 1 |
+
"""Tests for the one-function compress() API and integrations."""
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
|
| 5 |
+
import pytest
|
| 6 |
+
|
| 7 |
+
from headroom.compress import CompressResult, compress
|
| 8 |
+
from headroom.hooks import CompressionHooks
|
| 9 |
+
|
| 10 |
+
try:
|
| 11 |
+
from starlette.applications import Starlette
|
| 12 |
+
from starlette.requests import Request
|
| 13 |
+
from starlette.responses import JSONResponse
|
| 14 |
+
from starlette.routing import Route
|
| 15 |
+
from starlette.testclient import TestClient
|
| 16 |
+
|
| 17 |
+
from headroom.integrations.asgi import CompressionMiddleware
|
| 18 |
+
|
| 19 |
+
HAS_STARLETTE = True
|
| 20 |
+
except ImportError:
|
| 21 |
+
HAS_STARLETTE = False
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
# =============================================================================
|
| 25 |
+
# Tests: compress() function
|
| 26 |
+
# =============================================================================
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
class TestCompressFunction:
|
| 30 |
+
def test_empty_messages(self):
|
| 31 |
+
result = compress([], model="test")
|
| 32 |
+
assert result.messages == []
|
| 33 |
+
assert result.tokens_saved == 0
|
| 34 |
+
|
| 35 |
+
def test_small_messages_passthrough(self):
|
| 36 |
+
"""Small messages below compression threshold pass through unchanged."""
|
| 37 |
+
messages = [{"role": "user", "content": "hello"}]
|
| 38 |
+
result = compress(messages, model="gpt-4o")
|
| 39 |
+
assert result.messages[0]["content"] == "hello"
|
| 40 |
+
assert result.tokens_saved == 0
|
| 41 |
+
|
| 42 |
+
def test_returns_compress_result(self):
|
| 43 |
+
result = compress([{"role": "user", "content": "hi"}])
|
| 44 |
+
assert isinstance(result, CompressResult)
|
| 45 |
+
assert hasattr(result, "messages")
|
| 46 |
+
assert hasattr(result, "tokens_saved")
|
| 47 |
+
assert hasattr(result, "compression_ratio")
|
| 48 |
+
assert hasattr(result, "transforms_applied")
|
| 49 |
+
|
| 50 |
+
def test_large_tool_output_compressed(self):
|
| 51 |
+
"""Large JSON tool output should be compressed."""
|
| 52 |
+
big_data = json.dumps(
|
| 53 |
+
[
|
| 54 |
+
{"id": i, "status": "active", "name": f"item_{i}", "value": i * 17}
|
| 55 |
+
for i in range(200)
|
| 56 |
+
]
|
| 57 |
+
)
|
| 58 |
+
messages = [
|
| 59 |
+
{"role": "user", "content": "What are the top items?"},
|
| 60 |
+
{"role": "tool", "content": big_data, "tool_call_id": "call_1"},
|
| 61 |
+
]
|
| 62 |
+
result = compress(messages, model="gpt-4o")
|
| 63 |
+
assert result.tokens_after <= result.tokens_before
|
| 64 |
+
assert len(result.messages) == 2
|
| 65 |
+
|
| 66 |
+
def test_compact_json_counts_tokens_not_whitespace(self):
|
| 67 |
+
"""Compact JSON arrays should still compress under token thresholds."""
|
| 68 |
+
numbers = [42.0 + i * 0.1 for i in range(200)]
|
| 69 |
+
messages = [
|
| 70 |
+
{"role": "system", "content": "You are helpful."},
|
| 71 |
+
{"role": "user", "content": "Show metrics"},
|
| 72 |
+
{
|
| 73 |
+
"role": "assistant",
|
| 74 |
+
"content": None,
|
| 75 |
+
"tool_calls": [
|
| 76 |
+
{
|
| 77 |
+
"id": "call_1",
|
| 78 |
+
"type": "function",
|
| 79 |
+
"function": {"name": "get_metrics", "arguments": "{}"},
|
| 80 |
+
}
|
| 81 |
+
],
|
| 82 |
+
},
|
| 83 |
+
{"role": "tool", "tool_call_id": "call_1", "content": json.dumps(numbers)},
|
| 84 |
+
]
|
| 85 |
+
|
| 86 |
+
result = compress(messages, min_tokens_to_compress=250)
|
| 87 |
+
|
| 88 |
+
assert result.tokens_saved > 0
|
| 89 |
+
assert any(
|
| 90 |
+
transform.startswith("router:smart_crusher") for transform in result.transforms_applied
|
| 91 |
+
)
|
| 92 |
+
|
| 93 |
+
def test_optimize_false_passthrough(self):
|
| 94 |
+
"""optimize=False returns messages unchanged."""
|
| 95 |
+
messages = [{"role": "user", "content": "hello world " * 100}]
|
| 96 |
+
result = compress(messages, optimize=False)
|
| 97 |
+
assert result.messages is messages
|
| 98 |
+
assert result.tokens_saved == 0
|
| 99 |
+
|
| 100 |
+
def test_with_custom_hooks(self):
|
| 101 |
+
"""Hooks are called when provided."""
|
| 102 |
+
calls = []
|
| 103 |
+
|
| 104 |
+
class TrackingHooks(CompressionHooks):
|
| 105 |
+
def pre_compress(self, messages, ctx):
|
| 106 |
+
calls.append(("pre", len(messages)))
|
| 107 |
+
return messages
|
| 108 |
+
|
| 109 |
+
def compute_biases(self, messages, ctx):
|
| 110 |
+
calls.append(("biases", len(messages)))
|
| 111 |
+
return {}
|
| 112 |
+
|
| 113 |
+
def post_compress(self, event):
|
| 114 |
+
calls.append(("post", event.tokens_saved))
|
| 115 |
+
|
| 116 |
+
big_data = json.dumps([{"id": i, "status": "active"} for i in range(100)])
|
| 117 |
+
messages = [
|
| 118 |
+
{"role": "user", "content": "analyze"},
|
| 119 |
+
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
| 120 |
+
]
|
| 121 |
+
compress(messages, hooks=TrackingHooks())
|
| 122 |
+
|
| 123 |
+
assert any(c[0] == "pre" for c in calls)
|
| 124 |
+
assert any(c[0] == "biases" for c in calls)
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
class TestCompressResultFields:
|
| 128 |
+
def test_fields_populated(self):
|
| 129 |
+
big_data = json.dumps([{"id": i, "type": "log"} for i in range(100)])
|
| 130 |
+
messages = [
|
| 131 |
+
{"role": "user", "content": "summarize"},
|
| 132 |
+
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
| 133 |
+
]
|
| 134 |
+
result = compress(messages, model="claude-sonnet-4-5-20250929")
|
| 135 |
+
assert result.tokens_before > 0
|
| 136 |
+
assert result.tokens_after >= 0
|
| 137 |
+
assert result.tokens_saved >= 0
|
| 138 |
+
assert 0.0 <= result.compression_ratio <= 1.0
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
# =============================================================================
|
| 142 |
+
# Tests: ASGI CompressionMiddleware (requires starlette)
|
| 143 |
+
# =============================================================================
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def _make_asgi_app(middleware_kwargs=None):
|
| 147 |
+
"""Create a test ASGI app with CompressionMiddleware."""
|
| 148 |
+
|
| 149 |
+
async def chat_endpoint(request: Request) -> JSONResponse:
|
| 150 |
+
body = await request.json()
|
| 151 |
+
return JSONResponse(
|
| 152 |
+
{
|
| 153 |
+
"model": "gpt-4o",
|
| 154 |
+
"choices": [{"message": {"content": "response"}}],
|
| 155 |
+
"usage": {"prompt_tokens": 10, "completion_tokens": 5},
|
| 156 |
+
"_message_count": len(body.get("messages", [])),
|
| 157 |
+
}
|
| 158 |
+
)
|
| 159 |
+
|
| 160 |
+
async def health(request: Request) -> JSONResponse:
|
| 161 |
+
return JSONResponse({"status": "ok"})
|
| 162 |
+
|
| 163 |
+
app = Starlette(
|
| 164 |
+
routes=[
|
| 165 |
+
Route("/health", health),
|
| 166 |
+
Route("/v1/chat/completions", chat_endpoint, methods=["POST"]),
|
| 167 |
+
Route("/v1/messages", chat_endpoint, methods=["POST"]),
|
| 168 |
+
]
|
| 169 |
+
)
|
| 170 |
+
app.add_middleware(CompressionMiddleware, **(middleware_kwargs or {}))
|
| 171 |
+
return app
|
| 172 |
+
|
| 173 |
+
|
| 174 |
+
@pytest.mark.skipif(not HAS_STARLETTE, reason="starlette not installed")
|
| 175 |
+
class TestASGIMiddleware:
|
| 176 |
+
def test_non_llm_paths_passthrough(self):
|
| 177 |
+
app = _make_asgi_app()
|
| 178 |
+
client = TestClient(app)
|
| 179 |
+
resp = client.get("/health")
|
| 180 |
+
assert resp.status_code == 200
|
| 181 |
+
assert resp.json()["status"] == "ok"
|
| 182 |
+
|
| 183 |
+
def test_small_messages_passthrough(self):
|
| 184 |
+
app = _make_asgi_app()
|
| 185 |
+
client = TestClient(app)
|
| 186 |
+
resp = client.post(
|
| 187 |
+
"/v1/chat/completions",
|
| 188 |
+
json={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]},
|
| 189 |
+
)
|
| 190 |
+
assert resp.status_code == 200
|
| 191 |
+
|
| 192 |
+
def test_large_messages_compressed(self):
|
| 193 |
+
"""Large tool output should be compressed by middleware."""
|
| 194 |
+
app = _make_asgi_app()
|
| 195 |
+
client = TestClient(app)
|
| 196 |
+
|
| 197 |
+
big_data = json.dumps([{"id": i, "status": "active"} for i in range(200)])
|
| 198 |
+
resp = client.post(
|
| 199 |
+
"/v1/chat/completions",
|
| 200 |
+
json={
|
| 201 |
+
"model": "gpt-4o",
|
| 202 |
+
"messages": [
|
| 203 |
+
{"role": "user", "content": "analyze"},
|
| 204 |
+
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
| 205 |
+
],
|
| 206 |
+
},
|
| 207 |
+
)
|
| 208 |
+
assert resp.status_code == 200
|
| 209 |
+
|
| 210 |
+
def test_anthropic_path(self):
|
| 211 |
+
"""Works with Anthropic /v1/messages path."""
|
| 212 |
+
app = _make_asgi_app()
|
| 213 |
+
client = TestClient(app)
|
| 214 |
+
resp = client.post(
|
| 215 |
+
"/v1/messages",
|
| 216 |
+
json={
|
| 217 |
+
"model": "claude-sonnet-4-5-20250929",
|
| 218 |
+
"messages": [{"role": "user", "content": "hello"}],
|
| 219 |
+
},
|
| 220 |
+
)
|
| 221 |
+
assert resp.status_code == 200
|
| 222 |
+
|
| 223 |
+
def test_get_requests_passthrough(self):
|
| 224 |
+
"""GET requests to LLM paths pass through."""
|
| 225 |
+
app = _make_asgi_app()
|
| 226 |
+
client = TestClient(app)
|
| 227 |
+
resp = client.get("/v1/chat/completions")
|
| 228 |
+
assert resp.status_code in (200, 405)
|
| 229 |
+
|
| 230 |
+
|
| 231 |
+
# =============================================================================
|
| 232 |
+
# Tests: LiteLLM Callback
|
| 233 |
+
# =============================================================================
|
| 234 |
+
|
| 235 |
+
|
| 236 |
+
class TestLiteLLMCallback:
|
| 237 |
+
def test_callback_imports(self):
|
| 238 |
+
"""Verify the callback can be imported."""
|
| 239 |
+
from headroom.integrations.litellm_callback import HeadroomCallback
|
| 240 |
+
|
| 241 |
+
callback = HeadroomCallback()
|
| 242 |
+
assert callback.total_tokens_saved == 0
|
| 243 |
+
|
| 244 |
+
def test_callback_compresses_messages(self):
|
| 245 |
+
"""Callback compresses messages in pre_call_hook."""
|
| 246 |
+
import asyncio
|
| 247 |
+
|
| 248 |
+
from headroom.integrations.litellm_callback import HeadroomCallback
|
| 249 |
+
|
| 250 |
+
callback = HeadroomCallback()
|
| 251 |
+
|
| 252 |
+
big_data = json.dumps([{"id": i, "status": "active"} for i in range(200)])
|
| 253 |
+
data = {
|
| 254 |
+
"model": "gpt-4o",
|
| 255 |
+
"messages": [
|
| 256 |
+
{"role": "user", "content": "analyze"},
|
| 257 |
+
{"role": "tool", "content": big_data, "tool_call_id": "c1"},
|
| 258 |
+
],
|
| 259 |
+
}
|
| 260 |
+
|
| 261 |
+
result = asyncio.run(callback.async_pre_call_hook("key", data, "completion"))
|
| 262 |
+
assert result is data
|
| 263 |
+
|
| 264 |
+
def test_callback_ignores_non_completion(self):
|
| 265 |
+
"""Non-completion calls are passed through."""
|
| 266 |
+
import asyncio
|
| 267 |
+
|
| 268 |
+
from headroom.integrations.litellm_callback import HeadroomCallback
|
| 269 |
+
|
| 270 |
+
callback = HeadroomCallback()
|
| 271 |
+
data = {"messages": [{"role": "user", "content": "hi"}]}
|
| 272 |
+
|
| 273 |
+
result = asyncio.run(callback.async_pre_call_hook("key", data, "embedding"))
|
| 274 |
+
assert result is data
|
|
@@ -1,39 +1,39 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
import importlib
|
| 4 |
-
from types import SimpleNamespace
|
| 5 |
-
|
| 6 |
-
from headroom.compress import compress
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
class _FailingPipeline:
|
| 10 |
-
def apply(self, **kwargs): # noqa: ANN003, ANN201
|
| 11 |
-
raise RuntimeError("boom")
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
def test_compress_returns_original_messages_when_pipeline_fails(monkeypatch) -> None:
|
| 15 |
-
metrics: list[dict[str, str]] = []
|
| 16 |
-
compress_module = importlib.import_module("headroom.compress")
|
| 17 |
-
monkeypatch.setattr(compress_module, "_get_pipeline", lambda: _FailingPipeline())
|
| 18 |
-
monkeypatch.setattr(
|
| 19 |
-
compress_module,
|
| 20 |
-
"get_otel_metrics",
|
| 21 |
-
lambda: SimpleNamespace(
|
| 22 |
-
record_compression_failure=lambda **kwargs: metrics.append(kwargs),
|
| 23 |
-
),
|
| 24 |
-
)
|
| 25 |
-
|
| 26 |
-
messages = [{"role": "user", "content": "hello world " * 100}]
|
| 27 |
-
result = compress(messages, model="gpt-4o")
|
| 28 |
-
|
| 29 |
-
assert result.messages == messages
|
| 30 |
-
assert result.tokens_before == 0
|
| 31 |
-
assert result.tokens_after == 0
|
| 32 |
-
assert result.tokens_saved == 0
|
| 33 |
-
assert metrics == [
|
| 34 |
-
{
|
| 35 |
-
"model": "gpt-4o",
|
| 36 |
-
"operation": "compress",
|
| 37 |
-
"error_type": "RuntimeError",
|
| 38 |
-
}
|
| 39 |
-
]
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import importlib
|
| 4 |
+
from types import SimpleNamespace
|
| 5 |
+
|
| 6 |
+
from headroom.compress import compress
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
class _FailingPipeline:
|
| 10 |
+
def apply(self, **kwargs): # noqa: ANN003, ANN201
|
| 11 |
+
raise RuntimeError("boom")
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def test_compress_returns_original_messages_when_pipeline_fails(monkeypatch) -> None:
|
| 15 |
+
metrics: list[dict[str, str]] = []
|
| 16 |
+
compress_module = importlib.import_module("headroom.compress")
|
| 17 |
+
monkeypatch.setattr(compress_module, "_get_pipeline", lambda: _FailingPipeline())
|
| 18 |
+
monkeypatch.setattr(
|
| 19 |
+
compress_module,
|
| 20 |
+
"get_otel_metrics",
|
| 21 |
+
lambda: SimpleNamespace(
|
| 22 |
+
record_compression_failure=lambda **kwargs: metrics.append(kwargs),
|
| 23 |
+
),
|
| 24 |
+
)
|
| 25 |
+
|
| 26 |
+
messages = [{"role": "user", "content": "hello world " * 100}]
|
| 27 |
+
result = compress(messages, model="gpt-4o")
|
| 28 |
+
|
| 29 |
+
assert result.messages == messages
|
| 30 |
+
assert result.tokens_before == 0
|
| 31 |
+
assert result.tokens_after == 0
|
| 32 |
+
assert result.tokens_saved == 0
|
| 33 |
+
assert metrics == [
|
| 34 |
+
{
|
| 35 |
+
"model": "gpt-4o",
|
| 36 |
+
"operation": "compress",
|
| 37 |
+
"error_type": "RuntimeError",
|
| 38 |
+
}
|
| 39 |
+
]
|
|
@@ -1,331 +1,331 @@
|
|
| 1 |
-
"""Unit tests for headroom.subscription.copilot_quota."""
|
| 2 |
-
|
| 3 |
-
from __future__ import annotations
|
| 4 |
-
|
| 5 |
-
import time
|
| 6 |
-
|
| 7 |
-
import pytest
|
| 8 |
-
|
| 9 |
-
from headroom.subscription.copilot_quota import (
|
| 10 |
-
CopilotQuotaCategory,
|
| 11 |
-
CopilotQuotaSnapshot,
|
| 12 |
-
discover_github_token,
|
| 13 |
-
parse_copilot_quota,
|
| 14 |
-
)
|
| 15 |
-
|
| 16 |
-
# ---------------------------------------------------------------------------
|
| 17 |
-
# CopilotQuotaCategory helpers
|
| 18 |
-
# ---------------------------------------------------------------------------
|
| 19 |
-
|
| 20 |
-
|
| 21 |
-
class TestCopilotQuotaCategory:
|
| 22 |
-
def test_used_computed_from_entitlement_and_remaining(self):
|
| 23 |
-
cat = CopilotQuotaCategory(name="chat", entitlement=300, remaining=120)
|
| 24 |
-
assert cat.used == 180
|
| 25 |
-
|
| 26 |
-
def test_used_percent_computed(self):
|
| 27 |
-
cat = CopilotQuotaCategory(name="chat", entitlement=100, remaining=25)
|
| 28 |
-
assert cat.used_percent == pytest.approx(75.0)
|
| 29 |
-
|
| 30 |
-
def test_used_percent_from_percent_remaining(self):
|
| 31 |
-
cat = CopilotQuotaCategory(name="completions", percent_remaining=40.0)
|
| 32 |
-
assert cat.used_percent == pytest.approx(60.0)
|
| 33 |
-
|
| 34 |
-
def test_unlimited_used_percent_is_zero(self):
|
| 35 |
-
cat = CopilotQuotaCategory(name="premium_interactions", unlimited=True)
|
| 36 |
-
assert cat.used_percent == 0.0
|
| 37 |
-
|
| 38 |
-
def test_used_none_when_entitlement_missing(self):
|
| 39 |
-
cat = CopilotQuotaCategory(name="chat", remaining=50)
|
| 40 |
-
assert cat.used is None
|
| 41 |
-
|
| 42 |
-
def test_to_dict_keys(self):
|
| 43 |
-
cat = CopilotQuotaCategory(
|
| 44 |
-
name="chat",
|
| 45 |
-
entitlement=100,
|
| 46 |
-
remaining=60,
|
| 47 |
-
percent_remaining=60.0,
|
| 48 |
-
overage_count=2,
|
| 49 |
-
overage_permitted=True,
|
| 50 |
-
unlimited=False,
|
| 51 |
-
timestamp_utc="2025-01-01T00:00:00Z",
|
| 52 |
-
)
|
| 53 |
-
d = cat.to_dict()
|
| 54 |
-
assert d["name"] == "chat"
|
| 55 |
-
assert d["entitlement"] == 100
|
| 56 |
-
assert d["remaining"] == 60
|
| 57 |
-
assert d["used"] == 40
|
| 58 |
-
assert d["used_percent"] == pytest.approx(40.0)
|
| 59 |
-
assert d["overage_count"] == 2
|
| 60 |
-
assert d["overage_permitted"] is True
|
| 61 |
-
assert d["unlimited"] is False
|
| 62 |
-
|
| 63 |
-
def test_used_percent_clipped_at_zero(self):
|
| 64 |
-
# percent_remaining > 100 should not produce negative used_percent
|
| 65 |
-
cat = CopilotQuotaCategory(name="chat", percent_remaining=110.0)
|
| 66 |
-
assert cat.used_percent == pytest.approx(0.0)
|
| 67 |
-
|
| 68 |
-
|
| 69 |
-
# ---------------------------------------------------------------------------
|
| 70 |
-
# parse_copilot_quota
|
| 71 |
-
# ---------------------------------------------------------------------------
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
_SAMPLE_RESPONSE = {
|
| 75 |
-
"login": "octocat",
|
| 76 |
-
"copilot_plan": "individual",
|
| 77 |
-
"access_type_sku": "copilot_for_individuals",
|
| 78 |
-
"quota_reset_date_utc": "2025-02-01",
|
| 79 |
-
"quota_snapshots": {
|
| 80 |
-
"chat": {
|
| 81 |
-
"entitlement": 50,
|
| 82 |
-
"remaining": 30,
|
| 83 |
-
"quota_remaining": 30,
|
| 84 |
-
"percent_remaining": 60.0,
|
| 85 |
-
"overage_count": 0,
|
| 86 |
-
"overage_permitted": False,
|
| 87 |
-
"unlimited": False,
|
| 88 |
-
"timestamp_utc": "2025-01-15T10:00:00Z",
|
| 89 |
-
},
|
| 90 |
-
"completions": {
|
| 91 |
-
"entitlement": 2000,
|
| 92 |
-
"remaining": 1500,
|
| 93 |
-
"percent_remaining": 75.0,
|
| 94 |
-
"overage_count": 0,
|
| 95 |
-
"overage_permitted": True,
|
| 96 |
-
"unlimited": False,
|
| 97 |
-
"timestamp_utc": "2025-01-15T10:00:00Z",
|
| 98 |
-
},
|
| 99 |
-
"premium_interactions": {
|
| 100 |
-
"entitlement": 300,
|
| 101 |
-
"remaining": 298,
|
| 102 |
-
"percent_remaining": 99.3,
|
| 103 |
-
"overage_count": 2,
|
| 104 |
-
"overage_permitted": True,
|
| 105 |
-
"unlimited": False,
|
| 106 |
-
"timestamp_utc": "2025-01-15T10:00:00Z",
|
| 107 |
-
},
|
| 108 |
-
},
|
| 109 |
-
}
|
| 110 |
-
|
| 111 |
-
|
| 112 |
-
class TestParseCopilotQuota:
|
| 113 |
-
def test_basic_fields(self):
|
| 114 |
-
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
|
| 115 |
-
assert snap.login == "octocat"
|
| 116 |
-
assert snap.copilot_plan == "individual"
|
| 117 |
-
assert snap.access_type_sku == "copilot_for_individuals"
|
| 118 |
-
assert snap.quota_reset_date_utc == "2025-02-01"
|
| 119 |
-
|
| 120 |
-
def test_all_categories_parsed(self):
|
| 121 |
-
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
|
| 122 |
-
assert set(snap.categories.keys()) == {"chat", "completions", "premium_interactions"}
|
| 123 |
-
|
| 124 |
-
def test_chat_category(self):
|
| 125 |
-
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
|
| 126 |
-
chat = snap.categories["chat"]
|
| 127 |
-
assert chat.entitlement == 50
|
| 128 |
-
assert chat.remaining == 30
|
| 129 |
-
assert chat.percent_remaining == pytest.approx(60.0)
|
| 130 |
-
assert chat.unlimited is False
|
| 131 |
-
assert chat.overage_count == 0
|
| 132 |
-
|
| 133 |
-
def test_premium_interactions_overage(self):
|
| 134 |
-
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
|
| 135 |
-
prem = snap.categories["premium_interactions"]
|
| 136 |
-
assert prem.overage_count == 2
|
| 137 |
-
assert prem.overage_permitted is True
|
| 138 |
-
|
| 139 |
-
def test_quota_remaining_alias(self):
|
| 140 |
-
"""quota_remaining should be used when remaining is absent."""
|
| 141 |
-
data = {
|
| 142 |
-
"quota_snapshots": {
|
| 143 |
-
"chat": {
|
| 144 |
-
"entitlement": 100,
|
| 145 |
-
"quota_remaining": 75,
|
| 146 |
-
}
|
| 147 |
-
}
|
| 148 |
-
}
|
| 149 |
-
snap = parse_copilot_quota(data)
|
| 150 |
-
assert snap.categories["chat"].remaining == 75
|
| 151 |
-
|
| 152 |
-
def test_unlimited_category(self):
|
| 153 |
-
data = {"quota_snapshots": {"completions": {"unlimited": True}}}
|
| 154 |
-
snap = parse_copilot_quota(data)
|
| 155 |
-
assert snap.categories["completions"].unlimited is True
|
| 156 |
-
|
| 157 |
-
def test_empty_quota_snapshots(self):
|
| 158 |
-
snap = parse_copilot_quota({"login": "ghost"})
|
| 159 |
-
assert snap.login == "ghost"
|
| 160 |
-
assert snap.categories == {}
|
| 161 |
-
|
| 162 |
-
def test_quota_reset_date_fallback(self):
|
| 163 |
-
data = {"quota_reset_date": "2025-03-01"}
|
| 164 |
-
snap = parse_copilot_quota(data)
|
| 165 |
-
assert snap.quota_reset_date_utc == "2025-03-01"
|
| 166 |
-
|
| 167 |
-
def test_fetched_at_is_recent(self):
|
| 168 |
-
before = time.time()
|
| 169 |
-
snap = parse_copilot_quota({})
|
| 170 |
-
after = time.time()
|
| 171 |
-
assert before <= snap.fetched_at <= after
|
| 172 |
-
|
| 173 |
-
def test_to_dict_structure(self):
|
| 174 |
-
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
|
| 175 |
-
d = snap.to_dict()
|
| 176 |
-
assert "login" in d
|
| 177 |
-
assert "categories" in d
|
| 178 |
-
assert "chat" in d["categories"]
|
| 179 |
-
assert "used_percent" in d["categories"]["chat"]
|
| 180 |
-
|
| 181 |
-
def test_missing_categories_skipped(self):
|
| 182 |
-
data = {
|
| 183 |
-
"quota_snapshots": {
|
| 184 |
-
"chat": {"remaining": 10},
|
| 185 |
-
# completions and premium_interactions absent
|
| 186 |
-
}
|
| 187 |
-
}
|
| 188 |
-
snap = parse_copilot_quota(data)
|
| 189 |
-
assert "chat" in snap.categories
|
| 190 |
-
assert "completions" not in snap.categories
|
| 191 |
-
assert "premium_interactions" not in snap.categories
|
| 192 |
-
|
| 193 |
-
def test_free_plan(self):
|
| 194 |
-
data = {"copilot_plan": "free", "quota_snapshots": {}}
|
| 195 |
-
snap = parse_copilot_quota(data)
|
| 196 |
-
assert snap.copilot_plan == "free"
|
| 197 |
-
|
| 198 |
-
|
| 199 |
-
# ---------------------------------------------------------------------------
|
| 200 |
-
# discover_github_token
|
| 201 |
-
# ---------------------------------------------------------------------------
|
| 202 |
-
|
| 203 |
-
|
| 204 |
-
class TestDiscoverGithubToken:
|
| 205 |
-
def test_returns_none_when_no_env_vars(self, monkeypatch):
|
| 206 |
-
for var in [
|
| 207 |
-
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 208 |
-
"GITHUB_TOKEN",
|
| 209 |
-
"COPILOT_GITHUB_TOKEN",
|
| 210 |
-
"GITHUB_COPILOT_API_TOKEN",
|
| 211 |
-
]:
|
| 212 |
-
monkeypatch.delenv(var, raising=False)
|
| 213 |
-
assert discover_github_token() is None
|
| 214 |
-
|
| 215 |
-
def test_picks_up_github_token(self, monkeypatch):
|
| 216 |
-
for var in [
|
| 217 |
-
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 218 |
-
"GITHUB_TOKEN",
|
| 219 |
-
"COPILOT_GITHUB_TOKEN",
|
| 220 |
-
"GITHUB_COPILOT_API_TOKEN",
|
| 221 |
-
]:
|
| 222 |
-
monkeypatch.delenv(var, raising=False)
|
| 223 |
-
monkeypatch.setenv("GITHUB_TOKEN", "ghp_testtoken123")
|
| 224 |
-
assert discover_github_token() == "ghp_testtoken123"
|
| 225 |
-
|
| 226 |
-
def test_prefers_copilot_specific_token(self, monkeypatch):
|
| 227 |
-
for var in [
|
| 228 |
-
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 229 |
-
"GITHUB_TOKEN",
|
| 230 |
-
"COPILOT_GITHUB_TOKEN",
|
| 231 |
-
"GITHUB_COPILOT_API_TOKEN",
|
| 232 |
-
]:
|
| 233 |
-
monkeypatch.delenv(var, raising=False)
|
| 234 |
-
monkeypatch.setenv("GITHUB_COPILOT_GITHUB_TOKEN", "ghp_copilot_specific")
|
| 235 |
-
monkeypatch.setenv("GITHUB_TOKEN", "ghp_generic")
|
| 236 |
-
assert discover_github_token() == "ghp_copilot_specific"
|
| 237 |
-
|
| 238 |
-
def test_falls_through_to_next_env_var(self, monkeypatch):
|
| 239 |
-
for var in [
|
| 240 |
-
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 241 |
-
"GITHUB_TOKEN",
|
| 242 |
-
"COPILOT_GITHUB_TOKEN",
|
| 243 |
-
"GITHUB_COPILOT_API_TOKEN",
|
| 244 |
-
]:
|
| 245 |
-
monkeypatch.delenv(var, raising=False)
|
| 246 |
-
monkeypatch.setenv("COPILOT_GITHUB_TOKEN", "ghp_copilot")
|
| 247 |
-
assert discover_github_token() == "ghp_copilot"
|
| 248 |
-
|
| 249 |
-
def test_ignores_empty_strings(self, monkeypatch):
|
| 250 |
-
for var in [
|
| 251 |
-
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 252 |
-
"GITHUB_TOKEN",
|
| 253 |
-
"COPILOT_GITHUB_TOKEN",
|
| 254 |
-
"GITHUB_COPILOT_API_TOKEN",
|
| 255 |
-
]:
|
| 256 |
-
monkeypatch.delenv(var, raising=False)
|
| 257 |
-
monkeypatch.setenv("GITHUB_COPILOT_GITHUB_TOKEN", "")
|
| 258 |
-
monkeypatch.setenv("GITHUB_TOKEN", "ghp_valid")
|
| 259 |
-
assert discover_github_token() == "ghp_valid"
|
| 260 |
-
|
| 261 |
-
|
| 262 |
-
# ---------------------------------------------------------------------------
|
| 263 |
-
# CopilotQuotaSnapshot.to_dict
|
| 264 |
-
# ---------------------------------------------------------------------------
|
| 265 |
-
|
| 266 |
-
|
| 267 |
-
class TestCopilotQuotaSnapshot:
|
| 268 |
-
def test_to_dict_complete(self):
|
| 269 |
-
snap = CopilotQuotaSnapshot(
|
| 270 |
-
login="user1",
|
| 271 |
-
copilot_plan="business",
|
| 272 |
-
access_type_sku="copilot_enterprise",
|
| 273 |
-
quota_reset_date_utc="2025-02-01",
|
| 274 |
-
)
|
| 275 |
-
snap.categories["chat"] = CopilotQuotaCategory(name="chat", entitlement=50, remaining=25)
|
| 276 |
-
d = snap.to_dict()
|
| 277 |
-
assert d["login"] == "user1"
|
| 278 |
-
assert d["copilot_plan"] == "business"
|
| 279 |
-
assert "chat" in d["categories"]
|
| 280 |
-
assert d["categories"]["chat"]["entitlement"] == 50
|
| 281 |
-
|
| 282 |
-
|
| 283 |
-
# ---------------------------------------------------------------------------
|
| 284 |
-
# Poll-loop task-leak regression
|
| 285 |
-
# ---------------------------------------------------------------------------
|
| 286 |
-
|
| 287 |
-
|
| 288 |
-
class TestCopilotQuotaPollLoopLeak:
|
| 289 |
-
@pytest.mark.asyncio
|
| 290 |
-
async def test_poll_loop_does_not_leak_event_wait_tasks(self, monkeypatch):
|
| 291 |
-
"""Regression for the ``asyncio.shield(event.wait())`` pattern.
|
| 292 |
-
|
| 293 |
-
Matches the equivalent guard in ``tests/test_subscription_tracker.py``:
|
| 294 |
-
every poll interval the loop previously leaked one Event.wait
|
| 295 |
-
waiter because ``asyncio.shield`` prevented ``wait_for`` from
|
| 296 |
-
cancelling the inner wait on timeout.
|
| 297 |
-
"""
|
| 298 |
-
import asyncio
|
| 299 |
-
|
| 300 |
-
from headroom.subscription.copilot_quota import _CopilotQuotaTracker
|
| 301 |
-
|
| 302 |
-
# No token configured → _maybe_poll returns immediately each cycle.
|
| 303 |
-
for var in ("GITHUB_COPILOT_GITHUB_TOKEN", "GITHUB_TOKEN"):
|
| 304 |
-
monkeypatch.delenv(var, raising=False)
|
| 305 |
-
|
| 306 |
-
tracker = _CopilotQuotaTracker(poll_interval_s=0.05)
|
| 307 |
-
|
| 308 |
-
def _count_event_wait() -> int:
|
| 309 |
-
return sum(
|
| 310 |
-
1
|
| 311 |
-
for t in asyncio.all_tasks()
|
| 312 |
-
if (t.get_coro().__qualname__ if t.get_coro() else "") == "Event.wait"
|
| 313 |
-
)
|
| 314 |
-
|
| 315 |
-
baseline = _count_event_wait()
|
| 316 |
-
await tracker.start()
|
| 317 |
-
try:
|
| 318 |
-
await asyncio.sleep(0.3) # ~6 poll cycles
|
| 319 |
-
peak = _count_event_wait()
|
| 320 |
-
finally:
|
| 321 |
-
await tracker.stop()
|
| 322 |
-
|
| 323 |
-
await asyncio.sleep(0.05)
|
| 324 |
-
residual = _count_event_wait()
|
| 325 |
-
|
| 326 |
-
assert peak - baseline <= 1, (
|
| 327 |
-
f"CopilotQuotaTracker leaked Event.wait: baseline={baseline} peak={peak}"
|
| 328 |
-
)
|
| 329 |
-
assert residual <= baseline, (
|
| 330 |
-
f"CopilotQuotaTracker left residual Event.wait: baseline={baseline} residual={residual}"
|
| 331 |
-
)
|
|
|
|
| 1 |
+
"""Unit tests for headroom.subscription.copilot_quota."""
|
| 2 |
+
|
| 3 |
+
from __future__ import annotations
|
| 4 |
+
|
| 5 |
+
import time
|
| 6 |
+
|
| 7 |
+
import pytest
|
| 8 |
+
|
| 9 |
+
from headroom.subscription.copilot_quota import (
|
| 10 |
+
CopilotQuotaCategory,
|
| 11 |
+
CopilotQuotaSnapshot,
|
| 12 |
+
discover_github_token,
|
| 13 |
+
parse_copilot_quota,
|
| 14 |
+
)
|
| 15 |
+
|
| 16 |
+
# ---------------------------------------------------------------------------
|
| 17 |
+
# CopilotQuotaCategory helpers
|
| 18 |
+
# ---------------------------------------------------------------------------
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
class TestCopilotQuotaCategory:
|
| 22 |
+
def test_used_computed_from_entitlement_and_remaining(self):
|
| 23 |
+
cat = CopilotQuotaCategory(name="chat", entitlement=300, remaining=120)
|
| 24 |
+
assert cat.used == 180
|
| 25 |
+
|
| 26 |
+
def test_used_percent_computed(self):
|
| 27 |
+
cat = CopilotQuotaCategory(name="chat", entitlement=100, remaining=25)
|
| 28 |
+
assert cat.used_percent == pytest.approx(75.0)
|
| 29 |
+
|
| 30 |
+
def test_used_percent_from_percent_remaining(self):
|
| 31 |
+
cat = CopilotQuotaCategory(name="completions", percent_remaining=40.0)
|
| 32 |
+
assert cat.used_percent == pytest.approx(60.0)
|
| 33 |
+
|
| 34 |
+
def test_unlimited_used_percent_is_zero(self):
|
| 35 |
+
cat = CopilotQuotaCategory(name="premium_interactions", unlimited=True)
|
| 36 |
+
assert cat.used_percent == 0.0
|
| 37 |
+
|
| 38 |
+
def test_used_none_when_entitlement_missing(self):
|
| 39 |
+
cat = CopilotQuotaCategory(name="chat", remaining=50)
|
| 40 |
+
assert cat.used is None
|
| 41 |
+
|
| 42 |
+
def test_to_dict_keys(self):
|
| 43 |
+
cat = CopilotQuotaCategory(
|
| 44 |
+
name="chat",
|
| 45 |
+
entitlement=100,
|
| 46 |
+
remaining=60,
|
| 47 |
+
percent_remaining=60.0,
|
| 48 |
+
overage_count=2,
|
| 49 |
+
overage_permitted=True,
|
| 50 |
+
unlimited=False,
|
| 51 |
+
timestamp_utc="2025-01-01T00:00:00Z",
|
| 52 |
+
)
|
| 53 |
+
d = cat.to_dict()
|
| 54 |
+
assert d["name"] == "chat"
|
| 55 |
+
assert d["entitlement"] == 100
|
| 56 |
+
assert d["remaining"] == 60
|
| 57 |
+
assert d["used"] == 40
|
| 58 |
+
assert d["used_percent"] == pytest.approx(40.0)
|
| 59 |
+
assert d["overage_count"] == 2
|
| 60 |
+
assert d["overage_permitted"] is True
|
| 61 |
+
assert d["unlimited"] is False
|
| 62 |
+
|
| 63 |
+
def test_used_percent_clipped_at_zero(self):
|
| 64 |
+
# percent_remaining > 100 should not produce negative used_percent
|
| 65 |
+
cat = CopilotQuotaCategory(name="chat", percent_remaining=110.0)
|
| 66 |
+
assert cat.used_percent == pytest.approx(0.0)
|
| 67 |
+
|
| 68 |
+
|
| 69 |
+
# ---------------------------------------------------------------------------
|
| 70 |
+
# parse_copilot_quota
|
| 71 |
+
# ---------------------------------------------------------------------------
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
_SAMPLE_RESPONSE = {
|
| 75 |
+
"login": "octocat",
|
| 76 |
+
"copilot_plan": "individual",
|
| 77 |
+
"access_type_sku": "copilot_for_individuals",
|
| 78 |
+
"quota_reset_date_utc": "2025-02-01",
|
| 79 |
+
"quota_snapshots": {
|
| 80 |
+
"chat": {
|
| 81 |
+
"entitlement": 50,
|
| 82 |
+
"remaining": 30,
|
| 83 |
+
"quota_remaining": 30,
|
| 84 |
+
"percent_remaining": 60.0,
|
| 85 |
+
"overage_count": 0,
|
| 86 |
+
"overage_permitted": False,
|
| 87 |
+
"unlimited": False,
|
| 88 |
+
"timestamp_utc": "2025-01-15T10:00:00Z",
|
| 89 |
+
},
|
| 90 |
+
"completions": {
|
| 91 |
+
"entitlement": 2000,
|
| 92 |
+
"remaining": 1500,
|
| 93 |
+
"percent_remaining": 75.0,
|
| 94 |
+
"overage_count": 0,
|
| 95 |
+
"overage_permitted": True,
|
| 96 |
+
"unlimited": False,
|
| 97 |
+
"timestamp_utc": "2025-01-15T10:00:00Z",
|
| 98 |
+
},
|
| 99 |
+
"premium_interactions": {
|
| 100 |
+
"entitlement": 300,
|
| 101 |
+
"remaining": 298,
|
| 102 |
+
"percent_remaining": 99.3,
|
| 103 |
+
"overage_count": 2,
|
| 104 |
+
"overage_permitted": True,
|
| 105 |
+
"unlimited": False,
|
| 106 |
+
"timestamp_utc": "2025-01-15T10:00:00Z",
|
| 107 |
+
},
|
| 108 |
+
},
|
| 109 |
+
}
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
class TestParseCopilotQuota:
|
| 113 |
+
def test_basic_fields(self):
|
| 114 |
+
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
|
| 115 |
+
assert snap.login == "octocat"
|
| 116 |
+
assert snap.copilot_plan == "individual"
|
| 117 |
+
assert snap.access_type_sku == "copilot_for_individuals"
|
| 118 |
+
assert snap.quota_reset_date_utc == "2025-02-01"
|
| 119 |
+
|
| 120 |
+
def test_all_categories_parsed(self):
|
| 121 |
+
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
|
| 122 |
+
assert set(snap.categories.keys()) == {"chat", "completions", "premium_interactions"}
|
| 123 |
+
|
| 124 |
+
def test_chat_category(self):
|
| 125 |
+
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
|
| 126 |
+
chat = snap.categories["chat"]
|
| 127 |
+
assert chat.entitlement == 50
|
| 128 |
+
assert chat.remaining == 30
|
| 129 |
+
assert chat.percent_remaining == pytest.approx(60.0)
|
| 130 |
+
assert chat.unlimited is False
|
| 131 |
+
assert chat.overage_count == 0
|
| 132 |
+
|
| 133 |
+
def test_premium_interactions_overage(self):
|
| 134 |
+
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
|
| 135 |
+
prem = snap.categories["premium_interactions"]
|
| 136 |
+
assert prem.overage_count == 2
|
| 137 |
+
assert prem.overage_permitted is True
|
| 138 |
+
|
| 139 |
+
def test_quota_remaining_alias(self):
|
| 140 |
+
"""quota_remaining should be used when remaining is absent."""
|
| 141 |
+
data = {
|
| 142 |
+
"quota_snapshots": {
|
| 143 |
+
"chat": {
|
| 144 |
+
"entitlement": 100,
|
| 145 |
+
"quota_remaining": 75,
|
| 146 |
+
}
|
| 147 |
+
}
|
| 148 |
+
}
|
| 149 |
+
snap = parse_copilot_quota(data)
|
| 150 |
+
assert snap.categories["chat"].remaining == 75
|
| 151 |
+
|
| 152 |
+
def test_unlimited_category(self):
|
| 153 |
+
data = {"quota_snapshots": {"completions": {"unlimited": True}}}
|
| 154 |
+
snap = parse_copilot_quota(data)
|
| 155 |
+
assert snap.categories["completions"].unlimited is True
|
| 156 |
+
|
| 157 |
+
def test_empty_quota_snapshots(self):
|
| 158 |
+
snap = parse_copilot_quota({"login": "ghost"})
|
| 159 |
+
assert snap.login == "ghost"
|
| 160 |
+
assert snap.categories == {}
|
| 161 |
+
|
| 162 |
+
def test_quota_reset_date_fallback(self):
|
| 163 |
+
data = {"quota_reset_date": "2025-03-01"}
|
| 164 |
+
snap = parse_copilot_quota(data)
|
| 165 |
+
assert snap.quota_reset_date_utc == "2025-03-01"
|
| 166 |
+
|
| 167 |
+
def test_fetched_at_is_recent(self):
|
| 168 |
+
before = time.time()
|
| 169 |
+
snap = parse_copilot_quota({})
|
| 170 |
+
after = time.time()
|
| 171 |
+
assert before <= snap.fetched_at <= after
|
| 172 |
+
|
| 173 |
+
def test_to_dict_structure(self):
|
| 174 |
+
snap = parse_copilot_quota(_SAMPLE_RESPONSE)
|
| 175 |
+
d = snap.to_dict()
|
| 176 |
+
assert "login" in d
|
| 177 |
+
assert "categories" in d
|
| 178 |
+
assert "chat" in d["categories"]
|
| 179 |
+
assert "used_percent" in d["categories"]["chat"]
|
| 180 |
+
|
| 181 |
+
def test_missing_categories_skipped(self):
|
| 182 |
+
data = {
|
| 183 |
+
"quota_snapshots": {
|
| 184 |
+
"chat": {"remaining": 10},
|
| 185 |
+
# completions and premium_interactions absent
|
| 186 |
+
}
|
| 187 |
+
}
|
| 188 |
+
snap = parse_copilot_quota(data)
|
| 189 |
+
assert "chat" in snap.categories
|
| 190 |
+
assert "completions" not in snap.categories
|
| 191 |
+
assert "premium_interactions" not in snap.categories
|
| 192 |
+
|
| 193 |
+
def test_free_plan(self):
|
| 194 |
+
data = {"copilot_plan": "free", "quota_snapshots": {}}
|
| 195 |
+
snap = parse_copilot_quota(data)
|
| 196 |
+
assert snap.copilot_plan == "free"
|
| 197 |
+
|
| 198 |
+
|
| 199 |
+
# ---------------------------------------------------------------------------
|
| 200 |
+
# discover_github_token
|
| 201 |
+
# ---------------------------------------------------------------------------
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
class TestDiscoverGithubToken:
|
| 205 |
+
def test_returns_none_when_no_env_vars(self, monkeypatch):
|
| 206 |
+
for var in [
|
| 207 |
+
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 208 |
+
"GITHUB_TOKEN",
|
| 209 |
+
"COPILOT_GITHUB_TOKEN",
|
| 210 |
+
"GITHUB_COPILOT_API_TOKEN",
|
| 211 |
+
]:
|
| 212 |
+
monkeypatch.delenv(var, raising=False)
|
| 213 |
+
assert discover_github_token() is None
|
| 214 |
+
|
| 215 |
+
def test_picks_up_github_token(self, monkeypatch):
|
| 216 |
+
for var in [
|
| 217 |
+
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 218 |
+
"GITHUB_TOKEN",
|
| 219 |
+
"COPILOT_GITHUB_TOKEN",
|
| 220 |
+
"GITHUB_COPILOT_API_TOKEN",
|
| 221 |
+
]:
|
| 222 |
+
monkeypatch.delenv(var, raising=False)
|
| 223 |
+
monkeypatch.setenv("GITHUB_TOKEN", "ghp_testtoken123")
|
| 224 |
+
assert discover_github_token() == "ghp_testtoken123"
|
| 225 |
+
|
| 226 |
+
def test_prefers_copilot_specific_token(self, monkeypatch):
|
| 227 |
+
for var in [
|
| 228 |
+
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 229 |
+
"GITHUB_TOKEN",
|
| 230 |
+
"COPILOT_GITHUB_TOKEN",
|
| 231 |
+
"GITHUB_COPILOT_API_TOKEN",
|
| 232 |
+
]:
|
| 233 |
+
monkeypatch.delenv(var, raising=False)
|
| 234 |
+
monkeypatch.setenv("GITHUB_COPILOT_GITHUB_TOKEN", "ghp_copilot_specific")
|
| 235 |
+
monkeypatch.setenv("GITHUB_TOKEN", "ghp_generic")
|
| 236 |
+
assert discover_github_token() == "ghp_copilot_specific"
|
| 237 |
+
|
| 238 |
+
def test_falls_through_to_next_env_var(self, monkeypatch):
|
| 239 |
+
for var in [
|
| 240 |
+
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 241 |
+
"GITHUB_TOKEN",
|
| 242 |
+
"COPILOT_GITHUB_TOKEN",
|
| 243 |
+
"GITHUB_COPILOT_API_TOKEN",
|
| 244 |
+
]:
|
| 245 |
+
monkeypatch.delenv(var, raising=False)
|
| 246 |
+
monkeypatch.setenv("COPILOT_GITHUB_TOKEN", "ghp_copilot")
|
| 247 |
+
assert discover_github_token() == "ghp_copilot"
|
| 248 |
+
|
| 249 |
+
def test_ignores_empty_strings(self, monkeypatch):
|
| 250 |
+
for var in [
|
| 251 |
+
"GITHUB_COPILOT_GITHUB_TOKEN",
|
| 252 |
+
"GITHUB_TOKEN",
|
| 253 |
+
"COPILOT_GITHUB_TOKEN",
|
| 254 |
+
"GITHUB_COPILOT_API_TOKEN",
|
| 255 |
+
]:
|
| 256 |
+
monkeypatch.delenv(var, raising=False)
|
| 257 |
+
monkeypatch.setenv("GITHUB_COPILOT_GITHUB_TOKEN", "")
|
| 258 |
+
monkeypatch.setenv("GITHUB_TOKEN", "ghp_valid")
|
| 259 |
+
assert discover_github_token() == "ghp_valid"
|
| 260 |
+
|
| 261 |
+
|
| 262 |
+
# ---------------------------------------------------------------------------
|
| 263 |
+
# CopilotQuotaSnapshot.to_dict
|
| 264 |
+
# ---------------------------------------------------------------------------
|
| 265 |
+
|
| 266 |
+
|
| 267 |
+
class TestCopilotQuotaSnapshot:
|
| 268 |
+
def test_to_dict_complete(self):
|
| 269 |
+
snap = CopilotQuotaSnapshot(
|
| 270 |
+
login="user1",
|
| 271 |
+
copilot_plan="business",
|
| 272 |
+
access_type_sku="copilot_enterprise",
|
| 273 |
+
quota_reset_date_utc="2025-02-01",
|
| 274 |
+
)
|
| 275 |
+
snap.categories["chat"] = CopilotQuotaCategory(name="chat", entitlement=50, remaining=25)
|
| 276 |
+
d = snap.to_dict()
|
| 277 |
+
assert d["login"] == "user1"
|
| 278 |
+
assert d["copilot_plan"] == "business"
|
| 279 |
+
assert "chat" in d["categories"]
|
| 280 |
+
assert d["categories"]["chat"]["entitlement"] == 50
|
| 281 |
+
|
| 282 |
+
|
| 283 |
+
# ---------------------------------------------------------------------------
|
| 284 |
+
# Poll-loop task-leak regression
|
| 285 |
+
# ---------------------------------------------------------------------------
|
| 286 |
+
|
| 287 |
+
|
| 288 |
+
class TestCopilotQuotaPollLoopLeak:
|
| 289 |
+
@pytest.mark.asyncio
|
| 290 |
+
async def test_poll_loop_does_not_leak_event_wait_tasks(self, monkeypatch):
|
| 291 |
+
"""Regression for the ``asyncio.shield(event.wait())`` pattern.
|
| 292 |
+
|
| 293 |
+
Matches the equivalent guard in ``tests/test_subscription_tracker.py``:
|
| 294 |
+
every poll interval the loop previously leaked one Event.wait
|
| 295 |
+
waiter because ``asyncio.shield`` prevented ``wait_for`` from
|
| 296 |
+
cancelling the inner wait on timeout.
|
| 297 |
+
"""
|
| 298 |
+
import asyncio
|
| 299 |
+
|
| 300 |
+
from headroom.subscription.copilot_quota import _CopilotQuotaTracker
|
| 301 |
+
|
| 302 |
+
# No token configured → _maybe_poll returns immediately each cycle.
|
| 303 |
+
for var in ("GITHUB_COPILOT_GITHUB_TOKEN", "GITHUB_TOKEN"):
|
| 304 |
+
monkeypatch.delenv(var, raising=False)
|
| 305 |
+
|
| 306 |
+
tracker = _CopilotQuotaTracker(poll_interval_s=0.05)
|
| 307 |
+
|
| 308 |
+
def _count_event_wait() -> int:
|
| 309 |
+
return sum(
|
| 310 |
+
1
|
| 311 |
+
for t in asyncio.all_tasks()
|
| 312 |
+
if (t.get_coro().__qualname__ if t.get_coro() else "") == "Event.wait"
|
| 313 |
+
)
|
| 314 |
+
|
| 315 |
+
baseline = _count_event_wait()
|
| 316 |
+
await tracker.start()
|
| 317 |
+
try:
|
| 318 |
+
await asyncio.sleep(0.3) # ~6 poll cycles
|
| 319 |
+
peak = _count_event_wait()
|
| 320 |
+
finally:
|
| 321 |
+
await tracker.stop()
|
| 322 |
+
|
| 323 |
+
await asyncio.sleep(0.05)
|
| 324 |
+
residual = _count_event_wait()
|
| 325 |
+
|
| 326 |
+
assert peak - baseline <= 1, (
|
| 327 |
+
f"CopilotQuotaTracker leaked Event.wait: baseline={baseline} peak={peak}"
|
| 328 |
+
)
|
| 329 |
+
assert residual <= baseline, (
|
| 330 |
+
f"CopilotQuotaTracker left residual Event.wait: baseline={baseline} residual={residual}"
|
| 331 |
+
)
|
|
@@ -1,538 +1,538 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
import json
|
| 4 |
-
import sys
|
| 5 |
-
import urllib.request
|
| 6 |
-
from types import SimpleNamespace
|
| 7 |
-
from urllib.error import URLError
|
| 8 |
-
|
| 9 |
-
import pytest
|
| 10 |
-
|
| 11 |
-
from headroom.evals import datasets
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
def install_fake_datasets(
|
| 15 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 16 |
-
mapping: dict[tuple[str, str | None, str | None], list[dict[str, object]]],
|
| 17 |
-
) -> list[tuple[str, str | None, str | None]]:
|
| 18 |
-
calls: list[tuple[str, str | None, str | None]] = []
|
| 19 |
-
|
| 20 |
-
def fake_load_dataset(name: str, subset: str | None = None, split: str | None = None):
|
| 21 |
-
key = (name, subset, split)
|
| 22 |
-
calls.append(key)
|
| 23 |
-
return mapping[key]
|
| 24 |
-
|
| 25 |
-
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
|
| 26 |
-
return calls
|
| 27 |
-
|
| 28 |
-
|
| 29 |
-
def test_check_datasets_installed_errors_without_dependency(
|
| 30 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 31 |
-
) -> None:
|
| 32 |
-
monkeypatch.delitem(sys.modules, "datasets", raising=False)
|
| 33 |
-
|
| 34 |
-
import builtins
|
| 35 |
-
|
| 36 |
-
real_import = builtins.__import__
|
| 37 |
-
|
| 38 |
-
def fake_import(name, globals=None, locals=None, fromlist=(), level=0): # noqa: ANN001
|
| 39 |
-
if name == "datasets":
|
| 40 |
-
raise ImportError("missing")
|
| 41 |
-
return real_import(name, globals, locals, fromlist, level)
|
| 42 |
-
|
| 43 |
-
monkeypatch.setattr(builtins, "__import__", fake_import)
|
| 44 |
-
|
| 45 |
-
with pytest.raises(ImportError, match="HuggingFace datasets required"):
|
| 46 |
-
datasets._check_datasets_installed()
|
| 47 |
-
|
| 48 |
-
|
| 49 |
-
def test_load_hotpotqa_and_natural_questions(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 50 |
-
calls = install_fake_datasets(
|
| 51 |
-
monkeypatch,
|
| 52 |
-
{
|
| 53 |
-
("hotpotqa/hotpot_qa", "fullwiki", "validation"): [
|
| 54 |
-
{
|
| 55 |
-
"context": {"title": ["Page A"], "sentences": [["Line 1", "Line 2"]]},
|
| 56 |
-
"question": "Who?",
|
| 57 |
-
"answer": "Alice",
|
| 58 |
-
"type": "bridge",
|
| 59 |
-
"level": "easy",
|
| 60 |
-
}
|
| 61 |
-
],
|
| 62 |
-
("google-research-datasets/natural_questions", "default", "validation"): [
|
| 63 |
-
{"document": {}, "question": {"text": "skip me"}},
|
| 64 |
-
{
|
| 65 |
-
"document": {
|
| 66 |
-
"tokens": {
|
| 67 |
-
"token": ["<p>", "Ada", "Lovelace", "wrote", "notes"],
|
| 68 |
-
"is_html": [True, False, False, False, False],
|
| 69 |
-
}
|
| 70 |
-
},
|
| 71 |
-
"question": {"text": "Who wrote notes?"},
|
| 72 |
-
"annotations": {"short_answers": [[{"start_token": 1, "end_token": 3}]]},
|
| 73 |
-
},
|
| 74 |
-
],
|
| 75 |
-
},
|
| 76 |
-
)
|
| 77 |
-
|
| 78 |
-
hotpot = datasets.load_hotpotqa(n=1)
|
| 79 |
-
natural = datasets.load_natural_questions(n=1)
|
| 80 |
-
|
| 81 |
-
assert calls == [
|
| 82 |
-
("hotpotqa/hotpot_qa", "fullwiki", "validation"),
|
| 83 |
-
("google-research-datasets/natural_questions", "default", "validation"),
|
| 84 |
-
]
|
| 85 |
-
assert hotpot.name == "HotpotQA"
|
| 86 |
-
assert hotpot.cases[0].context == "## Page A\nLine 1\nLine 2"
|
| 87 |
-
assert hotpot.cases[0].metadata["type"] == "bridge"
|
| 88 |
-
assert natural.name == "Natural_Questions"
|
| 89 |
-
assert natural.cases[0].context == "Ada Lovelace wrote notes"
|
| 90 |
-
assert natural.cases[0].ground_truth == "Ada Lovelace"
|
| 91 |
-
|
| 92 |
-
|
| 93 |
-
def test_load_triviaqa_msmarco_and_squad(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 94 |
-
install_fake_datasets(
|
| 95 |
-
monkeypatch,
|
| 96 |
-
{
|
| 97 |
-
("trivia_qa", "rc", "validation"): [
|
| 98 |
-
{"question": "", "search_results": {"search_context": ["unused"]}},
|
| 99 |
-
{
|
| 100 |
-
"question": "Question 1",
|
| 101 |
-
"search_results": {"search_context": ["A", "B"]},
|
| 102 |
-
"answer": {"value": "Answer", "aliases": ["Alias"]},
|
| 103 |
-
},
|
| 104 |
-
{
|
| 105 |
-
"question": "Question 2",
|
| 106 |
-
"search_results": {"search_context": []},
|
| 107 |
-
"entity_pages": {"wiki_context": ["Wiki 1", "Wiki 2"]},
|
| 108 |
-
"answer": {"normalized_value": "Normalized"},
|
| 109 |
-
},
|
| 110 |
-
],
|
| 111 |
-
("microsoft/ms_marco", "v2.1", "validation"): [
|
| 112 |
-
{"query": "", "passages": {"passage_text": ["skip"], "is_selected": [True]}},
|
| 113 |
-
{
|
| 114 |
-
"query": "Find docs",
|
| 115 |
-
"passages": {"passage_text": ["Doc 1", "Doc 2"], "is_selected": [True, False]},
|
| 116 |
-
"answers": ["Primary answer"],
|
| 117 |
-
"query_type": "description",
|
| 118 |
-
},
|
| 119 |
-
],
|
| 120 |
-
("rajpurkar/squad_v2", None, "validation"): [
|
| 121 |
-
{"answers": {"text": []}, "context": "skip", "question": "skip"},
|
| 122 |
-
{
|
| 123 |
-
"context": "Context",
|
| 124 |
-
"question": "Question",
|
| 125 |
-
"answers": {"text": ["First answer"]},
|
| 126 |
-
"title": "Title",
|
| 127 |
-
},
|
| 128 |
-
],
|
| 129 |
-
},
|
| 130 |
-
)
|
| 131 |
-
|
| 132 |
-
trivia = datasets.load_triviaqa(n=2)
|
| 133 |
-
msmarco = datasets.load_msmarco(n=1)
|
| 134 |
-
squad = datasets.load_squad(n=1)
|
| 135 |
-
|
| 136 |
-
assert len(trivia.cases) == 2
|
| 137 |
-
assert trivia.cases[0].context == "A\n\nB"
|
| 138 |
-
assert trivia.cases[1].ground_truth == "Normalized"
|
| 139 |
-
assert trivia.cases[1].metadata["aliases"] == []
|
| 140 |
-
assert msmarco.cases[0].context.startswith("[RELEVANT] Passage 1: Doc 1")
|
| 141 |
-
assert msmarco.cases[0].metadata["num_passages"] == 2
|
| 142 |
-
assert squad.cases[0].ground_truth == "First answer"
|
| 143 |
-
assert squad.cases[0].metadata["title"] == "Title"
|
| 144 |
-
|
| 145 |
-
|
| 146 |
-
def test_load_longbench_narrativeqa_toolbench_codesearchnet_and_humaneval(
|
| 147 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 148 |
-
) -> None:
|
| 149 |
-
install_fake_datasets(
|
| 150 |
-
monkeypatch,
|
| 151 |
-
{
|
| 152 |
-
("THUDM/LongBench", "qasper", "test"): [
|
| 153 |
-
{"context": "", "input": "skip"},
|
| 154 |
-
{"context": "Long context", "input": "Question", "answers": ["Truth"]},
|
| 155 |
-
],
|
| 156 |
-
("deepmind/narrativeqa", None, "test"): [
|
| 157 |
-
{
|
| 158 |
-
"document": {"summary": {"text": "Story summary"}, "kind": "movie"},
|
| 159 |
-
"question": {"text": "What happened?"},
|
| 160 |
-
"answers": [{"text": "A"}, {"text": "B"}],
|
| 161 |
-
}
|
| 162 |
-
],
|
| 163 |
-
("ToolBench/ToolBench", "G1", "test"): [
|
| 164 |
-
{"api_list": [], "query": "skip"},
|
| 165 |
-
{
|
| 166 |
-
"api_list": [
|
| 167 |
-
{
|
| 168 |
-
"api_name": "weather",
|
| 169 |
-
"api_description": "Get weather",
|
| 170 |
-
"required_parameters": [{"name": "city"}],
|
| 171 |
-
"optional_parameters": [{"name": "unit"}],
|
| 172 |
-
}
|
| 173 |
-
],
|
| 174 |
-
"query": "Weather in SF?",
|
| 175 |
-
"answer": "Call weather",
|
| 176 |
-
},
|
| 177 |
-
],
|
| 178 |
-
("code_search_net", "python", "test"): [
|
| 179 |
-
{"func_code_string": "", "func_documentation_string": "skip"},
|
| 180 |
-
{
|
| 181 |
-
"func_code_string": "def add(a, b): return a + b",
|
| 182 |
-
"func_documentation_string": "Add two numbers.",
|
| 183 |
-
"func_name": "add",
|
| 184 |
-
"repository_name": "repo",
|
| 185 |
-
},
|
| 186 |
-
],
|
| 187 |
-
("openai_humaneval", None, "test"): [
|
| 188 |
-
{"prompt": "", "canonical_solution": "skip"},
|
| 189 |
-
{
|
| 190 |
-
"task_id": "HumanEval/1",
|
| 191 |
-
"prompt": "def solve(x):",
|
| 192 |
-
"canonical_solution": "return x",
|
| 193 |
-
"entry_point": "solve",
|
| 194 |
-
"test": "assert solve(1) == 1",
|
| 195 |
-
},
|
| 196 |
-
],
|
| 197 |
-
},
|
| 198 |
-
)
|
| 199 |
-
|
| 200 |
-
longbench = datasets.load_longbench(n=2, task="qasper")
|
| 201 |
-
narrative = datasets.load_narrativeqa(n=1)
|
| 202 |
-
toolbench = datasets.load_toolbench(n=1, category="G1")
|
| 203 |
-
codesearchnet = datasets.load_codesearchnet(n=1, language="python")
|
| 204 |
-
humaneval = datasets.load_humaneval(n=2)
|
| 205 |
-
|
| 206 |
-
assert longbench.name == "LongBench_qasper"
|
| 207 |
-
assert longbench.cases[0].metadata["context_length"] == len("Long context")
|
| 208 |
-
assert narrative.cases[0].metadata["all_answers"] == ["A", "B"]
|
| 209 |
-
assert toolbench.cases[0].metadata["num_tools"] == 1
|
| 210 |
-
assert '"name": "weather"' in toolbench.cases[0].context
|
| 211 |
-
assert codesearchnet.cases[0].ground_truth == "Add two numbers."
|
| 212 |
-
assert humaneval.cases[0].id == "humaneval_HumanEval/1"
|
| 213 |
-
assert humaneval.cases[0].metadata["entry_point"] == "solve"
|
| 214 |
-
|
| 215 |
-
|
| 216 |
-
def test_load_longbench_toolbench_and_codesearchnet_wrap_loader_errors(
|
| 217 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 218 |
-
) -> None:
|
| 219 |
-
def fake_load_dataset(name: str, subset: str | None = None, split: str | None = None): # noqa: ANN001
|
| 220 |
-
raise RuntimeError(f"broken {name}:{subset}:{split}")
|
| 221 |
-
|
| 222 |
-
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
|
| 223 |
-
|
| 224 |
-
with pytest.raises(ValueError, match="Failed to load LongBench task 'gov_report'"):
|
| 225 |
-
datasets.load_longbench(task="gov_report")
|
| 226 |
-
with pytest.raises(ValueError, match="Failed to load ToolBench category 'G2'"):
|
| 227 |
-
datasets.load_toolbench(category="G2")
|
| 228 |
-
with pytest.raises(ValueError, match="Failed to load CodeSearchNet for 'go'"):
|
| 229 |
-
datasets.load_codesearchnet(language="go")
|
| 230 |
-
|
| 231 |
-
|
| 232 |
-
def test_load_bfcl_success_and_download_failure(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 233 |
-
data_lines = "\n".join(
|
| 234 |
-
[
|
| 235 |
-
json.dumps(
|
| 236 |
-
{
|
| 237 |
-
"id": "case-1",
|
| 238 |
-
"question": [[{"role": "user", "content": "How is the weather?"}]],
|
| 239 |
-
"function": [{"name": "weather"}],
|
| 240 |
-
}
|
| 241 |
-
),
|
| 242 |
-
json.dumps({"question": [123], "function": []}),
|
| 243 |
-
]
|
| 244 |
-
)
|
| 245 |
-
gt_lines = json.dumps({"id": "case-1", "ground_truth": [{"name": "weather"}]})
|
| 246 |
-
|
| 247 |
-
def fake_urlopen(url: str): # noqa: ANN001
|
| 248 |
-
if "possible_answer/BFCL_v3_simple.json" in url:
|
| 249 |
-
return SimpleNamespace(read=lambda: gt_lines.encode("utf-8"))
|
| 250 |
-
if "BFCL_v3_simple.json" in url:
|
| 251 |
-
return SimpleNamespace(read=lambda: data_lines.encode("utf-8"))
|
| 252 |
-
raise URLError("missing")
|
| 253 |
-
|
| 254 |
-
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
|
| 255 |
-
|
| 256 |
-
suite = datasets.load_bfcl(n=2, category="simple")
|
| 257 |
-
assert suite.name == "BFCL_simple"
|
| 258 |
-
assert suite.cases[0].query == "How is the weather?"
|
| 259 |
-
assert suite.cases[0].ground_truth == '[{"name": "weather"}]'
|
| 260 |
-
assert suite.cases[0].metadata["num_functions"] == 1
|
| 261 |
-
|
| 262 |
-
def failing_urlopen(url: str): # noqa: ANN001
|
| 263 |
-
raise URLError("offline")
|
| 264 |
-
|
| 265 |
-
monkeypatch.setattr(urllib.request, "urlopen", failing_urlopen)
|
| 266 |
-
with pytest.raises(ValueError, match="Failed to download BFCL dataset 'BFCL_v3_parallel.json'"):
|
| 267 |
-
datasets.load_bfcl(category="parallel")
|
| 268 |
-
|
| 269 |
-
|
| 270 |
-
def test_tool_output_samples_custom_dataset_and_probe_generation(tmp_path) -> None:
|
| 271 |
-
tool_outputs = datasets.load_tool_output_samples()
|
| 272 |
-
assert tool_outputs.name == "ToolOutputSamples"
|
| 273 |
-
assert len(tool_outputs.cases) >= 8
|
| 274 |
-
assert tool_outputs.cases[0].ground_truth == "prompt-optimizer"
|
| 275 |
-
|
| 276 |
-
custom_path = tmp_path / "custom.jsonl"
|
| 277 |
-
custom_path.write_text(
|
| 278 |
-
json.dumps(
|
| 279 |
-
{"id": "case1", "context": "Context", "query": "Question", "ground_truth": "Answer"}
|
| 280 |
-
)
|
| 281 |
-
+ "\n",
|
| 282 |
-
encoding="utf-8",
|
| 283 |
-
)
|
| 284 |
-
custom_suite = datasets.load_custom_dataset(custom_path)
|
| 285 |
-
assert custom_suite.cases[0].id == "case1"
|
| 286 |
-
|
| 287 |
-
probes = datasets.generate_retrieval_probes(
|
| 288 |
-
'Alice Smith deployed API on 2024-01-15 at 99.9% confidence for "Launch Ready" and build_id',
|
| 289 |
-
n_probes=5,
|
| 290 |
-
)
|
| 291 |
-
assert "Alice Smith" in probes
|
| 292 |
-
assert "2024-01-15" in probes
|
| 293 |
-
assert "API" in probes
|
| 294 |
-
assert "99.9" in probes
|
| 295 |
-
assert "Launch Ready" in probes
|
| 296 |
-
|
| 297 |
-
|
| 298 |
-
def test_dataset_registry_helpers(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 299 |
-
categories = datasets.list_available_datasets()
|
| 300 |
-
assert "hotpotqa" in categories["rag"]
|
| 301 |
-
assert "tool_outputs" in categories["tool_use"]
|
| 302 |
-
|
| 303 |
-
seen: list[tuple[str, dict[str, object]]] = []
|
| 304 |
-
|
| 305 |
-
def fake_loader(*, n: int = 0, **kwargs): # noqa: ANN003
|
| 306 |
-
seen.append(("with-n", {"n": n, **kwargs}))
|
| 307 |
-
return "with-n-result"
|
| 308 |
-
|
| 309 |
-
def fixed_loader(**kwargs): # noqa: ANN003
|
| 310 |
-
seen.append(("fixed", kwargs))
|
| 311 |
-
return "fixed-result"
|
| 312 |
-
|
| 313 |
-
original_registry = dict(datasets.DATASET_REGISTRY)
|
| 314 |
-
monkeypatch.setattr(
|
| 315 |
-
datasets,
|
| 316 |
-
"DATASET_REGISTRY",
|
| 317 |
-
{
|
| 318 |
-
**original_registry,
|
| 319 |
-
"fake_n": {"loader": fake_loader, "category": "x", "description": "", "default_n": 3},
|
| 320 |
-
"fake_fixed": {
|
| 321 |
-
"loader": fixed_loader,
|
| 322 |
-
"category": "x",
|
| 323 |
-
"description": "",
|
| 324 |
-
"default_n": None,
|
| 325 |
-
},
|
| 326 |
-
},
|
| 327 |
-
)
|
| 328 |
-
|
| 329 |
-
assert datasets.load_dataset_by_name("fake_n") == "with-n-result"
|
| 330 |
-
assert datasets.load_dataset_by_name("fake_n", n=7, split="test") == "with-n-result"
|
| 331 |
-
assert datasets.load_dataset_by_name("fake_fixed", path="x") == "fixed-result"
|
| 332 |
-
assert seen == [
|
| 333 |
-
("with-n", {"n": 3}),
|
| 334 |
-
("with-n", {"n": 7, "split": "test"}),
|
| 335 |
-
("fixed", {"path": "x"}),
|
| 336 |
-
]
|
| 337 |
-
|
| 338 |
-
with pytest.raises(ValueError, match="Unknown dataset 'missing'"):
|
| 339 |
-
datasets.load_dataset_by_name("missing")
|
| 340 |
-
|
| 341 |
-
|
| 342 |
-
def test_dataset_loaders_cover_skip_and_limit_branches(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 343 |
-
install_fake_datasets(
|
| 344 |
-
monkeypatch,
|
| 345 |
-
{
|
| 346 |
-
("hotpotqa/hotpot_qa", "fullwiki", "validation"): [
|
| 347 |
-
{
|
| 348 |
-
"context": {"title": ["Page A"], "sentences": [["Line 1"]]},
|
| 349 |
-
"question": "Q1",
|
| 350 |
-
"answer": "A1",
|
| 351 |
-
},
|
| 352 |
-
{
|
| 353 |
-
"context": {"title": ["Page B"], "sentences": [["Line 2"]]},
|
| 354 |
-
"question": "Q2",
|
| 355 |
-
"answer": "A2",
|
| 356 |
-
},
|
| 357 |
-
],
|
| 358 |
-
("google-research-datasets/natural_questions", "default", "validation"): [
|
| 359 |
-
{
|
| 360 |
-
"document": {"tokens": {"token": ["x"], "is_html": [False]}},
|
| 361 |
-
"question": {"text": ""},
|
| 362 |
-
},
|
| 363 |
-
{
|
| 364 |
-
"document": {"tokens": {"token": ["<b>"], "is_html": [True]}},
|
| 365 |
-
"question": {"text": "blank context"},
|
| 366 |
-
},
|
| 367 |
-
{
|
| 368 |
-
"document": {"tokens": {"token": ["Ada", "wrote"], "is_html": [False, False]}},
|
| 369 |
-
"question": {"text": "Who?"},
|
| 370 |
-
"annotations": {"short_answers": [[{"start_token": 1, "end_token": 1}]]},
|
| 371 |
-
},
|
| 372 |
-
{
|
| 373 |
-
"document": {"tokens": {"token": ["Grace"], "is_html": [False]}},
|
| 374 |
-
"question": {"text": "Ignored by limit"},
|
| 375 |
-
},
|
| 376 |
-
],
|
| 377 |
-
("trivia_qa", "rc", "validation"): [
|
| 378 |
-
{"question": "skip", "search_results": {"search_context": []}, "entity_pages": {}},
|
| 379 |
-
{"question": "blank", "search_results": {"search_context": [""]}},
|
| 380 |
-
{
|
| 381 |
-
"question": "Good 1",
|
| 382 |
-
"search_results": {"search_context": ["Context 1"]},
|
| 383 |
-
"answer": {"value": "A1"},
|
| 384 |
-
},
|
| 385 |
-
{
|
| 386 |
-
"question": "Good 2",
|
| 387 |
-
"search_results": {"search_context": ["Context 2"]},
|
| 388 |
-
"answer": {"value": "A2"},
|
| 389 |
-
},
|
| 390 |
-
],
|
| 391 |
-
("microsoft/ms_marco", "v2.1", "validation"): [
|
| 392 |
-
{"query": "skip", "passages": {"passage_text": [], "is_selected": []}},
|
| 393 |
-
{
|
| 394 |
-
"query": "Find one",
|
| 395 |
-
"passages": {"passage_text": ["Doc 1"], "is_selected": [False]},
|
| 396 |
-
"answers": [],
|
| 397 |
-
},
|
| 398 |
-
{
|
| 399 |
-
"query": "Find two",
|
| 400 |
-
"passages": {"passage_text": ["Doc 2"], "is_selected": [True]},
|
| 401 |
-
"answers": ["A2"],
|
| 402 |
-
},
|
| 403 |
-
],
|
| 404 |
-
("rajpurkar/squad_v2", None, "validation"): [
|
| 405 |
-
{
|
| 406 |
-
"context": "Context 1",
|
| 407 |
-
"question": "Q1",
|
| 408 |
-
"answers": {"text": ["A1"]},
|
| 409 |
-
},
|
| 410 |
-
{
|
| 411 |
-
"context": "Context 2",
|
| 412 |
-
"question": "Q2",
|
| 413 |
-
"answers": {"text": ["A2"]},
|
| 414 |
-
},
|
| 415 |
-
],
|
| 416 |
-
("THUDM/LongBench", "qasper", "test"): [
|
| 417 |
-
{"context": "Context 1", "input": "Q1", "answers": ["A1"]},
|
| 418 |
-
{"context": "Has context", "input": ""},
|
| 419 |
-
{"context": "Context 2", "input": "Q2", "answers": ["A2"]},
|
| 420 |
-
],
|
| 421 |
-
("deepmind/narrativeqa", None, "test"): [
|
| 422 |
-
{"document": {"summary": {"text": ""}}, "question": {"text": "skip"}},
|
| 423 |
-
{"document": {"summary": {"text": "Story"}}, "question": {"text": ""}},
|
| 424 |
-
{
|
| 425 |
-
"document": {"summary": {"text": "Story 1"}, "kind": "book"},
|
| 426 |
-
"question": {"text": "Q1"},
|
| 427 |
-
"answers": [{"text": "A1"}],
|
| 428 |
-
},
|
| 429 |
-
{
|
| 430 |
-
"document": {"summary": {"text": "Story 2"}, "kind": "movie"},
|
| 431 |
-
"question": {"text": "Q2"},
|
| 432 |
-
"answers": [{"text": "A2"}],
|
| 433 |
-
},
|
| 434 |
-
],
|
| 435 |
-
("ToolBench/ToolBench", "G1", "test"): [
|
| 436 |
-
{"api_list": [], "query": "skip"},
|
| 437 |
-
{
|
| 438 |
-
"api_list": [
|
| 439 |
-
{
|
| 440 |
-
"api_name": "weather",
|
| 441 |
-
"required_parameters": [],
|
| 442 |
-
"optional_parameters": [],
|
| 443 |
-
}
|
| 444 |
-
],
|
| 445 |
-
"query": "",
|
| 446 |
-
},
|
| 447 |
-
{
|
| 448 |
-
"api_list": [
|
| 449 |
-
{"api_name": "calc", "required_parameters": [], "optional_parameters": []}
|
| 450 |
-
],
|
| 451 |
-
"query": "Good",
|
| 452 |
-
},
|
| 453 |
-
],
|
| 454 |
-
("code_search_net", "python", "test"): [
|
| 455 |
-
{
|
| 456 |
-
"func_code_string": "",
|
| 457 |
-
"whole_func_string": "",
|
| 458 |
-
"func_documentation_string": "skip",
|
| 459 |
-
},
|
| 460 |
-
{"whole_func_string": "def alt(): pass", "func_documentation_string": ""},
|
| 461 |
-
{
|
| 462 |
-
"whole_func_string": "def good(): pass",
|
| 463 |
-
"func_documentation_string": "Good doc",
|
| 464 |
-
"func_name": "good",
|
| 465 |
-
"repository_name": "repo",
|
| 466 |
-
},
|
| 467 |
-
{
|
| 468 |
-
"whole_func_string": "def ignored(): pass",
|
| 469 |
-
"func_documentation_string": "Ignored by limit",
|
| 470 |
-
},
|
| 471 |
-
],
|
| 472 |
-
("openai_humaneval", None, "test"): [
|
| 473 |
-
{
|
| 474 |
-
"task_id": "Task/1",
|
| 475 |
-
"prompt": "def solve():",
|
| 476 |
-
"canonical_solution": "return 1",
|
| 477 |
-
"test": "assert solve() == 1",
|
| 478 |
-
},
|
| 479 |
-
{
|
| 480 |
-
"task_id": "Task/2",
|
| 481 |
-
"prompt": "def other():",
|
| 482 |
-
"canonical_solution": "return 2",
|
| 483 |
-
"test": "assert other() == 2",
|
| 484 |
-
},
|
| 485 |
-
],
|
| 486 |
-
},
|
| 487 |
-
)
|
| 488 |
-
|
| 489 |
-
assert len(datasets.load_hotpotqa(n=1).cases) == 1
|
| 490 |
-
natural = datasets.load_natural_questions(n=1)
|
| 491 |
-
assert len(natural.cases) == 1
|
| 492 |
-
assert natural.cases[0].ground_truth is None
|
| 493 |
-
assert len(datasets.load_triviaqa(n=1).cases) == 1
|
| 494 |
-
msmarco = datasets.load_msmarco(n=1)
|
| 495 |
-
assert len(msmarco.cases) == 1
|
| 496 |
-
assert msmarco.cases[0].ground_truth is None
|
| 497 |
-
assert len(datasets.load_squad(n=1).cases) == 1
|
| 498 |
-
assert len(datasets.load_longbench(n=2, task="qasper").cases) == 1
|
| 499 |
-
assert len(datasets.load_narrativeqa(n=1).cases) == 1
|
| 500 |
-
assert len(datasets.load_toolbench(n=1, category="G1").cases) == 1
|
| 501 |
-
assert len(datasets.load_codesearchnet(n=1, language="python").cases) == 1
|
| 502 |
-
assert len(datasets.load_humaneval(n=1).cases) == 1
|
| 503 |
-
|
| 504 |
-
|
| 505 |
-
def test_load_bfcl_handles_optional_ground_truth_and_question_fallback(
|
| 506 |
-
monkeypatch: pytest.MonkeyPatch,
|
| 507 |
-
) -> None:
|
| 508 |
-
data_lines = "\n".join(
|
| 509 |
-
[
|
| 510 |
-
json.dumps(
|
| 511 |
-
{
|
| 512 |
-
"id": "case-1",
|
| 513 |
-
"question": [123],
|
| 514 |
-
"function": [{"name": "weather"}],
|
| 515 |
-
}
|
| 516 |
-
),
|
| 517 |
-
json.dumps({"id": "skip", "function": []}),
|
| 518 |
-
json.dumps(
|
| 519 |
-
{
|
| 520 |
-
"id": "case-2",
|
| 521 |
-
"question": [[{"role": "user", "content": "Ignored by limit"}]],
|
| 522 |
-
"function": [{"name": "time"}],
|
| 523 |
-
}
|
| 524 |
-
),
|
| 525 |
-
]
|
| 526 |
-
)
|
| 527 |
-
|
| 528 |
-
def fake_urlopen(url: str): # noqa: ANN001
|
| 529 |
-
if "possible_answer" in url:
|
| 530 |
-
raise URLError("missing ground truth")
|
| 531 |
-
return SimpleNamespace(read=lambda: data_lines.encode("utf-8"))
|
| 532 |
-
|
| 533 |
-
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
|
| 534 |
-
|
| 535 |
-
suite = datasets.load_bfcl(n=2, category="simple")
|
| 536 |
-
assert len(suite.cases) == 1
|
| 537 |
-
assert suite.cases[0].query == "[123]"
|
| 538 |
-
assert suite.cases[0].ground_truth is None
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
import sys
|
| 5 |
+
import urllib.request
|
| 6 |
+
from types import SimpleNamespace
|
| 7 |
+
from urllib.error import URLError
|
| 8 |
+
|
| 9 |
+
import pytest
|
| 10 |
+
|
| 11 |
+
from headroom.evals import datasets
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def install_fake_datasets(
|
| 15 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 16 |
+
mapping: dict[tuple[str, str | None, str | None], list[dict[str, object]]],
|
| 17 |
+
) -> list[tuple[str, str | None, str | None]]:
|
| 18 |
+
calls: list[tuple[str, str | None, str | None]] = []
|
| 19 |
+
|
| 20 |
+
def fake_load_dataset(name: str, subset: str | None = None, split: str | None = None):
|
| 21 |
+
key = (name, subset, split)
|
| 22 |
+
calls.append(key)
|
| 23 |
+
return mapping[key]
|
| 24 |
+
|
| 25 |
+
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
|
| 26 |
+
return calls
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
def test_check_datasets_installed_errors_without_dependency(
|
| 30 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 31 |
+
) -> None:
|
| 32 |
+
monkeypatch.delitem(sys.modules, "datasets", raising=False)
|
| 33 |
+
|
| 34 |
+
import builtins
|
| 35 |
+
|
| 36 |
+
real_import = builtins.__import__
|
| 37 |
+
|
| 38 |
+
def fake_import(name, globals=None, locals=None, fromlist=(), level=0): # noqa: ANN001
|
| 39 |
+
if name == "datasets":
|
| 40 |
+
raise ImportError("missing")
|
| 41 |
+
return real_import(name, globals, locals, fromlist, level)
|
| 42 |
+
|
| 43 |
+
monkeypatch.setattr(builtins, "__import__", fake_import)
|
| 44 |
+
|
| 45 |
+
with pytest.raises(ImportError, match="HuggingFace datasets required"):
|
| 46 |
+
datasets._check_datasets_installed()
|
| 47 |
+
|
| 48 |
+
|
| 49 |
+
def test_load_hotpotqa_and_natural_questions(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 50 |
+
calls = install_fake_datasets(
|
| 51 |
+
monkeypatch,
|
| 52 |
+
{
|
| 53 |
+
("hotpotqa/hotpot_qa", "fullwiki", "validation"): [
|
| 54 |
+
{
|
| 55 |
+
"context": {"title": ["Page A"], "sentences": [["Line 1", "Line 2"]]},
|
| 56 |
+
"question": "Who?",
|
| 57 |
+
"answer": "Alice",
|
| 58 |
+
"type": "bridge",
|
| 59 |
+
"level": "easy",
|
| 60 |
+
}
|
| 61 |
+
],
|
| 62 |
+
("google-research-datasets/natural_questions", "default", "validation"): [
|
| 63 |
+
{"document": {}, "question": {"text": "skip me"}},
|
| 64 |
+
{
|
| 65 |
+
"document": {
|
| 66 |
+
"tokens": {
|
| 67 |
+
"token": ["<p>", "Ada", "Lovelace", "wrote", "notes"],
|
| 68 |
+
"is_html": [True, False, False, False, False],
|
| 69 |
+
}
|
| 70 |
+
},
|
| 71 |
+
"question": {"text": "Who wrote notes?"},
|
| 72 |
+
"annotations": {"short_answers": [[{"start_token": 1, "end_token": 3}]]},
|
| 73 |
+
},
|
| 74 |
+
],
|
| 75 |
+
},
|
| 76 |
+
)
|
| 77 |
+
|
| 78 |
+
hotpot = datasets.load_hotpotqa(n=1)
|
| 79 |
+
natural = datasets.load_natural_questions(n=1)
|
| 80 |
+
|
| 81 |
+
assert calls == [
|
| 82 |
+
("hotpotqa/hotpot_qa", "fullwiki", "validation"),
|
| 83 |
+
("google-research-datasets/natural_questions", "default", "validation"),
|
| 84 |
+
]
|
| 85 |
+
assert hotpot.name == "HotpotQA"
|
| 86 |
+
assert hotpot.cases[0].context == "## Page A\nLine 1\nLine 2"
|
| 87 |
+
assert hotpot.cases[0].metadata["type"] == "bridge"
|
| 88 |
+
assert natural.name == "Natural_Questions"
|
| 89 |
+
assert natural.cases[0].context == "Ada Lovelace wrote notes"
|
| 90 |
+
assert natural.cases[0].ground_truth == "Ada Lovelace"
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def test_load_triviaqa_msmarco_and_squad(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 94 |
+
install_fake_datasets(
|
| 95 |
+
monkeypatch,
|
| 96 |
+
{
|
| 97 |
+
("trivia_qa", "rc", "validation"): [
|
| 98 |
+
{"question": "", "search_results": {"search_context": ["unused"]}},
|
| 99 |
+
{
|
| 100 |
+
"question": "Question 1",
|
| 101 |
+
"search_results": {"search_context": ["A", "B"]},
|
| 102 |
+
"answer": {"value": "Answer", "aliases": ["Alias"]},
|
| 103 |
+
},
|
| 104 |
+
{
|
| 105 |
+
"question": "Question 2",
|
| 106 |
+
"search_results": {"search_context": []},
|
| 107 |
+
"entity_pages": {"wiki_context": ["Wiki 1", "Wiki 2"]},
|
| 108 |
+
"answer": {"normalized_value": "Normalized"},
|
| 109 |
+
},
|
| 110 |
+
],
|
| 111 |
+
("microsoft/ms_marco", "v2.1", "validation"): [
|
| 112 |
+
{"query": "", "passages": {"passage_text": ["skip"], "is_selected": [True]}},
|
| 113 |
+
{
|
| 114 |
+
"query": "Find docs",
|
| 115 |
+
"passages": {"passage_text": ["Doc 1", "Doc 2"], "is_selected": [True, False]},
|
| 116 |
+
"answers": ["Primary answer"],
|
| 117 |
+
"query_type": "description",
|
| 118 |
+
},
|
| 119 |
+
],
|
| 120 |
+
("rajpurkar/squad_v2", None, "validation"): [
|
| 121 |
+
{"answers": {"text": []}, "context": "skip", "question": "skip"},
|
| 122 |
+
{
|
| 123 |
+
"context": "Context",
|
| 124 |
+
"question": "Question",
|
| 125 |
+
"answers": {"text": ["First answer"]},
|
| 126 |
+
"title": "Title",
|
| 127 |
+
},
|
| 128 |
+
],
|
| 129 |
+
},
|
| 130 |
+
)
|
| 131 |
+
|
| 132 |
+
trivia = datasets.load_triviaqa(n=2)
|
| 133 |
+
msmarco = datasets.load_msmarco(n=1)
|
| 134 |
+
squad = datasets.load_squad(n=1)
|
| 135 |
+
|
| 136 |
+
assert len(trivia.cases) == 2
|
| 137 |
+
assert trivia.cases[0].context == "A\n\nB"
|
| 138 |
+
assert trivia.cases[1].ground_truth == "Normalized"
|
| 139 |
+
assert trivia.cases[1].metadata["aliases"] == []
|
| 140 |
+
assert msmarco.cases[0].context.startswith("[RELEVANT] Passage 1: Doc 1")
|
| 141 |
+
assert msmarco.cases[0].metadata["num_passages"] == 2
|
| 142 |
+
assert squad.cases[0].ground_truth == "First answer"
|
| 143 |
+
assert squad.cases[0].metadata["title"] == "Title"
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def test_load_longbench_narrativeqa_toolbench_codesearchnet_and_humaneval(
|
| 147 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 148 |
+
) -> None:
|
| 149 |
+
install_fake_datasets(
|
| 150 |
+
monkeypatch,
|
| 151 |
+
{
|
| 152 |
+
("THUDM/LongBench", "qasper", "test"): [
|
| 153 |
+
{"context": "", "input": "skip"},
|
| 154 |
+
{"context": "Long context", "input": "Question", "answers": ["Truth"]},
|
| 155 |
+
],
|
| 156 |
+
("deepmind/narrativeqa", None, "test"): [
|
| 157 |
+
{
|
| 158 |
+
"document": {"summary": {"text": "Story summary"}, "kind": "movie"},
|
| 159 |
+
"question": {"text": "What happened?"},
|
| 160 |
+
"answers": [{"text": "A"}, {"text": "B"}],
|
| 161 |
+
}
|
| 162 |
+
],
|
| 163 |
+
("ToolBench/ToolBench", "G1", "test"): [
|
| 164 |
+
{"api_list": [], "query": "skip"},
|
| 165 |
+
{
|
| 166 |
+
"api_list": [
|
| 167 |
+
{
|
| 168 |
+
"api_name": "weather",
|
| 169 |
+
"api_description": "Get weather",
|
| 170 |
+
"required_parameters": [{"name": "city"}],
|
| 171 |
+
"optional_parameters": [{"name": "unit"}],
|
| 172 |
+
}
|
| 173 |
+
],
|
| 174 |
+
"query": "Weather in SF?",
|
| 175 |
+
"answer": "Call weather",
|
| 176 |
+
},
|
| 177 |
+
],
|
| 178 |
+
("code_search_net", "python", "test"): [
|
| 179 |
+
{"func_code_string": "", "func_documentation_string": "skip"},
|
| 180 |
+
{
|
| 181 |
+
"func_code_string": "def add(a, b): return a + b",
|
| 182 |
+
"func_documentation_string": "Add two numbers.",
|
| 183 |
+
"func_name": "add",
|
| 184 |
+
"repository_name": "repo",
|
| 185 |
+
},
|
| 186 |
+
],
|
| 187 |
+
("openai_humaneval", None, "test"): [
|
| 188 |
+
{"prompt": "", "canonical_solution": "skip"},
|
| 189 |
+
{
|
| 190 |
+
"task_id": "HumanEval/1",
|
| 191 |
+
"prompt": "def solve(x):",
|
| 192 |
+
"canonical_solution": "return x",
|
| 193 |
+
"entry_point": "solve",
|
| 194 |
+
"test": "assert solve(1) == 1",
|
| 195 |
+
},
|
| 196 |
+
],
|
| 197 |
+
},
|
| 198 |
+
)
|
| 199 |
+
|
| 200 |
+
longbench = datasets.load_longbench(n=2, task="qasper")
|
| 201 |
+
narrative = datasets.load_narrativeqa(n=1)
|
| 202 |
+
toolbench = datasets.load_toolbench(n=1, category="G1")
|
| 203 |
+
codesearchnet = datasets.load_codesearchnet(n=1, language="python")
|
| 204 |
+
humaneval = datasets.load_humaneval(n=2)
|
| 205 |
+
|
| 206 |
+
assert longbench.name == "LongBench_qasper"
|
| 207 |
+
assert longbench.cases[0].metadata["context_length"] == len("Long context")
|
| 208 |
+
assert narrative.cases[0].metadata["all_answers"] == ["A", "B"]
|
| 209 |
+
assert toolbench.cases[0].metadata["num_tools"] == 1
|
| 210 |
+
assert '"name": "weather"' in toolbench.cases[0].context
|
| 211 |
+
assert codesearchnet.cases[0].ground_truth == "Add two numbers."
|
| 212 |
+
assert humaneval.cases[0].id == "humaneval_HumanEval/1"
|
| 213 |
+
assert humaneval.cases[0].metadata["entry_point"] == "solve"
|
| 214 |
+
|
| 215 |
+
|
| 216 |
+
def test_load_longbench_toolbench_and_codesearchnet_wrap_loader_errors(
|
| 217 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 218 |
+
) -> None:
|
| 219 |
+
def fake_load_dataset(name: str, subset: str | None = None, split: str | None = None): # noqa: ANN001
|
| 220 |
+
raise RuntimeError(f"broken {name}:{subset}:{split}")
|
| 221 |
+
|
| 222 |
+
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
|
| 223 |
+
|
| 224 |
+
with pytest.raises(ValueError, match="Failed to load LongBench task 'gov_report'"):
|
| 225 |
+
datasets.load_longbench(task="gov_report")
|
| 226 |
+
with pytest.raises(ValueError, match="Failed to load ToolBench category 'G2'"):
|
| 227 |
+
datasets.load_toolbench(category="G2")
|
| 228 |
+
with pytest.raises(ValueError, match="Failed to load CodeSearchNet for 'go'"):
|
| 229 |
+
datasets.load_codesearchnet(language="go")
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
def test_load_bfcl_success_and_download_failure(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 233 |
+
data_lines = "\n".join(
|
| 234 |
+
[
|
| 235 |
+
json.dumps(
|
| 236 |
+
{
|
| 237 |
+
"id": "case-1",
|
| 238 |
+
"question": [[{"role": "user", "content": "How is the weather?"}]],
|
| 239 |
+
"function": [{"name": "weather"}],
|
| 240 |
+
}
|
| 241 |
+
),
|
| 242 |
+
json.dumps({"question": [123], "function": []}),
|
| 243 |
+
]
|
| 244 |
+
)
|
| 245 |
+
gt_lines = json.dumps({"id": "case-1", "ground_truth": [{"name": "weather"}]})
|
| 246 |
+
|
| 247 |
+
def fake_urlopen(url: str): # noqa: ANN001
|
| 248 |
+
if "possible_answer/BFCL_v3_simple.json" in url:
|
| 249 |
+
return SimpleNamespace(read=lambda: gt_lines.encode("utf-8"))
|
| 250 |
+
if "BFCL_v3_simple.json" in url:
|
| 251 |
+
return SimpleNamespace(read=lambda: data_lines.encode("utf-8"))
|
| 252 |
+
raise URLError("missing")
|
| 253 |
+
|
| 254 |
+
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
|
| 255 |
+
|
| 256 |
+
suite = datasets.load_bfcl(n=2, category="simple")
|
| 257 |
+
assert suite.name == "BFCL_simple"
|
| 258 |
+
assert suite.cases[0].query == "How is the weather?"
|
| 259 |
+
assert suite.cases[0].ground_truth == '[{"name": "weather"}]'
|
| 260 |
+
assert suite.cases[0].metadata["num_functions"] == 1
|
| 261 |
+
|
| 262 |
+
def failing_urlopen(url: str): # noqa: ANN001
|
| 263 |
+
raise URLError("offline")
|
| 264 |
+
|
| 265 |
+
monkeypatch.setattr(urllib.request, "urlopen", failing_urlopen)
|
| 266 |
+
with pytest.raises(ValueError, match="Failed to download BFCL dataset 'BFCL_v3_parallel.json'"):
|
| 267 |
+
datasets.load_bfcl(category="parallel")
|
| 268 |
+
|
| 269 |
+
|
| 270 |
+
def test_tool_output_samples_custom_dataset_and_probe_generation(tmp_path) -> None:
|
| 271 |
+
tool_outputs = datasets.load_tool_output_samples()
|
| 272 |
+
assert tool_outputs.name == "ToolOutputSamples"
|
| 273 |
+
assert len(tool_outputs.cases) >= 8
|
| 274 |
+
assert tool_outputs.cases[0].ground_truth == "prompt-optimizer"
|
| 275 |
+
|
| 276 |
+
custom_path = tmp_path / "custom.jsonl"
|
| 277 |
+
custom_path.write_text(
|
| 278 |
+
json.dumps(
|
| 279 |
+
{"id": "case1", "context": "Context", "query": "Question", "ground_truth": "Answer"}
|
| 280 |
+
)
|
| 281 |
+
+ "\n",
|
| 282 |
+
encoding="utf-8",
|
| 283 |
+
)
|
| 284 |
+
custom_suite = datasets.load_custom_dataset(custom_path)
|
| 285 |
+
assert custom_suite.cases[0].id == "case1"
|
| 286 |
+
|
| 287 |
+
probes = datasets.generate_retrieval_probes(
|
| 288 |
+
'Alice Smith deployed API on 2024-01-15 at 99.9% confidence for "Launch Ready" and build_id',
|
| 289 |
+
n_probes=5,
|
| 290 |
+
)
|
| 291 |
+
assert "Alice Smith" in probes
|
| 292 |
+
assert "2024-01-15" in probes
|
| 293 |
+
assert "API" in probes
|
| 294 |
+
assert "99.9" in probes
|
| 295 |
+
assert "Launch Ready" in probes
|
| 296 |
+
|
| 297 |
+
|
| 298 |
+
def test_dataset_registry_helpers(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 299 |
+
categories = datasets.list_available_datasets()
|
| 300 |
+
assert "hotpotqa" in categories["rag"]
|
| 301 |
+
assert "tool_outputs" in categories["tool_use"]
|
| 302 |
+
|
| 303 |
+
seen: list[tuple[str, dict[str, object]]] = []
|
| 304 |
+
|
| 305 |
+
def fake_loader(*, n: int = 0, **kwargs): # noqa: ANN003
|
| 306 |
+
seen.append(("with-n", {"n": n, **kwargs}))
|
| 307 |
+
return "with-n-result"
|
| 308 |
+
|
| 309 |
+
def fixed_loader(**kwargs): # noqa: ANN003
|
| 310 |
+
seen.append(("fixed", kwargs))
|
| 311 |
+
return "fixed-result"
|
| 312 |
+
|
| 313 |
+
original_registry = dict(datasets.DATASET_REGISTRY)
|
| 314 |
+
monkeypatch.setattr(
|
| 315 |
+
datasets,
|
| 316 |
+
"DATASET_REGISTRY",
|
| 317 |
+
{
|
| 318 |
+
**original_registry,
|
| 319 |
+
"fake_n": {"loader": fake_loader, "category": "x", "description": "", "default_n": 3},
|
| 320 |
+
"fake_fixed": {
|
| 321 |
+
"loader": fixed_loader,
|
| 322 |
+
"category": "x",
|
| 323 |
+
"description": "",
|
| 324 |
+
"default_n": None,
|
| 325 |
+
},
|
| 326 |
+
},
|
| 327 |
+
)
|
| 328 |
+
|
| 329 |
+
assert datasets.load_dataset_by_name("fake_n") == "with-n-result"
|
| 330 |
+
assert datasets.load_dataset_by_name("fake_n", n=7, split="test") == "with-n-result"
|
| 331 |
+
assert datasets.load_dataset_by_name("fake_fixed", path="x") == "fixed-result"
|
| 332 |
+
assert seen == [
|
| 333 |
+
("with-n", {"n": 3}),
|
| 334 |
+
("with-n", {"n": 7, "split": "test"}),
|
| 335 |
+
("fixed", {"path": "x"}),
|
| 336 |
+
]
|
| 337 |
+
|
| 338 |
+
with pytest.raises(ValueError, match="Unknown dataset 'missing'"):
|
| 339 |
+
datasets.load_dataset_by_name("missing")
|
| 340 |
+
|
| 341 |
+
|
| 342 |
+
def test_dataset_loaders_cover_skip_and_limit_branches(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 343 |
+
install_fake_datasets(
|
| 344 |
+
monkeypatch,
|
| 345 |
+
{
|
| 346 |
+
("hotpotqa/hotpot_qa", "fullwiki", "validation"): [
|
| 347 |
+
{
|
| 348 |
+
"context": {"title": ["Page A"], "sentences": [["Line 1"]]},
|
| 349 |
+
"question": "Q1",
|
| 350 |
+
"answer": "A1",
|
| 351 |
+
},
|
| 352 |
+
{
|
| 353 |
+
"context": {"title": ["Page B"], "sentences": [["Line 2"]]},
|
| 354 |
+
"question": "Q2",
|
| 355 |
+
"answer": "A2",
|
| 356 |
+
},
|
| 357 |
+
],
|
| 358 |
+
("google-research-datasets/natural_questions", "default", "validation"): [
|
| 359 |
+
{
|
| 360 |
+
"document": {"tokens": {"token": ["x"], "is_html": [False]}},
|
| 361 |
+
"question": {"text": ""},
|
| 362 |
+
},
|
| 363 |
+
{
|
| 364 |
+
"document": {"tokens": {"token": ["<b>"], "is_html": [True]}},
|
| 365 |
+
"question": {"text": "blank context"},
|
| 366 |
+
},
|
| 367 |
+
{
|
| 368 |
+
"document": {"tokens": {"token": ["Ada", "wrote"], "is_html": [False, False]}},
|
| 369 |
+
"question": {"text": "Who?"},
|
| 370 |
+
"annotations": {"short_answers": [[{"start_token": 1, "end_token": 1}]]},
|
| 371 |
+
},
|
| 372 |
+
{
|
| 373 |
+
"document": {"tokens": {"token": ["Grace"], "is_html": [False]}},
|
| 374 |
+
"question": {"text": "Ignored by limit"},
|
| 375 |
+
},
|
| 376 |
+
],
|
| 377 |
+
("trivia_qa", "rc", "validation"): [
|
| 378 |
+
{"question": "skip", "search_results": {"search_context": []}, "entity_pages": {}},
|
| 379 |
+
{"question": "blank", "search_results": {"search_context": [""]}},
|
| 380 |
+
{
|
| 381 |
+
"question": "Good 1",
|
| 382 |
+
"search_results": {"search_context": ["Context 1"]},
|
| 383 |
+
"answer": {"value": "A1"},
|
| 384 |
+
},
|
| 385 |
+
{
|
| 386 |
+
"question": "Good 2",
|
| 387 |
+
"search_results": {"search_context": ["Context 2"]},
|
| 388 |
+
"answer": {"value": "A2"},
|
| 389 |
+
},
|
| 390 |
+
],
|
| 391 |
+
("microsoft/ms_marco", "v2.1", "validation"): [
|
| 392 |
+
{"query": "skip", "passages": {"passage_text": [], "is_selected": []}},
|
| 393 |
+
{
|
| 394 |
+
"query": "Find one",
|
| 395 |
+
"passages": {"passage_text": ["Doc 1"], "is_selected": [False]},
|
| 396 |
+
"answers": [],
|
| 397 |
+
},
|
| 398 |
+
{
|
| 399 |
+
"query": "Find two",
|
| 400 |
+
"passages": {"passage_text": ["Doc 2"], "is_selected": [True]},
|
| 401 |
+
"answers": ["A2"],
|
| 402 |
+
},
|
| 403 |
+
],
|
| 404 |
+
("rajpurkar/squad_v2", None, "validation"): [
|
| 405 |
+
{
|
| 406 |
+
"context": "Context 1",
|
| 407 |
+
"question": "Q1",
|
| 408 |
+
"answers": {"text": ["A1"]},
|
| 409 |
+
},
|
| 410 |
+
{
|
| 411 |
+
"context": "Context 2",
|
| 412 |
+
"question": "Q2",
|
| 413 |
+
"answers": {"text": ["A2"]},
|
| 414 |
+
},
|
| 415 |
+
],
|
| 416 |
+
("THUDM/LongBench", "qasper", "test"): [
|
| 417 |
+
{"context": "Context 1", "input": "Q1", "answers": ["A1"]},
|
| 418 |
+
{"context": "Has context", "input": ""},
|
| 419 |
+
{"context": "Context 2", "input": "Q2", "answers": ["A2"]},
|
| 420 |
+
],
|
| 421 |
+
("deepmind/narrativeqa", None, "test"): [
|
| 422 |
+
{"document": {"summary": {"text": ""}}, "question": {"text": "skip"}},
|
| 423 |
+
{"document": {"summary": {"text": "Story"}}, "question": {"text": ""}},
|
| 424 |
+
{
|
| 425 |
+
"document": {"summary": {"text": "Story 1"}, "kind": "book"},
|
| 426 |
+
"question": {"text": "Q1"},
|
| 427 |
+
"answers": [{"text": "A1"}],
|
| 428 |
+
},
|
| 429 |
+
{
|
| 430 |
+
"document": {"summary": {"text": "Story 2"}, "kind": "movie"},
|
| 431 |
+
"question": {"text": "Q2"},
|
| 432 |
+
"answers": [{"text": "A2"}],
|
| 433 |
+
},
|
| 434 |
+
],
|
| 435 |
+
("ToolBench/ToolBench", "G1", "test"): [
|
| 436 |
+
{"api_list": [], "query": "skip"},
|
| 437 |
+
{
|
| 438 |
+
"api_list": [
|
| 439 |
+
{
|
| 440 |
+
"api_name": "weather",
|
| 441 |
+
"required_parameters": [],
|
| 442 |
+
"optional_parameters": [],
|
| 443 |
+
}
|
| 444 |
+
],
|
| 445 |
+
"query": "",
|
| 446 |
+
},
|
| 447 |
+
{
|
| 448 |
+
"api_list": [
|
| 449 |
+
{"api_name": "calc", "required_parameters": [], "optional_parameters": []}
|
| 450 |
+
],
|
| 451 |
+
"query": "Good",
|
| 452 |
+
},
|
| 453 |
+
],
|
| 454 |
+
("code_search_net", "python", "test"): [
|
| 455 |
+
{
|
| 456 |
+
"func_code_string": "",
|
| 457 |
+
"whole_func_string": "",
|
| 458 |
+
"func_documentation_string": "skip",
|
| 459 |
+
},
|
| 460 |
+
{"whole_func_string": "def alt(): pass", "func_documentation_string": ""},
|
| 461 |
+
{
|
| 462 |
+
"whole_func_string": "def good(): pass",
|
| 463 |
+
"func_documentation_string": "Good doc",
|
| 464 |
+
"func_name": "good",
|
| 465 |
+
"repository_name": "repo",
|
| 466 |
+
},
|
| 467 |
+
{
|
| 468 |
+
"whole_func_string": "def ignored(): pass",
|
| 469 |
+
"func_documentation_string": "Ignored by limit",
|
| 470 |
+
},
|
| 471 |
+
],
|
| 472 |
+
("openai_humaneval", None, "test"): [
|
| 473 |
+
{
|
| 474 |
+
"task_id": "Task/1",
|
| 475 |
+
"prompt": "def solve():",
|
| 476 |
+
"canonical_solution": "return 1",
|
| 477 |
+
"test": "assert solve() == 1",
|
| 478 |
+
},
|
| 479 |
+
{
|
| 480 |
+
"task_id": "Task/2",
|
| 481 |
+
"prompt": "def other():",
|
| 482 |
+
"canonical_solution": "return 2",
|
| 483 |
+
"test": "assert other() == 2",
|
| 484 |
+
},
|
| 485 |
+
],
|
| 486 |
+
},
|
| 487 |
+
)
|
| 488 |
+
|
| 489 |
+
assert len(datasets.load_hotpotqa(n=1).cases) == 1
|
| 490 |
+
natural = datasets.load_natural_questions(n=1)
|
| 491 |
+
assert len(natural.cases) == 1
|
| 492 |
+
assert natural.cases[0].ground_truth is None
|
| 493 |
+
assert len(datasets.load_triviaqa(n=1).cases) == 1
|
| 494 |
+
msmarco = datasets.load_msmarco(n=1)
|
| 495 |
+
assert len(msmarco.cases) == 1
|
| 496 |
+
assert msmarco.cases[0].ground_truth is None
|
| 497 |
+
assert len(datasets.load_squad(n=1).cases) == 1
|
| 498 |
+
assert len(datasets.load_longbench(n=2, task="qasper").cases) == 1
|
| 499 |
+
assert len(datasets.load_narrativeqa(n=1).cases) == 1
|
| 500 |
+
assert len(datasets.load_toolbench(n=1, category="G1").cases) == 1
|
| 501 |
+
assert len(datasets.load_codesearchnet(n=1, language="python").cases) == 1
|
| 502 |
+
assert len(datasets.load_humaneval(n=1).cases) == 1
|
| 503 |
+
|
| 504 |
+
|
| 505 |
+
def test_load_bfcl_handles_optional_ground_truth_and_question_fallback(
|
| 506 |
+
monkeypatch: pytest.MonkeyPatch,
|
| 507 |
+
) -> None:
|
| 508 |
+
data_lines = "\n".join(
|
| 509 |
+
[
|
| 510 |
+
json.dumps(
|
| 511 |
+
{
|
| 512 |
+
"id": "case-1",
|
| 513 |
+
"question": [123],
|
| 514 |
+
"function": [{"name": "weather"}],
|
| 515 |
+
}
|
| 516 |
+
),
|
| 517 |
+
json.dumps({"id": "skip", "function": []}),
|
| 518 |
+
json.dumps(
|
| 519 |
+
{
|
| 520 |
+
"id": "case-2",
|
| 521 |
+
"question": [[{"role": "user", "content": "Ignored by limit"}]],
|
| 522 |
+
"function": [{"name": "time"}],
|
| 523 |
+
}
|
| 524 |
+
),
|
| 525 |
+
]
|
| 526 |
+
)
|
| 527 |
+
|
| 528 |
+
def fake_urlopen(url: str): # noqa: ANN001
|
| 529 |
+
if "possible_answer" in url:
|
| 530 |
+
raise URLError("missing ground truth")
|
| 531 |
+
return SimpleNamespace(read=lambda: data_lines.encode("utf-8"))
|
| 532 |
+
|
| 533 |
+
monkeypatch.setattr(urllib.request, "urlopen", fake_urlopen)
|
| 534 |
+
|
| 535 |
+
suite = datasets.load_bfcl(n=2, category="simple")
|
| 536 |
+
assert len(suite.cases) == 1
|
| 537 |
+
assert suite.cases[0].query == "[123]"
|
| 538 |
+
assert suite.cases[0].ground_truth is None
|
|
@@ -1,132 +1,132 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
import math
|
| 4 |
-
import sys
|
| 5 |
-
from types import SimpleNamespace
|
| 6 |
-
|
| 7 |
-
import pytest
|
| 8 |
-
|
| 9 |
-
from headroom.evals import metrics
|
| 10 |
-
|
| 11 |
-
|
| 12 |
-
def test_normalize_tokenize_and_exact_match() -> None:
|
| 13 |
-
assert metrics.normalize_text(" Hello,\nWORLD ") == "hello, world"
|
| 14 |
-
assert metrics.tokenize("Hello, world! API_v2") == ["hello", "world", "api_v2"]
|
| 15 |
-
assert metrics.compute_exact_match(" Hello World ", "hello\nworld") is True
|
| 16 |
-
assert metrics.compute_exact_match("hello", "world") is False
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
def test_f1_bleu_and_rouge_cover_edge_cases() -> None:
|
| 20 |
-
assert metrics.compute_f1("", "value") == 0.0
|
| 21 |
-
assert metrics.compute_f1("alpha beta", "gamma delta") == 0.0
|
| 22 |
-
assert metrics.compute_f1("alpha beta gamma", "alpha gamma") == pytest.approx(0.8)
|
| 23 |
-
|
| 24 |
-
assert metrics.compute_bleu("", "value") == 0.0
|
| 25 |
-
assert metrics.compute_bleu("one", "one") == pytest.approx(1.0)
|
| 26 |
-
assert metrics.compute_bleu("alpha beta", "gamma delta") == 0.0
|
| 27 |
-
assert metrics.compute_bleu("alpha beta", "alpha beta gamma", max_n=4) == pytest.approx(1.0)
|
| 28 |
-
|
| 29 |
-
assert metrics.compute_rouge_l("", "value") == 0.0
|
| 30 |
-
assert metrics.compute_rouge_l("alpha beta", "gamma delta") == 0.0
|
| 31 |
-
assert metrics.compute_rouge_l("alpha beta gamma", "alpha gamma") == pytest.approx(0.8)
|
| 32 |
-
|
| 33 |
-
|
| 34 |
-
def test_compute_semantic_similarity_and_zero_norm(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 35 |
-
fake_numpy = SimpleNamespace(
|
| 36 |
-
dot=lambda a, b: sum(x * y for x, y in zip(a, b)),
|
| 37 |
-
linalg=SimpleNamespace(norm=lambda a: math.sqrt(sum(x * x for x in a))),
|
| 38 |
-
)
|
| 39 |
-
monkeypatch.setitem(sys.modules, "numpy", fake_numpy)
|
| 40 |
-
|
| 41 |
-
class FakeModel:
|
| 42 |
-
def __init__(self, embeddings: list[list[float]]) -> None:
|
| 43 |
-
self.embeddings = embeddings
|
| 44 |
-
|
| 45 |
-
def encode(self, values: list[str]) -> list[list[float]]:
|
| 46 |
-
assert values == ["first", "second"]
|
| 47 |
-
return self.embeddings
|
| 48 |
-
|
| 49 |
-
monkeypatch.setattr(
|
| 50 |
-
"headroom.models.ml_models.MLModelRegistry.get_sentence_transformer",
|
| 51 |
-
lambda model_name=None: FakeModel([[1.0, 0.0], [1.0, 0.0]]),
|
| 52 |
-
)
|
| 53 |
-
assert metrics.compute_semantic_similarity("first", "second") == 1.0
|
| 54 |
-
|
| 55 |
-
monkeypatch.setattr(
|
| 56 |
-
"headroom.models.ml_models.MLModelRegistry.get_sentence_transformer",
|
| 57 |
-
lambda model_name=None: FakeModel([[0.0, 0.0], [1.0, 0.0]]),
|
| 58 |
-
)
|
| 59 |
-
assert metrics.compute_semantic_similarity("first", "second") == 0.0
|
| 60 |
-
|
| 61 |
-
|
| 62 |
-
def test_compute_answer_equivalence_uses_multiple_paths(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 63 |
-
monkeypatch.setattr(metrics, "compute_semantic_similarity", lambda a, b: 0.2)
|
| 64 |
-
exact = metrics.compute_answer_equivalence("Answer", "answer", ground_truth="missing")
|
| 65 |
-
assert exact["exact_match"] is True
|
| 66 |
-
assert exact["equivalent"] is True
|
| 67 |
-
assert exact["ground_truth_in_a"] is False
|
| 68 |
-
assert exact["ground_truth_in_b"] is False
|
| 69 |
-
|
| 70 |
-
monkeypatch.setattr(metrics, "compute_semantic_similarity", lambda a, b: 0.1)
|
| 71 |
-
high_f1 = metrics.compute_answer_equivalence(
|
| 72 |
-
"alpha beta gamma",
|
| 73 |
-
"alpha gamma",
|
| 74 |
-
semantic_threshold=0.95,
|
| 75 |
-
f1_threshold=0.75,
|
| 76 |
-
)
|
| 77 |
-
assert high_f1["equivalent"] is True
|
| 78 |
-
assert high_f1["semantic_similarity"] == 0.1
|
| 79 |
-
|
| 80 |
-
monkeypatch.setattr(metrics, "compute_semantic_similarity", lambda a, b: 0.95)
|
| 81 |
-
semantic = metrics.compute_answer_equivalence(
|
| 82 |
-
"completely different",
|
| 83 |
-
"nothing in common",
|
| 84 |
-
semantic_threshold=0.9,
|
| 85 |
-
f1_threshold=0.99,
|
| 86 |
-
)
|
| 87 |
-
assert semantic["equivalent"] is True
|
| 88 |
-
assert semantic["semantic_similarity"] == 0.95
|
| 89 |
-
|
| 90 |
-
def raise_import_error(a: str, b: str) -> float:
|
| 91 |
-
raise ImportError("missing dependency")
|
| 92 |
-
|
| 93 |
-
monkeypatch.setattr(metrics, "compute_semantic_similarity", raise_import_error)
|
| 94 |
-
ground_truth = metrics.compute_answer_equivalence(
|
| 95 |
-
"The capital is Paris.",
|
| 96 |
-
"Paris is definitely the capital city.",
|
| 97 |
-
ground_truth="paris",
|
| 98 |
-
semantic_threshold=0.99,
|
| 99 |
-
f1_threshold=0.99,
|
| 100 |
-
)
|
| 101 |
-
assert ground_truth["semantic_similarity"] is None
|
| 102 |
-
assert ground_truth["ground_truth_in_a"] is True
|
| 103 |
-
assert ground_truth["ground_truth_in_b"] is True
|
| 104 |
-
assert ground_truth["equivalent"] is True
|
| 105 |
-
|
| 106 |
-
not_equivalent = metrics.compute_answer_equivalence(
|
| 107 |
-
"alpha beta",
|
| 108 |
-
"gamma delta",
|
| 109 |
-
ground_truth="omega",
|
| 110 |
-
semantic_threshold=0.99,
|
| 111 |
-
f1_threshold=0.99,
|
| 112 |
-
)
|
| 113 |
-
assert not_equivalent["equivalent"] is False
|
| 114 |
-
|
| 115 |
-
|
| 116 |
-
def test_information_recall_reports_preserved_and_missing_facts() -> None:
|
| 117 |
-
result = metrics.compute_information_recall(
|
| 118 |
-
"Alice likes pizza and Bob likes ramen.",
|
| 119 |
-
"Alice likes pizza.",
|
| 120 |
-
["Alice", "Bob", "ramen", "Carol"],
|
| 121 |
-
)
|
| 122 |
-
assert result == {
|
| 123 |
-
"total_probes": 4,
|
| 124 |
-
"facts_in_original": 3,
|
| 125 |
-
"facts_preserved": 1,
|
| 126 |
-
"facts_lost": ["Bob", "ramen"],
|
| 127 |
-
"recall": pytest.approx(1 / 3),
|
| 128 |
-
}
|
| 129 |
-
|
| 130 |
-
empty_original = metrics.compute_information_recall("No facts here", "Still none", ["Alice"])
|
| 131 |
-
assert empty_original["facts_in_original"] == 0
|
| 132 |
-
assert empty_original["recall"] == 1.0
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import math
|
| 4 |
+
import sys
|
| 5 |
+
from types import SimpleNamespace
|
| 6 |
+
|
| 7 |
+
import pytest
|
| 8 |
+
|
| 9 |
+
from headroom.evals import metrics
|
| 10 |
+
|
| 11 |
+
|
| 12 |
+
def test_normalize_tokenize_and_exact_match() -> None:
|
| 13 |
+
assert metrics.normalize_text(" Hello,\nWORLD ") == "hello, world"
|
| 14 |
+
assert metrics.tokenize("Hello, world! API_v2") == ["hello", "world", "api_v2"]
|
| 15 |
+
assert metrics.compute_exact_match(" Hello World ", "hello\nworld") is True
|
| 16 |
+
assert metrics.compute_exact_match("hello", "world") is False
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
def test_f1_bleu_and_rouge_cover_edge_cases() -> None:
|
| 20 |
+
assert metrics.compute_f1("", "value") == 0.0
|
| 21 |
+
assert metrics.compute_f1("alpha beta", "gamma delta") == 0.0
|
| 22 |
+
assert metrics.compute_f1("alpha beta gamma", "alpha gamma") == pytest.approx(0.8)
|
| 23 |
+
|
| 24 |
+
assert metrics.compute_bleu("", "value") == 0.0
|
| 25 |
+
assert metrics.compute_bleu("one", "one") == pytest.approx(1.0)
|
| 26 |
+
assert metrics.compute_bleu("alpha beta", "gamma delta") == 0.0
|
| 27 |
+
assert metrics.compute_bleu("alpha beta", "alpha beta gamma", max_n=4) == pytest.approx(1.0)
|
| 28 |
+
|
| 29 |
+
assert metrics.compute_rouge_l("", "value") == 0.0
|
| 30 |
+
assert metrics.compute_rouge_l("alpha beta", "gamma delta") == 0.0
|
| 31 |
+
assert metrics.compute_rouge_l("alpha beta gamma", "alpha gamma") == pytest.approx(0.8)
|
| 32 |
+
|
| 33 |
+
|
| 34 |
+
def test_compute_semantic_similarity_and_zero_norm(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 35 |
+
fake_numpy = SimpleNamespace(
|
| 36 |
+
dot=lambda a, b: sum(x * y for x, y in zip(a, b)),
|
| 37 |
+
linalg=SimpleNamespace(norm=lambda a: math.sqrt(sum(x * x for x in a))),
|
| 38 |
+
)
|
| 39 |
+
monkeypatch.setitem(sys.modules, "numpy", fake_numpy)
|
| 40 |
+
|
| 41 |
+
class FakeModel:
|
| 42 |
+
def __init__(self, embeddings: list[list[float]]) -> None:
|
| 43 |
+
self.embeddings = embeddings
|
| 44 |
+
|
| 45 |
+
def encode(self, values: list[str]) -> list[list[float]]:
|
| 46 |
+
assert values == ["first", "second"]
|
| 47 |
+
return self.embeddings
|
| 48 |
+
|
| 49 |
+
monkeypatch.setattr(
|
| 50 |
+
"headroom.models.ml_models.MLModelRegistry.get_sentence_transformer",
|
| 51 |
+
lambda model_name=None: FakeModel([[1.0, 0.0], [1.0, 0.0]]),
|
| 52 |
+
)
|
| 53 |
+
assert metrics.compute_semantic_similarity("first", "second") == 1.0
|
| 54 |
+
|
| 55 |
+
monkeypatch.setattr(
|
| 56 |
+
"headroom.models.ml_models.MLModelRegistry.get_sentence_transformer",
|
| 57 |
+
lambda model_name=None: FakeModel([[0.0, 0.0], [1.0, 0.0]]),
|
| 58 |
+
)
|
| 59 |
+
assert metrics.compute_semantic_similarity("first", "second") == 0.0
|
| 60 |
+
|
| 61 |
+
|
| 62 |
+
def test_compute_answer_equivalence_uses_multiple_paths(monkeypatch: pytest.MonkeyPatch) -> None:
|
| 63 |
+
monkeypatch.setattr(metrics, "compute_semantic_similarity", lambda a, b: 0.2)
|
| 64 |
+
exact = metrics.compute_answer_equivalence("Answer", "answer", ground_truth="missing")
|
| 65 |
+
assert exact["exact_match"] is True
|
| 66 |
+
assert exact["equivalent"] is True
|
| 67 |
+
assert exact["ground_truth_in_a"] is False
|
| 68 |
+
assert exact["ground_truth_in_b"] is False
|
| 69 |
+
|
| 70 |
+
monkeypatch.setattr(metrics, "compute_semantic_similarity", lambda a, b: 0.1)
|
| 71 |
+
high_f1 = metrics.compute_answer_equivalence(
|
| 72 |
+
"alpha beta gamma",
|
| 73 |
+
"alpha gamma",
|
| 74 |
+
semantic_threshold=0.95,
|
| 75 |
+
f1_threshold=0.75,
|
| 76 |
+
)
|
| 77 |
+
assert high_f1["equivalent"] is True
|
| 78 |
+
assert high_f1["semantic_similarity"] == 0.1
|
| 79 |
+
|
| 80 |
+
monkeypatch.setattr(metrics, "compute_semantic_similarity", lambda a, b: 0.95)
|
| 81 |
+
semantic = metrics.compute_answer_equivalence(
|
| 82 |
+
"completely different",
|
| 83 |
+
"nothing in common",
|
| 84 |
+
semantic_threshold=0.9,
|
| 85 |
+
f1_threshold=0.99,
|
| 86 |
+
)
|
| 87 |
+
assert semantic["equivalent"] is True
|
| 88 |
+
assert semantic["semantic_similarity"] == 0.95
|
| 89 |
+
|
| 90 |
+
def raise_import_error(a: str, b: str) -> float:
|
| 91 |
+
raise ImportError("missing dependency")
|
| 92 |
+
|
| 93 |
+
monkeypatch.setattr(metrics, "compute_semantic_similarity", raise_import_error)
|
| 94 |
+
ground_truth = metrics.compute_answer_equivalence(
|
| 95 |
+
"The capital is Paris.",
|
| 96 |
+
"Paris is definitely the capital city.",
|
| 97 |
+
ground_truth="paris",
|
| 98 |
+
semantic_threshold=0.99,
|
| 99 |
+
f1_threshold=0.99,
|
| 100 |
+
)
|
| 101 |
+
assert ground_truth["semantic_similarity"] is None
|
| 102 |
+
assert ground_truth["ground_truth_in_a"] is True
|
| 103 |
+
assert ground_truth["ground_truth_in_b"] is True
|
| 104 |
+
assert ground_truth["equivalent"] is True
|
| 105 |
+
|
| 106 |
+
not_equivalent = metrics.compute_answer_equivalence(
|
| 107 |
+
"alpha beta",
|
| 108 |
+
"gamma delta",
|
| 109 |
+
ground_truth="omega",
|
| 110 |
+
semantic_threshold=0.99,
|
| 111 |
+
f1_threshold=0.99,
|
| 112 |
+
)
|
| 113 |
+
assert not_equivalent["equivalent"] is False
|
| 114 |
+
|
| 115 |
+
|
| 116 |
+
def test_information_recall_reports_preserved_and_missing_facts() -> None:
|
| 117 |
+
result = metrics.compute_information_recall(
|
| 118 |
+
"Alice likes pizza and Bob likes ramen.",
|
| 119 |
+
"Alice likes pizza.",
|
| 120 |
+
["Alice", "Bob", "ramen", "Carol"],
|
| 121 |
+
)
|
| 122 |
+
assert result == {
|
| 123 |
+
"total_probes": 4,
|
| 124 |
+
"facts_in_original": 3,
|
| 125 |
+
"facts_preserved": 1,
|
| 126 |
+
"facts_lost": ["Bob", "ramen"],
|
| 127 |
+
"recall": pytest.approx(1 / 3),
|
| 128 |
+
}
|
| 129 |
+
|
| 130 |
+
empty_original = metrics.compute_information_recall("No facts here", "Still none", ["Alice"])
|
| 131 |
+
assert empty_original["facts_in_original"] == 0
|
| 132 |
+
assert empty_original["recall"] == 1.0
|
|
@@ -1,40 +1,40 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
from headroom.exceptions import (
|
| 4 |
-
CacheError,
|
| 5 |
-
CompressionError,
|
| 6 |
-
ConfigurationError,
|
| 7 |
-
HeadroomError,
|
| 8 |
-
ProviderError,
|
| 9 |
-
StorageError,
|
| 10 |
-
TokenizationError,
|
| 11 |
-
TransformError,
|
| 12 |
-
ValidationError,
|
| 13 |
-
)
|
| 14 |
-
|
| 15 |
-
|
| 16 |
-
def test_headroom_error_formats_details() -> None:
|
| 17 |
-
err = HeadroomError("bad config", details={"mode": "foo", "valid": "bar"})
|
| 18 |
-
assert err.message == "bad config"
|
| 19 |
-
assert err.details == {"mode": "foo", "valid": "bar"}
|
| 20 |
-
assert str(err) == "bad config (mode=foo, valid=bar)"
|
| 21 |
-
|
| 22 |
-
plain = HeadroomError("just bad")
|
| 23 |
-
assert plain.details == {}
|
| 24 |
-
assert str(plain) == "just bad"
|
| 25 |
-
|
| 26 |
-
|
| 27 |
-
def test_specialized_exceptions_inherit_headroom_error() -> None:
|
| 28 |
-
for exc_type in (
|
| 29 |
-
ConfigurationError,
|
| 30 |
-
ProviderError,
|
| 31 |
-
StorageError,
|
| 32 |
-
CompressionError,
|
| 33 |
-
TokenizationError,
|
| 34 |
-
CacheError,
|
| 35 |
-
ValidationError,
|
| 36 |
-
TransformError,
|
| 37 |
-
):
|
| 38 |
-
err = exc_type("problem", details={"kind": exc_type.__name__})
|
| 39 |
-
assert isinstance(err, HeadroomError)
|
| 40 |
-
assert str(err) == f"problem (kind={exc_type.__name__})"
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
from headroom.exceptions import (
|
| 4 |
+
CacheError,
|
| 5 |
+
CompressionError,
|
| 6 |
+
ConfigurationError,
|
| 7 |
+
HeadroomError,
|
| 8 |
+
ProviderError,
|
| 9 |
+
StorageError,
|
| 10 |
+
TokenizationError,
|
| 11 |
+
TransformError,
|
| 12 |
+
ValidationError,
|
| 13 |
+
)
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
def test_headroom_error_formats_details() -> None:
|
| 17 |
+
err = HeadroomError("bad config", details={"mode": "foo", "valid": "bar"})
|
| 18 |
+
assert err.message == "bad config"
|
| 19 |
+
assert err.details == {"mode": "foo", "valid": "bar"}
|
| 20 |
+
assert str(err) == "bad config (mode=foo, valid=bar)"
|
| 21 |
+
|
| 22 |
+
plain = HeadroomError("just bad")
|
| 23 |
+
assert plain.details == {}
|
| 24 |
+
assert str(plain) == "just bad"
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def test_specialized_exceptions_inherit_headroom_error() -> None:
|
| 28 |
+
for exc_type in (
|
| 29 |
+
ConfigurationError,
|
| 30 |
+
ProviderError,
|
| 31 |
+
StorageError,
|
| 32 |
+
CompressionError,
|
| 33 |
+
TokenizationError,
|
| 34 |
+
CacheError,
|
| 35 |
+
ValidationError,
|
| 36 |
+
TransformError,
|
| 37 |
+
):
|
| 38 |
+
err = exc_type("problem", details={"kind": exc_type.__name__})
|
| 39 |
+
assert isinstance(err, HeadroomError)
|
| 40 |
+
assert str(err) == f"problem (kind={exc_type.__name__})"
|
|
@@ -1,352 +1,352 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
import io
|
| 4 |
-
import json
|
| 5 |
-
import subprocess
|
| 6 |
-
import tarfile
|
| 7 |
-
from pathlib import Path
|
| 8 |
-
from types import ModuleType, SimpleNamespace
|
| 9 |
-
|
| 10 |
-
import pytest
|
| 11 |
-
|
| 12 |
-
from headroom.graph import installer, watcher
|
| 13 |
-
|
| 14 |
-
|
| 15 |
-
def _build_archive(member_name: str = installer.CBM_BIN_NAME) -> bytes:
|
| 16 |
-
payload = io.BytesIO()
|
| 17 |
-
with tarfile.open(fileobj=payload, mode="w:gz") as tar:
|
| 18 |
-
data = b"#!/bin/sh\necho version\n"
|
| 19 |
-
info = tarfile.TarInfo(name=member_name)
|
| 20 |
-
info.size = len(data)
|
| 21 |
-
tar.addfile(info, io.BytesIO(data))
|
| 22 |
-
return payload.getvalue()
|
| 23 |
-
|
| 24 |
-
|
| 25 |
-
class FakeResponse:
|
| 26 |
-
def __init__(self, data: bytes) -> None:
|
| 27 |
-
self._data = data
|
| 28 |
-
|
| 29 |
-
def __enter__(self):
|
| 30 |
-
return self
|
| 31 |
-
|
| 32 |
-
def __exit__(self, exc_type, exc, tb) -> None:
|
| 33 |
-
return None
|
| 34 |
-
|
| 35 |
-
def read(self) -> bytes:
|
| 36 |
-
return self._data
|
| 37 |
-
|
| 38 |
-
|
| 39 |
-
@pytest.mark.parametrize(
|
| 40 |
-
("system", "machine", "expected"),
|
| 41 |
-
[
|
| 42 |
-
("Darwin", "arm64", "darwin-arm64"),
|
| 43 |
-
("Darwin", "x86_64", "darwin-amd64"),
|
| 44 |
-
("Linux", "aarch64", "linux-arm64"),
|
| 45 |
-
("Linux", "arm64", "linux-arm64"),
|
| 46 |
-
("Linux", "x86_64", "linux-amd64"),
|
| 47 |
-
("Windows", "AMD64", "windows-amd64"),
|
| 48 |
-
],
|
| 49 |
-
)
|
| 50 |
-
def test_detect_platform_variants(monkeypatch, system: str, machine: str, expected: str) -> None:
|
| 51 |
-
monkeypatch.setattr(installer.platform, "system", lambda: system)
|
| 52 |
-
monkeypatch.setattr(installer.platform, "machine", lambda: machine)
|
| 53 |
-
assert installer._detect_platform() == expected
|
| 54 |
-
|
| 55 |
-
|
| 56 |
-
def test_detect_platform_rejects_unknown_system(monkeypatch) -> None:
|
| 57 |
-
monkeypatch.setattr(installer.platform, "system", lambda: "Solaris")
|
| 58 |
-
monkeypatch.setattr(installer.platform, "machine", lambda: "sparc")
|
| 59 |
-
with pytest.raises(RuntimeError, match="Unsupported platform"):
|
| 60 |
-
installer._detect_platform()
|
| 61 |
-
|
| 62 |
-
|
| 63 |
-
def test_get_cbm_path_prefers_path_then_install_dir(monkeypatch, tmp_path: Path) -> None:
|
| 64 |
-
on_path = tmp_path / "on-path"
|
| 65 |
-
installed = tmp_path / installer.CBM_BIN_NAME
|
| 66 |
-
installed.write_text("bin")
|
| 67 |
-
monkeypatch.setattr(installer, "CBM_BIN_DIR", tmp_path)
|
| 68 |
-
monkeypatch.setattr(installer.shutil, "which", lambda name: str(on_path))
|
| 69 |
-
assert installer.get_cbm_path() == on_path
|
| 70 |
-
|
| 71 |
-
monkeypatch.setattr(installer.shutil, "which", lambda name: None)
|
| 72 |
-
assert installer.get_cbm_path() == installed
|
| 73 |
-
|
| 74 |
-
installed.unlink()
|
| 75 |
-
assert installer.get_cbm_path() is None
|
| 76 |
-
|
| 77 |
-
|
| 78 |
-
def test_download_cbm_success_and_verification_paths(monkeypatch, tmp_path: Path) -> None:
|
| 79 |
-
monkeypatch.setattr(installer, "CBM_BIN_DIR", tmp_path)
|
| 80 |
-
monkeypatch.setattr(installer, "_detect_platform", lambda: "linux-amd64")
|
| 81 |
-
monkeypatch.setattr(
|
| 82 |
-
installer, "urlopen", lambda url, timeout=60: FakeResponse(_build_archive())
|
| 83 |
-
)
|
| 84 |
-
|
| 85 |
-
run_calls: list[list[str]] = []
|
| 86 |
-
|
| 87 |
-
def fake_run(command, **kwargs):
|
| 88 |
-
run_calls.append(command)
|
| 89 |
-
return SimpleNamespace(returncode=1, stdout="")
|
| 90 |
-
|
| 91 |
-
monkeypatch.setattr("subprocess.run", fake_run)
|
| 92 |
-
path = installer.download_cbm(version="v1.2.3")
|
| 93 |
-
assert path == tmp_path / installer.CBM_BIN_NAME
|
| 94 |
-
assert path.exists()
|
| 95 |
-
assert run_calls == [[str(path), "--version"]]
|
| 96 |
-
|
| 97 |
-
monkeypatch.setattr(
|
| 98 |
-
"subprocess.run", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("boom"))
|
| 99 |
-
)
|
| 100 |
-
assert installer.download_cbm(version="v1.2.3") == path
|
| 101 |
-
|
| 102 |
-
monkeypatch.setattr(
|
| 103 |
-
"subprocess.run",
|
| 104 |
-
lambda *args, **kwargs: SimpleNamespace(returncode=0, stdout="cbm v1.2.3\n"),
|
| 105 |
-
)
|
| 106 |
-
assert installer.download_cbm(version="v1.2.3") == path
|
| 107 |
-
|
| 108 |
-
|
| 109 |
-
def test_download_cbm_invalid_url_download_failure_and_extract_errors(
|
| 110 |
-
monkeypatch, tmp_path: Path
|
| 111 |
-
) -> None:
|
| 112 |
-
monkeypatch.setattr(installer, "CBM_BIN_DIR", tmp_path)
|
| 113 |
-
monkeypatch.setattr(installer, "_detect_platform", lambda: "linux-amd64")
|
| 114 |
-
|
| 115 |
-
monkeypatch.setattr(installer, "GITHUB_RELEASE_URL", "ftp://example.test/releases")
|
| 116 |
-
with pytest.raises(RuntimeError, match="Invalid URL"):
|
| 117 |
-
installer.download_cbm()
|
| 118 |
-
|
| 119 |
-
monkeypatch.setattr(installer, "GITHUB_RELEASE_URL", "https://example.test/releases")
|
| 120 |
-
monkeypatch.setattr(
|
| 121 |
-
installer,
|
| 122 |
-
"urlopen",
|
| 123 |
-
lambda url, timeout=60: (_ for _ in ()).throw(OSError("network down")),
|
| 124 |
-
)
|
| 125 |
-
with pytest.raises(RuntimeError, match="Failed to download codebase-memory-mcp"):
|
| 126 |
-
installer.download_cbm()
|
| 127 |
-
|
| 128 |
-
monkeypatch.setattr(
|
| 129 |
-
installer,
|
| 130 |
-
"urlopen",
|
| 131 |
-
lambda url, timeout=60: FakeResponse(_build_archive("some/other-binary")),
|
| 132 |
-
)
|
| 133 |
-
with pytest.raises(RuntimeError, match="binary not found in archive"):
|
| 134 |
-
installer.download_cbm()
|
| 135 |
-
|
| 136 |
-
monkeypatch.setattr(installer, "urlopen", lambda url, timeout=60: FakeResponse(b"not a tar"))
|
| 137 |
-
with pytest.raises(RuntimeError, match="Failed to extract archive"):
|
| 138 |
-
installer.download_cbm()
|
| 139 |
-
|
| 140 |
-
|
| 141 |
-
def test_ensure_cbm_uses_existing_or_returns_none_on_failure(monkeypatch, tmp_path: Path) -> None:
|
| 142 |
-
existing = tmp_path / installer.CBM_BIN_NAME
|
| 143 |
-
monkeypatch.setattr(installer, "get_cbm_path", lambda: existing)
|
| 144 |
-
assert installer.ensure_cbm() == existing
|
| 145 |
-
|
| 146 |
-
monkeypatch.setattr(installer, "get_cbm_path", lambda: None)
|
| 147 |
-
monkeypatch.setattr(
|
| 148 |
-
installer, "download_cbm", lambda: (_ for _ in ()).throw(RuntimeError("nope"))
|
| 149 |
-
)
|
| 150 |
-
assert installer.ensure_cbm() is None
|
| 151 |
-
|
| 152 |
-
|
| 153 |
-
def test_code_graph_watcher_init_start_stop_and_event_filtering(
|
| 154 |
-
monkeypatch, tmp_path: Path
|
| 155 |
-
) -> None:
|
| 156 |
-
monkeypatch.setattr("headroom.graph.installer.get_cbm_path", lambda: tmp_path / "cbm")
|
| 157 |
-
graph_watcher = watcher.CodeGraphWatcher(tmp_path)
|
| 158 |
-
assert graph_watcher.cbm_binary == str(tmp_path / "cbm")
|
| 159 |
-
|
| 160 |
-
explicit = watcher.CodeGraphWatcher(tmp_path, cbm_binary="explicit-cbm")
|
| 161 |
-
assert explicit.cbm_binary == "explicit-cbm"
|
| 162 |
-
|
| 163 |
-
missing = watcher.CodeGraphWatcher(tmp_path, cbm_binary=None)
|
| 164 |
-
missing.cbm_binary = None
|
| 165 |
-
assert missing.start() is False
|
| 166 |
-
|
| 167 |
-
watchdog_mod = ModuleType("watchdog")
|
| 168 |
-
events_mod = ModuleType("watchdog.events")
|
| 169 |
-
observers_mod = ModuleType("watchdog.observers")
|
| 170 |
-
|
| 171 |
-
class FileSystemEventHandler:
|
| 172 |
-
pass
|
| 173 |
-
|
| 174 |
-
class FakeObserver:
|
| 175 |
-
def __init__(self) -> None:
|
| 176 |
-
self.scheduled = None
|
| 177 |
-
self.daemon = False
|
| 178 |
-
self.started = False
|
| 179 |
-
self.stopped = False
|
| 180 |
-
self.join_timeout = None
|
| 181 |
-
|
| 182 |
-
def schedule(self, handler, project_dir, recursive=True) -> None:
|
| 183 |
-
self.scheduled = (handler, project_dir, recursive)
|
| 184 |
-
|
| 185 |
-
def start(self) -> None:
|
| 186 |
-
self.started = True
|
| 187 |
-
|
| 188 |
-
def stop(self) -> None:
|
| 189 |
-
self.stopped = True
|
| 190 |
-
|
| 191 |
-
def join(self, timeout=None) -> None:
|
| 192 |
-
self.join_timeout = timeout
|
| 193 |
-
|
| 194 |
-
events_mod.FileSystemEventHandler = FileSystemEventHandler
|
| 195 |
-
observers_mod.Observer = FakeObserver
|
| 196 |
-
monkeypatch.setitem(__import__("sys").modules, "watchdog", watchdog_mod)
|
| 197 |
-
monkeypatch.setitem(__import__("sys").modules, "watchdog.events", events_mod)
|
| 198 |
-
monkeypatch.setitem(__import__("sys").modules, "watchdog.observers", observers_mod)
|
| 199 |
-
|
| 200 |
-
scheduled: list[str] = []
|
| 201 |
-
monkeypatch.setattr(graph_watcher, "_schedule_reindex", lambda: scheduled.append("reindex"))
|
| 202 |
-
|
| 203 |
-
assert graph_watcher.start() is True
|
| 204 |
-
handler, project_dir, recursive = graph_watcher._observer.scheduled
|
| 205 |
-
assert project_dir == str(tmp_path)
|
| 206 |
-
assert recursive is True
|
| 207 |
-
|
| 208 |
-
handler.on_any_event(SimpleNamespace(src_path=""))
|
| 209 |
-
handler.on_any_event(SimpleNamespace(src_path=str(tmp_path / ".git" / "config")))
|
| 210 |
-
handler.on_any_event(SimpleNamespace(src_path=str(tmp_path / "notes.txt")))
|
| 211 |
-
handler.on_any_event(SimpleNamespace(src_path=str(tmp_path / ".temp.py")))
|
| 212 |
-
handler.on_any_event(SimpleNamespace(src_path=str(tmp_path / "main.py~")))
|
| 213 |
-
handler.on_any_event(SimpleNamespace(src_path=str(tmp_path / "main.py")))
|
| 214 |
-
assert scheduled == ["reindex"]
|
| 215 |
-
|
| 216 |
-
class FakeTimer:
|
| 217 |
-
def __init__(self) -> None:
|
| 218 |
-
self.cancelled = False
|
| 219 |
-
|
| 220 |
-
def cancel(self) -> None:
|
| 221 |
-
self.cancelled = True
|
| 222 |
-
|
| 223 |
-
timer = FakeTimer()
|
| 224 |
-
graph_watcher._debounce_timer = timer
|
| 225 |
-
graph_watcher._reindex_count = 1
|
| 226 |
-
graph_watcher.stop()
|
| 227 |
-
assert timer.cancelled is True
|
| 228 |
-
assert graph_watcher._observer is None
|
| 229 |
-
|
| 230 |
-
|
| 231 |
-
def test_code_graph_watcher_start_returns_false_without_watchdog(
|
| 232 |
-
monkeypatch, tmp_path: Path
|
| 233 |
-
) -> None:
|
| 234 |
-
graph_watcher = watcher.CodeGraphWatcher(tmp_path, cbm_binary="cbm")
|
| 235 |
-
|
| 236 |
-
import builtins
|
| 237 |
-
|
| 238 |
-
real_import = builtins.__import__
|
| 239 |
-
|
| 240 |
-
def fake_import(name, globals=None, locals=None, fromlist=(), level=0):
|
| 241 |
-
if name.startswith("watchdog"):
|
| 242 |
-
raise ImportError("missing watchdog")
|
| 243 |
-
return real_import(name, globals, locals, fromlist, level)
|
| 244 |
-
|
| 245 |
-
monkeypatch.setattr(builtins, "__import__", fake_import)
|
| 246 |
-
assert graph_watcher.start() is False
|
| 247 |
-
|
| 248 |
-
|
| 249 |
-
def test_code_graph_watcher_stop_handles_missing_timer_and_observer_methods(tmp_path: Path) -> None:
|
| 250 |
-
graph_watcher = watcher.CodeGraphWatcher(tmp_path, cbm_binary="cbm")
|
| 251 |
-
graph_watcher._observer = object()
|
| 252 |
-
graph_watcher.stop()
|
| 253 |
-
assert graph_watcher._observer is None
|
| 254 |
-
|
| 255 |
-
graph_watcher.stop()
|
| 256 |
-
|
| 257 |
-
|
| 258 |
-
def test_schedule_reindex_replaces_existing_timer(monkeypatch, tmp_path: Path) -> None:
|
| 259 |
-
graph_watcher = watcher.CodeGraphWatcher(tmp_path, debounce_seconds=3.5, cbm_binary="cbm")
|
| 260 |
-
timers: list[FakeTimer] = []
|
| 261 |
-
|
| 262 |
-
class FakeTimer:
|
| 263 |
-
def __init__(self, interval, callback) -> None:
|
| 264 |
-
self.interval = interval
|
| 265 |
-
self.callback = callback
|
| 266 |
-
self.daemon = False
|
| 267 |
-
self.started = False
|
| 268 |
-
self.cancelled = False
|
| 269 |
-
timers.append(self)
|
| 270 |
-
|
| 271 |
-
def start(self) -> None:
|
| 272 |
-
self.started = True
|
| 273 |
-
|
| 274 |
-
def cancel(self) -> None:
|
| 275 |
-
self.cancelled = True
|
| 276 |
-
|
| 277 |
-
monkeypatch.setattr(watcher.threading, "Timer", FakeTimer)
|
| 278 |
-
graph_watcher._schedule_reindex()
|
| 279 |
-
graph_watcher._schedule_reindex()
|
| 280 |
-
|
| 281 |
-
assert len(timers) == 2
|
| 282 |
-
assert timers[0].cancelled is True
|
| 283 |
-
assert timers[1].started is True
|
| 284 |
-
assert timers[1].daemon is True
|
| 285 |
-
assert timers[1].interval == 3.5
|
| 286 |
-
|
| 287 |
-
|
| 288 |
-
def test_do_reindex_success_failure_timeout_and_stats(monkeypatch, tmp_path: Path) -> None:
|
| 289 |
-
graph_watcher = watcher.CodeGraphWatcher(tmp_path, cbm_binary="cbm")
|
| 290 |
-
graph_watcher._running = True
|
| 291 |
-
|
| 292 |
-
monotonic_values = iter([10.0, 10.4, 20.0, 20.5, 30.0, 30.5, 40.0, 40.5])
|
| 293 |
-
monkeypatch.setattr(watcher.time, "monotonic", lambda: next(monotonic_values))
|
| 294 |
-
monkeypatch.setattr(watcher.time, "time", lambda: 1234.0)
|
| 295 |
-
|
| 296 |
-
run_calls: list[list[str]] = []
|
| 297 |
-
|
| 298 |
-
def success_run(command, **kwargs):
|
| 299 |
-
run_calls.append(command)
|
| 300 |
-
return SimpleNamespace(returncode=0, stderr="indexed\nchanged=7 files\n")
|
| 301 |
-
|
| 302 |
-
monkeypatch.setattr(watcher.subprocess, "run", success_run)
|
| 303 |
-
graph_watcher._do_reindex()
|
| 304 |
-
assert graph_watcher.stats == {
|
| 305 |
-
"running": True,
|
| 306 |
-
"project_dir": str(tmp_path),
|
| 307 |
-
"reindex_count": 1,
|
| 308 |
-
"last_reindex": 1234.0,
|
| 309 |
-
"debounce_seconds": 2.0,
|
| 310 |
-
}
|
| 311 |
-
assert run_calls == [
|
| 312 |
-
["cbm", "cli", "index_repository", json.dumps({"repo_path": str(tmp_path), "mode": "fast"})]
|
| 313 |
-
]
|
| 314 |
-
|
| 315 |
-
monkeypatch.setattr(
|
| 316 |
-
watcher.subprocess,
|
| 317 |
-
"run",
|
| 318 |
-
lambda command, **kwargs: SimpleNamespace(returncode=1, stderr="failed"),
|
| 319 |
-
)
|
| 320 |
-
graph_watcher._do_reindex()
|
| 321 |
-
assert graph_watcher._reindex_count == 2
|
| 322 |
-
|
| 323 |
-
monkeypatch.setattr(
|
| 324 |
-
watcher.subprocess,
|
| 325 |
-
"run",
|
| 326 |
-
lambda command, **kwargs: SimpleNamespace(
|
| 327 |
-
returncode=0, stderr="indexed\nchanged=oops\nstill running\n"
|
| 328 |
-
),
|
| 329 |
-
)
|
| 330 |
-
graph_watcher._do_reindex()
|
| 331 |
-
assert graph_watcher._reindex_count == 3
|
| 332 |
-
|
| 333 |
-
monkeypatch.setattr(
|
| 334 |
-
watcher.subprocess,
|
| 335 |
-
"run",
|
| 336 |
-
lambda command, **kwargs: (_ for _ in ()).throw(subprocess.TimeoutExpired(command, 30)),
|
| 337 |
-
)
|
| 338 |
-
graph_watcher._do_reindex()
|
| 339 |
-
|
| 340 |
-
monkeypatch.setattr(
|
| 341 |
-
watcher.subprocess,
|
| 342 |
-
"run",
|
| 343 |
-
lambda command, **kwargs: (_ for _ in ()).throw(RuntimeError("boom")),
|
| 344 |
-
)
|
| 345 |
-
graph_watcher._do_reindex()
|
| 346 |
-
|
| 347 |
-
graph_watcher._running = False
|
| 348 |
-
graph_watcher._do_reindex()
|
| 349 |
-
|
| 350 |
-
graph_watcher._running = True
|
| 351 |
-
graph_watcher.cbm_binary = None
|
| 352 |
-
graph_watcher._do_reindex()
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import io
|
| 4 |
+
import json
|
| 5 |
+
import subprocess
|
| 6 |
+
import tarfile
|
| 7 |
+
from pathlib import Path
|
| 8 |
+
from types import ModuleType, SimpleNamespace
|
| 9 |
+
|
| 10 |
+
import pytest
|
| 11 |
+
|
| 12 |
+
from headroom.graph import installer, watcher
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
def _build_archive(member_name: str = installer.CBM_BIN_NAME) -> bytes:
|
| 16 |
+
payload = io.BytesIO()
|
| 17 |
+
with tarfile.open(fileobj=payload, mode="w:gz") as tar:
|
| 18 |
+
data = b"#!/bin/sh\necho version\n"
|
| 19 |
+
info = tarfile.TarInfo(name=member_name)
|
| 20 |
+
info.size = len(data)
|
| 21 |
+
tar.addfile(info, io.BytesIO(data))
|
| 22 |
+
return payload.getvalue()
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
class FakeResponse:
|
| 26 |
+
def __init__(self, data: bytes) -> None:
|
| 27 |
+
self._data = data
|
| 28 |
+
|
| 29 |
+
def __enter__(self):
|
| 30 |
+
return self
|
| 31 |
+
|
| 32 |
+
def __exit__(self, exc_type, exc, tb) -> None:
|
| 33 |
+
return None
|
| 34 |
+
|
| 35 |
+
def read(self) -> bytes:
|
| 36 |
+
return self._data
|
| 37 |
+
|
| 38 |
+
|
| 39 |
+
@pytest.mark.parametrize(
|
| 40 |
+
("system", "machine", "expected"),
|
| 41 |
+
[
|
| 42 |
+
("Darwin", "arm64", "darwin-arm64"),
|
| 43 |
+
("Darwin", "x86_64", "darwin-amd64"),
|
| 44 |
+
("Linux", "aarch64", "linux-arm64"),
|
| 45 |
+
("Linux", "arm64", "linux-arm64"),
|
| 46 |
+
("Linux", "x86_64", "linux-amd64"),
|
| 47 |
+
("Windows", "AMD64", "windows-amd64"),
|
| 48 |
+
],
|
| 49 |
+
)
|
| 50 |
+
def test_detect_platform_variants(monkeypatch, system: str, machine: str, expected: str) -> None:
|
| 51 |
+
monkeypatch.setattr(installer.platform, "system", lambda: system)
|
| 52 |
+
monkeypatch.setattr(installer.platform, "machine", lambda: machine)
|
| 53 |
+
assert installer._detect_platform() == expected
|
| 54 |
+
|
| 55 |
+
|
| 56 |
+
def test_detect_platform_rejects_unknown_system(monkeypatch) -> None:
|
| 57 |
+
monkeypatch.setattr(installer.platform, "system", lambda: "Solaris")
|
| 58 |
+
monkeypatch.setattr(installer.platform, "machine", lambda: "sparc")
|
| 59 |
+
with pytest.raises(RuntimeError, match="Unsupported platform"):
|
| 60 |
+
installer._detect_platform()
|
| 61 |
+
|
| 62 |
+
|
| 63 |
+
def test_get_cbm_path_prefers_path_then_install_dir(monkeypatch, tmp_path: Path) -> None:
|
| 64 |
+
on_path = tmp_path / "on-path"
|
| 65 |
+
installed = tmp_path / installer.CBM_BIN_NAME
|
| 66 |
+
installed.write_text("bin")
|
| 67 |
+
monkeypatch.setattr(installer, "CBM_BIN_DIR", tmp_path)
|
| 68 |
+
monkeypatch.setattr(installer.shutil, "which", lambda name: str(on_path))
|
| 69 |
+
assert installer.get_cbm_path() == on_path
|
| 70 |
+
|
| 71 |
+
monkeypatch.setattr(installer.shutil, "which", lambda name: None)
|
| 72 |
+
assert installer.get_cbm_path() == installed
|
| 73 |
+
|
| 74 |
+
installed.unlink()
|
| 75 |
+
assert installer.get_cbm_path() is None
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
def test_download_cbm_success_and_verification_paths(monkeypatch, tmp_path: Path) -> None:
|
| 79 |
+
monkeypatch.setattr(installer, "CBM_BIN_DIR", tmp_path)
|
| 80 |
+
monkeypatch.setattr(installer, "_detect_platform", lambda: "linux-amd64")
|
| 81 |
+
monkeypatch.setattr(
|
| 82 |
+
installer, "urlopen", lambda url, timeout=60: FakeResponse(_build_archive())
|
| 83 |
+
)
|
| 84 |
+
|
| 85 |
+
run_calls: list[list[str]] = []
|
| 86 |
+
|
| 87 |
+
def fake_run(command, **kwargs):
|
| 88 |
+
run_calls.append(command)
|
| 89 |
+
return SimpleNamespace(returncode=1, stdout="")
|
| 90 |
+
|
| 91 |
+
monkeypatch.setattr("subprocess.run", fake_run)
|
| 92 |
+
path = installer.download_cbm(version="v1.2.3")
|
| 93 |
+
assert path == tmp_path / installer.CBM_BIN_NAME
|
| 94 |
+
assert path.exists()
|
| 95 |
+
assert run_calls == [[str(path), "--version"]]
|
| 96 |
+
|
| 97 |
+
monkeypatch.setattr(
|
| 98 |
+
"subprocess.run", lambda *args, **kwargs: (_ for _ in ()).throw(RuntimeError("boom"))
|
| 99 |
+
)
|
| 100 |
+
assert installer.download_cbm(version="v1.2.3") == path
|
| 101 |
+
|
| 102 |
+
monkeypatch.setattr(
|
| 103 |
+
"subprocess.run",
|
| 104 |
+
lambda *args, **kwargs: SimpleNamespace(returncode=0, stdout="cbm v1.2.3\n"),
|
| 105 |
+
)
|
| 106 |
+
assert installer.download_cbm(version="v1.2.3") == path
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
def test_download_cbm_invalid_url_download_failure_and_extract_errors(
|
| 110 |
+
monkeypatch, tmp_path: Path
|
| 111 |
+
) -> None:
|
| 112 |
+
monkeypatch.setattr(installer, "CBM_BIN_DIR", tmp_path)
|
| 113 |
+
monkeypatch.setattr(installer, "_detect_platform", lambda: "linux-amd64")
|
| 114 |
+
|
| 115 |
+
monkeypatch.setattr(installer, "GITHUB_RELEASE_URL", "ftp://example.test/releases")
|
| 116 |
+
with pytest.raises(RuntimeError, match="Invalid URL"):
|
| 117 |
+
installer.download_cbm()
|
| 118 |
+
|
| 119 |
+
monkeypatch.setattr(installer, "GITHUB_RELEASE_URL", "https://example.test/releases")
|
| 120 |
+
monkeypatch.setattr(
|
| 121 |
+
installer,
|
| 122 |
+
"urlopen",
|
| 123 |
+
lambda url, timeout=60: (_ for _ in ()).throw(OSError("network down")),
|
| 124 |
+
)
|
| 125 |
+
with pytest.raises(RuntimeError, match="Failed to download codebase-memory-mcp"):
|
| 126 |
+
installer.download_cbm()
|
| 127 |
+
|
| 128 |
+
monkeypatch.setattr(
|
| 129 |
+
installer,
|
| 130 |
+
"urlopen",
|
| 131 |
+
lambda url, timeout=60: FakeResponse(_build_archive("some/other-binary")),
|
| 132 |
+
)
|
| 133 |
+
with pytest.raises(RuntimeError, match="binary not found in archive"):
|
| 134 |
+
installer.download_cbm()
|
| 135 |
+
|
| 136 |
+
monkeypatch.setattr(installer, "urlopen", lambda url, timeout=60: FakeResponse(b"not a tar"))
|
| 137 |
+
with pytest.raises(RuntimeError, match="Failed to extract archive"):
|
| 138 |
+
installer.download_cbm()
|
| 139 |
+
|
| 140 |
+
|
| 141 |
+
def test_ensure_cbm_uses_existing_or_returns_none_on_failure(monkeypatch, tmp_path: Path) -> None:
|
| 142 |
+
existing = tmp_path / installer.CBM_BIN_NAME
|
| 143 |
+
monkeypatch.setattr(installer, "get_cbm_path", lambda: existing)
|
| 144 |
+
assert installer.ensure_cbm() == existing
|
| 145 |
+
|
| 146 |
+
monkeypatch.setattr(installer, "get_cbm_path", lambda: None)
|
| 147 |
+
monkeypatch.setattr(
|
| 148 |
+
installer, "download_cbm", lambda: (_ for _ in ()).throw(RuntimeError("nope"))
|
| 149 |
+
)
|
| 150 |
+
assert installer.ensure_cbm() is None
|
| 151 |
+
|
| 152 |
+
|
| 153 |
+
def test_code_graph_watcher_init_start_stop_and_event_filtering(
|
| 154 |
+
monkeypatch, tmp_path: Path
|
| 155 |
+
) -> None:
|
| 156 |
+
monkeypatch.setattr("headroom.graph.installer.get_cbm_path", lambda: tmp_path / "cbm")
|
| 157 |
+
graph_watcher = watcher.CodeGraphWatcher(tmp_path)
|
| 158 |
+
assert graph_watcher.cbm_binary == str(tmp_path / "cbm")
|
| 159 |
+
|
| 160 |
+
explicit = watcher.CodeGraphWatcher(tmp_path, cbm_binary="explicit-cbm")
|
| 161 |
+
assert explicit.cbm_binary == "explicit-cbm"
|
| 162 |
+
|
| 163 |
+
missing = watcher.CodeGraphWatcher(tmp_path, cbm_binary=None)
|
| 164 |
+
missing.cbm_binary = None
|
| 165 |
+
assert missing.start() is False
|
| 166 |
+
|
| 167 |
+
watchdog_mod = ModuleType("watchdog")
|
| 168 |
+
events_mod = ModuleType("watchdog.events")
|
| 169 |
+
observers_mod = ModuleType("watchdog.observers")
|
| 170 |
+
|
| 171 |
+
class FileSystemEventHandler:
|
| 172 |
+
pass
|
| 173 |
+
|
| 174 |
+
class FakeObserver:
|
| 175 |
+
def __init__(self) -> None:
|
| 176 |
+
self.scheduled = None
|
| 177 |
+
self.daemon = False
|
| 178 |
+
self.started = False
|
| 179 |
+
self.stopped = False
|
| 180 |
+
self.join_timeout = None
|
| 181 |
+
|
| 182 |
+
def schedule(self, handler, project_dir, recursive=True) -> None:
|
| 183 |
+
self.scheduled = (handler, project_dir, recursive)
|
| 184 |
+
|
| 185 |
+
def start(self) -> None:
|
| 186 |
+
self.started = True
|
| 187 |
+
|
| 188 |
+
def stop(self) -> None:
|
| 189 |
+
self.stopped = True
|
| 190 |
+
|
| 191 |
+
def join(self, timeout=None) -> None:
|
| 192 |
+
self.join_timeout = timeout
|
| 193 |
+
|
| 194 |
+
events_mod.FileSystemEventHandler = FileSystemEventHandler
|
| 195 |
+
observers_mod.Observer = FakeObserver
|
| 196 |
+
monkeypatch.setitem(__import__("sys").modules, "watchdog", watchdog_mod)
|
| 197 |
+
monkeypatch.setitem(__import__("sys").modules, "watchdog.events", events_mod)
|
| 198 |
+
monkeypatch.setitem(__import__("sys").modules, "watchdog.observers", observers_mod)
|
| 199 |
+
|
| 200 |
+
scheduled: list[str] = []
|
| 201 |
+
monkeypatch.setattr(graph_watcher, "_schedule_reindex", lambda: scheduled.append("reindex"))
|
| 202 |
+
|
| 203 |
+
assert graph_watcher.start() is True
|
| 204 |
+
handler, project_dir, recursive = graph_watcher._observer.scheduled
|
| 205 |
+
assert project_dir == str(tmp_path)
|
| 206 |
+
assert recursive is True
|
| 207 |
+
|
| 208 |
+
handler.on_any_event(SimpleNamespace(src_path=""))
|
| 209 |
+
handler.on_any_event(SimpleNamespace(src_path=str(tmp_path / ".git" / "config")))
|
| 210 |
+
handler.on_any_event(SimpleNamespace(src_path=str(tmp_path / "notes.txt")))
|
| 211 |
+
handler.on_any_event(SimpleNamespace(src_path=str(tmp_path / ".temp.py")))
|
| 212 |
+
handler.on_any_event(SimpleNamespace(src_path=str(tmp_path / "main.py~")))
|
| 213 |
+
handler.on_any_event(SimpleNamespace(src_path=str(tmp_path / "main.py")))
|
| 214 |
+
assert scheduled == ["reindex"]
|
| 215 |
+
|
| 216 |
+
class FakeTimer:
|
| 217 |
+
def __init__(self) -> None:
|
| 218 |
+
self.cancelled = False
|
| 219 |
+
|
| 220 |
+
def cancel(self) -> None:
|
| 221 |
+
self.cancelled = True
|
| 222 |
+
|
| 223 |
+
timer = FakeTimer()
|
| 224 |
+
graph_watcher._debounce_timer = timer
|
| 225 |
+
graph_watcher._reindex_count = 1
|
| 226 |
+
graph_watcher.stop()
|
| 227 |
+
assert timer.cancelled is True
|
| 228 |
+
assert graph_watcher._observer is None
|
| 229 |
+
|
| 230 |
+
|
| 231 |
+
def test_code_graph_watcher_start_returns_false_without_watchdog(
|
| 232 |
+
monkeypatch, tmp_path: Path
|
| 233 |
+
) -> None:
|
| 234 |
+
graph_watcher = watcher.CodeGraphWatcher(tmp_path, cbm_binary="cbm")
|
| 235 |
+
|
| 236 |
+
import builtins
|
| 237 |
+
|
| 238 |
+
real_import = builtins.__import__
|
| 239 |
+
|
| 240 |
+
def fake_import(name, globals=None, locals=None, fromlist=(), level=0):
|
| 241 |
+
if name.startswith("watchdog"):
|
| 242 |
+
raise ImportError("missing watchdog")
|
| 243 |
+
return real_import(name, globals, locals, fromlist, level)
|
| 244 |
+
|
| 245 |
+
monkeypatch.setattr(builtins, "__import__", fake_import)
|
| 246 |
+
assert graph_watcher.start() is False
|
| 247 |
+
|
| 248 |
+
|
| 249 |
+
def test_code_graph_watcher_stop_handles_missing_timer_and_observer_methods(tmp_path: Path) -> None:
|
| 250 |
+
graph_watcher = watcher.CodeGraphWatcher(tmp_path, cbm_binary="cbm")
|
| 251 |
+
graph_watcher._observer = object()
|
| 252 |
+
graph_watcher.stop()
|
| 253 |
+
assert graph_watcher._observer is None
|
| 254 |
+
|
| 255 |
+
graph_watcher.stop()
|
| 256 |
+
|
| 257 |
+
|
| 258 |
+
def test_schedule_reindex_replaces_existing_timer(monkeypatch, tmp_path: Path) -> None:
|
| 259 |
+
graph_watcher = watcher.CodeGraphWatcher(tmp_path, debounce_seconds=3.5, cbm_binary="cbm")
|
| 260 |
+
timers: list[FakeTimer] = []
|
| 261 |
+
|
| 262 |
+
class FakeTimer:
|
| 263 |
+
def __init__(self, interval, callback) -> None:
|
| 264 |
+
self.interval = interval
|
| 265 |
+
self.callback = callback
|
| 266 |
+
self.daemon = False
|
| 267 |
+
self.started = False
|
| 268 |
+
self.cancelled = False
|
| 269 |
+
timers.append(self)
|
| 270 |
+
|
| 271 |
+
def start(self) -> None:
|
| 272 |
+
self.started = True
|
| 273 |
+
|
| 274 |
+
def cancel(self) -> None:
|
| 275 |
+
self.cancelled = True
|
| 276 |
+
|
| 277 |
+
monkeypatch.setattr(watcher.threading, "Timer", FakeTimer)
|
| 278 |
+
graph_watcher._schedule_reindex()
|
| 279 |
+
graph_watcher._schedule_reindex()
|
| 280 |
+
|
| 281 |
+
assert len(timers) == 2
|
| 282 |
+
assert timers[0].cancelled is True
|
| 283 |
+
assert timers[1].started is True
|
| 284 |
+
assert timers[1].daemon is True
|
| 285 |
+
assert timers[1].interval == 3.5
|
| 286 |
+
|
| 287 |
+
|
| 288 |
+
def test_do_reindex_success_failure_timeout_and_stats(monkeypatch, tmp_path: Path) -> None:
|
| 289 |
+
graph_watcher = watcher.CodeGraphWatcher(tmp_path, cbm_binary="cbm")
|
| 290 |
+
graph_watcher._running = True
|
| 291 |
+
|
| 292 |
+
monotonic_values = iter([10.0, 10.4, 20.0, 20.5, 30.0, 30.5, 40.0, 40.5])
|
| 293 |
+
monkeypatch.setattr(watcher.time, "monotonic", lambda: next(monotonic_values))
|
| 294 |
+
monkeypatch.setattr(watcher.time, "time", lambda: 1234.0)
|
| 295 |
+
|
| 296 |
+
run_calls: list[list[str]] = []
|
| 297 |
+
|
| 298 |
+
def success_run(command, **kwargs):
|
| 299 |
+
run_calls.append(command)
|
| 300 |
+
return SimpleNamespace(returncode=0, stderr="indexed\nchanged=7 files\n")
|
| 301 |
+
|
| 302 |
+
monkeypatch.setattr(watcher.subprocess, "run", success_run)
|
| 303 |
+
graph_watcher._do_reindex()
|
| 304 |
+
assert graph_watcher.stats == {
|
| 305 |
+
"running": True,
|
| 306 |
+
"project_dir": str(tmp_path),
|
| 307 |
+
"reindex_count": 1,
|
| 308 |
+
"last_reindex": 1234.0,
|
| 309 |
+
"debounce_seconds": 2.0,
|
| 310 |
+
}
|
| 311 |
+
assert run_calls == [
|
| 312 |
+
["cbm", "cli", "index_repository", json.dumps({"repo_path": str(tmp_path), "mode": "fast"})]
|
| 313 |
+
]
|
| 314 |
+
|
| 315 |
+
monkeypatch.setattr(
|
| 316 |
+
watcher.subprocess,
|
| 317 |
+
"run",
|
| 318 |
+
lambda command, **kwargs: SimpleNamespace(returncode=1, stderr="failed"),
|
| 319 |
+
)
|
| 320 |
+
graph_watcher._do_reindex()
|
| 321 |
+
assert graph_watcher._reindex_count == 2
|
| 322 |
+
|
| 323 |
+
monkeypatch.setattr(
|
| 324 |
+
watcher.subprocess,
|
| 325 |
+
"run",
|
| 326 |
+
lambda command, **kwargs: SimpleNamespace(
|
| 327 |
+
returncode=0, stderr="indexed\nchanged=oops\nstill running\n"
|
| 328 |
+
),
|
| 329 |
+
)
|
| 330 |
+
graph_watcher._do_reindex()
|
| 331 |
+
assert graph_watcher._reindex_count == 3
|
| 332 |
+
|
| 333 |
+
monkeypatch.setattr(
|
| 334 |
+
watcher.subprocess,
|
| 335 |
+
"run",
|
| 336 |
+
lambda command, **kwargs: (_ for _ in ()).throw(subprocess.TimeoutExpired(command, 30)),
|
| 337 |
+
)
|
| 338 |
+
graph_watcher._do_reindex()
|
| 339 |
+
|
| 340 |
+
monkeypatch.setattr(
|
| 341 |
+
watcher.subprocess,
|
| 342 |
+
"run",
|
| 343 |
+
lambda command, **kwargs: (_ for _ in ()).throw(RuntimeError("boom")),
|
| 344 |
+
)
|
| 345 |
+
graph_watcher._do_reindex()
|
| 346 |
+
|
| 347 |
+
graph_watcher._running = False
|
| 348 |
+
graph_watcher._do_reindex()
|
| 349 |
+
|
| 350 |
+
graph_watcher._running = True
|
| 351 |
+
graph_watcher.cbm_binary = None
|
| 352 |
+
graph_watcher._do_reindex()
|
|
@@ -1,68 +1,68 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
from pathlib import Path
|
| 4 |
-
|
| 5 |
-
import click
|
| 6 |
-
import pytest
|
| 7 |
-
|
| 8 |
-
from headroom.install import paths as install_paths
|
| 9 |
-
|
| 10 |
-
|
| 11 |
-
def test_validate_profile_name_accepts_and_rejects_values() -> None:
|
| 12 |
-
assert install_paths.validate_profile_name("good.profile-1_2") == "good.profile-1_2"
|
| 13 |
-
|
| 14 |
-
for value in (".", "..", "bad/name", "bad space", ""):
|
| 15 |
-
with pytest.raises(click.ClickException, match="Invalid profile name"):
|
| 16 |
-
install_paths.validate_profile_name(value)
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
def test_profile_and_artifact_paths(monkeypatch, tmp_path: Path) -> None:
|
| 20 |
-
monkeypatch.setattr("headroom.install.paths._paths.deploy_root", lambda: tmp_path / "deploy")
|
| 21 |
-
|
| 22 |
-
assert install_paths.deploy_root() == tmp_path / "deploy"
|
| 23 |
-
assert install_paths.profile_root("demo") == tmp_path / "deploy" / "demo"
|
| 24 |
-
assert install_paths.manifest_path("demo") == tmp_path / "deploy" / "demo" / "manifest.json"
|
| 25 |
-
assert install_paths.log_path("demo") == tmp_path / "deploy" / "demo" / "runner.log"
|
| 26 |
-
assert install_paths.pid_path("demo") == tmp_path / "deploy" / "demo" / "runner.pid"
|
| 27 |
-
assert (
|
| 28 |
-
install_paths.unix_run_script_path("demo")
|
| 29 |
-
== tmp_path / "deploy" / "demo" / "run-headroom.sh"
|
| 30 |
-
)
|
| 31 |
-
assert install_paths.unix_ensure_script_path("demo") == (
|
| 32 |
-
tmp_path / "deploy" / "demo" / "ensure-headroom.sh"
|
| 33 |
-
)
|
| 34 |
-
assert install_paths.windows_run_script_path("demo") == (
|
| 35 |
-
tmp_path / "deploy" / "demo" / "run-headroom.ps1"
|
| 36 |
-
)
|
| 37 |
-
assert install_paths.windows_run_cmd_path("demo") == (
|
| 38 |
-
tmp_path / "deploy" / "demo" / "run-headroom.cmd"
|
| 39 |
-
)
|
| 40 |
-
assert install_paths.windows_ensure_script_path("demo") == (
|
| 41 |
-
tmp_path / "deploy" / "demo" / "ensure-headroom.ps1"
|
| 42 |
-
)
|
| 43 |
-
assert install_paths.windows_ensure_cmd_path("demo") == (
|
| 44 |
-
tmp_path / "deploy" / "demo" / "ensure-headroom.cmd"
|
| 45 |
-
)
|
| 46 |
-
|
| 47 |
-
|
| 48 |
-
def test_env_target_and_config_paths(monkeypatch, tmp_path: Path) -> None:
|
| 49 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 50 |
-
monkeypatch.setattr("headroom.install.paths.sys.platform", "linux")
|
| 51 |
-
|
| 52 |
-
assert install_paths.unix_user_env_targets() == [
|
| 53 |
-
tmp_path / ".bashrc",
|
| 54 |
-
tmp_path / ".zshrc",
|
| 55 |
-
tmp_path / ".profile",
|
| 56 |
-
]
|
| 57 |
-
assert install_paths.unix_system_env_targets() == [Path("/etc/profile.d/headroom.sh")]
|
| 58 |
-
|
| 59 |
-
monkeypatch.setattr("headroom.install.paths.sys.platform", "darwin")
|
| 60 |
-
assert install_paths.unix_system_env_targets() == [
|
| 61 |
-
Path("/etc/profile"),
|
| 62 |
-
Path("/etc/zprofile"),
|
| 63 |
-
Path("/etc/bashrc"),
|
| 64 |
-
]
|
| 65 |
-
|
| 66 |
-
assert install_paths.claude_settings_path() == tmp_path / ".claude" / "settings.json"
|
| 67 |
-
assert install_paths.codex_config_path() == tmp_path / ".codex" / "config.toml"
|
| 68 |
-
assert install_paths.openclaw_config_path() == tmp_path / ".openclaw" / "openclaw.json"
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
|
| 5 |
+
import click
|
| 6 |
+
import pytest
|
| 7 |
+
|
| 8 |
+
from headroom.install import paths as install_paths
|
| 9 |
+
|
| 10 |
+
|
| 11 |
+
def test_validate_profile_name_accepts_and_rejects_values() -> None:
|
| 12 |
+
assert install_paths.validate_profile_name("good.profile-1_2") == "good.profile-1_2"
|
| 13 |
+
|
| 14 |
+
for value in (".", "..", "bad/name", "bad space", ""):
|
| 15 |
+
with pytest.raises(click.ClickException, match="Invalid profile name"):
|
| 16 |
+
install_paths.validate_profile_name(value)
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
def test_profile_and_artifact_paths(monkeypatch, tmp_path: Path) -> None:
|
| 20 |
+
monkeypatch.setattr("headroom.install.paths._paths.deploy_root", lambda: tmp_path / "deploy")
|
| 21 |
+
|
| 22 |
+
assert install_paths.deploy_root() == tmp_path / "deploy"
|
| 23 |
+
assert install_paths.profile_root("demo") == tmp_path / "deploy" / "demo"
|
| 24 |
+
assert install_paths.manifest_path("demo") == tmp_path / "deploy" / "demo" / "manifest.json"
|
| 25 |
+
assert install_paths.log_path("demo") == tmp_path / "deploy" / "demo" / "runner.log"
|
| 26 |
+
assert install_paths.pid_path("demo") == tmp_path / "deploy" / "demo" / "runner.pid"
|
| 27 |
+
assert (
|
| 28 |
+
install_paths.unix_run_script_path("demo")
|
| 29 |
+
== tmp_path / "deploy" / "demo" / "run-headroom.sh"
|
| 30 |
+
)
|
| 31 |
+
assert install_paths.unix_ensure_script_path("demo") == (
|
| 32 |
+
tmp_path / "deploy" / "demo" / "ensure-headroom.sh"
|
| 33 |
+
)
|
| 34 |
+
assert install_paths.windows_run_script_path("demo") == (
|
| 35 |
+
tmp_path / "deploy" / "demo" / "run-headroom.ps1"
|
| 36 |
+
)
|
| 37 |
+
assert install_paths.windows_run_cmd_path("demo") == (
|
| 38 |
+
tmp_path / "deploy" / "demo" / "run-headroom.cmd"
|
| 39 |
+
)
|
| 40 |
+
assert install_paths.windows_ensure_script_path("demo") == (
|
| 41 |
+
tmp_path / "deploy" / "demo" / "ensure-headroom.ps1"
|
| 42 |
+
)
|
| 43 |
+
assert install_paths.windows_ensure_cmd_path("demo") == (
|
| 44 |
+
tmp_path / "deploy" / "demo" / "ensure-headroom.cmd"
|
| 45 |
+
)
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def test_env_target_and_config_paths(monkeypatch, tmp_path: Path) -> None:
|
| 49 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 50 |
+
monkeypatch.setattr("headroom.install.paths.sys.platform", "linux")
|
| 51 |
+
|
| 52 |
+
assert install_paths.unix_user_env_targets() == [
|
| 53 |
+
tmp_path / ".bashrc",
|
| 54 |
+
tmp_path / ".zshrc",
|
| 55 |
+
tmp_path / ".profile",
|
| 56 |
+
]
|
| 57 |
+
assert install_paths.unix_system_env_targets() == [Path("/etc/profile.d/headroom.sh")]
|
| 58 |
+
|
| 59 |
+
monkeypatch.setattr("headroom.install.paths.sys.platform", "darwin")
|
| 60 |
+
assert install_paths.unix_system_env_targets() == [
|
| 61 |
+
Path("/etc/profile"),
|
| 62 |
+
Path("/etc/zprofile"),
|
| 63 |
+
Path("/etc/bashrc"),
|
| 64 |
+
]
|
| 65 |
+
|
| 66 |
+
assert install_paths.claude_settings_path() == tmp_path / ".claude" / "settings.json"
|
| 67 |
+
assert install_paths.codex_config_path() == tmp_path / ".codex" / "config.toml"
|
| 68 |
+
assert install_paths.openclaw_config_path() == tmp_path / ".openclaw" / "openclaw.json"
|
|
@@ -1,468 +1,468 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
import signal
|
| 4 |
-
from pathlib import Path
|
| 5 |
-
|
| 6 |
-
from headroom.install.models import DeploymentManifest, InstallPreset
|
| 7 |
-
from headroom.install.runtime import (
|
| 8 |
-
_clear_pid,
|
| 9 |
-
_deployment_env,
|
| 10 |
-
_mount_source,
|
| 11 |
-
_read_pid,
|
| 12 |
-
_runtime_env,
|
| 13 |
-
_write_pid,
|
| 14 |
-
build_runtime_command,
|
| 15 |
-
resolve_headroom_command,
|
| 16 |
-
run_foreground,
|
| 17 |
-
runtime_status,
|
| 18 |
-
start_detached_agent,
|
| 19 |
-
start_persistent_docker,
|
| 20 |
-
stop_runtime,
|
| 21 |
-
wait_ready,
|
| 22 |
-
)
|
| 23 |
-
|
| 24 |
-
|
| 25 |
-
def test_build_runtime_command_for_docker_includes_deployment_env(
|
| 26 |
-
monkeypatch, tmp_path: Path
|
| 27 |
-
) -> None:
|
| 28 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 29 |
-
manifest = DeploymentManifest(
|
| 30 |
-
profile="default",
|
| 31 |
-
preset="persistent-docker",
|
| 32 |
-
runtime_kind="docker",
|
| 33 |
-
supervisor_kind="none",
|
| 34 |
-
scope="user",
|
| 35 |
-
provider_mode="manual",
|
| 36 |
-
targets=["claude"],
|
| 37 |
-
port=8787,
|
| 38 |
-
host="127.0.0.1",
|
| 39 |
-
backend="anthropic",
|
| 40 |
-
image="ghcr.io/chopratejas/headroom:latest",
|
| 41 |
-
base_env={"HEADROOM_PORT": "8787"},
|
| 42 |
-
proxy_args=["--host", "127.0.0.1", "--port", "8787"],
|
| 43 |
-
)
|
| 44 |
-
|
| 45 |
-
command = build_runtime_command(manifest)
|
| 46 |
-
|
| 47 |
-
joined = " ".join(command)
|
| 48 |
-
assert command[:3] == ["docker", "run", "--rm"]
|
| 49 |
-
assert "HEADROOM_DEPLOYMENT_PROFILE=default" in joined
|
| 50 |
-
assert "HEADROOM_DEPLOYMENT_PRESET=persistent-docker" in joined
|
| 51 |
-
assert "127.0.0.1:8787:8787" in joined
|
| 52 |
-
assert "ghcr.io/chopratejas/headroom:latest" in command
|
| 53 |
-
# Canonical Headroom filesystem contract (issue #175) forwarded into
|
| 54 |
-
# the container.
|
| 55 |
-
assert "HEADROOM_WORKSPACE_DIR=/tmp/headroom-home/.headroom" in command
|
| 56 |
-
assert "HEADROOM_CONFIG_DIR=/tmp/headroom-home/.headroom/config" in command
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
def test_build_runtime_command_for_docker_matches_wrapper_parity(
|
| 60 |
-
monkeypatch, tmp_path: Path
|
| 61 |
-
) -> None:
|
| 62 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 63 |
-
monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key")
|
| 64 |
-
monkeypatch.setenv("OPENAI_API_KEY", "test-openai")
|
| 65 |
-
manifest = DeploymentManifest(
|
| 66 |
-
profile="default",
|
| 67 |
-
preset="persistent-docker",
|
| 68 |
-
runtime_kind="docker",
|
| 69 |
-
supervisor_kind="none",
|
| 70 |
-
scope="user",
|
| 71 |
-
provider_mode="manual",
|
| 72 |
-
targets=["claude"],
|
| 73 |
-
port=8787,
|
| 74 |
-
host="127.0.0.1",
|
| 75 |
-
backend="anthropic",
|
| 76 |
-
image="ghcr.io/chopratejas/headroom:latest",
|
| 77 |
-
base_env={"HEADROOM_PORT": "8787"},
|
| 78 |
-
proxy_args=["--host", "127.0.0.1", "--port", "8787"],
|
| 79 |
-
)
|
| 80 |
-
|
| 81 |
-
command = build_runtime_command(manifest)
|
| 82 |
-
|
| 83 |
-
assert (tmp_path / ".headroom").is_dir()
|
| 84 |
-
assert (tmp_path / ".claude").is_dir()
|
| 85 |
-
assert (tmp_path / ".codex").is_dir()
|
| 86 |
-
assert (tmp_path / ".gemini").is_dir()
|
| 87 |
-
assert "--env" in command
|
| 88 |
-
joined = " ".join(command)
|
| 89 |
-
assert "ANTHROPIC_API_KEY" in joined
|
| 90 |
-
assert "OPENAI_API_KEY" in joined
|
| 91 |
-
|
| 92 |
-
|
| 93 |
-
def test_resolve_headroom_command_prefers_headroom_binary(monkeypatch) -> None:
|
| 94 |
-
monkeypatch.setattr(
|
| 95 |
-
"shutil.which", lambda name: "/usr/bin/headroom" if name == "headroom" else None
|
| 96 |
-
)
|
| 97 |
-
|
| 98 |
-
assert resolve_headroom_command() == ["/usr/bin/headroom"]
|
| 99 |
-
|
| 100 |
-
|
| 101 |
-
def test_resolve_headroom_command_falls_back_to_python_module(monkeypatch) -> None:
|
| 102 |
-
monkeypatch.setattr("shutil.which", lambda name: None)
|
| 103 |
-
monkeypatch.setattr("headroom.install.runtime.sys.executable", "/usr/bin/python")
|
| 104 |
-
assert resolve_headroom_command() == ["/usr/bin/python", "-m", "headroom.cli"]
|
| 105 |
-
|
| 106 |
-
|
| 107 |
-
def test_runtime_env_and_mount_source(monkeypatch) -> None:
|
| 108 |
-
manifest = DeploymentManifest(
|
| 109 |
-
profile="default",
|
| 110 |
-
preset="persistent-service",
|
| 111 |
-
runtime_kind="python",
|
| 112 |
-
supervisor_kind="service",
|
| 113 |
-
scope="user",
|
| 114 |
-
provider_mode="manual",
|
| 115 |
-
targets=[],
|
| 116 |
-
port=8787,
|
| 117 |
-
host="127.0.0.1",
|
| 118 |
-
backend="anthropic",
|
| 119 |
-
base_env={"EXTRA": "1"},
|
| 120 |
-
)
|
| 121 |
-
monkeypatch.setattr("headroom.install.runtime.os.environ", {"BASE": "x"})
|
| 122 |
-
|
| 123 |
-
assert _deployment_env(manifest) == {
|
| 124 |
-
"HEADROOM_DEPLOYMENT_PROFILE": "default",
|
| 125 |
-
"HEADROOM_DEPLOYMENT_PRESET": "persistent-service",
|
| 126 |
-
"HEADROOM_DEPLOYMENT_RUNTIME": "python",
|
| 127 |
-
"HEADROOM_DEPLOYMENT_SUPERVISOR": "service",
|
| 128 |
-
"HEADROOM_DEPLOYMENT_SCOPE": "user",
|
| 129 |
-
}
|
| 130 |
-
assert _runtime_env(manifest)["BASE"] == "x"
|
| 131 |
-
assert _runtime_env(manifest)["EXTRA"] == "1"
|
| 132 |
-
assert _runtime_env(manifest)["HEADROOM_DEPLOYMENT_PROFILE"] == "default"
|
| 133 |
-
|
| 134 |
-
monkeypatch.setattr("headroom.install.runtime.sys.platform", "win32")
|
| 135 |
-
assert _mount_source("C:\\Users\\me", ".headroom") == "C:\\Users\\me\\.headroom"
|
| 136 |
-
monkeypatch.setattr("headroom.install.runtime.sys.platform", "linux")
|
| 137 |
-
assert _mount_source("/home/me", ".headroom") == "/home/me/.headroom"
|
| 138 |
-
|
| 139 |
-
|
| 140 |
-
def test_build_runtime_command_python_and_docker_user(monkeypatch, tmp_path: Path) -> None:
|
| 141 |
-
monkeypatch.setattr("headroom.install.runtime.sys.executable", "/usr/bin/python")
|
| 142 |
-
manifest = DeploymentManifest(
|
| 143 |
-
profile="default",
|
| 144 |
-
preset="persistent-service",
|
| 145 |
-
runtime_kind="python",
|
| 146 |
-
supervisor_kind="service",
|
| 147 |
-
scope="user",
|
| 148 |
-
provider_mode="manual",
|
| 149 |
-
targets=[],
|
| 150 |
-
port=8787,
|
| 151 |
-
host="127.0.0.1",
|
| 152 |
-
backend="anthropic",
|
| 153 |
-
proxy_args=["--host", "127.0.0.1", "--port", "8787"],
|
| 154 |
-
)
|
| 155 |
-
assert build_runtime_command(manifest) == [
|
| 156 |
-
"/usr/bin/python",
|
| 157 |
-
"-m",
|
| 158 |
-
"headroom.cli",
|
| 159 |
-
"proxy",
|
| 160 |
-
"--host",
|
| 161 |
-
"127.0.0.1",
|
| 162 |
-
"--port",
|
| 163 |
-
"8787",
|
| 164 |
-
]
|
| 165 |
-
|
| 166 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 167 |
-
monkeypatch.setattr("headroom.install.runtime.sys.platform", "linux")
|
| 168 |
-
monkeypatch.setattr("headroom.install.runtime.os.getuid", lambda: 1000, raising=False)
|
| 169 |
-
monkeypatch.setattr("headroom.install.runtime.os.getgid", lambda: 1001, raising=False)
|
| 170 |
-
docker_manifest = DeploymentManifest(
|
| 171 |
-
profile="default",
|
| 172 |
-
preset="persistent-docker",
|
| 173 |
-
runtime_kind="docker",
|
| 174 |
-
supervisor_kind="none",
|
| 175 |
-
scope="user",
|
| 176 |
-
provider_mode="manual",
|
| 177 |
-
targets=[],
|
| 178 |
-
port=8787,
|
| 179 |
-
host="127.0.0.1",
|
| 180 |
-
backend="anthropic",
|
| 181 |
-
image="ghcr.io/chopratejas/headroom:latest",
|
| 182 |
-
base_env={"HEADROOM_PORT": "8787"},
|
| 183 |
-
proxy_args=["--host", "127.0.0.1", "--port", "8787"],
|
| 184 |
-
)
|
| 185 |
-
command = build_runtime_command(docker_manifest)
|
| 186 |
-
assert "--user" in command
|
| 187 |
-
assert "1000:1001" in command
|
| 188 |
-
|
| 189 |
-
|
| 190 |
-
def test_read_pid_handles_invalid_content(monkeypatch, tmp_path: Path) -> None:
|
| 191 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 192 |
-
pid_file = tmp_path / ".headroom" / "deploy" / "default" / "runner.pid"
|
| 193 |
-
pid_file.parent.mkdir(parents=True)
|
| 194 |
-
pid_file.write_text("not-a-pid", encoding="utf-8")
|
| 195 |
-
|
| 196 |
-
assert _read_pid("default") is None
|
| 197 |
-
_clear_pid("default")
|
| 198 |
-
assert not pid_file.exists()
|
| 199 |
-
|
| 200 |
-
|
| 201 |
-
def test_write_read_and_clear_pid(monkeypatch, tmp_path: Path) -> None:
|
| 202 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 203 |
-
_write_pid("default", 456)
|
| 204 |
-
assert _read_pid("default") == 456
|
| 205 |
-
_clear_pid("default")
|
| 206 |
-
assert _read_pid("default") is None
|
| 207 |
-
|
| 208 |
-
|
| 209 |
-
def test_run_foreground_and_detached_helpers(monkeypatch, tmp_path: Path) -> None:
|
| 210 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 211 |
-
monkeypatch.setattr(
|
| 212 |
-
"headroom.install.runtime.build_runtime_command", lambda manifest: ["headroom", "proxy"]
|
| 213 |
-
)
|
| 214 |
-
monkeypatch.setattr("headroom.install.runtime._runtime_env", lambda manifest: {"ENV": "1"})
|
| 215 |
-
signal_calls: list[int] = []
|
| 216 |
-
monkeypatch.setattr(
|
| 217 |
-
"headroom.install.runtime.signal.signal", lambda sig, fn: signal_calls.append(sig)
|
| 218 |
-
)
|
| 219 |
-
|
| 220 |
-
class FakeProc:
|
| 221 |
-
def __init__(self, returncode: int = 0, pid: int = 321) -> None:
|
| 222 |
-
self.returncode = returncode
|
| 223 |
-
self.pid = pid
|
| 224 |
-
self.terminated = False
|
| 225 |
-
self.killed = False
|
| 226 |
-
|
| 227 |
-
def wait(self, timeout: int | None = None) -> int:
|
| 228 |
-
return self.returncode
|
| 229 |
-
|
| 230 |
-
def poll(self):
|
| 231 |
-
return None if not self.terminated else self.returncode
|
| 232 |
-
|
| 233 |
-
def terminate(self) -> None:
|
| 234 |
-
self.terminated = True
|
| 235 |
-
|
| 236 |
-
def kill(self) -> None:
|
| 237 |
-
self.killed = True
|
| 238 |
-
|
| 239 |
-
fake_proc = FakeProc(returncode=7)
|
| 240 |
-
popen_calls: list[tuple[list[str], dict]] = []
|
| 241 |
-
|
| 242 |
-
def fake_popen(command: list[str], **kwargs):
|
| 243 |
-
popen_calls.append((command, kwargs))
|
| 244 |
-
return fake_proc
|
| 245 |
-
|
| 246 |
-
monkeypatch.setattr("headroom.install.runtime.subprocess.Popen", fake_popen)
|
| 247 |
-
manifest = DeploymentManifest(
|
| 248 |
-
profile="default",
|
| 249 |
-
preset="persistent-service",
|
| 250 |
-
runtime_kind="python",
|
| 251 |
-
supervisor_kind="service",
|
| 252 |
-
scope="user",
|
| 253 |
-
provider_mode="manual",
|
| 254 |
-
targets=[],
|
| 255 |
-
port=8787,
|
| 256 |
-
host="127.0.0.1",
|
| 257 |
-
backend="anthropic",
|
| 258 |
-
)
|
| 259 |
-
assert run_foreground(manifest) == 7
|
| 260 |
-
assert popen_calls[0][0] == ["headroom", "proxy"]
|
| 261 |
-
assert signal.SIGINT in signal_calls
|
| 262 |
-
assert signal.SIGTERM in signal_calls
|
| 263 |
-
assert _read_pid("default") is None
|
| 264 |
-
|
| 265 |
-
monkeypatch.setattr("headroom.install.runtime.resolve_headroom_command", lambda: ["headroom"])
|
| 266 |
-
monkeypatch.setattr("headroom.install.runtime.sys.platform", "win32")
|
| 267 |
-
monkeypatch.setattr("headroom.install.runtime.subprocess.DETACHED_PROCESS", 1, raising=False)
|
| 268 |
-
monkeypatch.setattr(
|
| 269 |
-
"headroom.install.runtime.subprocess.CREATE_NEW_PROCESS_GROUP", 2, raising=False
|
| 270 |
-
)
|
| 271 |
-
fake_proc_nt = FakeProc()
|
| 272 |
-
monkeypatch.setattr(
|
| 273 |
-
"headroom.install.runtime.subprocess.Popen", lambda command, **kwargs: fake_proc_nt
|
| 274 |
-
)
|
| 275 |
-
assert start_detached_agent("demo") is fake_proc_nt
|
| 276 |
-
|
| 277 |
-
monkeypatch.setattr("headroom.install.runtime.sys.platform", "linux")
|
| 278 |
-
fake_proc_posix = FakeProc()
|
| 279 |
-
monkeypatch.setattr(
|
| 280 |
-
"headroom.install.runtime.subprocess.Popen", lambda command, **kwargs: fake_proc_posix
|
| 281 |
-
)
|
| 282 |
-
assert start_detached_agent("demo") is fake_proc_posix
|
| 283 |
-
|
| 284 |
-
|
| 285 |
-
def test_start_stop_wait_and_runtime_status_branches(monkeypatch, tmp_path: Path) -> None:
|
| 286 |
-
calls: list[list[str]] = []
|
| 287 |
-
monkeypatch.setattr(
|
| 288 |
-
"headroom.install.runtime.subprocess.run",
|
| 289 |
-
lambda command, **kwargs: calls.append(command) or type("Result", (), {"stdout": ""})(),
|
| 290 |
-
)
|
| 291 |
-
monkeypatch.setattr(
|
| 292 |
-
"headroom.install.runtime.build_runtime_command",
|
| 293 |
-
lambda manifest: [
|
| 294 |
-
"docker",
|
| 295 |
-
"run",
|
| 296 |
-
"--rm",
|
| 297 |
-
"--name",
|
| 298 |
-
"demo",
|
| 299 |
-
"-p",
|
| 300 |
-
"127.0.0.1:8787:8787",
|
| 301 |
-
"image",
|
| 302 |
-
],
|
| 303 |
-
)
|
| 304 |
-
manifest = DeploymentManifest(
|
| 305 |
-
profile="default",
|
| 306 |
-
preset=InstallPreset.PERSISTENT_DOCKER.value,
|
| 307 |
-
runtime_kind="docker",
|
| 308 |
-
supervisor_kind="none",
|
| 309 |
-
scope="user",
|
| 310 |
-
provider_mode="manual",
|
| 311 |
-
targets=[],
|
| 312 |
-
port=8787,
|
| 313 |
-
host="127.0.0.1",
|
| 314 |
-
backend="anthropic",
|
| 315 |
-
container_name="headroom-default",
|
| 316 |
-
)
|
| 317 |
-
start_persistent_docker(manifest)
|
| 318 |
-
assert calls == [
|
| 319 |
-
["docker", "rm", "-f", "headroom-default"],
|
| 320 |
-
[
|
| 321 |
-
"docker",
|
| 322 |
-
"run",
|
| 323 |
-
"-d",
|
| 324 |
-
"--restart",
|
| 325 |
-
"unless-stopped",
|
| 326 |
-
"--name",
|
| 327 |
-
"headroom-default",
|
| 328 |
-
"-p",
|
| 329 |
-
"127.0.0.1:8787:8787",
|
| 330 |
-
"image",
|
| 331 |
-
],
|
| 332 |
-
]
|
| 333 |
-
|
| 334 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 335 |
-
python_manifest = DeploymentManifest(
|
| 336 |
-
profile="default",
|
| 337 |
-
preset="persistent-service",
|
| 338 |
-
runtime_kind="python",
|
| 339 |
-
supervisor_kind="service",
|
| 340 |
-
scope="user",
|
| 341 |
-
provider_mode="manual",
|
| 342 |
-
targets=[],
|
| 343 |
-
port=8787,
|
| 344 |
-
host="127.0.0.1",
|
| 345 |
-
backend="anthropic",
|
| 346 |
-
health_url="http://127.0.0.1:8787/health",
|
| 347 |
-
)
|
| 348 |
-
_write_pid("default", 123)
|
| 349 |
-
killed: list[tuple[int, int]] = []
|
| 350 |
-
monkeypatch.setattr(
|
| 351 |
-
"headroom.install.runtime.os.kill", lambda pid, sig: killed.append((pid, sig))
|
| 352 |
-
)
|
| 353 |
-
stop_runtime(python_manifest)
|
| 354 |
-
assert killed == [(123, signal.SIGTERM)]
|
| 355 |
-
assert _read_pid("default") is None
|
| 356 |
-
|
| 357 |
-
_write_pid("default", 124)
|
| 358 |
-
monkeypatch.setattr(
|
| 359 |
-
"headroom.install.runtime.os.kill",
|
| 360 |
-
lambda pid, sig: (_ for _ in ()).throw(OSError("gone")),
|
| 361 |
-
)
|
| 362 |
-
stop_runtime(python_manifest)
|
| 363 |
-
assert _read_pid("default") is None
|
| 364 |
-
|
| 365 |
-
probe_results = iter([False, False, True])
|
| 366 |
-
sleeps: list[int] = []
|
| 367 |
-
monkeypatch.setattr("headroom.install.runtime.probe_ready", lambda url: next(probe_results))
|
| 368 |
-
monkeypatch.setattr(
|
| 369 |
-
"headroom.install.runtime.time.sleep", lambda seconds: sleeps.append(seconds)
|
| 370 |
-
)
|
| 371 |
-
assert wait_ready(python_manifest, timeout_seconds=3) is True
|
| 372 |
-
assert sleeps == [1, 1]
|
| 373 |
-
|
| 374 |
-
monkeypatch.setattr("headroom.install.runtime.probe_ready", lambda url: False)
|
| 375 |
-
sleeps.clear()
|
| 376 |
-
assert wait_ready(python_manifest, timeout_seconds=2) is False
|
| 377 |
-
assert sleeps == [1, 1]
|
| 378 |
-
|
| 379 |
-
class Result:
|
| 380 |
-
def __init__(self, stdout: str = "") -> None:
|
| 381 |
-
self.stdout = stdout
|
| 382 |
-
|
| 383 |
-
monkeypatch.setattr(
|
| 384 |
-
"headroom.install.runtime.subprocess.run",
|
| 385 |
-
lambda command, **kwargs: Result(stdout=""),
|
| 386 |
-
)
|
| 387 |
-
assert runtime_status(manifest) == "stopped"
|
| 388 |
-
assert runtime_status(python_manifest) == "stopped"
|
| 389 |
-
|
| 390 |
-
_write_pid("default", 125)
|
| 391 |
-
monkeypatch.setattr(
|
| 392 |
-
"headroom.install.runtime.os.kill", lambda pid, sig: (_ for _ in ()).throw(OSError())
|
| 393 |
-
)
|
| 394 |
-
assert runtime_status(python_manifest) == "stopped"
|
| 395 |
-
|
| 396 |
-
|
| 397 |
-
def test_stop_runtime_for_docker_stops_and_removes_container(monkeypatch) -> None:
|
| 398 |
-
calls: list[list[str]] = []
|
| 399 |
-
manifest = DeploymentManifest(
|
| 400 |
-
profile="default",
|
| 401 |
-
preset="persistent-docker",
|
| 402 |
-
runtime_kind="docker",
|
| 403 |
-
supervisor_kind="none",
|
| 404 |
-
scope="user",
|
| 405 |
-
provider_mode="manual",
|
| 406 |
-
targets=[],
|
| 407 |
-
port=8787,
|
| 408 |
-
host="127.0.0.1",
|
| 409 |
-
backend="anthropic",
|
| 410 |
-
container_name="headroom-default",
|
| 411 |
-
)
|
| 412 |
-
|
| 413 |
-
monkeypatch.setattr(
|
| 414 |
-
"headroom.install.runtime.subprocess.run",
|
| 415 |
-
lambda command, **kwargs: calls.append(command),
|
| 416 |
-
)
|
| 417 |
-
|
| 418 |
-
stop_runtime(manifest)
|
| 419 |
-
|
| 420 |
-
assert calls == [
|
| 421 |
-
["docker", "stop", "headroom-default"],
|
| 422 |
-
["docker", "rm", "-f", "headroom-default"],
|
| 423 |
-
]
|
| 424 |
-
|
| 425 |
-
|
| 426 |
-
def test_runtime_status_reads_container_and_pid_state(monkeypatch, tmp_path: Path) -> None:
|
| 427 |
-
docker_manifest = DeploymentManifest(
|
| 428 |
-
profile="default",
|
| 429 |
-
preset="persistent-docker",
|
| 430 |
-
runtime_kind="docker",
|
| 431 |
-
supervisor_kind="none",
|
| 432 |
-
scope="user",
|
| 433 |
-
provider_mode="manual",
|
| 434 |
-
targets=[],
|
| 435 |
-
port=8787,
|
| 436 |
-
host="127.0.0.1",
|
| 437 |
-
backend="anthropic",
|
| 438 |
-
container_name="headroom-default",
|
| 439 |
-
)
|
| 440 |
-
|
| 441 |
-
class Result:
|
| 442 |
-
def __init__(self, stdout: str = "") -> None:
|
| 443 |
-
self.stdout = stdout
|
| 444 |
-
|
| 445 |
-
monkeypatch.setattr(
|
| 446 |
-
"headroom.install.runtime.subprocess.run",
|
| 447 |
-
lambda command, **kwargs: Result(stdout="headroom-default\n"),
|
| 448 |
-
)
|
| 449 |
-
assert runtime_status(docker_manifest) == "running"
|
| 450 |
-
|
| 451 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 452 |
-
pid_file = tmp_path / ".headroom" / "deploy" / "default" / "runner.pid"
|
| 453 |
-
pid_file.parent.mkdir(parents=True)
|
| 454 |
-
pid_file.write_text("123", encoding="utf-8")
|
| 455 |
-
monkeypatch.setattr("headroom.install.runtime.os.kill", lambda pid, sig: None)
|
| 456 |
-
python_manifest = DeploymentManifest(
|
| 457 |
-
profile="default",
|
| 458 |
-
preset="persistent-service",
|
| 459 |
-
runtime_kind="python",
|
| 460 |
-
supervisor_kind="service",
|
| 461 |
-
scope="user",
|
| 462 |
-
provider_mode="manual",
|
| 463 |
-
targets=[],
|
| 464 |
-
port=8787,
|
| 465 |
-
host="127.0.0.1",
|
| 466 |
-
backend="anthropic",
|
| 467 |
-
)
|
| 468 |
-
assert runtime_status(python_manifest) == "running"
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import signal
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
from headroom.install.models import DeploymentManifest, InstallPreset
|
| 7 |
+
from headroom.install.runtime import (
|
| 8 |
+
_clear_pid,
|
| 9 |
+
_deployment_env,
|
| 10 |
+
_mount_source,
|
| 11 |
+
_read_pid,
|
| 12 |
+
_runtime_env,
|
| 13 |
+
_write_pid,
|
| 14 |
+
build_runtime_command,
|
| 15 |
+
resolve_headroom_command,
|
| 16 |
+
run_foreground,
|
| 17 |
+
runtime_status,
|
| 18 |
+
start_detached_agent,
|
| 19 |
+
start_persistent_docker,
|
| 20 |
+
stop_runtime,
|
| 21 |
+
wait_ready,
|
| 22 |
+
)
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def test_build_runtime_command_for_docker_includes_deployment_env(
|
| 26 |
+
monkeypatch, tmp_path: Path
|
| 27 |
+
) -> None:
|
| 28 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 29 |
+
manifest = DeploymentManifest(
|
| 30 |
+
profile="default",
|
| 31 |
+
preset="persistent-docker",
|
| 32 |
+
runtime_kind="docker",
|
| 33 |
+
supervisor_kind="none",
|
| 34 |
+
scope="user",
|
| 35 |
+
provider_mode="manual",
|
| 36 |
+
targets=["claude"],
|
| 37 |
+
port=8787,
|
| 38 |
+
host="127.0.0.1",
|
| 39 |
+
backend="anthropic",
|
| 40 |
+
image="ghcr.io/chopratejas/headroom:latest",
|
| 41 |
+
base_env={"HEADROOM_PORT": "8787"},
|
| 42 |
+
proxy_args=["--host", "127.0.0.1", "--port", "8787"],
|
| 43 |
+
)
|
| 44 |
+
|
| 45 |
+
command = build_runtime_command(manifest)
|
| 46 |
+
|
| 47 |
+
joined = " ".join(command)
|
| 48 |
+
assert command[:3] == ["docker", "run", "--rm"]
|
| 49 |
+
assert "HEADROOM_DEPLOYMENT_PROFILE=default" in joined
|
| 50 |
+
assert "HEADROOM_DEPLOYMENT_PRESET=persistent-docker" in joined
|
| 51 |
+
assert "127.0.0.1:8787:8787" in joined
|
| 52 |
+
assert "ghcr.io/chopratejas/headroom:latest" in command
|
| 53 |
+
# Canonical Headroom filesystem contract (issue #175) forwarded into
|
| 54 |
+
# the container.
|
| 55 |
+
assert "HEADROOM_WORKSPACE_DIR=/tmp/headroom-home/.headroom" in command
|
| 56 |
+
assert "HEADROOM_CONFIG_DIR=/tmp/headroom-home/.headroom/config" in command
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def test_build_runtime_command_for_docker_matches_wrapper_parity(
|
| 60 |
+
monkeypatch, tmp_path: Path
|
| 61 |
+
) -> None:
|
| 62 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 63 |
+
monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key")
|
| 64 |
+
monkeypatch.setenv("OPENAI_API_KEY", "test-openai")
|
| 65 |
+
manifest = DeploymentManifest(
|
| 66 |
+
profile="default",
|
| 67 |
+
preset="persistent-docker",
|
| 68 |
+
runtime_kind="docker",
|
| 69 |
+
supervisor_kind="none",
|
| 70 |
+
scope="user",
|
| 71 |
+
provider_mode="manual",
|
| 72 |
+
targets=["claude"],
|
| 73 |
+
port=8787,
|
| 74 |
+
host="127.0.0.1",
|
| 75 |
+
backend="anthropic",
|
| 76 |
+
image="ghcr.io/chopratejas/headroom:latest",
|
| 77 |
+
base_env={"HEADROOM_PORT": "8787"},
|
| 78 |
+
proxy_args=["--host", "127.0.0.1", "--port", "8787"],
|
| 79 |
+
)
|
| 80 |
+
|
| 81 |
+
command = build_runtime_command(manifest)
|
| 82 |
+
|
| 83 |
+
assert (tmp_path / ".headroom").is_dir()
|
| 84 |
+
assert (tmp_path / ".claude").is_dir()
|
| 85 |
+
assert (tmp_path / ".codex").is_dir()
|
| 86 |
+
assert (tmp_path / ".gemini").is_dir()
|
| 87 |
+
assert "--env" in command
|
| 88 |
+
joined = " ".join(command)
|
| 89 |
+
assert "ANTHROPIC_API_KEY" in joined
|
| 90 |
+
assert "OPENAI_API_KEY" in joined
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def test_resolve_headroom_command_prefers_headroom_binary(monkeypatch) -> None:
|
| 94 |
+
monkeypatch.setattr(
|
| 95 |
+
"shutil.which", lambda name: "/usr/bin/headroom" if name == "headroom" else None
|
| 96 |
+
)
|
| 97 |
+
|
| 98 |
+
assert resolve_headroom_command() == ["/usr/bin/headroom"]
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def test_resolve_headroom_command_falls_back_to_python_module(monkeypatch) -> None:
|
| 102 |
+
monkeypatch.setattr("shutil.which", lambda name: None)
|
| 103 |
+
monkeypatch.setattr("headroom.install.runtime.sys.executable", "/usr/bin/python")
|
| 104 |
+
assert resolve_headroom_command() == ["/usr/bin/python", "-m", "headroom.cli"]
|
| 105 |
+
|
| 106 |
+
|
| 107 |
+
def test_runtime_env_and_mount_source(monkeypatch) -> None:
|
| 108 |
+
manifest = DeploymentManifest(
|
| 109 |
+
profile="default",
|
| 110 |
+
preset="persistent-service",
|
| 111 |
+
runtime_kind="python",
|
| 112 |
+
supervisor_kind="service",
|
| 113 |
+
scope="user",
|
| 114 |
+
provider_mode="manual",
|
| 115 |
+
targets=[],
|
| 116 |
+
port=8787,
|
| 117 |
+
host="127.0.0.1",
|
| 118 |
+
backend="anthropic",
|
| 119 |
+
base_env={"EXTRA": "1"},
|
| 120 |
+
)
|
| 121 |
+
monkeypatch.setattr("headroom.install.runtime.os.environ", {"BASE": "x"})
|
| 122 |
+
|
| 123 |
+
assert _deployment_env(manifest) == {
|
| 124 |
+
"HEADROOM_DEPLOYMENT_PROFILE": "default",
|
| 125 |
+
"HEADROOM_DEPLOYMENT_PRESET": "persistent-service",
|
| 126 |
+
"HEADROOM_DEPLOYMENT_RUNTIME": "python",
|
| 127 |
+
"HEADROOM_DEPLOYMENT_SUPERVISOR": "service",
|
| 128 |
+
"HEADROOM_DEPLOYMENT_SCOPE": "user",
|
| 129 |
+
}
|
| 130 |
+
assert _runtime_env(manifest)["BASE"] == "x"
|
| 131 |
+
assert _runtime_env(manifest)["EXTRA"] == "1"
|
| 132 |
+
assert _runtime_env(manifest)["HEADROOM_DEPLOYMENT_PROFILE"] == "default"
|
| 133 |
+
|
| 134 |
+
monkeypatch.setattr("headroom.install.runtime.sys.platform", "win32")
|
| 135 |
+
assert _mount_source("C:\\Users\\me", ".headroom") == "C:\\Users\\me\\.headroom"
|
| 136 |
+
monkeypatch.setattr("headroom.install.runtime.sys.platform", "linux")
|
| 137 |
+
assert _mount_source("/home/me", ".headroom") == "/home/me/.headroom"
|
| 138 |
+
|
| 139 |
+
|
| 140 |
+
def test_build_runtime_command_python_and_docker_user(monkeypatch, tmp_path: Path) -> None:
|
| 141 |
+
monkeypatch.setattr("headroom.install.runtime.sys.executable", "/usr/bin/python")
|
| 142 |
+
manifest = DeploymentManifest(
|
| 143 |
+
profile="default",
|
| 144 |
+
preset="persistent-service",
|
| 145 |
+
runtime_kind="python",
|
| 146 |
+
supervisor_kind="service",
|
| 147 |
+
scope="user",
|
| 148 |
+
provider_mode="manual",
|
| 149 |
+
targets=[],
|
| 150 |
+
port=8787,
|
| 151 |
+
host="127.0.0.1",
|
| 152 |
+
backend="anthropic",
|
| 153 |
+
proxy_args=["--host", "127.0.0.1", "--port", "8787"],
|
| 154 |
+
)
|
| 155 |
+
assert build_runtime_command(manifest) == [
|
| 156 |
+
"/usr/bin/python",
|
| 157 |
+
"-m",
|
| 158 |
+
"headroom.cli",
|
| 159 |
+
"proxy",
|
| 160 |
+
"--host",
|
| 161 |
+
"127.0.0.1",
|
| 162 |
+
"--port",
|
| 163 |
+
"8787",
|
| 164 |
+
]
|
| 165 |
+
|
| 166 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 167 |
+
monkeypatch.setattr("headroom.install.runtime.sys.platform", "linux")
|
| 168 |
+
monkeypatch.setattr("headroom.install.runtime.os.getuid", lambda: 1000, raising=False)
|
| 169 |
+
monkeypatch.setattr("headroom.install.runtime.os.getgid", lambda: 1001, raising=False)
|
| 170 |
+
docker_manifest = DeploymentManifest(
|
| 171 |
+
profile="default",
|
| 172 |
+
preset="persistent-docker",
|
| 173 |
+
runtime_kind="docker",
|
| 174 |
+
supervisor_kind="none",
|
| 175 |
+
scope="user",
|
| 176 |
+
provider_mode="manual",
|
| 177 |
+
targets=[],
|
| 178 |
+
port=8787,
|
| 179 |
+
host="127.0.0.1",
|
| 180 |
+
backend="anthropic",
|
| 181 |
+
image="ghcr.io/chopratejas/headroom:latest",
|
| 182 |
+
base_env={"HEADROOM_PORT": "8787"},
|
| 183 |
+
proxy_args=["--host", "127.0.0.1", "--port", "8787"],
|
| 184 |
+
)
|
| 185 |
+
command = build_runtime_command(docker_manifest)
|
| 186 |
+
assert "--user" in command
|
| 187 |
+
assert "1000:1001" in command
|
| 188 |
+
|
| 189 |
+
|
| 190 |
+
def test_read_pid_handles_invalid_content(monkeypatch, tmp_path: Path) -> None:
|
| 191 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 192 |
+
pid_file = tmp_path / ".headroom" / "deploy" / "default" / "runner.pid"
|
| 193 |
+
pid_file.parent.mkdir(parents=True)
|
| 194 |
+
pid_file.write_text("not-a-pid", encoding="utf-8")
|
| 195 |
+
|
| 196 |
+
assert _read_pid("default") is None
|
| 197 |
+
_clear_pid("default")
|
| 198 |
+
assert not pid_file.exists()
|
| 199 |
+
|
| 200 |
+
|
| 201 |
+
def test_write_read_and_clear_pid(monkeypatch, tmp_path: Path) -> None:
|
| 202 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 203 |
+
_write_pid("default", 456)
|
| 204 |
+
assert _read_pid("default") == 456
|
| 205 |
+
_clear_pid("default")
|
| 206 |
+
assert _read_pid("default") is None
|
| 207 |
+
|
| 208 |
+
|
| 209 |
+
def test_run_foreground_and_detached_helpers(monkeypatch, tmp_path: Path) -> None:
|
| 210 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 211 |
+
monkeypatch.setattr(
|
| 212 |
+
"headroom.install.runtime.build_runtime_command", lambda manifest: ["headroom", "proxy"]
|
| 213 |
+
)
|
| 214 |
+
monkeypatch.setattr("headroom.install.runtime._runtime_env", lambda manifest: {"ENV": "1"})
|
| 215 |
+
signal_calls: list[int] = []
|
| 216 |
+
monkeypatch.setattr(
|
| 217 |
+
"headroom.install.runtime.signal.signal", lambda sig, fn: signal_calls.append(sig)
|
| 218 |
+
)
|
| 219 |
+
|
| 220 |
+
class FakeProc:
|
| 221 |
+
def __init__(self, returncode: int = 0, pid: int = 321) -> None:
|
| 222 |
+
self.returncode = returncode
|
| 223 |
+
self.pid = pid
|
| 224 |
+
self.terminated = False
|
| 225 |
+
self.killed = False
|
| 226 |
+
|
| 227 |
+
def wait(self, timeout: int | None = None) -> int:
|
| 228 |
+
return self.returncode
|
| 229 |
+
|
| 230 |
+
def poll(self):
|
| 231 |
+
return None if not self.terminated else self.returncode
|
| 232 |
+
|
| 233 |
+
def terminate(self) -> None:
|
| 234 |
+
self.terminated = True
|
| 235 |
+
|
| 236 |
+
def kill(self) -> None:
|
| 237 |
+
self.killed = True
|
| 238 |
+
|
| 239 |
+
fake_proc = FakeProc(returncode=7)
|
| 240 |
+
popen_calls: list[tuple[list[str], dict]] = []
|
| 241 |
+
|
| 242 |
+
def fake_popen(command: list[str], **kwargs):
|
| 243 |
+
popen_calls.append((command, kwargs))
|
| 244 |
+
return fake_proc
|
| 245 |
+
|
| 246 |
+
monkeypatch.setattr("headroom.install.runtime.subprocess.Popen", fake_popen)
|
| 247 |
+
manifest = DeploymentManifest(
|
| 248 |
+
profile="default",
|
| 249 |
+
preset="persistent-service",
|
| 250 |
+
runtime_kind="python",
|
| 251 |
+
supervisor_kind="service",
|
| 252 |
+
scope="user",
|
| 253 |
+
provider_mode="manual",
|
| 254 |
+
targets=[],
|
| 255 |
+
port=8787,
|
| 256 |
+
host="127.0.0.1",
|
| 257 |
+
backend="anthropic",
|
| 258 |
+
)
|
| 259 |
+
assert run_foreground(manifest) == 7
|
| 260 |
+
assert popen_calls[0][0] == ["headroom", "proxy"]
|
| 261 |
+
assert signal.SIGINT in signal_calls
|
| 262 |
+
assert signal.SIGTERM in signal_calls
|
| 263 |
+
assert _read_pid("default") is None
|
| 264 |
+
|
| 265 |
+
monkeypatch.setattr("headroom.install.runtime.resolve_headroom_command", lambda: ["headroom"])
|
| 266 |
+
monkeypatch.setattr("headroom.install.runtime.sys.platform", "win32")
|
| 267 |
+
monkeypatch.setattr("headroom.install.runtime.subprocess.DETACHED_PROCESS", 1, raising=False)
|
| 268 |
+
monkeypatch.setattr(
|
| 269 |
+
"headroom.install.runtime.subprocess.CREATE_NEW_PROCESS_GROUP", 2, raising=False
|
| 270 |
+
)
|
| 271 |
+
fake_proc_nt = FakeProc()
|
| 272 |
+
monkeypatch.setattr(
|
| 273 |
+
"headroom.install.runtime.subprocess.Popen", lambda command, **kwargs: fake_proc_nt
|
| 274 |
+
)
|
| 275 |
+
assert start_detached_agent("demo") is fake_proc_nt
|
| 276 |
+
|
| 277 |
+
monkeypatch.setattr("headroom.install.runtime.sys.platform", "linux")
|
| 278 |
+
fake_proc_posix = FakeProc()
|
| 279 |
+
monkeypatch.setattr(
|
| 280 |
+
"headroom.install.runtime.subprocess.Popen", lambda command, **kwargs: fake_proc_posix
|
| 281 |
+
)
|
| 282 |
+
assert start_detached_agent("demo") is fake_proc_posix
|
| 283 |
+
|
| 284 |
+
|
| 285 |
+
def test_start_stop_wait_and_runtime_status_branches(monkeypatch, tmp_path: Path) -> None:
|
| 286 |
+
calls: list[list[str]] = []
|
| 287 |
+
monkeypatch.setattr(
|
| 288 |
+
"headroom.install.runtime.subprocess.run",
|
| 289 |
+
lambda command, **kwargs: calls.append(command) or type("Result", (), {"stdout": ""})(),
|
| 290 |
+
)
|
| 291 |
+
monkeypatch.setattr(
|
| 292 |
+
"headroom.install.runtime.build_runtime_command",
|
| 293 |
+
lambda manifest: [
|
| 294 |
+
"docker",
|
| 295 |
+
"run",
|
| 296 |
+
"--rm",
|
| 297 |
+
"--name",
|
| 298 |
+
"demo",
|
| 299 |
+
"-p",
|
| 300 |
+
"127.0.0.1:8787:8787",
|
| 301 |
+
"image",
|
| 302 |
+
],
|
| 303 |
+
)
|
| 304 |
+
manifest = DeploymentManifest(
|
| 305 |
+
profile="default",
|
| 306 |
+
preset=InstallPreset.PERSISTENT_DOCKER.value,
|
| 307 |
+
runtime_kind="docker",
|
| 308 |
+
supervisor_kind="none",
|
| 309 |
+
scope="user",
|
| 310 |
+
provider_mode="manual",
|
| 311 |
+
targets=[],
|
| 312 |
+
port=8787,
|
| 313 |
+
host="127.0.0.1",
|
| 314 |
+
backend="anthropic",
|
| 315 |
+
container_name="headroom-default",
|
| 316 |
+
)
|
| 317 |
+
start_persistent_docker(manifest)
|
| 318 |
+
assert calls == [
|
| 319 |
+
["docker", "rm", "-f", "headroom-default"],
|
| 320 |
+
[
|
| 321 |
+
"docker",
|
| 322 |
+
"run",
|
| 323 |
+
"-d",
|
| 324 |
+
"--restart",
|
| 325 |
+
"unless-stopped",
|
| 326 |
+
"--name",
|
| 327 |
+
"headroom-default",
|
| 328 |
+
"-p",
|
| 329 |
+
"127.0.0.1:8787:8787",
|
| 330 |
+
"image",
|
| 331 |
+
],
|
| 332 |
+
]
|
| 333 |
+
|
| 334 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 335 |
+
python_manifest = DeploymentManifest(
|
| 336 |
+
profile="default",
|
| 337 |
+
preset="persistent-service",
|
| 338 |
+
runtime_kind="python",
|
| 339 |
+
supervisor_kind="service",
|
| 340 |
+
scope="user",
|
| 341 |
+
provider_mode="manual",
|
| 342 |
+
targets=[],
|
| 343 |
+
port=8787,
|
| 344 |
+
host="127.0.0.1",
|
| 345 |
+
backend="anthropic",
|
| 346 |
+
health_url="http://127.0.0.1:8787/health",
|
| 347 |
+
)
|
| 348 |
+
_write_pid("default", 123)
|
| 349 |
+
killed: list[tuple[int, int]] = []
|
| 350 |
+
monkeypatch.setattr(
|
| 351 |
+
"headroom.install.runtime.os.kill", lambda pid, sig: killed.append((pid, sig))
|
| 352 |
+
)
|
| 353 |
+
stop_runtime(python_manifest)
|
| 354 |
+
assert killed == [(123, signal.SIGTERM)]
|
| 355 |
+
assert _read_pid("default") is None
|
| 356 |
+
|
| 357 |
+
_write_pid("default", 124)
|
| 358 |
+
monkeypatch.setattr(
|
| 359 |
+
"headroom.install.runtime.os.kill",
|
| 360 |
+
lambda pid, sig: (_ for _ in ()).throw(OSError("gone")),
|
| 361 |
+
)
|
| 362 |
+
stop_runtime(python_manifest)
|
| 363 |
+
assert _read_pid("default") is None
|
| 364 |
+
|
| 365 |
+
probe_results = iter([False, False, True])
|
| 366 |
+
sleeps: list[int] = []
|
| 367 |
+
monkeypatch.setattr("headroom.install.runtime.probe_ready", lambda url: next(probe_results))
|
| 368 |
+
monkeypatch.setattr(
|
| 369 |
+
"headroom.install.runtime.time.sleep", lambda seconds: sleeps.append(seconds)
|
| 370 |
+
)
|
| 371 |
+
assert wait_ready(python_manifest, timeout_seconds=3) is True
|
| 372 |
+
assert sleeps == [1, 1]
|
| 373 |
+
|
| 374 |
+
monkeypatch.setattr("headroom.install.runtime.probe_ready", lambda url: False)
|
| 375 |
+
sleeps.clear()
|
| 376 |
+
assert wait_ready(python_manifest, timeout_seconds=2) is False
|
| 377 |
+
assert sleeps == [1, 1]
|
| 378 |
+
|
| 379 |
+
class Result:
|
| 380 |
+
def __init__(self, stdout: str = "") -> None:
|
| 381 |
+
self.stdout = stdout
|
| 382 |
+
|
| 383 |
+
monkeypatch.setattr(
|
| 384 |
+
"headroom.install.runtime.subprocess.run",
|
| 385 |
+
lambda command, **kwargs: Result(stdout=""),
|
| 386 |
+
)
|
| 387 |
+
assert runtime_status(manifest) == "stopped"
|
| 388 |
+
assert runtime_status(python_manifest) == "stopped"
|
| 389 |
+
|
| 390 |
+
_write_pid("default", 125)
|
| 391 |
+
monkeypatch.setattr(
|
| 392 |
+
"headroom.install.runtime.os.kill", lambda pid, sig: (_ for _ in ()).throw(OSError())
|
| 393 |
+
)
|
| 394 |
+
assert runtime_status(python_manifest) == "stopped"
|
| 395 |
+
|
| 396 |
+
|
| 397 |
+
def test_stop_runtime_for_docker_stops_and_removes_container(monkeypatch) -> None:
|
| 398 |
+
calls: list[list[str]] = []
|
| 399 |
+
manifest = DeploymentManifest(
|
| 400 |
+
profile="default",
|
| 401 |
+
preset="persistent-docker",
|
| 402 |
+
runtime_kind="docker",
|
| 403 |
+
supervisor_kind="none",
|
| 404 |
+
scope="user",
|
| 405 |
+
provider_mode="manual",
|
| 406 |
+
targets=[],
|
| 407 |
+
port=8787,
|
| 408 |
+
host="127.0.0.1",
|
| 409 |
+
backend="anthropic",
|
| 410 |
+
container_name="headroom-default",
|
| 411 |
+
)
|
| 412 |
+
|
| 413 |
+
monkeypatch.setattr(
|
| 414 |
+
"headroom.install.runtime.subprocess.run",
|
| 415 |
+
lambda command, **kwargs: calls.append(command),
|
| 416 |
+
)
|
| 417 |
+
|
| 418 |
+
stop_runtime(manifest)
|
| 419 |
+
|
| 420 |
+
assert calls == [
|
| 421 |
+
["docker", "stop", "headroom-default"],
|
| 422 |
+
["docker", "rm", "-f", "headroom-default"],
|
| 423 |
+
]
|
| 424 |
+
|
| 425 |
+
|
| 426 |
+
def test_runtime_status_reads_container_and_pid_state(monkeypatch, tmp_path: Path) -> None:
|
| 427 |
+
docker_manifest = DeploymentManifest(
|
| 428 |
+
profile="default",
|
| 429 |
+
preset="persistent-docker",
|
| 430 |
+
runtime_kind="docker",
|
| 431 |
+
supervisor_kind="none",
|
| 432 |
+
scope="user",
|
| 433 |
+
provider_mode="manual",
|
| 434 |
+
targets=[],
|
| 435 |
+
port=8787,
|
| 436 |
+
host="127.0.0.1",
|
| 437 |
+
backend="anthropic",
|
| 438 |
+
container_name="headroom-default",
|
| 439 |
+
)
|
| 440 |
+
|
| 441 |
+
class Result:
|
| 442 |
+
def __init__(self, stdout: str = "") -> None:
|
| 443 |
+
self.stdout = stdout
|
| 444 |
+
|
| 445 |
+
monkeypatch.setattr(
|
| 446 |
+
"headroom.install.runtime.subprocess.run",
|
| 447 |
+
lambda command, **kwargs: Result(stdout="headroom-default\n"),
|
| 448 |
+
)
|
| 449 |
+
assert runtime_status(docker_manifest) == "running"
|
| 450 |
+
|
| 451 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 452 |
+
pid_file = tmp_path / ".headroom" / "deploy" / "default" / "runner.pid"
|
| 453 |
+
pid_file.parent.mkdir(parents=True)
|
| 454 |
+
pid_file.write_text("123", encoding="utf-8")
|
| 455 |
+
monkeypatch.setattr("headroom.install.runtime.os.kill", lambda pid, sig: None)
|
| 456 |
+
python_manifest = DeploymentManifest(
|
| 457 |
+
profile="default",
|
| 458 |
+
preset="persistent-service",
|
| 459 |
+
runtime_kind="python",
|
| 460 |
+
supervisor_kind="service",
|
| 461 |
+
scope="user",
|
| 462 |
+
provider_mode="manual",
|
| 463 |
+
targets=[],
|
| 464 |
+
port=8787,
|
| 465 |
+
host="127.0.0.1",
|
| 466 |
+
backend="anthropic",
|
| 467 |
+
)
|
| 468 |
+
assert runtime_status(python_manifest) == "running"
|
|
@@ -1,471 +1,471 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
from pathlib import Path
|
| 4 |
-
|
| 5 |
-
import click
|
| 6 |
-
import pytest
|
| 7 |
-
|
| 8 |
-
from headroom.install.models import DeploymentManifest, SupervisorKind
|
| 9 |
-
from headroom.install.supervisors import (
|
| 10 |
-
_command_for_script,
|
| 11 |
-
_linux_service_unit,
|
| 12 |
-
_linux_task_spec,
|
| 13 |
-
_macos_launchd_plist,
|
| 14 |
-
_render_unix_runner,
|
| 15 |
-
_render_windows_runner,
|
| 16 |
-
install_supervisor,
|
| 17 |
-
remove_supervisor,
|
| 18 |
-
render_runner_scripts,
|
| 19 |
-
start_supervisor,
|
| 20 |
-
stop_supervisor,
|
| 21 |
-
)
|
| 22 |
-
|
| 23 |
-
|
| 24 |
-
def _manifest(
|
| 25 |
-
*, profile: str = "default", scope: str = "user", supervisor: str = "service"
|
| 26 |
-
) -> DeploymentManifest:
|
| 27 |
-
return DeploymentManifest(
|
| 28 |
-
profile=profile,
|
| 29 |
-
preset="persistent-service",
|
| 30 |
-
runtime_kind="python",
|
| 31 |
-
supervisor_kind=supervisor,
|
| 32 |
-
scope=scope,
|
| 33 |
-
provider_mode="manual",
|
| 34 |
-
targets=[],
|
| 35 |
-
port=8787,
|
| 36 |
-
host="127.0.0.1",
|
| 37 |
-
backend="anthropic",
|
| 38 |
-
service_name=f"headroom-{profile}",
|
| 39 |
-
)
|
| 40 |
-
|
| 41 |
-
|
| 42 |
-
def test_linux_service_unit_uses_user_systemd_path(monkeypatch, tmp_path: Path) -> None:
|
| 43 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 44 |
-
manifest = _manifest()
|
| 45 |
-
|
| 46 |
-
unit_path, content = _linux_service_unit(manifest, tmp_path / "run-headroom.sh")
|
| 47 |
-
|
| 48 |
-
assert unit_path == tmp_path / ".config" / "systemd" / "user" / "headroom-default.service"
|
| 49 |
-
assert "ExecStart=" + str(tmp_path / "run-headroom.sh") in content
|
| 50 |
-
assert "Restart=on-failure" in content
|
| 51 |
-
|
| 52 |
-
|
| 53 |
-
def test_command_for_script_and_unix_runner(monkeypatch, tmp_path: Path) -> None:
|
| 54 |
-
monkeypatch.setattr(
|
| 55 |
-
"headroom.install.supervisors.resolve_headroom_command",
|
| 56 |
-
lambda: ["python", "-m", "headroom"],
|
| 57 |
-
)
|
| 58 |
-
|
| 59 |
-
assert _command_for_script("install", "agent", "run") == [
|
| 60 |
-
"python",
|
| 61 |
-
"-m",
|
| 62 |
-
"headroom",
|
| 63 |
-
"install",
|
| 64 |
-
"agent",
|
| 65 |
-
"run",
|
| 66 |
-
]
|
| 67 |
-
|
| 68 |
-
record = _render_unix_runner(
|
| 69 |
-
tmp_path / "scripts" / "run-headroom.sh", ["headroom", "run", "--flag"]
|
| 70 |
-
)
|
| 71 |
-
assert record.kind == "script"
|
| 72 |
-
content = Path(record.path).read_text(encoding="utf-8")
|
| 73 |
-
assert content.startswith("#!/usr/bin/env bash")
|
| 74 |
-
assert "exec headroom run --flag" in content
|
| 75 |
-
|
| 76 |
-
|
| 77 |
-
def test_linux_task_spec_for_user_scope_includes_crontab_markers(tmp_path: Path) -> None:
|
| 78 |
-
manifest = _manifest(profile="smoke", supervisor=SupervisorKind.TASK.value)
|
| 79 |
-
|
| 80 |
-
cron_path, content = _linux_task_spec(manifest, tmp_path / "ensure-headroom.sh")
|
| 81 |
-
|
| 82 |
-
assert cron_path is None
|
| 83 |
-
assert "# >>> headroom smoke >>>" in content
|
| 84 |
-
assert "# <<< headroom smoke <<<" in content
|
| 85 |
-
assert "@reboot" in content
|
| 86 |
-
assert "*/5 * * * *" in content
|
| 87 |
-
|
| 88 |
-
|
| 89 |
-
def test_macos_launchd_plist_switches_between_keepalive_and_interval(
|
| 90 |
-
monkeypatch, tmp_path: Path
|
| 91 |
-
) -> None:
|
| 92 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 93 |
-
|
| 94 |
-
service_manifest = _manifest(supervisor=SupervisorKind.SERVICE.value)
|
| 95 |
-
service_path, service_content = _macos_launchd_plist(
|
| 96 |
-
service_manifest, tmp_path / "run-headroom.sh"
|
| 97 |
-
)
|
| 98 |
-
assert service_path == tmp_path / "Library" / "LaunchAgents" / "com.headroom.default.plist"
|
| 99 |
-
assert "<key>KeepAlive</key>" in service_content
|
| 100 |
-
assert "<key>StartInterval</key>" not in service_content
|
| 101 |
-
|
| 102 |
-
task_manifest = _manifest(profile="tasky", supervisor=SupervisorKind.TASK.value)
|
| 103 |
-
task_path, task_content = _macos_launchd_plist(
|
| 104 |
-
task_manifest, tmp_path / "ensure-headroom.sh", interval=300
|
| 105 |
-
)
|
| 106 |
-
assert task_path == tmp_path / "Library" / "LaunchAgents" / "com.headroom.tasky.plist"
|
| 107 |
-
assert "<key>StartInterval</key>" in task_content
|
| 108 |
-
assert "<integer>300</integer>" in task_content
|
| 109 |
-
|
| 110 |
-
|
| 111 |
-
def test_render_windows_runner_writes_ps1_and_cmd_wrappers(tmp_path: Path) -> None:
|
| 112 |
-
ps1_path = tmp_path / "run-headroom.ps1"
|
| 113 |
-
cmd_path = tmp_path / "run-headroom.cmd"
|
| 114 |
-
|
| 115 |
-
records = _render_windows_runner(
|
| 116 |
-
ps1_path,
|
| 117 |
-
cmd_path,
|
| 118 |
-
["C:\\Program Files\\Python\\python.exe", "headroom", "install", "agent", "run"],
|
| 119 |
-
)
|
| 120 |
-
|
| 121 |
-
assert [record.path for record in records] == [str(ps1_path), str(cmd_path)]
|
| 122 |
-
ps1_content = ps1_path.read_text(encoding="utf-8")
|
| 123 |
-
cmd_content = cmd_path.read_text(encoding="utf-8")
|
| 124 |
-
assert '& "C:\\Program Files\\Python\\python.exe" headroom install agent run' in ps1_content
|
| 125 |
-
assert (
|
| 126 |
-
'powershell -NoProfile -ExecutionPolicy Bypass -File "%~dp0run-headroom.ps1" %*'
|
| 127 |
-
in cmd_content
|
| 128 |
-
)
|
| 129 |
-
|
| 130 |
-
|
| 131 |
-
def test_render_runner_scripts_writes_unix_scripts(monkeypatch, tmp_path: Path) -> None:
|
| 132 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 133 |
-
monkeypatch.setattr(
|
| 134 |
-
"headroom.install.supervisors.resolve_headroom_command", lambda: ["headroom"]
|
| 135 |
-
)
|
| 136 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 137 |
-
manifest = _manifest()
|
| 138 |
-
|
| 139 |
-
records = render_runner_scripts(manifest)
|
| 140 |
-
|
| 141 |
-
assert {record.path.split("\\")[-1].split("/")[-1] for record in records} == {
|
| 142 |
-
"run-headroom.sh",
|
| 143 |
-
"ensure-headroom.sh",
|
| 144 |
-
}
|
| 145 |
-
|
| 146 |
-
|
| 147 |
-
def test_render_runner_scripts_writes_windows_scripts(monkeypatch, tmp_path: Path) -> None:
|
| 148 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 149 |
-
monkeypatch.setattr(
|
| 150 |
-
"headroom.install.supervisors.resolve_headroom_command", lambda: ["headroom.exe"]
|
| 151 |
-
)
|
| 152 |
-
monkeypatch.setattr(
|
| 153 |
-
"headroom.install.supervisors.windows_run_script_path",
|
| 154 |
-
lambda profile: tmp_path / "run-headroom.ps1",
|
| 155 |
-
)
|
| 156 |
-
monkeypatch.setattr(
|
| 157 |
-
"headroom.install.supervisors.windows_run_cmd_path",
|
| 158 |
-
lambda profile: tmp_path / "run-headroom.cmd",
|
| 159 |
-
)
|
| 160 |
-
monkeypatch.setattr(
|
| 161 |
-
"headroom.install.supervisors.windows_ensure_script_path",
|
| 162 |
-
lambda profile: tmp_path / "ensure-headroom.ps1",
|
| 163 |
-
)
|
| 164 |
-
monkeypatch.setattr(
|
| 165 |
-
"headroom.install.supervisors.windows_ensure_cmd_path",
|
| 166 |
-
lambda profile: tmp_path / "ensure-headroom.cmd",
|
| 167 |
-
)
|
| 168 |
-
|
| 169 |
-
records = render_runner_scripts(_manifest(profile="win"))
|
| 170 |
-
|
| 171 |
-
assert [Path(record.path).name for record in records] == [
|
| 172 |
-
"run-headroom.ps1",
|
| 173 |
-
"run-headroom.cmd",
|
| 174 |
-
"ensure-headroom.ps1",
|
| 175 |
-
"ensure-headroom.cmd",
|
| 176 |
-
]
|
| 177 |
-
|
| 178 |
-
|
| 179 |
-
def test_install_supervisor_none_returns_runner_records(monkeypatch, tmp_path: Path) -> None:
|
| 180 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 181 |
-
monkeypatch.setattr(
|
| 182 |
-
"headroom.install.supervisors.resolve_headroom_command", lambda: ["headroom"]
|
| 183 |
-
)
|
| 184 |
-
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 185 |
-
manifest = _manifest(supervisor=SupervisorKind.NONE.value)
|
| 186 |
-
|
| 187 |
-
records = install_supervisor(manifest)
|
| 188 |
-
|
| 189 |
-
assert len(records) == 2
|
| 190 |
-
assert all(record.kind == "script" for record in records)
|
| 191 |
-
|
| 192 |
-
|
| 193 |
-
def test_start_and_stop_supervisor_use_linux_systemctl(monkeypatch) -> None:
|
| 194 |
-
calls: list[list[str]] = []
|
| 195 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 196 |
-
monkeypatch.setattr(
|
| 197 |
-
"headroom.install.supervisors.subprocess.run",
|
| 198 |
-
lambda command, **kwargs: calls.append(command),
|
| 199 |
-
)
|
| 200 |
-
manifest = _manifest()
|
| 201 |
-
|
| 202 |
-
start_supervisor(manifest)
|
| 203 |
-
stop_supervisor(manifest)
|
| 204 |
-
|
| 205 |
-
assert calls == [
|
| 206 |
-
["systemctl", "--user", "restart", "headroom-default"],
|
| 207 |
-
["systemctl", "--user", "stop", "headroom-default"],
|
| 208 |
-
]
|
| 209 |
-
|
| 210 |
-
|
| 211 |
-
def test_install_supervisor_linux_service_and_tasks(monkeypatch, tmp_path: Path) -> None:
|
| 212 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 213 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 214 |
-
run_script = tmp_path / "run-headroom.sh"
|
| 215 |
-
ensure_script = tmp_path / "ensure-headroom.sh"
|
| 216 |
-
monkeypatch.setattr(
|
| 217 |
-
"headroom.install.supervisors.render_runner_scripts",
|
| 218 |
-
lambda manifest: [
|
| 219 |
-
type("Record", (), {"kind": "script", "path": run_script.as_posix()})(),
|
| 220 |
-
type("Record", (), {"kind": "script", "path": ensure_script.as_posix()})(),
|
| 221 |
-
],
|
| 222 |
-
)
|
| 223 |
-
unit_path = tmp_path / "headroom-default.service"
|
| 224 |
-
monkeypatch.setattr(
|
| 225 |
-
"headroom.install.supervisors._linux_service_unit",
|
| 226 |
-
lambda manifest, script: (unit_path, "UNIT"),
|
| 227 |
-
)
|
| 228 |
-
calls: list[tuple[list[str], dict]] = []
|
| 229 |
-
|
| 230 |
-
def fake_run(command: list[str], **kwargs):
|
| 231 |
-
calls.append((command, kwargs))
|
| 232 |
-
return type("Result", (), {"returncode": 0, "stdout": "# old cron\n"})()
|
| 233 |
-
|
| 234 |
-
monkeypatch.setattr("headroom.install.supervisors.subprocess.run", fake_run)
|
| 235 |
-
|
| 236 |
-
service_records = install_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 237 |
-
assert unit_path.read_text(encoding="utf-8") == "UNIT"
|
| 238 |
-
assert ["systemctl", "--user", "daemon-reload"] in [call[0] for call in calls]
|
| 239 |
-
assert ["systemctl", "--user", "enable", "headroom-default"] in [call[0] for call in calls]
|
| 240 |
-
assert service_records[-1].kind == "service-unit"
|
| 241 |
-
|
| 242 |
-
cron_path = tmp_path / "headroom-system"
|
| 243 |
-
monkeypatch.setattr(
|
| 244 |
-
"headroom.install.supervisors._linux_task_spec",
|
| 245 |
-
lambda manifest, script: (cron_path, "@reboot root ensure\n"),
|
| 246 |
-
)
|
| 247 |
-
system_task_records = install_supervisor(
|
| 248 |
-
_manifest(profile="system-task", scope="system", supervisor=SupervisorKind.TASK.value)
|
| 249 |
-
)
|
| 250 |
-
assert cron_path.read_text(encoding="utf-8") == "@reboot root ensure\n"
|
| 251 |
-
assert system_task_records[-1].kind == "cron"
|
| 252 |
-
|
| 253 |
-
monkeypatch.setattr(
|
| 254 |
-
"headroom.install.supervisors._linux_task_spec",
|
| 255 |
-
lambda manifest, script: (
|
| 256 |
-
None,
|
| 257 |
-
"# >>> headroom default >>>\n@reboot ensure\n# <<< headroom default <<<\n",
|
| 258 |
-
),
|
| 259 |
-
)
|
| 260 |
-
user_task_records = install_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 261 |
-
assert user_task_records[-1].kind == "crontab"
|
| 262 |
-
assert calls[-1][0] == ["crontab", "-"]
|
| 263 |
-
assert "@reboot ensure" in calls[-1][1]["input"]
|
| 264 |
-
|
| 265 |
-
|
| 266 |
-
def test_install_supervisor_darwin_windows_and_unsupported(monkeypatch, tmp_path: Path) -> None:
|
| 267 |
-
run_script = tmp_path / "run-headroom.sh"
|
| 268 |
-
ensure_script = tmp_path / "ensure-headroom.sh"
|
| 269 |
-
monkeypatch.setattr(
|
| 270 |
-
"headroom.install.supervisors.render_runner_scripts",
|
| 271 |
-
lambda manifest: [
|
| 272 |
-
type("Record", (), {"kind": "script", "path": run_script.as_posix()})(),
|
| 273 |
-
type("Record", (), {"kind": "script", "path": ensure_script.as_posix()})(),
|
| 274 |
-
],
|
| 275 |
-
)
|
| 276 |
-
calls: list[list[str]] = []
|
| 277 |
-
monkeypatch.setattr(
|
| 278 |
-
"headroom.install.supervisors.subprocess.run",
|
| 279 |
-
lambda command, **kwargs: calls.append(command),
|
| 280 |
-
)
|
| 281 |
-
monkeypatch.setattr("headroom.install.supervisors.os.getuid", lambda: 123, raising=False)
|
| 282 |
-
|
| 283 |
-
plist_path = tmp_path / "com.headroom.default.plist"
|
| 284 |
-
monkeypatch.setattr(
|
| 285 |
-
"headroom.install.supervisors._macos_launchd_plist",
|
| 286 |
-
lambda manifest, script, interval=None: (plist_path, f"plist-{interval}"),
|
| 287 |
-
)
|
| 288 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "darwin")
|
| 289 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "darwin")
|
| 290 |
-
service_records = install_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 291 |
-
task_records = install_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 292 |
-
assert plist_path.read_text(encoding="utf-8") == "plist-300"
|
| 293 |
-
assert service_records[-1].kind == "plist"
|
| 294 |
-
assert task_records[-1].kind == "plist"
|
| 295 |
-
assert ["launchctl", "bootstrap", "gui/123", str(plist_path)] in calls
|
| 296 |
-
|
| 297 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 298 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 299 |
-
monkeypatch.setattr(
|
| 300 |
-
"headroom.install.supervisors.windows_run_cmd_path",
|
| 301 |
-
lambda profile: Path(f"C:\\tmp\\{profile}\\run-headroom.cmd"),
|
| 302 |
-
)
|
| 303 |
-
monkeypatch.setattr(
|
| 304 |
-
"headroom.install.supervisors.windows_ensure_cmd_path",
|
| 305 |
-
lambda profile: Path(f"C:\\tmp\\{profile}\\ensure-headroom.cmd"),
|
| 306 |
-
)
|
| 307 |
-
win_service = install_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 308 |
-
win_task = install_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 309 |
-
assert win_service[-1].kind == "windows-service"
|
| 310 |
-
assert win_task[-2].path.endswith("-startup")
|
| 311 |
-
assert [
|
| 312 |
-
"sc.exe",
|
| 313 |
-
"create",
|
| 314 |
-
"headroom-default",
|
| 315 |
-
'binPath= cmd.exe /c "C:\\tmp\\default\\run-headroom.cmd"',
|
| 316 |
-
"start= auto",
|
| 317 |
-
] in calls
|
| 318 |
-
assert [
|
| 319 |
-
"schtasks",
|
| 320 |
-
"/Create",
|
| 321 |
-
"/TN",
|
| 322 |
-
"headroom-default-health",
|
| 323 |
-
"/TR",
|
| 324 |
-
"C:\\tmp\\default\\ensure-headroom.cmd",
|
| 325 |
-
"/SC",
|
| 326 |
-
"MINUTE",
|
| 327 |
-
"/MO",
|
| 328 |
-
"5",
|
| 329 |
-
"/F",
|
| 330 |
-
] in calls
|
| 331 |
-
|
| 332 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "plan9")
|
| 333 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "plan9")
|
| 334 |
-
with pytest.raises(click.ClickException, match="not supported"):
|
| 335 |
-
install_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 336 |
-
|
| 337 |
-
|
| 338 |
-
def test_start_and_stop_supervisor_darwin_windows_and_none(monkeypatch) -> None:
|
| 339 |
-
calls: list[list[str]] = []
|
| 340 |
-
monkeypatch.setattr(
|
| 341 |
-
"headroom.install.supervisors.subprocess.run",
|
| 342 |
-
lambda command, **kwargs: calls.append(command),
|
| 343 |
-
)
|
| 344 |
-
monkeypatch.setattr("headroom.install.supervisors.os.getuid", lambda: 77, raising=False)
|
| 345 |
-
|
| 346 |
-
start_supervisor(_manifest(supervisor=SupervisorKind.NONE.value))
|
| 347 |
-
stop_supervisor(_manifest(supervisor=SupervisorKind.NONE.value))
|
| 348 |
-
assert calls == []
|
| 349 |
-
|
| 350 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "darwin")
|
| 351 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "darwin")
|
| 352 |
-
start_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 353 |
-
stop_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 354 |
-
assert calls == [
|
| 355 |
-
["launchctl", "kickstart", "-k", "gui/77/com.headroom.default"],
|
| 356 |
-
["launchctl", "bootout", "gui/77/com.headroom.default"],
|
| 357 |
-
]
|
| 358 |
-
|
| 359 |
-
calls.clear()
|
| 360 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 361 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 362 |
-
start_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 363 |
-
stop_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 364 |
-
assert calls == [
|
| 365 |
-
["sc.exe", "start", "headroom-default"],
|
| 366 |
-
["sc.exe", "stop", "headroom-default"],
|
| 367 |
-
]
|
| 368 |
-
|
| 369 |
-
|
| 370 |
-
def test_remove_supervisor_removes_user_crontab_block(monkeypatch) -> None:
|
| 371 |
-
calls: list[tuple[list[str], str | None]] = []
|
| 372 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 373 |
-
|
| 374 |
-
class Result:
|
| 375 |
-
def __init__(self, returncode: int = 0, stdout: str = "") -> None:
|
| 376 |
-
self.returncode = returncode
|
| 377 |
-
self.stdout = stdout
|
| 378 |
-
|
| 379 |
-
def fake_run(command: list[str], **kwargs):
|
| 380 |
-
calls.append((command, kwargs.get("input")))
|
| 381 |
-
if command == ["crontab", "-l"]:
|
| 382 |
-
return Result(
|
| 383 |
-
stdout="# >>> headroom default >>>\n@reboot /tmp/ensure\n# <<< headroom default <<<\n"
|
| 384 |
-
)
|
| 385 |
-
return Result()
|
| 386 |
-
|
| 387 |
-
monkeypatch.setattr("headroom.install.supervisors.subprocess.run", fake_run)
|
| 388 |
-
manifest = _manifest(supervisor=SupervisorKind.TASK.value)
|
| 389 |
-
|
| 390 |
-
remove_supervisor(manifest)
|
| 391 |
-
|
| 392 |
-
assert calls[0][0] == ["crontab", "-l"]
|
| 393 |
-
assert calls[1][0] == ["crontab", "-"]
|
| 394 |
-
|
| 395 |
-
|
| 396 |
-
def test_remove_supervisor_linux_service_cron_path_and_missing_crontab(
|
| 397 |
-
monkeypatch, tmp_path: Path
|
| 398 |
-
) -> None:
|
| 399 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 400 |
-
calls: list[list[str]] = []
|
| 401 |
-
|
| 402 |
-
def fake_run(command: list[str], **kwargs):
|
| 403 |
-
calls.append(command)
|
| 404 |
-
return type("Result", (), {"returncode": 1, "stdout": ""})()
|
| 405 |
-
|
| 406 |
-
monkeypatch.setattr("headroom.install.supervisors.subprocess.run", fake_run)
|
| 407 |
-
unit_path = tmp_path / "headroom-default.service"
|
| 408 |
-
unit_path.write_text("unit", encoding="utf-8")
|
| 409 |
-
monkeypatch.setattr(
|
| 410 |
-
"headroom.install.supervisors._linux_service_unit",
|
| 411 |
-
lambda manifest, script: (unit_path, "unit"),
|
| 412 |
-
)
|
| 413 |
-
remove_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 414 |
-
assert not unit_path.exists()
|
| 415 |
-
assert ["systemctl", "--user", "disable", "--now", "headroom-default"] in calls
|
| 416 |
-
assert ["systemctl", "--user", "daemon-reload"] in calls
|
| 417 |
-
|
| 418 |
-
cron_path = tmp_path / "headroom-task"
|
| 419 |
-
cron_path.write_text("cron", encoding="utf-8")
|
| 420 |
-
monkeypatch.setattr(
|
| 421 |
-
"headroom.install.supervisors._linux_task_spec",
|
| 422 |
-
lambda manifest, script: (cron_path, "cron"),
|
| 423 |
-
)
|
| 424 |
-
remove_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 425 |
-
assert not cron_path.exists()
|
| 426 |
-
|
| 427 |
-
monkeypatch.setattr(
|
| 428 |
-
"headroom.install.supervisors._linux_task_spec",
|
| 429 |
-
lambda manifest, script: (None, "cron"),
|
| 430 |
-
)
|
| 431 |
-
remove_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 432 |
-
assert calls[-1] == ["crontab", "-l"]
|
| 433 |
-
|
| 434 |
-
|
| 435 |
-
def test_remove_supervisor_darwin_and_windows(monkeypatch, tmp_path: Path) -> None:
|
| 436 |
-
calls: list[list[str]] = []
|
| 437 |
-
monkeypatch.setattr(
|
| 438 |
-
"headroom.install.supervisors.subprocess.run",
|
| 439 |
-
lambda command, **kwargs: calls.append(command),
|
| 440 |
-
)
|
| 441 |
-
monkeypatch.setattr("headroom.install.supervisors.os.getuid", lambda: 55, raising=False)
|
| 442 |
-
|
| 443 |
-
plist_path = tmp_path / "com.headroom.default.plist"
|
| 444 |
-
plist_path.write_text("plist", encoding="utf-8")
|
| 445 |
-
monkeypatch.setattr(
|
| 446 |
-
"headroom.install.supervisors.unix_run_script_path",
|
| 447 |
-
lambda profile: tmp_path / "run-headroom.sh",
|
| 448 |
-
)
|
| 449 |
-
monkeypatch.setattr(
|
| 450 |
-
"headroom.install.supervisors.unix_ensure_script_path",
|
| 451 |
-
lambda profile: tmp_path / "ensure-headroom.sh",
|
| 452 |
-
)
|
| 453 |
-
monkeypatch.setattr(
|
| 454 |
-
"headroom.install.supervisors._macos_launchd_plist",
|
| 455 |
-
lambda manifest, script, interval=None: (plist_path, "plist"),
|
| 456 |
-
)
|
| 457 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "darwin")
|
| 458 |
-
remove_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 459 |
-
assert not plist_path.exists()
|
| 460 |
-
assert calls[0] == ["launchctl", "bootout", "gui/55/com.headroom.default"]
|
| 461 |
-
|
| 462 |
-
calls.clear()
|
| 463 |
-
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 464 |
-
remove_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 465 |
-
remove_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 466 |
-
assert calls == [
|
| 467 |
-
["sc.exe", "stop", "headroom-default"],
|
| 468 |
-
["sc.exe", "delete", "headroom-default"],
|
| 469 |
-
["schtasks", "/Delete", "/TN", "headroom-default-startup", "/F"],
|
| 470 |
-
["schtasks", "/Delete", "/TN", "headroom-default-health", "/F"],
|
| 471 |
-
]
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
from pathlib import Path
|
| 4 |
+
|
| 5 |
+
import click
|
| 6 |
+
import pytest
|
| 7 |
+
|
| 8 |
+
from headroom.install.models import DeploymentManifest, SupervisorKind
|
| 9 |
+
from headroom.install.supervisors import (
|
| 10 |
+
_command_for_script,
|
| 11 |
+
_linux_service_unit,
|
| 12 |
+
_linux_task_spec,
|
| 13 |
+
_macos_launchd_plist,
|
| 14 |
+
_render_unix_runner,
|
| 15 |
+
_render_windows_runner,
|
| 16 |
+
install_supervisor,
|
| 17 |
+
remove_supervisor,
|
| 18 |
+
render_runner_scripts,
|
| 19 |
+
start_supervisor,
|
| 20 |
+
stop_supervisor,
|
| 21 |
+
)
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
def _manifest(
|
| 25 |
+
*, profile: str = "default", scope: str = "user", supervisor: str = "service"
|
| 26 |
+
) -> DeploymentManifest:
|
| 27 |
+
return DeploymentManifest(
|
| 28 |
+
profile=profile,
|
| 29 |
+
preset="persistent-service",
|
| 30 |
+
runtime_kind="python",
|
| 31 |
+
supervisor_kind=supervisor,
|
| 32 |
+
scope=scope,
|
| 33 |
+
provider_mode="manual",
|
| 34 |
+
targets=[],
|
| 35 |
+
port=8787,
|
| 36 |
+
host="127.0.0.1",
|
| 37 |
+
backend="anthropic",
|
| 38 |
+
service_name=f"headroom-{profile}",
|
| 39 |
+
)
|
| 40 |
+
|
| 41 |
+
|
| 42 |
+
def test_linux_service_unit_uses_user_systemd_path(monkeypatch, tmp_path: Path) -> None:
|
| 43 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 44 |
+
manifest = _manifest()
|
| 45 |
+
|
| 46 |
+
unit_path, content = _linux_service_unit(manifest, tmp_path / "run-headroom.sh")
|
| 47 |
+
|
| 48 |
+
assert unit_path == tmp_path / ".config" / "systemd" / "user" / "headroom-default.service"
|
| 49 |
+
assert "ExecStart=" + str(tmp_path / "run-headroom.sh") in content
|
| 50 |
+
assert "Restart=on-failure" in content
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
def test_command_for_script_and_unix_runner(monkeypatch, tmp_path: Path) -> None:
|
| 54 |
+
monkeypatch.setattr(
|
| 55 |
+
"headroom.install.supervisors.resolve_headroom_command",
|
| 56 |
+
lambda: ["python", "-m", "headroom"],
|
| 57 |
+
)
|
| 58 |
+
|
| 59 |
+
assert _command_for_script("install", "agent", "run") == [
|
| 60 |
+
"python",
|
| 61 |
+
"-m",
|
| 62 |
+
"headroom",
|
| 63 |
+
"install",
|
| 64 |
+
"agent",
|
| 65 |
+
"run",
|
| 66 |
+
]
|
| 67 |
+
|
| 68 |
+
record = _render_unix_runner(
|
| 69 |
+
tmp_path / "scripts" / "run-headroom.sh", ["headroom", "run", "--flag"]
|
| 70 |
+
)
|
| 71 |
+
assert record.kind == "script"
|
| 72 |
+
content = Path(record.path).read_text(encoding="utf-8")
|
| 73 |
+
assert content.startswith("#!/usr/bin/env bash")
|
| 74 |
+
assert "exec headroom run --flag" in content
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
def test_linux_task_spec_for_user_scope_includes_crontab_markers(tmp_path: Path) -> None:
|
| 78 |
+
manifest = _manifest(profile="smoke", supervisor=SupervisorKind.TASK.value)
|
| 79 |
+
|
| 80 |
+
cron_path, content = _linux_task_spec(manifest, tmp_path / "ensure-headroom.sh")
|
| 81 |
+
|
| 82 |
+
assert cron_path is None
|
| 83 |
+
assert "# >>> headroom smoke >>>" in content
|
| 84 |
+
assert "# <<< headroom smoke <<<" in content
|
| 85 |
+
assert "@reboot" in content
|
| 86 |
+
assert "*/5 * * * *" in content
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def test_macos_launchd_plist_switches_between_keepalive_and_interval(
|
| 90 |
+
monkeypatch, tmp_path: Path
|
| 91 |
+
) -> None:
|
| 92 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 93 |
+
|
| 94 |
+
service_manifest = _manifest(supervisor=SupervisorKind.SERVICE.value)
|
| 95 |
+
service_path, service_content = _macos_launchd_plist(
|
| 96 |
+
service_manifest, tmp_path / "run-headroom.sh"
|
| 97 |
+
)
|
| 98 |
+
assert service_path == tmp_path / "Library" / "LaunchAgents" / "com.headroom.default.plist"
|
| 99 |
+
assert "<key>KeepAlive</key>" in service_content
|
| 100 |
+
assert "<key>StartInterval</key>" not in service_content
|
| 101 |
+
|
| 102 |
+
task_manifest = _manifest(profile="tasky", supervisor=SupervisorKind.TASK.value)
|
| 103 |
+
task_path, task_content = _macos_launchd_plist(
|
| 104 |
+
task_manifest, tmp_path / "ensure-headroom.sh", interval=300
|
| 105 |
+
)
|
| 106 |
+
assert task_path == tmp_path / "Library" / "LaunchAgents" / "com.headroom.tasky.plist"
|
| 107 |
+
assert "<key>StartInterval</key>" in task_content
|
| 108 |
+
assert "<integer>300</integer>" in task_content
|
| 109 |
+
|
| 110 |
+
|
| 111 |
+
def test_render_windows_runner_writes_ps1_and_cmd_wrappers(tmp_path: Path) -> None:
|
| 112 |
+
ps1_path = tmp_path / "run-headroom.ps1"
|
| 113 |
+
cmd_path = tmp_path / "run-headroom.cmd"
|
| 114 |
+
|
| 115 |
+
records = _render_windows_runner(
|
| 116 |
+
ps1_path,
|
| 117 |
+
cmd_path,
|
| 118 |
+
["C:\\Program Files\\Python\\python.exe", "headroom", "install", "agent", "run"],
|
| 119 |
+
)
|
| 120 |
+
|
| 121 |
+
assert [record.path for record in records] == [str(ps1_path), str(cmd_path)]
|
| 122 |
+
ps1_content = ps1_path.read_text(encoding="utf-8")
|
| 123 |
+
cmd_content = cmd_path.read_text(encoding="utf-8")
|
| 124 |
+
assert '& "C:\\Program Files\\Python\\python.exe" headroom install agent run' in ps1_content
|
| 125 |
+
assert (
|
| 126 |
+
'powershell -NoProfile -ExecutionPolicy Bypass -File "%~dp0run-headroom.ps1" %*'
|
| 127 |
+
in cmd_content
|
| 128 |
+
)
|
| 129 |
+
|
| 130 |
+
|
| 131 |
+
def test_render_runner_scripts_writes_unix_scripts(monkeypatch, tmp_path: Path) -> None:
|
| 132 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 133 |
+
monkeypatch.setattr(
|
| 134 |
+
"headroom.install.supervisors.resolve_headroom_command", lambda: ["headroom"]
|
| 135 |
+
)
|
| 136 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 137 |
+
manifest = _manifest()
|
| 138 |
+
|
| 139 |
+
records = render_runner_scripts(manifest)
|
| 140 |
+
|
| 141 |
+
assert {record.path.split("\\")[-1].split("/")[-1] for record in records} == {
|
| 142 |
+
"run-headroom.sh",
|
| 143 |
+
"ensure-headroom.sh",
|
| 144 |
+
}
|
| 145 |
+
|
| 146 |
+
|
| 147 |
+
def test_render_runner_scripts_writes_windows_scripts(monkeypatch, tmp_path: Path) -> None:
|
| 148 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 149 |
+
monkeypatch.setattr(
|
| 150 |
+
"headroom.install.supervisors.resolve_headroom_command", lambda: ["headroom.exe"]
|
| 151 |
+
)
|
| 152 |
+
monkeypatch.setattr(
|
| 153 |
+
"headroom.install.supervisors.windows_run_script_path",
|
| 154 |
+
lambda profile: tmp_path / "run-headroom.ps1",
|
| 155 |
+
)
|
| 156 |
+
monkeypatch.setattr(
|
| 157 |
+
"headroom.install.supervisors.windows_run_cmd_path",
|
| 158 |
+
lambda profile: tmp_path / "run-headroom.cmd",
|
| 159 |
+
)
|
| 160 |
+
monkeypatch.setattr(
|
| 161 |
+
"headroom.install.supervisors.windows_ensure_script_path",
|
| 162 |
+
lambda profile: tmp_path / "ensure-headroom.ps1",
|
| 163 |
+
)
|
| 164 |
+
monkeypatch.setattr(
|
| 165 |
+
"headroom.install.supervisors.windows_ensure_cmd_path",
|
| 166 |
+
lambda profile: tmp_path / "ensure-headroom.cmd",
|
| 167 |
+
)
|
| 168 |
+
|
| 169 |
+
records = render_runner_scripts(_manifest(profile="win"))
|
| 170 |
+
|
| 171 |
+
assert [Path(record.path).name for record in records] == [
|
| 172 |
+
"run-headroom.ps1",
|
| 173 |
+
"run-headroom.cmd",
|
| 174 |
+
"ensure-headroom.ps1",
|
| 175 |
+
"ensure-headroom.cmd",
|
| 176 |
+
]
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
def test_install_supervisor_none_returns_runner_records(monkeypatch, tmp_path: Path) -> None:
|
| 180 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 181 |
+
monkeypatch.setattr(
|
| 182 |
+
"headroom.install.supervisors.resolve_headroom_command", lambda: ["headroom"]
|
| 183 |
+
)
|
| 184 |
+
monkeypatch.setattr(Path, "home", lambda: tmp_path)
|
| 185 |
+
manifest = _manifest(supervisor=SupervisorKind.NONE.value)
|
| 186 |
+
|
| 187 |
+
records = install_supervisor(manifest)
|
| 188 |
+
|
| 189 |
+
assert len(records) == 2
|
| 190 |
+
assert all(record.kind == "script" for record in records)
|
| 191 |
+
|
| 192 |
+
|
| 193 |
+
def test_start_and_stop_supervisor_use_linux_systemctl(monkeypatch) -> None:
|
| 194 |
+
calls: list[list[str]] = []
|
| 195 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 196 |
+
monkeypatch.setattr(
|
| 197 |
+
"headroom.install.supervisors.subprocess.run",
|
| 198 |
+
lambda command, **kwargs: calls.append(command),
|
| 199 |
+
)
|
| 200 |
+
manifest = _manifest()
|
| 201 |
+
|
| 202 |
+
start_supervisor(manifest)
|
| 203 |
+
stop_supervisor(manifest)
|
| 204 |
+
|
| 205 |
+
assert calls == [
|
| 206 |
+
["systemctl", "--user", "restart", "headroom-default"],
|
| 207 |
+
["systemctl", "--user", "stop", "headroom-default"],
|
| 208 |
+
]
|
| 209 |
+
|
| 210 |
+
|
| 211 |
+
def test_install_supervisor_linux_service_and_tasks(monkeypatch, tmp_path: Path) -> None:
|
| 212 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 213 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 214 |
+
run_script = tmp_path / "run-headroom.sh"
|
| 215 |
+
ensure_script = tmp_path / "ensure-headroom.sh"
|
| 216 |
+
monkeypatch.setattr(
|
| 217 |
+
"headroom.install.supervisors.render_runner_scripts",
|
| 218 |
+
lambda manifest: [
|
| 219 |
+
type("Record", (), {"kind": "script", "path": run_script.as_posix()})(),
|
| 220 |
+
type("Record", (), {"kind": "script", "path": ensure_script.as_posix()})(),
|
| 221 |
+
],
|
| 222 |
+
)
|
| 223 |
+
unit_path = tmp_path / "headroom-default.service"
|
| 224 |
+
monkeypatch.setattr(
|
| 225 |
+
"headroom.install.supervisors._linux_service_unit",
|
| 226 |
+
lambda manifest, script: (unit_path, "UNIT"),
|
| 227 |
+
)
|
| 228 |
+
calls: list[tuple[list[str], dict]] = []
|
| 229 |
+
|
| 230 |
+
def fake_run(command: list[str], **kwargs):
|
| 231 |
+
calls.append((command, kwargs))
|
| 232 |
+
return type("Result", (), {"returncode": 0, "stdout": "# old cron\n"})()
|
| 233 |
+
|
| 234 |
+
monkeypatch.setattr("headroom.install.supervisors.subprocess.run", fake_run)
|
| 235 |
+
|
| 236 |
+
service_records = install_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 237 |
+
assert unit_path.read_text(encoding="utf-8") == "UNIT"
|
| 238 |
+
assert ["systemctl", "--user", "daemon-reload"] in [call[0] for call in calls]
|
| 239 |
+
assert ["systemctl", "--user", "enable", "headroom-default"] in [call[0] for call in calls]
|
| 240 |
+
assert service_records[-1].kind == "service-unit"
|
| 241 |
+
|
| 242 |
+
cron_path = tmp_path / "headroom-system"
|
| 243 |
+
monkeypatch.setattr(
|
| 244 |
+
"headroom.install.supervisors._linux_task_spec",
|
| 245 |
+
lambda manifest, script: (cron_path, "@reboot root ensure\n"),
|
| 246 |
+
)
|
| 247 |
+
system_task_records = install_supervisor(
|
| 248 |
+
_manifest(profile="system-task", scope="system", supervisor=SupervisorKind.TASK.value)
|
| 249 |
+
)
|
| 250 |
+
assert cron_path.read_text(encoding="utf-8") == "@reboot root ensure\n"
|
| 251 |
+
assert system_task_records[-1].kind == "cron"
|
| 252 |
+
|
| 253 |
+
monkeypatch.setattr(
|
| 254 |
+
"headroom.install.supervisors._linux_task_spec",
|
| 255 |
+
lambda manifest, script: (
|
| 256 |
+
None,
|
| 257 |
+
"# >>> headroom default >>>\n@reboot ensure\n# <<< headroom default <<<\n",
|
| 258 |
+
),
|
| 259 |
+
)
|
| 260 |
+
user_task_records = install_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 261 |
+
assert user_task_records[-1].kind == "crontab"
|
| 262 |
+
assert calls[-1][0] == ["crontab", "-"]
|
| 263 |
+
assert "@reboot ensure" in calls[-1][1]["input"]
|
| 264 |
+
|
| 265 |
+
|
| 266 |
+
def test_install_supervisor_darwin_windows_and_unsupported(monkeypatch, tmp_path: Path) -> None:
|
| 267 |
+
run_script = tmp_path / "run-headroom.sh"
|
| 268 |
+
ensure_script = tmp_path / "ensure-headroom.sh"
|
| 269 |
+
monkeypatch.setattr(
|
| 270 |
+
"headroom.install.supervisors.render_runner_scripts",
|
| 271 |
+
lambda manifest: [
|
| 272 |
+
type("Record", (), {"kind": "script", "path": run_script.as_posix()})(),
|
| 273 |
+
type("Record", (), {"kind": "script", "path": ensure_script.as_posix()})(),
|
| 274 |
+
],
|
| 275 |
+
)
|
| 276 |
+
calls: list[list[str]] = []
|
| 277 |
+
monkeypatch.setattr(
|
| 278 |
+
"headroom.install.supervisors.subprocess.run",
|
| 279 |
+
lambda command, **kwargs: calls.append(command),
|
| 280 |
+
)
|
| 281 |
+
monkeypatch.setattr("headroom.install.supervisors.os.getuid", lambda: 123, raising=False)
|
| 282 |
+
|
| 283 |
+
plist_path = tmp_path / "com.headroom.default.plist"
|
| 284 |
+
monkeypatch.setattr(
|
| 285 |
+
"headroom.install.supervisors._macos_launchd_plist",
|
| 286 |
+
lambda manifest, script, interval=None: (plist_path, f"plist-{interval}"),
|
| 287 |
+
)
|
| 288 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "darwin")
|
| 289 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "darwin")
|
| 290 |
+
service_records = install_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 291 |
+
task_records = install_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 292 |
+
assert plist_path.read_text(encoding="utf-8") == "plist-300"
|
| 293 |
+
assert service_records[-1].kind == "plist"
|
| 294 |
+
assert task_records[-1].kind == "plist"
|
| 295 |
+
assert ["launchctl", "bootstrap", "gui/123", str(plist_path)] in calls
|
| 296 |
+
|
| 297 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 298 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 299 |
+
monkeypatch.setattr(
|
| 300 |
+
"headroom.install.supervisors.windows_run_cmd_path",
|
| 301 |
+
lambda profile: Path(f"C:\\tmp\\{profile}\\run-headroom.cmd"),
|
| 302 |
+
)
|
| 303 |
+
monkeypatch.setattr(
|
| 304 |
+
"headroom.install.supervisors.windows_ensure_cmd_path",
|
| 305 |
+
lambda profile: Path(f"C:\\tmp\\{profile}\\ensure-headroom.cmd"),
|
| 306 |
+
)
|
| 307 |
+
win_service = install_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 308 |
+
win_task = install_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 309 |
+
assert win_service[-1].kind == "windows-service"
|
| 310 |
+
assert win_task[-2].path.endswith("-startup")
|
| 311 |
+
assert [
|
| 312 |
+
"sc.exe",
|
| 313 |
+
"create",
|
| 314 |
+
"headroom-default",
|
| 315 |
+
'binPath= cmd.exe /c "C:\\tmp\\default\\run-headroom.cmd"',
|
| 316 |
+
"start= auto",
|
| 317 |
+
] in calls
|
| 318 |
+
assert [
|
| 319 |
+
"schtasks",
|
| 320 |
+
"/Create",
|
| 321 |
+
"/TN",
|
| 322 |
+
"headroom-default-health",
|
| 323 |
+
"/TR",
|
| 324 |
+
"C:\\tmp\\default\\ensure-headroom.cmd",
|
| 325 |
+
"/SC",
|
| 326 |
+
"MINUTE",
|
| 327 |
+
"/MO",
|
| 328 |
+
"5",
|
| 329 |
+
"/F",
|
| 330 |
+
] in calls
|
| 331 |
+
|
| 332 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "plan9")
|
| 333 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "plan9")
|
| 334 |
+
with pytest.raises(click.ClickException, match="not supported"):
|
| 335 |
+
install_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 336 |
+
|
| 337 |
+
|
| 338 |
+
def test_start_and_stop_supervisor_darwin_windows_and_none(monkeypatch) -> None:
|
| 339 |
+
calls: list[list[str]] = []
|
| 340 |
+
monkeypatch.setattr(
|
| 341 |
+
"headroom.install.supervisors.subprocess.run",
|
| 342 |
+
lambda command, **kwargs: calls.append(command),
|
| 343 |
+
)
|
| 344 |
+
monkeypatch.setattr("headroom.install.supervisors.os.getuid", lambda: 77, raising=False)
|
| 345 |
+
|
| 346 |
+
start_supervisor(_manifest(supervisor=SupervisorKind.NONE.value))
|
| 347 |
+
stop_supervisor(_manifest(supervisor=SupervisorKind.NONE.value))
|
| 348 |
+
assert calls == []
|
| 349 |
+
|
| 350 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "darwin")
|
| 351 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "darwin")
|
| 352 |
+
start_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 353 |
+
stop_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 354 |
+
assert calls == [
|
| 355 |
+
["launchctl", "kickstart", "-k", "gui/77/com.headroom.default"],
|
| 356 |
+
["launchctl", "bootout", "gui/77/com.headroom.default"],
|
| 357 |
+
]
|
| 358 |
+
|
| 359 |
+
calls.clear()
|
| 360 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 361 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 362 |
+
start_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 363 |
+
stop_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 364 |
+
assert calls == [
|
| 365 |
+
["sc.exe", "start", "headroom-default"],
|
| 366 |
+
["sc.exe", "stop", "headroom-default"],
|
| 367 |
+
]
|
| 368 |
+
|
| 369 |
+
|
| 370 |
+
def test_remove_supervisor_removes_user_crontab_block(monkeypatch) -> None:
|
| 371 |
+
calls: list[tuple[list[str], str | None]] = []
|
| 372 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 373 |
+
|
| 374 |
+
class Result:
|
| 375 |
+
def __init__(self, returncode: int = 0, stdout: str = "") -> None:
|
| 376 |
+
self.returncode = returncode
|
| 377 |
+
self.stdout = stdout
|
| 378 |
+
|
| 379 |
+
def fake_run(command: list[str], **kwargs):
|
| 380 |
+
calls.append((command, kwargs.get("input")))
|
| 381 |
+
if command == ["crontab", "-l"]:
|
| 382 |
+
return Result(
|
| 383 |
+
stdout="# >>> headroom default >>>\n@reboot /tmp/ensure\n# <<< headroom default <<<\n"
|
| 384 |
+
)
|
| 385 |
+
return Result()
|
| 386 |
+
|
| 387 |
+
monkeypatch.setattr("headroom.install.supervisors.subprocess.run", fake_run)
|
| 388 |
+
manifest = _manifest(supervisor=SupervisorKind.TASK.value)
|
| 389 |
+
|
| 390 |
+
remove_supervisor(manifest)
|
| 391 |
+
|
| 392 |
+
assert calls[0][0] == ["crontab", "-l"]
|
| 393 |
+
assert calls[1][0] == ["crontab", "-"]
|
| 394 |
+
|
| 395 |
+
|
| 396 |
+
def test_remove_supervisor_linux_service_cron_path_and_missing_crontab(
|
| 397 |
+
monkeypatch, tmp_path: Path
|
| 398 |
+
) -> None:
|
| 399 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "linux")
|
| 400 |
+
calls: list[list[str]] = []
|
| 401 |
+
|
| 402 |
+
def fake_run(command: list[str], **kwargs):
|
| 403 |
+
calls.append(command)
|
| 404 |
+
return type("Result", (), {"returncode": 1, "stdout": ""})()
|
| 405 |
+
|
| 406 |
+
monkeypatch.setattr("headroom.install.supervisors.subprocess.run", fake_run)
|
| 407 |
+
unit_path = tmp_path / "headroom-default.service"
|
| 408 |
+
unit_path.write_text("unit", encoding="utf-8")
|
| 409 |
+
monkeypatch.setattr(
|
| 410 |
+
"headroom.install.supervisors._linux_service_unit",
|
| 411 |
+
lambda manifest, script: (unit_path, "unit"),
|
| 412 |
+
)
|
| 413 |
+
remove_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 414 |
+
assert not unit_path.exists()
|
| 415 |
+
assert ["systemctl", "--user", "disable", "--now", "headroom-default"] in calls
|
| 416 |
+
assert ["systemctl", "--user", "daemon-reload"] in calls
|
| 417 |
+
|
| 418 |
+
cron_path = tmp_path / "headroom-task"
|
| 419 |
+
cron_path.write_text("cron", encoding="utf-8")
|
| 420 |
+
monkeypatch.setattr(
|
| 421 |
+
"headroom.install.supervisors._linux_task_spec",
|
| 422 |
+
lambda manifest, script: (cron_path, "cron"),
|
| 423 |
+
)
|
| 424 |
+
remove_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 425 |
+
assert not cron_path.exists()
|
| 426 |
+
|
| 427 |
+
monkeypatch.setattr(
|
| 428 |
+
"headroom.install.supervisors._linux_task_spec",
|
| 429 |
+
lambda manifest, script: (None, "cron"),
|
| 430 |
+
)
|
| 431 |
+
remove_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 432 |
+
assert calls[-1] == ["crontab", "-l"]
|
| 433 |
+
|
| 434 |
+
|
| 435 |
+
def test_remove_supervisor_darwin_and_windows(monkeypatch, tmp_path: Path) -> None:
|
| 436 |
+
calls: list[list[str]] = []
|
| 437 |
+
monkeypatch.setattr(
|
| 438 |
+
"headroom.install.supervisors.subprocess.run",
|
| 439 |
+
lambda command, **kwargs: calls.append(command),
|
| 440 |
+
)
|
| 441 |
+
monkeypatch.setattr("headroom.install.supervisors.os.getuid", lambda: 55, raising=False)
|
| 442 |
+
|
| 443 |
+
plist_path = tmp_path / "com.headroom.default.plist"
|
| 444 |
+
plist_path.write_text("plist", encoding="utf-8")
|
| 445 |
+
monkeypatch.setattr(
|
| 446 |
+
"headroom.install.supervisors.unix_run_script_path",
|
| 447 |
+
lambda profile: tmp_path / "run-headroom.sh",
|
| 448 |
+
)
|
| 449 |
+
monkeypatch.setattr(
|
| 450 |
+
"headroom.install.supervisors.unix_ensure_script_path",
|
| 451 |
+
lambda profile: tmp_path / "ensure-headroom.sh",
|
| 452 |
+
)
|
| 453 |
+
monkeypatch.setattr(
|
| 454 |
+
"headroom.install.supervisors._macos_launchd_plist",
|
| 455 |
+
lambda manifest, script, interval=None: (plist_path, "plist"),
|
| 456 |
+
)
|
| 457 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "darwin")
|
| 458 |
+
remove_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 459 |
+
assert not plist_path.exists()
|
| 460 |
+
assert calls[0] == ["launchctl", "bootout", "gui/55/com.headroom.default"]
|
| 461 |
+
|
| 462 |
+
calls.clear()
|
| 463 |
+
monkeypatch.setattr("headroom.install.supervisors.sys.platform", "win32")
|
| 464 |
+
remove_supervisor(_manifest(supervisor=SupervisorKind.SERVICE.value))
|
| 465 |
+
remove_supervisor(_manifest(supervisor=SupervisorKind.TASK.value))
|
| 466 |
+
assert calls == [
|
| 467 |
+
["sc.exe", "stop", "headroom-default"],
|
| 468 |
+
["sc.exe", "delete", "headroom-default"],
|
| 469 |
+
["schtasks", "/Delete", "/TN", "headroom-default-startup", "/F"],
|
| 470 |
+
["schtasks", "/Delete", "/TN", "headroom-default-health", "/F"],
|
| 471 |
+
]
|
|
@@ -1,55 +1,55 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
import json
|
| 4 |
-
from pathlib import Path
|
| 5 |
-
|
| 6 |
-
REPO_ROOT = Path(__file__).resolve().parents[1]
|
| 7 |
-
|
| 8 |
-
|
| 9 |
-
def _load_json(relative_path: str) -> object:
|
| 10 |
-
return json.loads((REPO_ROOT / relative_path).read_text(encoding="utf-8"))
|
| 11 |
-
|
| 12 |
-
|
| 13 |
-
def test_marketplace_manifests_match() -> None:
|
| 14 |
-
assert _load_json(".claude-plugin/marketplace.json") == _load_json(
|
| 15 |
-
".github/plugin/marketplace.json"
|
| 16 |
-
)
|
| 17 |
-
|
| 18 |
-
|
| 19 |
-
def test_plugin_manifests_share_core_metadata() -> None:
|
| 20 |
-
claude = _load_json("plugins/headroom-agent-hooks/.claude-plugin/plugin.json")
|
| 21 |
-
copilot = _load_json("plugins/headroom-agent-hooks/.github/plugin/plugin.json")
|
| 22 |
-
assert isinstance(claude, dict)
|
| 23 |
-
assert isinstance(copilot, dict)
|
| 24 |
-
for key in ("name", "version", "description", "author", "homepage", "repository", "keywords"):
|
| 25 |
-
assert claude[key] == copilot[key]
|
| 26 |
-
assert "hooks" not in claude
|
| 27 |
-
assert copilot["hooks"] == "./hooks"
|
| 28 |
-
|
| 29 |
-
|
| 30 |
-
def test_marketplace_entry_points_to_plugin_root() -> None:
|
| 31 |
-
marketplace = _load_json(".claude-plugin/marketplace.json")
|
| 32 |
-
assert isinstance(marketplace, dict)
|
| 33 |
-
plugins = marketplace["plugins"]
|
| 34 |
-
assert isinstance(plugins, list)
|
| 35 |
-
plugin = plugins[0]
|
| 36 |
-
assert plugin["name"] == "headroom"
|
| 37 |
-
plugin_root = (REPO_ROOT / plugin["source"]).resolve()
|
| 38 |
-
assert plugin_root.is_dir()
|
| 39 |
-
assert (plugin_root / ".claude-plugin" / "plugin.json").is_file()
|
| 40 |
-
assert (plugin_root / "hooks" / "hooks.json").is_file()
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
def test_plugin_metadata_points_to_upstream_repo() -> None:
|
| 44 |
-
expected_repo = "https://github.com/chopratejas/headroom"
|
| 45 |
-
marketplace = _load_json(".claude-plugin/marketplace.json")
|
| 46 |
-
claude = _load_json("plugins/headroom-agent-hooks/.claude-plugin/plugin.json")
|
| 47 |
-
assert isinstance(marketplace, dict)
|
| 48 |
-
assert isinstance(claude, dict)
|
| 49 |
-
plugin = marketplace["plugins"][0]
|
| 50 |
-
assert plugin["author"]["url"] == expected_repo
|
| 51 |
-
assert plugin["homepage"] == expected_repo
|
| 52 |
-
assert plugin["repository"] == expected_repo
|
| 53 |
-
assert claude["author"]["url"] == expected_repo
|
| 54 |
-
assert claude["homepage"] == expected_repo
|
| 55 |
-
assert claude["repository"] == expected_repo
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
from pathlib import Path
|
| 5 |
+
|
| 6 |
+
REPO_ROOT = Path(__file__).resolve().parents[1]
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
def _load_json(relative_path: str) -> object:
|
| 10 |
+
return json.loads((REPO_ROOT / relative_path).read_text(encoding="utf-8"))
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def test_marketplace_manifests_match() -> None:
|
| 14 |
+
assert _load_json(".claude-plugin/marketplace.json") == _load_json(
|
| 15 |
+
".github/plugin/marketplace.json"
|
| 16 |
+
)
|
| 17 |
+
|
| 18 |
+
|
| 19 |
+
def test_plugin_manifests_share_core_metadata() -> None:
|
| 20 |
+
claude = _load_json("plugins/headroom-agent-hooks/.claude-plugin/plugin.json")
|
| 21 |
+
copilot = _load_json("plugins/headroom-agent-hooks/.github/plugin/plugin.json")
|
| 22 |
+
assert isinstance(claude, dict)
|
| 23 |
+
assert isinstance(copilot, dict)
|
| 24 |
+
for key in ("name", "version", "description", "author", "homepage", "repository", "keywords"):
|
| 25 |
+
assert claude[key] == copilot[key]
|
| 26 |
+
assert "hooks" not in claude
|
| 27 |
+
assert copilot["hooks"] == "./hooks"
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
def test_marketplace_entry_points_to_plugin_root() -> None:
|
| 31 |
+
marketplace = _load_json(".claude-plugin/marketplace.json")
|
| 32 |
+
assert isinstance(marketplace, dict)
|
| 33 |
+
plugins = marketplace["plugins"]
|
| 34 |
+
assert isinstance(plugins, list)
|
| 35 |
+
plugin = plugins[0]
|
| 36 |
+
assert plugin["name"] == "headroom"
|
| 37 |
+
plugin_root = (REPO_ROOT / plugin["source"]).resolve()
|
| 38 |
+
assert plugin_root.is_dir()
|
| 39 |
+
assert (plugin_root / ".claude-plugin" / "plugin.json").is_file()
|
| 40 |
+
assert (plugin_root / "hooks" / "hooks.json").is_file()
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def test_plugin_metadata_points_to_upstream_repo() -> None:
|
| 44 |
+
expected_repo = "https://github.com/chopratejas/headroom"
|
| 45 |
+
marketplace = _load_json(".claude-plugin/marketplace.json")
|
| 46 |
+
claude = _load_json("plugins/headroom-agent-hooks/.claude-plugin/plugin.json")
|
| 47 |
+
assert isinstance(marketplace, dict)
|
| 48 |
+
assert isinstance(claude, dict)
|
| 49 |
+
plugin = marketplace["plugins"][0]
|
| 50 |
+
assert plugin["author"]["url"] == expected_repo
|
| 51 |
+
assert plugin["homepage"] == expected_repo
|
| 52 |
+
assert plugin["repository"] == expected_repo
|
| 53 |
+
assert claude["author"]["url"] == expected_repo
|
| 54 |
+
assert claude["homepage"] == expected_repo
|
| 55 |
+
assert claude["repository"] == expected_repo
|
|
@@ -1,130 +1,130 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
from dataclasses import FrozenInstanceError
|
| 4 |
-
from datetime import date, timedelta
|
| 5 |
-
|
| 6 |
-
import pytest
|
| 7 |
-
|
| 8 |
-
import headroom.pricing as pricing
|
| 9 |
-
from headroom.pricing.anthropic_prices import ANTHROPIC_PRICES, get_anthropic_registry
|
| 10 |
-
from headroom.pricing.openai_prices import OPENAI_PRICES, get_openai_registry
|
| 11 |
-
from headroom.pricing.registry import ModelPricing, PricingRegistry
|
| 12 |
-
|
| 13 |
-
|
| 14 |
-
def test_pricing_public_exports_and_provider_registries() -> None:
|
| 15 |
-
assert pricing.ModelPricing is ModelPricing
|
| 16 |
-
assert pricing.PricingRegistry is PricingRegistry
|
| 17 |
-
assert "get_openai_registry" in pricing.__all__
|
| 18 |
-
assert "get_anthropic_registry" in pricing.__all__
|
| 19 |
-
assert "estimate_cost" in pricing.__all__
|
| 20 |
-
|
| 21 |
-
openai_registry = get_openai_registry()
|
| 22 |
-
anthropic_registry = get_anthropic_registry()
|
| 23 |
-
assert openai_registry.source_url == "https://openai.com/api/pricing/"
|
| 24 |
-
assert anthropic_registry.source_url == "https://www.anthropic.com/pricing"
|
| 25 |
-
assert openai_registry.prices["gpt-4o"] == OPENAI_PRICES["gpt-4o"]
|
| 26 |
-
assert (
|
| 27 |
-
anthropic_registry.prices["claude-3-5-sonnet-20241022"]
|
| 28 |
-
== ANTHROPIC_PRICES["claude-3-5-sonnet-20241022"]
|
| 29 |
-
)
|
| 30 |
-
|
| 31 |
-
openai_registry.prices.pop("gpt-4o")
|
| 32 |
-
anthropic_registry.prices.pop("claude-3-5-sonnet-20241022")
|
| 33 |
-
assert "gpt-4o" in OPENAI_PRICES
|
| 34 |
-
assert "claude-3-5-sonnet-20241022" in ANTHROPIC_PRICES
|
| 35 |
-
|
| 36 |
-
|
| 37 |
-
def test_model_pricing_is_frozen() -> None:
|
| 38 |
-
model = ModelPricing(model="demo", provider="test", input_per_1m=1.5, output_per_1m=2.5)
|
| 39 |
-
with pytest.raises(FrozenInstanceError):
|
| 40 |
-
model.model = "other" # type: ignore[misc]
|
| 41 |
-
|
| 42 |
-
|
| 43 |
-
def test_registry_staleness_and_warning() -> None:
|
| 44 |
-
fresh = PricingRegistry(last_updated=date.today() - timedelta(days=30))
|
| 45 |
-
assert fresh.is_stale() is False
|
| 46 |
-
assert fresh.staleness_warning() is None
|
| 47 |
-
|
| 48 |
-
stale = PricingRegistry(
|
| 49 |
-
last_updated=date.today() - timedelta(days=31),
|
| 50 |
-
source_url="https://example.test/pricing",
|
| 51 |
-
)
|
| 52 |
-
assert stale.is_stale() is True
|
| 53 |
-
assert stale.staleness_warning() == (
|
| 54 |
-
f"Pricing data is 31 days old (last updated: {stale.last_updated})."
|
| 55 |
-
" Please verify at: https://example.test/pricing"
|
| 56 |
-
)
|
| 57 |
-
|
| 58 |
-
|
| 59 |
-
def test_registry_estimate_cost_with_all_token_types() -> None:
|
| 60 |
-
registry = PricingRegistry(
|
| 61 |
-
last_updated=date.today() - timedelta(days=31),
|
| 62 |
-
prices={
|
| 63 |
-
"demo": ModelPricing(
|
| 64 |
-
model="demo",
|
| 65 |
-
provider="test",
|
| 66 |
-
input_per_1m=2.0,
|
| 67 |
-
output_per_1m=4.0,
|
| 68 |
-
cached_input_per_1m=1.0,
|
| 69 |
-
batch_input_per_1m=0.5,
|
| 70 |
-
batch_output_per_1m=0.25,
|
| 71 |
-
)
|
| 72 |
-
},
|
| 73 |
-
)
|
| 74 |
-
|
| 75 |
-
estimate = registry.estimate_cost(
|
| 76 |
-
"demo",
|
| 77 |
-
input_tokens=1_000_000,
|
| 78 |
-
output_tokens=500_000,
|
| 79 |
-
cached_input_tokens=250_000,
|
| 80 |
-
batch_input_tokens=200_000,
|
| 81 |
-
batch_output_tokens=100_000,
|
| 82 |
-
)
|
| 83 |
-
|
| 84 |
-
assert estimate.cost_usd == pytest.approx(4.375)
|
| 85 |
-
assert estimate.breakdown == {
|
| 86 |
-
"input": {"tokens": 1_000_000, "rate_per_1m": 2.0, "cost_usd": 2.0},
|
| 87 |
-
"output": {"tokens": 500_000, "rate_per_1m": 4.0, "cost_usd": 2.0},
|
| 88 |
-
"cached_input": {"tokens": 250_000, "rate_per_1m": 1.0, "cost_usd": 0.25},
|
| 89 |
-
"batch_input": {"tokens": 200_000, "rate_per_1m": 0.5, "cost_usd": 0.1},
|
| 90 |
-
"batch_output": {"tokens": 100_000, "rate_per_1m": 0.25, "cost_usd": 0.025},
|
| 91 |
-
}
|
| 92 |
-
assert estimate.pricing_date == registry.last_updated
|
| 93 |
-
assert estimate.is_stale is True
|
| 94 |
-
assert estimate.warning == (
|
| 95 |
-
f"Pricing data is 31 days old (last updated: {registry.last_updated})."
|
| 96 |
-
)
|
| 97 |
-
|
| 98 |
-
|
| 99 |
-
def test_registry_estimate_cost_zero_usage_returns_empty_breakdown() -> None:
|
| 100 |
-
registry = PricingRegistry(
|
| 101 |
-
last_updated=date.today(),
|
| 102 |
-
prices={
|
| 103 |
-
"demo": ModelPricing(model="demo", provider="test", input_per_1m=1.0, output_per_1m=2.0)
|
| 104 |
-
},
|
| 105 |
-
)
|
| 106 |
-
estimate = registry.estimate_cost("demo")
|
| 107 |
-
assert estimate.cost_usd == 0.0
|
| 108 |
-
assert estimate.breakdown == {}
|
| 109 |
-
assert estimate.is_stale is False
|
| 110 |
-
assert estimate.warning is None
|
| 111 |
-
|
| 112 |
-
|
| 113 |
-
@pytest.mark.parametrize(
|
| 114 |
-
("kwargs", "message"),
|
| 115 |
-
[
|
| 116 |
-
({}, "Model 'missing' not found in registry"),
|
| 117 |
-
({"cached_input_tokens": 1}, "Model 'demo' does not have cached input pricing"),
|
| 118 |
-
({"batch_input_tokens": 1}, "Model 'demo' does not have batch input pricing"),
|
| 119 |
-
({"batch_output_tokens": 1}, "Model 'demo' does not have batch output pricing"),
|
| 120 |
-
],
|
| 121 |
-
)
|
| 122 |
-
def test_registry_estimate_cost_error_paths(kwargs: dict[str, int], message: str) -> None:
|
| 123 |
-
registry = PricingRegistry(
|
| 124 |
-
last_updated=date.today(),
|
| 125 |
-
prices={
|
| 126 |
-
"demo": ModelPricing(model="demo", provider="test", input_per_1m=1.0, output_per_1m=2.0)
|
| 127 |
-
},
|
| 128 |
-
)
|
| 129 |
-
with pytest.raises(ValueError, match=message):
|
| 130 |
-
registry.estimate_cost("missing" if not kwargs else "demo", **kwargs)
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
from dataclasses import FrozenInstanceError
|
| 4 |
+
from datetime import date, timedelta
|
| 5 |
+
|
| 6 |
+
import pytest
|
| 7 |
+
|
| 8 |
+
import headroom.pricing as pricing
|
| 9 |
+
from headroom.pricing.anthropic_prices import ANTHROPIC_PRICES, get_anthropic_registry
|
| 10 |
+
from headroom.pricing.openai_prices import OPENAI_PRICES, get_openai_registry
|
| 11 |
+
from headroom.pricing.registry import ModelPricing, PricingRegistry
|
| 12 |
+
|
| 13 |
+
|
| 14 |
+
def test_pricing_public_exports_and_provider_registries() -> None:
|
| 15 |
+
assert pricing.ModelPricing is ModelPricing
|
| 16 |
+
assert pricing.PricingRegistry is PricingRegistry
|
| 17 |
+
assert "get_openai_registry" in pricing.__all__
|
| 18 |
+
assert "get_anthropic_registry" in pricing.__all__
|
| 19 |
+
assert "estimate_cost" in pricing.__all__
|
| 20 |
+
|
| 21 |
+
openai_registry = get_openai_registry()
|
| 22 |
+
anthropic_registry = get_anthropic_registry()
|
| 23 |
+
assert openai_registry.source_url == "https://openai.com/api/pricing/"
|
| 24 |
+
assert anthropic_registry.source_url == "https://www.anthropic.com/pricing"
|
| 25 |
+
assert openai_registry.prices["gpt-4o"] == OPENAI_PRICES["gpt-4o"]
|
| 26 |
+
assert (
|
| 27 |
+
anthropic_registry.prices["claude-3-5-sonnet-20241022"]
|
| 28 |
+
== ANTHROPIC_PRICES["claude-3-5-sonnet-20241022"]
|
| 29 |
+
)
|
| 30 |
+
|
| 31 |
+
openai_registry.prices.pop("gpt-4o")
|
| 32 |
+
anthropic_registry.prices.pop("claude-3-5-sonnet-20241022")
|
| 33 |
+
assert "gpt-4o" in OPENAI_PRICES
|
| 34 |
+
assert "claude-3-5-sonnet-20241022" in ANTHROPIC_PRICES
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
def test_model_pricing_is_frozen() -> None:
|
| 38 |
+
model = ModelPricing(model="demo", provider="test", input_per_1m=1.5, output_per_1m=2.5)
|
| 39 |
+
with pytest.raises(FrozenInstanceError):
|
| 40 |
+
model.model = "other" # type: ignore[misc]
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
def test_registry_staleness_and_warning() -> None:
|
| 44 |
+
fresh = PricingRegistry(last_updated=date.today() - timedelta(days=30))
|
| 45 |
+
assert fresh.is_stale() is False
|
| 46 |
+
assert fresh.staleness_warning() is None
|
| 47 |
+
|
| 48 |
+
stale = PricingRegistry(
|
| 49 |
+
last_updated=date.today() - timedelta(days=31),
|
| 50 |
+
source_url="https://example.test/pricing",
|
| 51 |
+
)
|
| 52 |
+
assert stale.is_stale() is True
|
| 53 |
+
assert stale.staleness_warning() == (
|
| 54 |
+
f"Pricing data is 31 days old (last updated: {stale.last_updated})."
|
| 55 |
+
" Please verify at: https://example.test/pricing"
|
| 56 |
+
)
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def test_registry_estimate_cost_with_all_token_types() -> None:
|
| 60 |
+
registry = PricingRegistry(
|
| 61 |
+
last_updated=date.today() - timedelta(days=31),
|
| 62 |
+
prices={
|
| 63 |
+
"demo": ModelPricing(
|
| 64 |
+
model="demo",
|
| 65 |
+
provider="test",
|
| 66 |
+
input_per_1m=2.0,
|
| 67 |
+
output_per_1m=4.0,
|
| 68 |
+
cached_input_per_1m=1.0,
|
| 69 |
+
batch_input_per_1m=0.5,
|
| 70 |
+
batch_output_per_1m=0.25,
|
| 71 |
+
)
|
| 72 |
+
},
|
| 73 |
+
)
|
| 74 |
+
|
| 75 |
+
estimate = registry.estimate_cost(
|
| 76 |
+
"demo",
|
| 77 |
+
input_tokens=1_000_000,
|
| 78 |
+
output_tokens=500_000,
|
| 79 |
+
cached_input_tokens=250_000,
|
| 80 |
+
batch_input_tokens=200_000,
|
| 81 |
+
batch_output_tokens=100_000,
|
| 82 |
+
)
|
| 83 |
+
|
| 84 |
+
assert estimate.cost_usd == pytest.approx(4.375)
|
| 85 |
+
assert estimate.breakdown == {
|
| 86 |
+
"input": {"tokens": 1_000_000, "rate_per_1m": 2.0, "cost_usd": 2.0},
|
| 87 |
+
"output": {"tokens": 500_000, "rate_per_1m": 4.0, "cost_usd": 2.0},
|
| 88 |
+
"cached_input": {"tokens": 250_000, "rate_per_1m": 1.0, "cost_usd": 0.25},
|
| 89 |
+
"batch_input": {"tokens": 200_000, "rate_per_1m": 0.5, "cost_usd": 0.1},
|
| 90 |
+
"batch_output": {"tokens": 100_000, "rate_per_1m": 0.25, "cost_usd": 0.025},
|
| 91 |
+
}
|
| 92 |
+
assert estimate.pricing_date == registry.last_updated
|
| 93 |
+
assert estimate.is_stale is True
|
| 94 |
+
assert estimate.warning == (
|
| 95 |
+
f"Pricing data is 31 days old (last updated: {registry.last_updated})."
|
| 96 |
+
)
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
def test_registry_estimate_cost_zero_usage_returns_empty_breakdown() -> None:
|
| 100 |
+
registry = PricingRegistry(
|
| 101 |
+
last_updated=date.today(),
|
| 102 |
+
prices={
|
| 103 |
+
"demo": ModelPricing(model="demo", provider="test", input_per_1m=1.0, output_per_1m=2.0)
|
| 104 |
+
},
|
| 105 |
+
)
|
| 106 |
+
estimate = registry.estimate_cost("demo")
|
| 107 |
+
assert estimate.cost_usd == 0.0
|
| 108 |
+
assert estimate.breakdown == {}
|
| 109 |
+
assert estimate.is_stale is False
|
| 110 |
+
assert estimate.warning is None
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
@pytest.mark.parametrize(
|
| 114 |
+
("kwargs", "message"),
|
| 115 |
+
[
|
| 116 |
+
({}, "Model 'missing' not found in registry"),
|
| 117 |
+
({"cached_input_tokens": 1}, "Model 'demo' does not have cached input pricing"),
|
| 118 |
+
({"batch_input_tokens": 1}, "Model 'demo' does not have batch input pricing"),
|
| 119 |
+
({"batch_output_tokens": 1}, "Model 'demo' does not have batch output pricing"),
|
| 120 |
+
],
|
| 121 |
+
)
|
| 122 |
+
def test_registry_estimate_cost_error_paths(kwargs: dict[str, int], message: str) -> None:
|
| 123 |
+
registry = PricingRegistry(
|
| 124 |
+
last_updated=date.today(),
|
| 125 |
+
prices={
|
| 126 |
+
"demo": ModelPricing(model="demo", provider="test", input_per_1m=1.0, output_per_1m=2.0)
|
| 127 |
+
},
|
| 128 |
+
)
|
| 129 |
+
with pytest.raises(ValueError, match=message):
|
| 130 |
+
registry.estimate_cost("missing" if not kwargs else "demo", **kwargs)
|
|
@@ -1,97 +1,97 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
from types import SimpleNamespace
|
| 4 |
-
|
| 5 |
-
from headroom.pricing import litellm_pricing
|
| 6 |
-
|
| 7 |
-
|
| 8 |
-
def test_litellm_helpers_when_dependency_is_unavailable(monkeypatch) -> None:
|
| 9 |
-
monkeypatch.setattr(litellm_pricing, "LITELLM_AVAILABLE", False)
|
| 10 |
-
monkeypatch.setattr(litellm_pricing, "litellm", None)
|
| 11 |
-
|
| 12 |
-
assert litellm_pricing.get_litellm_model_cost() == {}
|
| 13 |
-
assert litellm_pricing.get_model_pricing("gpt-4o") is None
|
| 14 |
-
assert litellm_pricing.estimate_cost("gpt-4o", input_tokens=1, output_tokens=1) is None
|
| 15 |
-
assert litellm_pricing.list_available_models() == []
|
| 16 |
-
|
| 17 |
-
|
| 18 |
-
def test_litellm_model_pricing_exact_match_and_defaults(monkeypatch) -> None:
|
| 19 |
-
fake_litellm = SimpleNamespace(
|
| 20 |
-
model_cost={
|
| 21 |
-
"gpt-4o": {
|
| 22 |
-
"input_cost_per_token": 0.0000025,
|
| 23 |
-
"output_cost_per_token": 0.00001,
|
| 24 |
-
"max_tokens": 128000,
|
| 25 |
-
}
|
| 26 |
-
}
|
| 27 |
-
)
|
| 28 |
-
monkeypatch.setattr(litellm_pricing, "LITELLM_AVAILABLE", True)
|
| 29 |
-
monkeypatch.setattr(litellm_pricing, "litellm", fake_litellm)
|
| 30 |
-
|
| 31 |
-
assert litellm_pricing.get_litellm_model_cost() == fake_litellm.model_cost
|
| 32 |
-
pricing = litellm_pricing.get_model_pricing("gpt-4o")
|
| 33 |
-
assert pricing is not None
|
| 34 |
-
assert pricing.model == "gpt-4o"
|
| 35 |
-
assert pricing.input_cost_per_1m == 2.5
|
| 36 |
-
assert pricing.output_cost_per_1m == 10.0
|
| 37 |
-
assert pricing.max_tokens == 128000
|
| 38 |
-
assert pricing.max_input_tokens is None
|
| 39 |
-
assert pricing.max_output_tokens is None
|
| 40 |
-
assert pricing.supports_vision is False
|
| 41 |
-
assert pricing.supports_function_calling is False
|
| 42 |
-
assert (
|
| 43 |
-
litellm_pricing.estimate_cost("gpt-4o", input_tokens=200_000, output_tokens=300_000) == 3.5
|
| 44 |
-
)
|
| 45 |
-
assert litellm_pricing.list_available_models() == ["gpt-4o"]
|
| 46 |
-
|
| 47 |
-
|
| 48 |
-
def test_litellm_model_pricing_uses_provider_prefixes(monkeypatch) -> None:
|
| 49 |
-
fake_litellm = SimpleNamespace(
|
| 50 |
-
model_cost={
|
| 51 |
-
"openai/gpt-4o-mini": {
|
| 52 |
-
"input_cost_per_token": 0.00000015,
|
| 53 |
-
"output_cost_per_token": 0.0000006,
|
| 54 |
-
"supports_vision": True,
|
| 55 |
-
"supports_function_calling": True,
|
| 56 |
-
"max_input_tokens": 64000,
|
| 57 |
-
"max_output_tokens": 16000,
|
| 58 |
-
}
|
| 59 |
-
}
|
| 60 |
-
)
|
| 61 |
-
monkeypatch.setattr(litellm_pricing, "LITELLM_AVAILABLE", True)
|
| 62 |
-
monkeypatch.setattr(litellm_pricing, "litellm", fake_litellm)
|
| 63 |
-
|
| 64 |
-
pricing = litellm_pricing.get_model_pricing("gpt-4o-mini")
|
| 65 |
-
assert pricing is not None
|
| 66 |
-
assert pricing.input_cost_per_1m == 0.15
|
| 67 |
-
assert pricing.output_cost_per_1m == 0.6
|
| 68 |
-
assert pricing.max_input_tokens == 64000
|
| 69 |
-
assert pricing.max_output_tokens == 16000
|
| 70 |
-
assert pricing.supports_vision is True
|
| 71 |
-
assert pricing.supports_function_calling is True
|
| 72 |
-
|
| 73 |
-
|
| 74 |
-
def test_litellm_model_pricing_uses_aliases_and_zero_cost_defaults(monkeypatch) -> None:
|
| 75 |
-
fake_litellm = SimpleNamespace(
|
| 76 |
-
model_cost={
|
| 77 |
-
"claude-sonnet-4-20250514": {
|
| 78 |
-
"input_cost_per_token": None,
|
| 79 |
-
"output_cost_per_token": None,
|
| 80 |
-
}
|
| 81 |
-
}
|
| 82 |
-
)
|
| 83 |
-
monkeypatch.setattr(litellm_pricing, "LITELLM_AVAILABLE", True)
|
| 84 |
-
monkeypatch.setattr(litellm_pricing, "litellm", fake_litellm)
|
| 85 |
-
|
| 86 |
-
pricing = litellm_pricing.get_model_pricing("claude-3-5-sonnet-20241022")
|
| 87 |
-
assert pricing is not None
|
| 88 |
-
assert pricing.model == "claude-3-5-sonnet-20241022"
|
| 89 |
-
assert pricing.input_cost_per_1m == 0
|
| 90 |
-
assert pricing.output_cost_per_1m == 0
|
| 91 |
-
assert litellm_pricing.estimate_cost("claude-3-5-sonnet-20241022", input_tokens=1) == 0
|
| 92 |
-
|
| 93 |
-
|
| 94 |
-
def test_litellm_model_pricing_returns_none_for_unknown_models(monkeypatch) -> None:
|
| 95 |
-
monkeypatch.setattr(litellm_pricing, "LITELLM_AVAILABLE", True)
|
| 96 |
-
monkeypatch.setattr(litellm_pricing, "litellm", SimpleNamespace(model_cost={}))
|
| 97 |
-
assert litellm_pricing.get_model_pricing("missing") is None
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
from types import SimpleNamespace
|
| 4 |
+
|
| 5 |
+
from headroom.pricing import litellm_pricing
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
def test_litellm_helpers_when_dependency_is_unavailable(monkeypatch) -> None:
|
| 9 |
+
monkeypatch.setattr(litellm_pricing, "LITELLM_AVAILABLE", False)
|
| 10 |
+
monkeypatch.setattr(litellm_pricing, "litellm", None)
|
| 11 |
+
|
| 12 |
+
assert litellm_pricing.get_litellm_model_cost() == {}
|
| 13 |
+
assert litellm_pricing.get_model_pricing("gpt-4o") is None
|
| 14 |
+
assert litellm_pricing.estimate_cost("gpt-4o", input_tokens=1, output_tokens=1) is None
|
| 15 |
+
assert litellm_pricing.list_available_models() == []
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
def test_litellm_model_pricing_exact_match_and_defaults(monkeypatch) -> None:
|
| 19 |
+
fake_litellm = SimpleNamespace(
|
| 20 |
+
model_cost={
|
| 21 |
+
"gpt-4o": {
|
| 22 |
+
"input_cost_per_token": 0.0000025,
|
| 23 |
+
"output_cost_per_token": 0.00001,
|
| 24 |
+
"max_tokens": 128000,
|
| 25 |
+
}
|
| 26 |
+
}
|
| 27 |
+
)
|
| 28 |
+
monkeypatch.setattr(litellm_pricing, "LITELLM_AVAILABLE", True)
|
| 29 |
+
monkeypatch.setattr(litellm_pricing, "litellm", fake_litellm)
|
| 30 |
+
|
| 31 |
+
assert litellm_pricing.get_litellm_model_cost() == fake_litellm.model_cost
|
| 32 |
+
pricing = litellm_pricing.get_model_pricing("gpt-4o")
|
| 33 |
+
assert pricing is not None
|
| 34 |
+
assert pricing.model == "gpt-4o"
|
| 35 |
+
assert pricing.input_cost_per_1m == 2.5
|
| 36 |
+
assert pricing.output_cost_per_1m == 10.0
|
| 37 |
+
assert pricing.max_tokens == 128000
|
| 38 |
+
assert pricing.max_input_tokens is None
|
| 39 |
+
assert pricing.max_output_tokens is None
|
| 40 |
+
assert pricing.supports_vision is False
|
| 41 |
+
assert pricing.supports_function_calling is False
|
| 42 |
+
assert (
|
| 43 |
+
litellm_pricing.estimate_cost("gpt-4o", input_tokens=200_000, output_tokens=300_000) == 3.5
|
| 44 |
+
)
|
| 45 |
+
assert litellm_pricing.list_available_models() == ["gpt-4o"]
|
| 46 |
+
|
| 47 |
+
|
| 48 |
+
def test_litellm_model_pricing_uses_provider_prefixes(monkeypatch) -> None:
|
| 49 |
+
fake_litellm = SimpleNamespace(
|
| 50 |
+
model_cost={
|
| 51 |
+
"openai/gpt-4o-mini": {
|
| 52 |
+
"input_cost_per_token": 0.00000015,
|
| 53 |
+
"output_cost_per_token": 0.0000006,
|
| 54 |
+
"supports_vision": True,
|
| 55 |
+
"supports_function_calling": True,
|
| 56 |
+
"max_input_tokens": 64000,
|
| 57 |
+
"max_output_tokens": 16000,
|
| 58 |
+
}
|
| 59 |
+
}
|
| 60 |
+
)
|
| 61 |
+
monkeypatch.setattr(litellm_pricing, "LITELLM_AVAILABLE", True)
|
| 62 |
+
monkeypatch.setattr(litellm_pricing, "litellm", fake_litellm)
|
| 63 |
+
|
| 64 |
+
pricing = litellm_pricing.get_model_pricing("gpt-4o-mini")
|
| 65 |
+
assert pricing is not None
|
| 66 |
+
assert pricing.input_cost_per_1m == 0.15
|
| 67 |
+
assert pricing.output_cost_per_1m == 0.6
|
| 68 |
+
assert pricing.max_input_tokens == 64000
|
| 69 |
+
assert pricing.max_output_tokens == 16000
|
| 70 |
+
assert pricing.supports_vision is True
|
| 71 |
+
assert pricing.supports_function_calling is True
|
| 72 |
+
|
| 73 |
+
|
| 74 |
+
def test_litellm_model_pricing_uses_aliases_and_zero_cost_defaults(monkeypatch) -> None:
|
| 75 |
+
fake_litellm = SimpleNamespace(
|
| 76 |
+
model_cost={
|
| 77 |
+
"claude-sonnet-4-20250514": {
|
| 78 |
+
"input_cost_per_token": None,
|
| 79 |
+
"output_cost_per_token": None,
|
| 80 |
+
}
|
| 81 |
+
}
|
| 82 |
+
)
|
| 83 |
+
monkeypatch.setattr(litellm_pricing, "LITELLM_AVAILABLE", True)
|
| 84 |
+
monkeypatch.setattr(litellm_pricing, "litellm", fake_litellm)
|
| 85 |
+
|
| 86 |
+
pricing = litellm_pricing.get_model_pricing("claude-3-5-sonnet-20241022")
|
| 87 |
+
assert pricing is not None
|
| 88 |
+
assert pricing.model == "claude-3-5-sonnet-20241022"
|
| 89 |
+
assert pricing.input_cost_per_1m == 0
|
| 90 |
+
assert pricing.output_cost_per_1m == 0
|
| 91 |
+
assert litellm_pricing.estimate_cost("claude-3-5-sonnet-20241022", input_tokens=1) == 0
|
| 92 |
+
|
| 93 |
+
|
| 94 |
+
def test_litellm_model_pricing_returns_none_for_unknown_models(monkeypatch) -> None:
|
| 95 |
+
monkeypatch.setattr(litellm_pricing, "LITELLM_AVAILABLE", True)
|
| 96 |
+
monkeypatch.setattr(litellm_pricing, "litellm", SimpleNamespace(model_cost={}))
|
| 97 |
+
assert litellm_pricing.get_model_pricing("missing") is None
|
|
@@ -1,33 +1,33 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
from headroom.providers.aider.install import build_install_env
|
| 4 |
-
from headroom.providers.aider.runtime import build_launch_env
|
| 5 |
-
|
| 6 |
-
|
| 7 |
-
def test_aider_build_launch_env_sets_proxy_urls_without_mutating_input() -> None:
|
| 8 |
-
# Arrange
|
| 9 |
-
source_env = {"EXISTING": "value"}
|
| 10 |
-
|
| 11 |
-
# Act
|
| 12 |
-
env, lines = build_launch_env(port=9999, environ=source_env)
|
| 13 |
-
|
| 14 |
-
# Assert
|
| 15 |
-
assert source_env == {"EXISTING": "value"}
|
| 16 |
-
assert env["EXISTING"] == "value"
|
| 17 |
-
assert env["OPENAI_API_BASE"] == "http://127.0.0.1:9999/v1"
|
| 18 |
-
assert env["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:9999"
|
| 19 |
-
assert lines == [
|
| 20 |
-
"OPENAI_API_BASE=http://127.0.0.1:9999/v1",
|
| 21 |
-
"ANTHROPIC_BASE_URL=http://127.0.0.1:9999",
|
| 22 |
-
]
|
| 23 |
-
|
| 24 |
-
|
| 25 |
-
def test_aider_build_install_env_returns_only_persistent_proxy_variables() -> None:
|
| 26 |
-
# Arrange / Act
|
| 27 |
-
env = build_install_env(port=8787, backend="ignored")
|
| 28 |
-
|
| 29 |
-
# Assert
|
| 30 |
-
assert env == {
|
| 31 |
-
"OPENAI_API_BASE": "http://127.0.0.1:8787/v1",
|
| 32 |
-
"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787",
|
| 33 |
-
}
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
from headroom.providers.aider.install import build_install_env
|
| 4 |
+
from headroom.providers.aider.runtime import build_launch_env
|
| 5 |
+
|
| 6 |
+
|
| 7 |
+
def test_aider_build_launch_env_sets_proxy_urls_without_mutating_input() -> None:
|
| 8 |
+
# Arrange
|
| 9 |
+
source_env = {"EXISTING": "value"}
|
| 10 |
+
|
| 11 |
+
# Act
|
| 12 |
+
env, lines = build_launch_env(port=9999, environ=source_env)
|
| 13 |
+
|
| 14 |
+
# Assert
|
| 15 |
+
assert source_env == {"EXISTING": "value"}
|
| 16 |
+
assert env["EXISTING"] == "value"
|
| 17 |
+
assert env["OPENAI_API_BASE"] == "http://127.0.0.1:9999/v1"
|
| 18 |
+
assert env["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:9999"
|
| 19 |
+
assert lines == [
|
| 20 |
+
"OPENAI_API_BASE=http://127.0.0.1:9999/v1",
|
| 21 |
+
"ANTHROPIC_BASE_URL=http://127.0.0.1:9999",
|
| 22 |
+
]
|
| 23 |
+
|
| 24 |
+
|
| 25 |
+
def test_aider_build_install_env_returns_only_persistent_proxy_variables() -> None:
|
| 26 |
+
# Arrange / Act
|
| 27 |
+
env = build_install_env(port=8787, backend="ignored")
|
| 28 |
+
|
| 29 |
+
# Assert
|
| 30 |
+
assert env == {
|
| 31 |
+
"OPENAI_API_BASE": "http://127.0.0.1:8787/v1",
|
| 32 |
+
"ANTHROPIC_BASE_URL": "http://127.0.0.1:8787",
|
| 33 |
+
}
|
|
@@ -1,9 +1,9 @@
|
|
| 1 |
-
from __future__ import annotations
|
| 2 |
-
|
| 3 |
-
from headroom.providers.claude import DEFAULT_API_URL, proxy_base_url
|
| 4 |
-
|
| 5 |
-
|
| 6 |
-
def test_claude_runtime_exposes_default_api_and_local_proxy_url() -> None:
|
| 7 |
-
# Arrange / Act / Assert
|
| 8 |
-
assert DEFAULT_API_URL == "https://api.anthropic.com"
|
| 9 |
-
assert proxy_base_url(4321) == "http://127.0.0.1:4321"
|
|
|
|
| 1 |
+
from __future__ import annotations
|
| 2 |
+
|
| 3 |
+
from headroom.providers.claude import DEFAULT_API_URL, proxy_base_url
|
| 4 |
+
|
| 5 |
+
|
| 6 |
+
def test_claude_runtime_exposes_default_api_and_local_proxy_url() -> None:
|
| 7 |
+
# Arrange / Act / Assert
|
| 8 |
+
assert DEFAULT_API_URL == "https://api.anthropic.com"
|
| 9 |
+
assert proxy_base_url(4321) == "http://127.0.0.1:4321"
|