stravinsky 0.2.52__py3-none-any.whl → 0.4.18__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of stravinsky might be problematic. Click here for more details.

Files changed (58) hide show
  1. mcp_bridge/__init__.py +1 -1
  2. mcp_bridge/auth/token_store.py +113 -11
  3. mcp_bridge/cli/__init__.py +6 -0
  4. mcp_bridge/cli/install_hooks.py +1265 -0
  5. mcp_bridge/cli/session_report.py +585 -0
  6. mcp_bridge/config/MANIFEST_SCHEMA.md +305 -0
  7. mcp_bridge/config/README.md +276 -0
  8. mcp_bridge/config/hook_config.py +249 -0
  9. mcp_bridge/config/hooks_manifest.json +138 -0
  10. mcp_bridge/config/rate_limits.py +222 -0
  11. mcp_bridge/config/skills_manifest.json +128 -0
  12. mcp_bridge/hooks/HOOKS_SETTINGS.json +175 -0
  13. mcp_bridge/hooks/README.md +215 -0
  14. mcp_bridge/hooks/__init__.py +119 -60
  15. mcp_bridge/hooks/edit_recovery.py +42 -37
  16. mcp_bridge/hooks/git_noninteractive.py +89 -0
  17. mcp_bridge/hooks/keyword_detector.py +30 -0
  18. mcp_bridge/hooks/manager.py +8 -0
  19. mcp_bridge/hooks/notification_hook.py +103 -0
  20. mcp_bridge/hooks/parallel_execution.py +111 -0
  21. mcp_bridge/hooks/pre_compact.py +82 -183
  22. mcp_bridge/hooks/rules_injector.py +507 -0
  23. mcp_bridge/hooks/session_notifier.py +125 -0
  24. mcp_bridge/{native_hooks → hooks}/stravinsky_mode.py +51 -16
  25. mcp_bridge/hooks/subagent_stop.py +98 -0
  26. mcp_bridge/hooks/task_validator.py +73 -0
  27. mcp_bridge/hooks/tmux_manager.py +141 -0
  28. mcp_bridge/hooks/todo_continuation.py +90 -0
  29. mcp_bridge/hooks/todo_delegation.py +88 -0
  30. mcp_bridge/hooks/tool_messaging.py +267 -0
  31. mcp_bridge/hooks/truncator.py +21 -17
  32. mcp_bridge/notifications.py +151 -0
  33. mcp_bridge/prompts/multimodal.py +24 -3
  34. mcp_bridge/server.py +214 -49
  35. mcp_bridge/server_tools.py +445 -0
  36. mcp_bridge/tools/__init__.py +22 -18
  37. mcp_bridge/tools/agent_manager.py +220 -32
  38. mcp_bridge/tools/code_search.py +97 -11
  39. mcp_bridge/tools/lsp/__init__.py +7 -0
  40. mcp_bridge/tools/lsp/manager.py +448 -0
  41. mcp_bridge/tools/lsp/tools.py +637 -150
  42. mcp_bridge/tools/model_invoke.py +208 -106
  43. mcp_bridge/tools/query_classifier.py +323 -0
  44. mcp_bridge/tools/semantic_search.py +3042 -0
  45. mcp_bridge/tools/templates.py +32 -18
  46. mcp_bridge/update_manager.py +589 -0
  47. mcp_bridge/update_manager_pypi.py +299 -0
  48. stravinsky-0.4.18.dist-info/METADATA +468 -0
  49. stravinsky-0.4.18.dist-info/RECORD +88 -0
  50. stravinsky-0.4.18.dist-info/entry_points.txt +5 -0
  51. mcp_bridge/native_hooks/edit_recovery.py +0 -46
  52. mcp_bridge/native_hooks/todo_delegation.py +0 -54
  53. mcp_bridge/native_hooks/truncator.py +0 -23
  54. stravinsky-0.2.52.dist-info/METADATA +0 -204
  55. stravinsky-0.2.52.dist-info/RECORD +0 -63
  56. stravinsky-0.2.52.dist-info/entry_points.txt +0 -3
  57. /mcp_bridge/{native_hooks → hooks}/context.py +0 -0
  58. {stravinsky-0.2.52.dist-info → stravinsky-0.4.18.dist-info}/WHEEL +0 -0
