Garm Claude Opus 4.7 (1M context) commited on
Commit
9d5bc33
·
1 Parent(s): 3c23d80

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>

This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. headroom/cli/wrap.py +0 -0
  2. headroom/compress.py +347 -347
  3. headroom/copilot_auth.py +444 -444
  4. headroom/install/health.py +28 -28
  5. headroom/install/providers.py +174 -174
  6. headroom/install/runtime.py +279 -279
  7. headroom/providers/aider/install.py +12 -12
  8. headroom/providers/claude/install.py +63 -63
  9. headroom/providers/codex/install.py +68 -68
  10. headroom/providers/copilot/install.py +25 -25
  11. headroom/providers/cursor/install.py +15 -15
  12. headroom/providers/install_registry.py +86 -86
  13. headroom/providers/openclaw/install.py +50 -50
  14. headroom/proxy/handlers/openai.py +0 -0
  15. headroom/proxy/server.py +0 -0
  16. headroom/release_version.py +310 -310
  17. headroom/subscription/__init__.py +72 -72
  18. headroom/subscription/base.py +230 -230
  19. headroom/subscription/client.py +131 -131
  20. headroom/subscription/codex_rate_limits.py +247 -247
  21. headroom/subscription/copilot_quota.py +366 -366
  22. headroom/subscription/models.py +395 -395
  23. headroom/subscription/session_tracking.py +189 -189
  24. headroom/subscription/tracker.py +464 -464
  25. headroom/transforms/content_router.py +0 -0
  26. scripts/changelog-gen.py +203 -203
  27. scripts/sync-plugin-versions.py +55 -55
  28. scripts/tests/test_changelog_gen.py +302 -302
  29. scripts/tests/test_sync_plugin_versions.py +68 -68
  30. tests/test_backend_anyllm.py +384 -384
  31. tests/test_ccr_batch_store.py +126 -126
  32. tests/test_ccr_response_handler_extra.py +372 -372
  33. tests/test_cli/test_wrap_copilot.py +335 -335
  34. tests/test_cli_learn.py +266 -266
  35. tests/test_codex_rate_limits.py +227 -227
  36. tests/test_compress_api.py +274 -274
  37. tests/test_compress_failure.py +39 -39
  38. tests/test_copilot_quota.py +331 -331
  39. tests/test_evals_datasets.py +538 -538
  40. tests/test_evals_metrics.py +132 -132
  41. tests/test_exceptions.py +40 -40
  42. tests/test_graph.py +352 -352
  43. tests/test_install/test_paths.py +68 -68
  44. tests/test_install/test_runtime.py +468 -468
  45. tests/test_install/test_supervisors.py +471 -471
  46. tests/test_plugin_manifests.py +55 -55
  47. tests/test_pricing.py +130 -130
  48. tests/test_pricing_litellm.py +97 -97
  49. tests/test_provider_aider.py +33 -33
  50. tests/test_provider_claude.py +9 -9
headroom/cli/wrap.py CHANGED
The diff for this file is too large to render. See raw diff
 
headroom/compress.py CHANGED
@@ -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
headroom/copilot_auth.py CHANGED
@@ -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
headroom/install/health.py CHANGED
@@ -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")
headroom/install/providers.py CHANGED
@@ -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)
headroom/install/runtime.py CHANGED
@@ -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"
headroom/providers/aider/install.py CHANGED
@@ -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")}
headroom/providers/claude/install.py CHANGED
@@ -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")
headroom/providers/codex/install.py CHANGED
@@ -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")
headroom/providers/copilot/install.py CHANGED
@@ -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
+ }
headroom/providers/cursor/install.py CHANGED
@@ -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
+ }
headroom/providers/install_registry.py CHANGED
@@ -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)
headroom/providers/openclaw/install.py CHANGED
@@ -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)
headroom/proxy/handlers/openai.py CHANGED
The diff for this file is too large to render. See raw diff
 
headroom/proxy/server.py CHANGED
The diff for this file is too large to render. See raw diff
 
headroom/release_version.py CHANGED
@@ -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()
headroom/subscription/__init__.py CHANGED
@@ -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
+ ]
headroom/subscription/base.py CHANGED
@@ -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()
headroom/subscription/client.py CHANGED
@@ -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
headroom/subscription/codex_rate_limits.py CHANGED
@@ -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
headroom/subscription/copilot_quota.py CHANGED
@@ -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
headroom/subscription/models.py CHANGED
@@ -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
headroom/subscription/session_tracking.py CHANGED
@@ -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
+ }
headroom/subscription/tracker.py CHANGED
@@ -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()
headroom/transforms/content_router.py CHANGED
The diff for this file is too large to render. See raw diff
 
scripts/changelog-gen.py CHANGED
@@ -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()
scripts/sync-plugin-versions.py CHANGED
@@ -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()
scripts/tests/test_changelog_gen.py CHANGED
@@ -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"])
scripts/tests/test_sync_plugin_versions.py CHANGED
@@ -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
+ ]
tests/test_backend_anyllm.py CHANGED
@@ -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)
tests/test_ccr_batch_store.py CHANGED
@@ -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
tests/test_ccr_response_handler_extra.py CHANGED
@@ -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"]
tests/test_cli/test_wrap_copilot.py CHANGED
@@ -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
tests/test_cli_learn.py CHANGED
@@ -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
tests/test_codex_rate_limits.py CHANGED
@@ -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
tests/test_compress_api.py CHANGED
@@ -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
tests/test_compress_failure.py CHANGED
@@ -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
+ ]
tests/test_copilot_quota.py CHANGED
@@ -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
+ )
tests/test_evals_datasets.py CHANGED
@@ -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
tests/test_evals_metrics.py CHANGED
@@ -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
tests/test_exceptions.py CHANGED
@@ -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__})"
tests/test_graph.py CHANGED
@@ -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()
tests/test_install/test_paths.py CHANGED
@@ -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"
tests/test_install/test_runtime.py CHANGED
@@ -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"
tests/test_install/test_supervisors.py CHANGED
@@ -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
+ ]
tests/test_plugin_manifests.py CHANGED
@@ -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
tests/test_pricing.py CHANGED
@@ -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)
tests/test_pricing_litellm.py CHANGED
@@ -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
tests/test_provider_aider.py CHANGED
@@ -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
+ }
tests/test_provider_claude.py CHANGED
@@ -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"