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.
- mcp_bridge/__init__.py +1 -1
- mcp_bridge/auth/token_store.py +113 -11
- mcp_bridge/cli/__init__.py +6 -0
- mcp_bridge/cli/install_hooks.py +1265 -0
- mcp_bridge/cli/session_report.py +585 -0
- mcp_bridge/config/MANIFEST_SCHEMA.md +305 -0
- mcp_bridge/config/README.md +276 -0
- mcp_bridge/config/hook_config.py +249 -0
- mcp_bridge/config/hooks_manifest.json +138 -0
- mcp_bridge/config/rate_limits.py +222 -0
- mcp_bridge/config/skills_manifest.json +128 -0
- mcp_bridge/hooks/HOOKS_SETTINGS.json +175 -0
- mcp_bridge/hooks/README.md +215 -0
- mcp_bridge/hooks/__init__.py +119 -60
- mcp_bridge/hooks/edit_recovery.py +42 -37
- mcp_bridge/hooks/git_noninteractive.py +89 -0
- mcp_bridge/hooks/keyword_detector.py +30 -0
- mcp_bridge/hooks/manager.py +8 -0
- mcp_bridge/hooks/notification_hook.py +103 -0
- mcp_bridge/hooks/parallel_execution.py +111 -0
- mcp_bridge/hooks/pre_compact.py +82 -183
- mcp_bridge/hooks/rules_injector.py +507 -0
- mcp_bridge/hooks/session_notifier.py +125 -0
- mcp_bridge/{native_hooks → hooks}/stravinsky_mode.py +51 -16
- mcp_bridge/hooks/subagent_stop.py +98 -0
- mcp_bridge/hooks/task_validator.py +73 -0
- mcp_bridge/hooks/tmux_manager.py +141 -0
- mcp_bridge/hooks/todo_continuation.py +90 -0
- mcp_bridge/hooks/todo_delegation.py +88 -0
- mcp_bridge/hooks/tool_messaging.py +267 -0
- mcp_bridge/hooks/truncator.py +21 -17
- mcp_bridge/notifications.py +151 -0
- mcp_bridge/prompts/multimodal.py +24 -3
- mcp_bridge/server.py +214 -49
- mcp_bridge/server_tools.py +445 -0
- mcp_bridge/tools/__init__.py +22 -18
- mcp_bridge/tools/agent_manager.py +220 -32
- mcp_bridge/tools/code_search.py +97 -11
- mcp_bridge/tools/lsp/__init__.py +7 -0
- mcp_bridge/tools/lsp/manager.py +448 -0
- mcp_bridge/tools/lsp/tools.py +637 -150
- mcp_bridge/tools/model_invoke.py +208 -106
- mcp_bridge/tools/query_classifier.py +323 -0
- mcp_bridge/tools/semantic_search.py +3042 -0
- mcp_bridge/tools/templates.py +32 -18
- mcp_bridge/update_manager.py +589 -0
- mcp_bridge/update_manager_pypi.py +299 -0
- stravinsky-0.4.18.dist-info/METADATA +468 -0
- stravinsky-0.4.18.dist-info/RECORD +88 -0
- stravinsky-0.4.18.dist-info/entry_points.txt +5 -0
- mcp_bridge/native_hooks/edit_recovery.py +0 -46
- mcp_bridge/native_hooks/todo_delegation.py +0 -54
- mcp_bridge/native_hooks/truncator.py +0 -23
- stravinsky-0.2.52.dist-info/METADATA +0 -204
- stravinsky-0.2.52.dist-info/RECORD +0 -63
- stravinsky-0.2.52.dist-info/entry_points.txt +0 -3
- /mcp_bridge/{native_hooks → hooks}/context.py +0 -0
- {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)
|