@@ -0,0 +1,249 @@
1
+ """
2
+ Hook configuration with selective disabling support.
3
+
4
+ Provides batteries-included defaults with user-configurable overrides.
5
+ Users can disable specific hooks via ~/.stravinsky/disable_hooks.txt
6
+ """
7
+
8
+ import os
9
+ from pathlib import Path
10
+ from typing import Set, Optional
11
+ import logging
12
+
13
+ logger = logging.getLogger(__name__)
14
+
15
+ # Default locations for disable hooks config
16
+ DISABLE_HOOKS_PATHS = [
17
+ Path.home() / ".stravinsky" / "disable_hooks.txt",
18
+ Path(".stravinsky") / "disable_hooks.txt",
19
+ Path(".claude") / "disable_hooks.txt",
20
+ ]
21
+
22
+
23
+ def get_disabled_hooks() -> Set[str]:
24
+ """
25
+ Load disabled hooks from config files.
26
+
27
+ Checks (in order):
28
+ 1. ~/.stravinsky/disable_hooks.txt (user global)
29
+ 2. .stravinsky/disable_hooks.txt (project local)
30
+ 3. .claude/disable_hooks.txt (claude project local)
31
+
32
+ Returns:
33
+ Set of hook names that should be disabled.
34
+ """
35
+ disabled = set()
36
+
37
+ for path in DISABLE_HOOKS_PATHS:
38
+ if path.exists():
39
+ try:
40
+ content = path.read_text()
41
+ for line in content.splitlines():
42
+ line = line.strip()
43
+ # Skip comments and empty lines
44
+ if line and not line.startswith("#"):
45
+ disabled.add(line)
46
+ logger.debug(f"Loaded disabled hooks from {path}: {disabled}")
47
+ except Exception as e:
48
+ logger.warning(f"Failed to read {path}: {e}")
49
+
50
+ return disabled
51
+
52
+
53
+ def is_hook_enabled(hook_name: str) -> bool:
54
+ """
55
+ Check if a specific hook is enabled.
56
+
57
+ Args:
58
+ hook_name: Name of the hook (e.g., 'comment_checker', 'session_recovery')
59
+
60
+ Returns:
61
+ True if the hook is enabled (not in disable list), False otherwise.
62
+ """
63
+ disabled = get_disabled_hooks()
64
+ return hook_name not in disabled
65
+
66
+
67
+ def get_hook_config_path() -> Path:
68
+ """
69
+ Get the path to the user's hook config directory.
70
+ Creates it if it doesn't exist.
71
+ """
72
+ config_dir = Path.home() / ".stravinsky"
73
+ config_dir.mkdir(parents=True, exist_ok=True)
74
+ return config_dir
75
+
76
+
77
+ def create_sample_disable_hooks() -> Optional[Path]:
78
+ """
79
+ Create a sample disable_hooks.txt file with documentation.
80
+
81
+ Returns:
82
+ Path to the created file, or None if it already exists.
83
+ """
84
+ config_dir = get_hook_config_path()
85
+ disable_file = config_dir / "disable_hooks.txt"
86
+
87
+ if disable_file.exists():
88
+ return None
89
+
90
+ sample_content = """# Stravinsky Hook Disabling Configuration
91
+ # Add hook names (one per line) to disable them.
92
+ # Lines starting with # are comments.
93
+ #
94
+ # Available hooks:
95
+ # ================
96
+ #
97
+ # PreToolUse Hooks:
98
+ # - comment_checker (checks git commit comments for quality)
99
+ # - stravinsky_mode (blocks direct tool calls, forces delegation)
100
+ # - notification_hook (displays agent spawn notifications)
101
+ #
102
+ # PostToolUse Hooks:
103
+ # - session_recovery (detects API errors and logs recovery info)
104
+ # - parallel_execution (injects parallel execution instructions)
105
+ # - todo_delegation (enforces parallel Task spawning for todos)
106
+ # - tool_messaging (user-friendly MCP tool messages)
107
+ # - edit_recovery (suggests recovery for Edit failures)
108
+ # - truncator (truncates long responses)
109
+ # - subagent_stop (handles subagent completion)
110
+ #
111
+ # UserPromptSubmit Hooks:
112
+ # - context (injects CLAUDE.md content)
113
+ # - todo_continuation (reminds about incomplete todos)
114
+ #
115
+ # PreCompact Hooks:
116
+ # - pre_compact (preserves critical context before compaction)
117
+ #
118
+ # Example - to disable the comment checker:
119
+ # comment_checker
120
+ #
121
+ # Example - to disable ultrawork mode detection:
122
+ # parallel_execution
123
+ """
124
+
125
+ disable_file.write_text(sample_content)
126
+ logger.info(f"Created sample disable_hooks.txt at {disable_file}")
127
+ return disable_file
128
+
129
+
130
+ # Hook metadata for batteries-included config
131
+ HOOK_DEFAULTS = {
132
+ # PreToolUse hooks
133
+ "comment_checker": {
134
+ "type": "PreToolUse",
135
+ "description": "Checks git commit comments for quality issues",
136
+ "default_enabled": True,
137
+ "exit_on_block": 0, # Warn but don't block
138
+ },
139
+ "stravinsky_mode": {
140
+ "type": "PreToolUse",
141
+ "description": "Blocks direct tool calls, forces Task delegation",
142
+ "default_enabled": True,
143
+ "exit_on_block": 2, # Hard block
144
+ },
145
+ "notification_hook": {
146
+ "type": "PreToolUse",
147
+ "description": "Displays agent spawn notifications",
148
+ "default_enabled": True,
149
+ "exit_on_block": 0,
150
+ },
151
+ # PostToolUse hooks
152
+ "session_recovery": {
153
+ "type": "PostToolUse",
154
+ "description": "Detects API errors and logs recovery suggestions",
155
+ "default_enabled": True,
156
+ "exit_on_block": 0,
157
+ },
158
+ "parallel_execution": {
159
+ "type": "PostToolUse",
160
+ "description": "Injects parallel execution and ULTRAWORK mode",
161
+ "default_enabled": True,
162
+ "exit_on_block": 0,
163
+ },
164
+ "todo_delegation": {
165
+ "type": "PostToolUse",
166
+ "description": "Enforces parallel Task spawning for 2+ todos",
167
+ "default_enabled": True,
168
+ "exit_on_block": 2, # Hard block in stravinsky mode
169
+ },
170
+ "tool_messaging": {
171
+ "type": "PostToolUse",
172
+ "description": "User-friendly messages for MCP tools",
173
+ "default_enabled": True,
174
+ "exit_on_block": 0,
175
+ },
176
+ "edit_recovery": {
177
+ "type": "PostToolUse",
178
+ "description": "Suggests recovery for Edit failures",
179
+ "default_enabled": True,
180
+ "exit_on_block": 0,
181
+ },
182
+ "truncator": {
183
+ "type": "PostToolUse",
184
+ "description": "Truncates responses longer than 30k chars",
185
+ "default_enabled": True,
186
+ "exit_on_block": 0,
187
+ },
188
+ "subagent_stop": {
189
+ "type": "SubagentStop",
190
+ "description": "Handles subagent completion events",
191
+ "default_enabled": True,
192
+ "exit_on_block": 0,
193
+ },
194
+ # UserPromptSubmit hooks
195
+ "context": {
196
+ "type": "UserPromptSubmit",
197
+ "description": "Injects CLAUDE.md content to prompts",
198
+ "default_enabled": True,
199
+ "exit_on_block": 0,
200
+ },
201
+ "todo_continuation": {
202
+ "type": "UserPromptSubmit",
203
+ "description": "Reminds about incomplete todos",
204
+ "default_enabled": True,
205
+ "exit_on_block": 0,
206
+ },
207
+ # PreCompact hooks
208
+ "pre_compact": {
209
+ "type": "PreCompact",
210
+ "description": "Preserves critical context before compaction",
211
+ "default_enabled": True,
212
+ "exit_on_block": 0,
213
+ },
214
+ }
215
+
216
+
217
+ def get_enabled_hooks() -> dict:
218
+ """
219
+ Get all enabled hooks with their configuration.
220
+
221
+ Returns:
222
+ Dict of hook_name -> hook_config for enabled hooks only.
223
+ """
224
+ disabled = get_disabled_hooks()
225
+ enabled = {}
226
+
227
+ for hook_name, config in HOOK_DEFAULTS.items():
228
+ if hook_name not in disabled and config.get("default_enabled", True):
229
+ enabled[hook_name] = config
230
+
231
+ return enabled
232
+
233
+
234
+ def list_hooks() -> str:
235
+ """
236
+ List all hooks with their status.
237
+
238
+ Returns:
239
+ Formatted string showing hook status.
240
+ """
241
+ disabled = get_disabled_hooks()
242
+ lines = ["# Stravinsky Hooks Status", ""]
243
+
244
+ for hook_name, config in sorted(HOOK_DEFAULTS.items()):
245
+ status = "DISABLED" if hook_name in disabled else "enabled"
246
+ icon = "" if hook_name in disabled else ""
247
+ lines.append(f"{icon} {hook_name}: {status} - {config['description']}")
248
+
249
+ return "\n".join(lines)
@@ -0,0 +1,138 @@
1
+ {
2
+ "schema_version": "1.0.0",
3
+ "manifest_version": "0.3.9",
4
+ "description": "Stravinsky hooks for Claude Code integration",
5
+ "generated_date": "2026-01-08T23:53:08.836595Z",
6
+ "hooks": {
7
+ "context.py": {
8
+ "version": "0.3.9",
9
+ "source": "mcp_bridge/hooks/context.py",
10
+ "description": "Project context injection from local files",
11
+ "checksum": "2411ea9d7ef9",
12
+ "lines_of_code": 38,
13
+ "updatable": true,
14
+ "priority": "high",
15
+ "required": true
16
+ },
17
+ "context_monitor.py": {
18
+ "version": "0.3.9",
19
+ "source": "mcp_bridge/hooks/context_monitor.py",
20
+ "description": "Pre-emptive context optimization at 70% threshold",
21
+ "checksum": "7a8d0615af4f",
22
+ "lines_of_code": 153,
23
+ "updatable": true,
24
+ "priority": "high",
25
+ "required": true
26
+ },
27
+ "edit_recovery.py": {
28
+ "version": "0.3.9",
29
+ "source": "mcp_bridge/hooks/edit_recovery.py",
30
+ "description": "Edit/MultiEdit error recovery helper",
31
+ "checksum": "d4e5a96f7bfc",
32
+ "lines_of_code": 46,
33
+ "updatable": true,
34
+ "priority": "high",
35
+ "required": true
36
+ },
37
+ "notification_hook.py": {
38
+ "version": "0.3.9",
39
+ "source": "mcp_bridge/hooks/notification_hook.py",
40
+ "description": "Agent spawn message formatting",
41
+ "checksum": "184947c5a227",
42
+ "lines_of_code": 103,
43
+ "updatable": true,
44
+ "priority": "high",
45
+ "required": true
46
+ },
47
+ "parallel_execution.py": {
48
+ "version": "0.3.9",
49
+ "source": "mcp_bridge/hooks/parallel_execution.py",
50
+ "description": "Pre-emptive parallel execution enforcement",
51
+ "checksum": "9c820d3d19be",
52
+ "lines_of_code": 111,
53
+ "updatable": true,
54
+ "priority": "high",
55
+ "required": true
56
+ },
57
+ "pre_compact.py": {
58
+ "version": "0.3.9",
59
+ "source": "mcp_bridge/hooks/pre_compact.py",
60
+ "description": "Context preservation before compaction",
61
+ "checksum": "4177023bd901",
62
+ "lines_of_code": 123,
63
+ "updatable": true,
64
+ "priority": "high",
65
+ "required": true
66
+ },
67
+ "stop_hook.py": {
68
+ "version": "0.3.9",
69
+ "source": "mcp_bridge/hooks/stop_hook.py",
70
+ "description": "Continuation loop handler",
71
+ "checksum": "820aef797e2e",
72
+ "lines_of_code": 234,
73
+ "updatable": true,
74
+ "priority": "high",
75
+ "required": true
76
+ },
77
+ "stravinsky_mode.py": {
78
+ "version": "0.3.9",
79
+ "source": "mcp_bridge/hooks/stravinsky_mode.py",
80
+ "description": "Hard blocking of native tools",
81
+ "checksum": "5968a95ebcbe",
82
+ "lines_of_code": 146,
83
+ "updatable": true,
84
+ "priority": "high",
85
+ "required": true
86
+ },
87
+ "subagent_stop.py": {
88
+ "version": "0.3.9",
89
+ "source": "mcp_bridge/hooks/subagent_stop.py",
90
+ "description": "Subagent completion handler",
91
+ "checksum": "1943d8dc5355",
92
+ "lines_of_code": 98,
93
+ "updatable": true,
94
+ "priority": "high",
95
+ "required": true
96
+ },
97
+ "todo_continuation.py": {
98
+ "version": "0.3.9",
99
+ "source": "mcp_bridge/hooks/todo_continuation.py",
100
+ "description": "Todo continuation enforcer",
101
+ "checksum": "b6685355f319",
102
+ "lines_of_code": 90,
103
+ "updatable": true,
104
+ "priority": "high",
105
+ "required": true
106
+ },
107
+ "todo_delegation.py": {
108
+ "version": "0.3.9",
109
+ "source": "mcp_bridge/hooks/todo_delegation.py",
110
+ "description": "Parallel task spawning enforcement",
111
+ "checksum": "b4e004d51600",
112
+ "lines_of_code": 88,
113
+ "updatable": true,
114
+ "priority": "high",
115
+ "required": true
116
+ },
117
+ "tool_messaging.py": {
118
+ "version": "0.3.9",
119
+ "source": "mcp_bridge/hooks/tool_messaging.py",
120
+ "description": "User-friendly tool messaging",
121
+ "checksum": "04a10e76f890",
122
+ "lines_of_code": 263,
123
+ "updatable": true,
124
+ "priority": "high",
125
+ "required": true
126
+ },
127
+ "truncator.py": {
128
+ "version": "0.3.9",
129
+ "source": "mcp_bridge/hooks/truncator.py",
130
+ "description": "Tool response truncation at 30k chars",
131
+ "checksum": "87785bf2c657",
132
+ "lines_of_code": 23,
133
+ "updatable": true,
134
+ "priority": "high",
135
+ "required": true
136
+ }
137
+ }
138
+ }
@@ -0,0 +1,222 @@
1
+ """
2
+ Rate Limiting Configuration for Stravinsky Agent Manager.
3
+
4
+ Provides per-model concurrency limits to prevent API overload.
5
+ Implements semaphore-based rate limiting with configurable limits
6
+ per model family.
7
+
8
+ Configuration file: ~/.stravinsky/config.json
9
+ {
10
+ "rate_limits": {
11
+ "claude-opus-4": 2,
12
+ "claude-sonnet-4.5": 5,
13
+ "gemini-3-flash": 10,
14
+ "gemini-3-pro-high": 5,
15
+ "gpt-5.2": 3
16
+ }
17
+ }
18
+ """
19
+
20
+ import asyncio
21
+ import json
22
+ import threading
23
+ from pathlib import Path
24
+ from typing import Dict, Optional
25
+ from collections import defaultdict
26
+ from datetime import datetime
27
+ import logging
28
+
29
+ logger = logging.getLogger(__name__)
30
+
31
+ # Default rate limits per model (conservative defaults)
32
+ DEFAULT_RATE_LIMITS = {
33
+ # Claude models via CLI
34
+ "opus": 2, # Expensive, limit parallel calls
35
+ "sonnet": 5, # Moderate cost
36
+ "haiku": 10, # Cheap, allow more
37
+ # Gemini models via MCP
38
+ "gemini-3-flash": 10, # Free/cheap, allow many
39
+ "gemini-3-pro-high": 5, # Medium cost
40
+ # OpenAI models via MCP
41
+ "gpt-5.2": 3, # Expensive
42
+ # Default for unknown models
43
+ "_default": 5,
44
+ }
45
+
46
+ # Config file location
47
+ CONFIG_FILE = Path.home() / ".stravinsky" / "config.json"
48
+
49
+
50
+ class RateLimiter:
51
+ """
52
+ Semaphore-based rate limiter for model concurrency.
53
+
54
+ Thread-safe implementation that limits concurrent requests
55
+ per model family to prevent API overload.
56
+ """
57
+
58
+ def __init__(self):
59
+ self._semaphores: Dict[str, threading.Semaphore] = {}
60
+ self._lock = threading.Lock()
61
+ self._limits = self._load_limits()
62
+ self._active_counts: Dict[str, int] = defaultdict(int)
63
+ self._queue_counts: Dict[str, int] = defaultdict(int)
64
+
65
+ def _load_limits(self) -> Dict[str, int]:
66
+ """Load rate limits from config file or use defaults."""
67
+ limits = DEFAULT_RATE_LIMITS.copy()
68
+
69
+ if CONFIG_FILE.exists():
70
+ try:
71
+ with open(CONFIG_FILE) as f:
72
+ config = json.load(f)
73
+ if "rate_limits" in config:
74
+ limits.update(config["rate_limits"])
75
+ logger.info(f"[RateLimiter] Loaded custom limits from {CONFIG_FILE}")
76
+ except (json.JSONDecodeError, IOError) as e:
77
+ logger.warning(f"[RateLimiter] Failed to load config: {e}")
78
+
79
+ return limits
80
+
81
+ def _get_semaphore(self, model: str) -> threading.Semaphore:
82
+ """Get or create a semaphore for a model."""
83
+ with self._lock:
84
+ if model not in self._semaphores:
85
+ limit = self._limits.get(model, self._limits.get("_default", 5))
86
+ self._semaphores[model] = threading.Semaphore(limit)
87
+ logger.debug(f"[RateLimiter] Created semaphore for {model} with limit {limit}")
88
+ return self._semaphores[model]
89
+
90
+ def _normalize_model(self, model: str) -> str:
91
+ """Normalize model name to match config keys."""
92
+ model_lower = model.lower()
93
+
94
+ # Match known patterns
95
+ if "opus" in model_lower:
96
+ return "opus"
97
+ elif "sonnet" in model_lower:
98
+ return "sonnet"
99
+ elif "haiku" in model_lower:
100
+ return "haiku"
101
+ elif "gemini" in model_lower and "flash" in model_lower:
102
+ return "gemini-3-flash"
103
+ elif "gemini" in model_lower and ("pro" in model_lower or "high" in model_lower):
104
+ return "gemini-3-pro-high"
105
+ elif "gpt" in model_lower:
106
+ return "gpt-5.2"
107
+
108
+ return model_lower
109
+
110
+ def acquire(self, model: str, timeout: float = 60.0) -> bool:
111
+ """
112
+ Acquire a slot for the given model.
113
+
114
+ Args:
115
+ model: Model name to acquire slot for
116
+ timeout: Maximum time to wait in seconds
117
+
118
+ Returns:
119
+ True if slot acquired, False if timed out
120
+ """
121
+ normalized = self._normalize_model(model)
122
+ semaphore = self._get_semaphore(normalized)
123
+
124
+ with self._lock:
125
+ self._queue_counts[normalized] += 1
126
+
127
+ logger.debug(f"[RateLimiter] Acquiring slot for {normalized}")
128
+ acquired = semaphore.acquire(blocking=True, timeout=timeout)
129
+
130
+ with self._lock:
131
+ self._queue_counts[normalized] -= 1
132
+ if acquired:
133
+ self._active_counts[normalized] += 1
134
+
135
+ if acquired:
136
+ logger.debug(f"[RateLimiter] Acquired slot for {normalized}")
137
+ else:
138
+ logger.warning(f"[RateLimiter] Timeout waiting for slot for {normalized}")
139
+
140
+ return acquired
141
+
142
+ def release(self, model: str):
143
+ """Release a slot for the given model."""
144
+ normalized = self._normalize_model(model)
145
+ semaphore = self._get_semaphore(normalized)
146
+
147
+ with self._lock:
148
+ self._active_counts[normalized] = max(0, self._active_counts[normalized] - 1)
149
+
150
+ semaphore.release()
151
+ logger.debug(f"[RateLimiter] Released slot for {normalized}")
152
+
153
+ def get_status(self) -> Dict[str, Dict[str, int]]:
154
+ """Get current rate limiter status."""
155
+ with self._lock:
156
+ return {
157
+ model: {
158
+ "limit": self._limits.get(model, self._limits.get("_default", 5)),
159
+ "active": self._active_counts[model],
160
+ "queued": self._queue_counts[model],
161
+ }
162
+ for model in set(list(self._active_counts.keys()) + list(self._queue_counts.keys()))
163
+ }
164
+
165
+ def update_limits(self, new_limits: Dict[str, int]):
166
+ """
167
+ Update rate limits dynamically.
168
+
169
+ Note: This only affects new semaphores. Existing ones
170
+ will continue with their original limits until recreated.
171
+ """
172
+ with self._lock:
173
+ self._limits.update(new_limits)
174
+ logger.info(f"[RateLimiter] Updated limits: {new_limits}")
175
+
176
+
177
+ class RateLimitContext:
178
+ """Context manager for rate-limited model access."""
179
+
180
+ def __init__(self, limiter: RateLimiter, model: str, timeout: float = 60.0):
181
+ self.limiter = limiter
182
+ self.model = model
183
+ self.timeout = timeout
184
+ self.acquired = False
185
+
186
+ def __enter__(self):
187
+ self.acquired = self.limiter.acquire(self.model, self.timeout)
188
+ if not self.acquired:
189
+ raise TimeoutError(f"Rate limit timeout for model {self.model}")
190
+ return self
191
+
192
+ def __exit__(self, exc_type, exc_val, exc_tb):
193
+ if self.acquired:
194
+ self.limiter.release(self.model)
195
+ return False
196
+
197
+
198
+ # Global rate limiter instance
199
+ _rate_limiter: Optional[RateLimiter] = None
200
+ _rate_limiter_lock = threading.Lock()
201
+
202
+
203
+ def get_rate_limiter() -> RateLimiter:
204
+ """Get or create the global RateLimiter instance."""
205
+ global _rate_limiter
206
+ if _rate_limiter is None:
207
+ with _rate_limiter_lock:
208
+ if _rate_limiter is None:
209
+ _rate_limiter = RateLimiter()
210
+ return _rate_limiter
211
+
212
+
213
+ def rate_limited(model: str, timeout: float = 60.0) -> RateLimitContext:
214
+ """
215
+ Get a rate-limited context for a model.
216
+
217
+ Usage:
218
+ with rate_limited("gemini-3-flash") as ctx:
219
+ # Make API call
220
+ pass
221
+ """
222
+ return RateLimitContext(get_rate_limiter(), model, timeout)