tool-throttle 0.1.0__tar.gz
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.
- tool_throttle-0.1.0/LICENSE +21 -0
- tool_throttle-0.1.0/PKG-INFO +132 -0
- tool_throttle-0.1.0/README.md +119 -0
- tool_throttle-0.1.0/pyproject.toml +22 -0
- tool_throttle-0.1.0/setup.cfg +4 -0
- tool_throttle-0.1.0/tests/test_smoke.py +188 -0
- tool_throttle-0.1.0/tool_throttle/__init__.py +17 -0
- tool_throttle-0.1.0/tool_throttle/bucket.py +69 -0
- tool_throttle-0.1.0/tool_throttle/cli.py +117 -0
- tool_throttle-0.1.0/tool_throttle/config.py +102 -0
- tool_throttle-0.1.0/tool_throttle/throttler.py +162 -0
- tool_throttle-0.1.0/tool_throttle.egg-info/PKG-INFO +132 -0
- tool_throttle-0.1.0/tool_throttle.egg-info/SOURCES.txt +14 -0
- tool_throttle-0.1.0/tool_throttle.egg-info/dependency_links.txt +1 -0
- tool_throttle-0.1.0/tool_throttle.egg-info/entry_points.txt +2 -0
- tool_throttle-0.1.0/tool_throttle.egg-info/top_level.txt +1 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 hao li
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: tool-throttle
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: MCP tool-call throttling: token-bucket rate limits per tool name, a fuse against runaway agents
|
|
5
|
+
Author: hao li
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/hahahahahahahahah6/tool-throttle
|
|
8
|
+
Keywords: mcp,ai-agents,rate-limit,throttling,token-bucket
|
|
9
|
+
Requires-Python: >=3.9
|
|
10
|
+
Description-Content-Type: text/markdown
|
|
11
|
+
License-File: LICENSE
|
|
12
|
+
Dynamic: license-file
|
|
13
|
+
|
|
14
|
+
# tool-throttle
|
|
15
|
+
|
|
16
|
+
[](LICENSE)
|
|
17
|
+
[](pyproject.toml)
|
|
18
|
+
[](pyproject.toml)
|
|
19
|
+
[](https://pypi.org/project/tool-throttle/)
|
|
20
|
+
|
|
21
|
+
A token-bucket rate limiter for **MCP tool calls** — a fuse against runaway
|
|
22
|
+
agents. When an agent gets stuck in a loop hammering `web_search` 400 times a
|
|
23
|
+
minute, tool-throttle slows it down or stops it, per tool name, with per-tool
|
|
24
|
+
quotas you configure in JSON.
|
|
25
|
+
|
|
26
|
+
Zero dependencies. Standard library only. Thread-safe.
|
|
27
|
+
|
|
28
|
+
```bash
|
|
29
|
+
pip install tool-throttle
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
```python
|
|
33
|
+
from tool_throttle import Throttler, ThrottlerConfig, Quota, REJECT
|
|
34
|
+
|
|
35
|
+
th = Throttler(ThrottlerConfig(
|
|
36
|
+
default=Quota(calls=60, per_seconds=60), # 60 calls/min for everything
|
|
37
|
+
per_tool={
|
|
38
|
+
"web_search": Quota(calls=10, per_seconds=60), # tight leash
|
|
39
|
+
"expensive_*": Quota(calls=2, per_seconds=60), # glob patterns work
|
|
40
|
+
},
|
|
41
|
+
on_exceed=REJECT, # or "queue": wait for a token, up to max_queue_wait
|
|
42
|
+
))
|
|
43
|
+
|
|
44
|
+
# wrap any callable -- your MCP tool dispatcher, a client method, anything:
|
|
45
|
+
safe_search = th.wrap("web_search", mcp_client.call_tool)
|
|
46
|
+
safe_search("web_search", {"query": "..."}) # -> ThrottledError past quota
|
|
47
|
+
|
|
48
|
+
# or decide explicitly without consuming a token:
|
|
49
|
+
d = th.decide("web_search")
|
|
50
|
+
# Decision(tool='web_search', allowed=False, wait_seconds=42.1, ...)
|
|
51
|
+
|
|
52
|
+
th.stats() # per-tool allowed / rejected / queued / waited_seconds
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
Config round-trips through JSON, so quotas can live in a file:
|
|
56
|
+
|
|
57
|
+
```python
|
|
58
|
+
cfg = ThrottlerConfig(
|
|
59
|
+
default=Quota(calls=60, per_seconds=60),
|
|
60
|
+
per_tool={"web_*": Quota(calls=10, per_seconds=60, burst=3)},
|
|
61
|
+
on_exceed="queue",
|
|
62
|
+
max_queue_wait=30.0,
|
|
63
|
+
)
|
|
64
|
+
cfg.save("throttle.json")
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
## CLI
|
|
68
|
+
|
|
69
|
+
```bash
|
|
70
|
+
# dry-run: print a decision per {"tool": ...} JSON line on stdin
|
|
71
|
+
echo '{"tool": "web_search"}' | tool-throttle check --config throttle.json
|
|
72
|
+
|
|
73
|
+
# enforce: consume tokens (queues or rejects over-quota calls)
|
|
74
|
+
echo '{"tool": "web_search"}' | tool-throttle acquire --config throttle.json
|
|
75
|
+
|
|
76
|
+
# built-in demo of a looping agent being fused
|
|
77
|
+
tool-throttle demo
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
## How it works
|
|
81
|
+
|
|
82
|
+
Each distinct tool name gets its own `TokenBucket`, sized by the first
|
|
83
|
+
matching quota in `per_tool` (glob patterns, insertion order) or `default`.
|
|
84
|
+
Buckets refill steadily over time, so a `Quota(calls=10, per_seconds=60)`
|
|
85
|
+
allows 10 calls per minute sustained, with a burst of up to `burst`
|
|
86
|
+
(default: `calls`).
|
|
87
|
+
|
|
88
|
+
On quota exhaustion:
|
|
89
|
+
- `on_exceed="reject"` (default in the demo): raise `ThrottledError` with a
|
|
90
|
+
`retry_after` hint — fail fast.
|
|
91
|
+
- `on_exceed="queue"`: block until a token frees up, bounded by
|
|
92
|
+
`max_queue_wait` seconds; if the wait would exceed the bound, raise
|
|
93
|
+
`ThrottledError(reason="queue_wait_exceeded")`.
|
|
94
|
+
|
|
95
|
+
## Differentiation: why not just rate-limit at the API gateway?
|
|
96
|
+
|
|
97
|
+
API gateways throttle **HTTP requests**. That is the wrong layer for agent
|
|
98
|
+
safety:
|
|
99
|
+
|
|
100
|
+
- One HTTP request can fan out into many tool calls; one tool call pattern
|
|
101
|
+
can span many HTTP endpoints. Gateway counters never see *tool-call
|
|
102
|
+
semantics* — they see bytes and paths.
|
|
103
|
+
- Gateway limits are per API key / per IP. An agent looping on a *single
|
|
104
|
+
expensive tool* (say, a search tool that costs $0.02 per call) looks
|
|
105
|
+
perfectly normal at the HTTP layer while burning money at the tool layer.
|
|
106
|
+
- Quotas here are expressed in the language of the agent: *tool names and
|
|
107
|
+
glob patterns* (`expensive_*`), with per-tool bursts and queue-or-reject
|
|
108
|
+
policies — the thing an agent developer actually wants to tune.
|
|
109
|
+
|
|
110
|
+
tool-throttle is the **fuse at the MCP tool-call semantic layer**: it sits
|
|
111
|
+
between the agent's tool-call intent and execution, trips before the loop
|
|
112
|
+
does real damage, and reports per-tool stats so you can see which tool the
|
|
113
|
+
agent was abusing. It complements gateway rate limiting; it does not
|
|
114
|
+
replace it.
|
|
115
|
+
|
|
116
|
+
## API reference
|
|
117
|
+
|
|
118
|
+
| Symbol | Description |
|
|
119
|
+
|---|---|
|
|
120
|
+
| `Throttler(config)` | the interceptor; one per agent/process |
|
|
121
|
+
| `ThrottlerConfig(default, per_tool, on_exceed, max_queue_wait)` | quotas + policy |
|
|
122
|
+
| `Quota(calls, per_seconds=60, burst=None)` | rate quota for one pattern |
|
|
123
|
+
| `th.acquire(tool)` | consume a token; queues or raises `ThrottledError` |
|
|
124
|
+
| `th.decide(tool)` | dry-run check, does not consume |
|
|
125
|
+
| `th.wrap(tool, fn)` / `th.call(tool, fn, *a, **k)` | throttle a callable |
|
|
126
|
+
| `th.stats()` / `th.reset(tool?)` / `th.clear()` | observability and reset |
|
|
127
|
+
| `ThrottledError` | `.tool`, `.retry_after`, `.reason` |
|
|
128
|
+
| `TokenBucket(capacity, refill_per_second, clock=...)` | the primitive underneath |
|
|
129
|
+
|
|
130
|
+
## License
|
|
131
|
+
|
|
132
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,119 @@
|
|
|
1
|
+
# tool-throttle
|
|
2
|
+
|
|
3
|
+
[](LICENSE)
|
|
4
|
+
[](pyproject.toml)
|
|
5
|
+
[](pyproject.toml)
|
|
6
|
+
[](https://pypi.org/project/tool-throttle/)
|
|
7
|
+
|
|
8
|
+
A token-bucket rate limiter for **MCP tool calls** — a fuse against runaway
|
|
9
|
+
agents. When an agent gets stuck in a loop hammering `web_search` 400 times a
|
|
10
|
+
minute, tool-throttle slows it down or stops it, per tool name, with per-tool
|
|
11
|
+
quotas you configure in JSON.
|
|
12
|
+
|
|
13
|
+
Zero dependencies. Standard library only. Thread-safe.
|
|
14
|
+
|
|
15
|
+
```bash
|
|
16
|
+
pip install tool-throttle
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
```python
|
|
20
|
+
from tool_throttle import Throttler, ThrottlerConfig, Quota, REJECT
|
|
21
|
+
|
|
22
|
+
th = Throttler(ThrottlerConfig(
|
|
23
|
+
default=Quota(calls=60, per_seconds=60), # 60 calls/min for everything
|
|
24
|
+
per_tool={
|
|
25
|
+
"web_search": Quota(calls=10, per_seconds=60), # tight leash
|
|
26
|
+
"expensive_*": Quota(calls=2, per_seconds=60), # glob patterns work
|
|
27
|
+
},
|
|
28
|
+
on_exceed=REJECT, # or "queue": wait for a token, up to max_queue_wait
|
|
29
|
+
))
|
|
30
|
+
|
|
31
|
+
# wrap any callable -- your MCP tool dispatcher, a client method, anything:
|
|
32
|
+
safe_search = th.wrap("web_search", mcp_client.call_tool)
|
|
33
|
+
safe_search("web_search", {"query": "..."}) # -> ThrottledError past quota
|
|
34
|
+
|
|
35
|
+
# or decide explicitly without consuming a token:
|
|
36
|
+
d = th.decide("web_search")
|
|
37
|
+
# Decision(tool='web_search', allowed=False, wait_seconds=42.1, ...)
|
|
38
|
+
|
|
39
|
+
th.stats() # per-tool allowed / rejected / queued / waited_seconds
|
|
40
|
+
```
|
|
41
|
+
|
|
42
|
+
Config round-trips through JSON, so quotas can live in a file:
|
|
43
|
+
|
|
44
|
+
```python
|
|
45
|
+
cfg = ThrottlerConfig(
|
|
46
|
+
default=Quota(calls=60, per_seconds=60),
|
|
47
|
+
per_tool={"web_*": Quota(calls=10, per_seconds=60, burst=3)},
|
|
48
|
+
on_exceed="queue",
|
|
49
|
+
max_queue_wait=30.0,
|
|
50
|
+
)
|
|
51
|
+
cfg.save("throttle.json")
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
## CLI
|
|
55
|
+
|
|
56
|
+
```bash
|
|
57
|
+
# dry-run: print a decision per {"tool": ...} JSON line on stdin
|
|
58
|
+
echo '{"tool": "web_search"}' | tool-throttle check --config throttle.json
|
|
59
|
+
|
|
60
|
+
# enforce: consume tokens (queues or rejects over-quota calls)
|
|
61
|
+
echo '{"tool": "web_search"}' | tool-throttle acquire --config throttle.json
|
|
62
|
+
|
|
63
|
+
# built-in demo of a looping agent being fused
|
|
64
|
+
tool-throttle demo
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
## How it works
|
|
68
|
+
|
|
69
|
+
Each distinct tool name gets its own `TokenBucket`, sized by the first
|
|
70
|
+
matching quota in `per_tool` (glob patterns, insertion order) or `default`.
|
|
71
|
+
Buckets refill steadily over time, so a `Quota(calls=10, per_seconds=60)`
|
|
72
|
+
allows 10 calls per minute sustained, with a burst of up to `burst`
|
|
73
|
+
(default: `calls`).
|
|
74
|
+
|
|
75
|
+
On quota exhaustion:
|
|
76
|
+
- `on_exceed="reject"` (default in the demo): raise `ThrottledError` with a
|
|
77
|
+
`retry_after` hint — fail fast.
|
|
78
|
+
- `on_exceed="queue"`: block until a token frees up, bounded by
|
|
79
|
+
`max_queue_wait` seconds; if the wait would exceed the bound, raise
|
|
80
|
+
`ThrottledError(reason="queue_wait_exceeded")`.
|
|
81
|
+
|
|
82
|
+
## Differentiation: why not just rate-limit at the API gateway?
|
|
83
|
+
|
|
84
|
+
API gateways throttle **HTTP requests**. That is the wrong layer for agent
|
|
85
|
+
safety:
|
|
86
|
+
|
|
87
|
+
- One HTTP request can fan out into many tool calls; one tool call pattern
|
|
88
|
+
can span many HTTP endpoints. Gateway counters never see *tool-call
|
|
89
|
+
semantics* — they see bytes and paths.
|
|
90
|
+
- Gateway limits are per API key / per IP. An agent looping on a *single
|
|
91
|
+
expensive tool* (say, a search tool that costs $0.02 per call) looks
|
|
92
|
+
perfectly normal at the HTTP layer while burning money at the tool layer.
|
|
93
|
+
- Quotas here are expressed in the language of the agent: *tool names and
|
|
94
|
+
glob patterns* (`expensive_*`), with per-tool bursts and queue-or-reject
|
|
95
|
+
policies — the thing an agent developer actually wants to tune.
|
|
96
|
+
|
|
97
|
+
tool-throttle is the **fuse at the MCP tool-call semantic layer**: it sits
|
|
98
|
+
between the agent's tool-call intent and execution, trips before the loop
|
|
99
|
+
does real damage, and reports per-tool stats so you can see which tool the
|
|
100
|
+
agent was abusing. It complements gateway rate limiting; it does not
|
|
101
|
+
replace it.
|
|
102
|
+
|
|
103
|
+
## API reference
|
|
104
|
+
|
|
105
|
+
| Symbol | Description |
|
|
106
|
+
|---|---|
|
|
107
|
+
| `Throttler(config)` | the interceptor; one per agent/process |
|
|
108
|
+
| `ThrottlerConfig(default, per_tool, on_exceed, max_queue_wait)` | quotas + policy |
|
|
109
|
+
| `Quota(calls, per_seconds=60, burst=None)` | rate quota for one pattern |
|
|
110
|
+
| `th.acquire(tool)` | consume a token; queues or raises `ThrottledError` |
|
|
111
|
+
| `th.decide(tool)` | dry-run check, does not consume |
|
|
112
|
+
| `th.wrap(tool, fn)` / `th.call(tool, fn, *a, **k)` | throttle a callable |
|
|
113
|
+
| `th.stats()` / `th.reset(tool?)` / `th.clear()` | observability and reset |
|
|
114
|
+
| `ThrottledError` | `.tool`, `.retry_after`, `.reason` |
|
|
115
|
+
| `TokenBucket(capacity, refill_per_second, clock=...)` | the primitive underneath |
|
|
116
|
+
|
|
117
|
+
## License
|
|
118
|
+
|
|
119
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=61"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "tool-throttle"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "MCP tool-call throttling: token-bucket rate limits per tool name, a fuse against runaway agents"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = {text = "MIT"}
|
|
12
|
+
authors = [{name = "hao li"}]
|
|
13
|
+
keywords = ["mcp", "ai-agents", "rate-limit", "throttling", "token-bucket"]
|
|
14
|
+
|
|
15
|
+
[project.scripts]
|
|
16
|
+
tool-throttle = "tool_throttle.cli:main"
|
|
17
|
+
|
|
18
|
+
[project.urls]
|
|
19
|
+
Homepage = "https://github.com/hahahahahahahahah6/tool-throttle"
|
|
20
|
+
|
|
21
|
+
[tool.setuptools.packages.find]
|
|
22
|
+
include = ["tool_throttle*"]
|
|
@@ -0,0 +1,188 @@
|
|
|
1
|
+
"""Smoke tests for tool-throttle. All must be green before shipping."""
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import os
|
|
5
|
+
import tempfile
|
|
6
|
+
import time
|
|
7
|
+
|
|
8
|
+
import pytest
|
|
9
|
+
|
|
10
|
+
from tool_throttle import (
|
|
11
|
+
QUEUE, REJECT, Quota, ThrottledError, Throttler, ThrottlerConfig,
|
|
12
|
+
TokenBucket,
|
|
13
|
+
)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
class FakeClock:
|
|
17
|
+
def __init__(self):
|
|
18
|
+
self.t = 1000.0
|
|
19
|
+
|
|
20
|
+
def __call__(self):
|
|
21
|
+
return self.t
|
|
22
|
+
|
|
23
|
+
def advance(self, dt):
|
|
24
|
+
self.t += dt
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
# 1. tokens refill at the configured rate (deterministic, fake clock)
|
|
28
|
+
def test_refill_math():
|
|
29
|
+
c = FakeClock()
|
|
30
|
+
b = TokenBucket(capacity=5, refill_per_second=2.0, clock=c)
|
|
31
|
+
assert b.try_consume(5) is True
|
|
32
|
+
assert b.try_consume(1) is False
|
|
33
|
+
c.advance(1.0)
|
|
34
|
+
assert b.peek() == pytest.approx(2.0)
|
|
35
|
+
assert b.try_consume(2) is True
|
|
36
|
+
assert b.try_consume(1) is False
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
# 2. idle buckets never exceed capacity (burst cap)
|
|
40
|
+
def test_burst_cap():
|
|
41
|
+
c = FakeClock()
|
|
42
|
+
b = TokenBucket(capacity=3, refill_per_second=100.0, clock=c)
|
|
43
|
+
b.try_consume(3)
|
|
44
|
+
c.advance(3600) # idle for an hour
|
|
45
|
+
assert b.peek() == pytest.approx(3.0)
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
# 3. per-tool quotas override the default; unmatched tools use default
|
|
49
|
+
def test_per_tool_override():
|
|
50
|
+
c = FakeClock()
|
|
51
|
+
th = Throttler(ThrottlerConfig(
|
|
52
|
+
default=Quota(calls=60, per_seconds=60),
|
|
53
|
+
per_tool={"expensive_tool": Quota(calls=2, per_seconds=60)},
|
|
54
|
+
on_exceed=REJECT), clock=c)
|
|
55
|
+
th.acquire("expensive_tool")
|
|
56
|
+
th.acquire("expensive_tool")
|
|
57
|
+
with pytest.raises(ThrottledError):
|
|
58
|
+
th.acquire("expensive_tool")
|
|
59
|
+
# default tool is unaffected
|
|
60
|
+
th.acquire("cheap_tool")
|
|
61
|
+
d = th.decide("cheap_tool")
|
|
62
|
+
assert d.pattern is None # default has no pattern
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
# 4. first matching pattern wins (insertion order)
|
|
66
|
+
def test_first_match_wins():
|
|
67
|
+
c = FakeClock()
|
|
68
|
+
th = Throttler(ThrottlerConfig(
|
|
69
|
+
per_tool={"web_*": Quota(calls=2, per_seconds=60),
|
|
70
|
+
"web_search": Quota(calls=50, per_seconds=60)},
|
|
71
|
+
on_exceed=REJECT), clock=c)
|
|
72
|
+
d = th.decide("web_search")
|
|
73
|
+
assert d.pattern == "web_*"
|
|
74
|
+
th.acquire("web_search")
|
|
75
|
+
th.acquire("web_search")
|
|
76
|
+
with pytest.raises(ThrottledError):
|
|
77
|
+
th.acquire("web_search")
|
|
78
|
+
|
|
79
|
+
|
|
80
|
+
# 5. reject policy fails fast and reports retry_after
|
|
81
|
+
def test_reject_reports_retry_after():
|
|
82
|
+
c = FakeClock()
|
|
83
|
+
th = Throttler(ThrottlerConfig(
|
|
84
|
+
default=Quota(calls=1, per_seconds=60), on_exceed=REJECT), clock=c)
|
|
85
|
+
th.acquire("t")
|
|
86
|
+
t0 = time.monotonic()
|
|
87
|
+
with pytest.raises(ThrottledError) as ei:
|
|
88
|
+
th.acquire("t")
|
|
89
|
+
assert time.monotonic() - t0 < 0.5 # fast, no blocking
|
|
90
|
+
assert ei.value.retry_after == pytest.approx(60.0, abs=1.0)
|
|
91
|
+
assert ei.value.tool == "t"
|
|
92
|
+
assert th.stats()["t"]["rejected"] == 1
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
# 6. queue policy waits for a token then proceeds (real clock)
|
|
96
|
+
def test_queue_waits_then_proceeds():
|
|
97
|
+
th = Throttler(ThrottlerConfig(
|
|
98
|
+
default=Quota(calls=2, per_seconds=1, burst=1), # 1 token every 0.5s
|
|
99
|
+
on_exceed=QUEUE, max_queue_wait=5.0))
|
|
100
|
+
th.acquire("t") # consumes the 1-token burst
|
|
101
|
+
t0 = time.monotonic()
|
|
102
|
+
d = th.acquire("t")
|
|
103
|
+
elapsed = time.monotonic() - t0
|
|
104
|
+
assert elapsed >= 0.4 # actually waited for refill
|
|
105
|
+
assert d.allowed and d.wait_seconds >= 0.4
|
|
106
|
+
assert th.stats()["t"]["queued"] == 1
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
# 7. queue policy is bounded: a call needing an hour-long wait fails fast
|
|
110
|
+
def test_queue_bound_never_blocks():
|
|
111
|
+
th = Throttler(ThrottlerConfig(
|
|
112
|
+
default=Quota(calls=1, per_seconds=3600), # ~1/hour refill
|
|
113
|
+
on_exceed=QUEUE, max_queue_wait=0.2))
|
|
114
|
+
th.acquire("t")
|
|
115
|
+
t0 = time.monotonic()
|
|
116
|
+
with pytest.raises(ThrottledError) as ei:
|
|
117
|
+
th.acquire("t")
|
|
118
|
+
elapsed = time.monotonic() - t0
|
|
119
|
+
assert elapsed < 1.0 # bounded, never blocks for the full hour
|
|
120
|
+
assert ei.value.reason == "quota_exceeded"
|
|
121
|
+
assert ei.value.retry_after > 3000
|
|
122
|
+
assert th.stats()["t"]["rejected"] == 1
|
|
123
|
+
|
|
124
|
+
|
|
125
|
+
# 8. wrap() passes return values through and enforces the quota
|
|
126
|
+
def test_wrap_integration():
|
|
127
|
+
c = FakeClock()
|
|
128
|
+
th = Throttler(ThrottlerConfig(
|
|
129
|
+
default=Quota(calls=3, per_seconds=60), on_exceed=REJECT), clock=c)
|
|
130
|
+
|
|
131
|
+
def add(a, b):
|
|
132
|
+
"""adds two numbers"""
|
|
133
|
+
return a + b
|
|
134
|
+
|
|
135
|
+
safe_add = th.wrap("math", add)
|
|
136
|
+
assert safe_add(2, 3) == 5
|
|
137
|
+
assert safe_add.__name__ == "add"
|
|
138
|
+
safe_add(1, 1)
|
|
139
|
+
safe_add(1, 1)
|
|
140
|
+
with pytest.raises(ThrottledError):
|
|
141
|
+
safe_add(1, 1)
|
|
142
|
+
s = th.stats()["math"]
|
|
143
|
+
assert s["allowed"] == 3 and s["rejected"] == 1
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
# 9. stats() exposes per-tool counters
|
|
147
|
+
def test_stats_shape():
|
|
148
|
+
c = FakeClock()
|
|
149
|
+
th = Throttler(ThrottlerConfig(
|
|
150
|
+
default=Quota(calls=1, per_seconds=60),
|
|
151
|
+
per_tool={"web_*": Quota(calls=5, per_seconds=60)},
|
|
152
|
+
on_exceed=REJECT), clock=c)
|
|
153
|
+
th.acquire("web_search")
|
|
154
|
+
try:
|
|
155
|
+
th.acquire("other")
|
|
156
|
+
th.acquire("other")
|
|
157
|
+
except ThrottledError:
|
|
158
|
+
pass
|
|
159
|
+
s = th.stats()
|
|
160
|
+
assert set(s) == {"web_search", "other"}
|
|
161
|
+
for tool, counters in s.items():
|
|
162
|
+
assert set(counters) == {"allowed", "rejected", "queued",
|
|
163
|
+
"waited_seconds", "quota", "pattern"}
|
|
164
|
+
assert s["web_search"]["allowed"] == 1
|
|
165
|
+
assert s["other"]["rejected"] == 1
|
|
166
|
+
assert s["web_search"]["pattern"] == "web_*"
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
# 10. config survives a JSON round-trip
|
|
170
|
+
def test_config_json_roundtrip():
|
|
171
|
+
cfg = ThrottlerConfig(
|
|
172
|
+
default=Quota(calls=30, per_seconds=60),
|
|
173
|
+
per_tool={"expensive_*": Quota(calls=2, per_seconds=60, burst=4)},
|
|
174
|
+
on_exceed=REJECT, max_queue_wait=7.5)
|
|
175
|
+
with tempfile.TemporaryDirectory() as d:
|
|
176
|
+
p = os.path.join(d, "throttle.json")
|
|
177
|
+
cfg.save(p)
|
|
178
|
+
loaded = ThrottlerConfig.load(p)
|
|
179
|
+
assert loaded == cfg
|
|
180
|
+
q, pat = loaded.quota_for("expensive_llm")
|
|
181
|
+
assert (q.calls, q.per_seconds, q.burst, pat) == (2, 60.0, 4, "expensive_*")
|
|
182
|
+
# round-tripped config behaves identically (burst=4 -> 4 tokens)
|
|
183
|
+
c = FakeClock()
|
|
184
|
+
th = Throttler(loaded, clock=c)
|
|
185
|
+
for _ in range(4):
|
|
186
|
+
th.acquire("expensive_llm")
|
|
187
|
+
with pytest.raises(ThrottledError):
|
|
188
|
+
th.acquire("expensive_llm")
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
"""tool-throttle: MCP tool-call throttling (token bucket per tool name)."""
|
|
2
|
+
|
|
3
|
+
from .bucket import TokenBucket
|
|
4
|
+
from .config import QUEUE, REJECT, Quota, ThrottlerConfig
|
|
5
|
+
from .throttler import Decision, ThrottledError, Throttler
|
|
6
|
+
|
|
7
|
+
__version__ = "0.1.0"
|
|
8
|
+
__all__ = [
|
|
9
|
+
"TokenBucket",
|
|
10
|
+
"Quota",
|
|
11
|
+
"ThrottlerConfig",
|
|
12
|
+
"QUEUE",
|
|
13
|
+
"REJECT",
|
|
14
|
+
"Decision",
|
|
15
|
+
"ThrottledError",
|
|
16
|
+
"Throttler",
|
|
17
|
+
]
|
|
@@ -0,0 +1,69 @@
|
|
|
1
|
+
"""Thread-safe token bucket used by tool-throttle.
|
|
2
|
+
|
|
3
|
+
The bucket is deliberately simple: a fixed ``capacity`` (burst size) and a
|
|
4
|
+
steady ``refill_per_second`` rate. A ``clock`` callable can be injected so
|
|
5
|
+
tests can control time deterministically.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import threading
|
|
11
|
+
import time
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class TokenBucket:
|
|
15
|
+
"""A classic token bucket.
|
|
16
|
+
|
|
17
|
+
* ``capacity`` -- maximum burst size (tokens).
|
|
18
|
+
* ``refill_per_second`` -- steady refill rate (tokens per second).
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
def __init__(self, capacity: float, refill_per_second: float,
|
|
22
|
+
clock=time.monotonic) -> None:
|
|
23
|
+
if capacity <= 0:
|
|
24
|
+
raise ValueError("capacity must be > 0")
|
|
25
|
+
if refill_per_second <= 0:
|
|
26
|
+
raise ValueError("refill_per_second must be > 0")
|
|
27
|
+
self.capacity = float(capacity)
|
|
28
|
+
self.refill_per_second = float(refill_per_second)
|
|
29
|
+
self._clock = clock
|
|
30
|
+
self._lock = threading.Lock()
|
|
31
|
+
self._tokens = float(capacity)
|
|
32
|
+
self._last = float(clock())
|
|
33
|
+
|
|
34
|
+
def _refill(self, now: float) -> None:
|
|
35
|
+
elapsed = now - self._last
|
|
36
|
+
if elapsed > 0:
|
|
37
|
+
self._tokens = min(self.capacity,
|
|
38
|
+
self._tokens + elapsed * self.refill_per_second)
|
|
39
|
+
self._last = now
|
|
40
|
+
|
|
41
|
+
def try_consume(self, n: float = 1.0) -> bool:
|
|
42
|
+
"""Consume ``n`` tokens if available; return True, else leave state
|
|
43
|
+
untouched and return False."""
|
|
44
|
+
with self._lock:
|
|
45
|
+
self._refill(self._clock())
|
|
46
|
+
if self._tokens >= n:
|
|
47
|
+
self._tokens -= n
|
|
48
|
+
return True
|
|
49
|
+
return False
|
|
50
|
+
|
|
51
|
+
def peek(self) -> float:
|
|
52
|
+
"""Current token count after applying refill (does not consume)."""
|
|
53
|
+
with self._lock:
|
|
54
|
+
self._refill(self._clock())
|
|
55
|
+
return self._tokens
|
|
56
|
+
|
|
57
|
+
def deficit_wait(self, n: float = 1.0) -> float:
|
|
58
|
+
"""Seconds until ``n`` tokens are available (>= 0)."""
|
|
59
|
+
with self._lock:
|
|
60
|
+
self._refill(self._clock())
|
|
61
|
+
if self._tokens >= n:
|
|
62
|
+
return 0.0
|
|
63
|
+
return (n - self._tokens) / self.refill_per_second
|
|
64
|
+
|
|
65
|
+
def reset(self) -> None:
|
|
66
|
+
"""Refill the bucket to full capacity."""
|
|
67
|
+
with self._lock:
|
|
68
|
+
self._tokens = self.capacity
|
|
69
|
+
self._last = float(self._clock())
|
|
@@ -0,0 +1,117 @@
|
|
|
1
|
+
"""CLI for tool-throttle.
|
|
2
|
+
|
|
3
|
+
Subcommands:
|
|
4
|
+
check read {"tool": ...} JSON lines from stdin, print a throttling
|
|
5
|
+
decision per line (dry-run: does NOT consume tokens).
|
|
6
|
+
acquire read {"tool": ...} JSON lines from stdin, apply the throttle
|
|
7
|
+
(consumes tokens; queues or rejects), print the decision.
|
|
8
|
+
demo run a built-in scenario showing queue vs reject behaviour.
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from __future__ import annotations
|
|
12
|
+
|
|
13
|
+
import argparse
|
|
14
|
+
import json
|
|
15
|
+
import sys
|
|
16
|
+
|
|
17
|
+
from . import __version__
|
|
18
|
+
from .config import QUEUE, REJECT, Quota, ThrottlerConfig
|
|
19
|
+
from .throttler import ThrottledError, Throttler
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def demo_config() -> ThrottlerConfig:
|
|
23
|
+
return ThrottlerConfig(
|
|
24
|
+
default=Quota(calls=60, per_seconds=60),
|
|
25
|
+
per_tool={
|
|
26
|
+
"web_search": Quota(calls=10, per_seconds=60),
|
|
27
|
+
"expensive_*": Quota(calls=2, per_seconds=60),
|
|
28
|
+
},
|
|
29
|
+
on_exceed=REJECT,
|
|
30
|
+
max_queue_wait=5.0,
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _iter_stdin():
|
|
35
|
+
for line in sys.stdin:
|
|
36
|
+
line = line.strip()
|
|
37
|
+
if not line:
|
|
38
|
+
continue
|
|
39
|
+
try:
|
|
40
|
+
yield json.loads(line)
|
|
41
|
+
except json.JSONDecodeError as e:
|
|
42
|
+
print(json.dumps({"error": f"bad JSON: {e}"}), flush=True)
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def cmd_check(args) -> int:
|
|
46
|
+
cfg = ThrottlerConfig.load(args.config) if args.config else demo_config()
|
|
47
|
+
th = Throttler(cfg)
|
|
48
|
+
for obj in _iter_stdin():
|
|
49
|
+
tool = obj.get("tool", "")
|
|
50
|
+
d = th.decide(tool)
|
|
51
|
+
print(json.dumps({"tool": d.tool, "allowed": d.allowed,
|
|
52
|
+
"wait_seconds": round(d.wait_seconds, 3),
|
|
53
|
+
"reason": d.reason, "pattern": d.pattern}),
|
|
54
|
+
flush=True)
|
|
55
|
+
return 0
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def cmd_acquire(args) -> int:
|
|
59
|
+
cfg = ThrottlerConfig.load(args.config) if args.config else demo_config()
|
|
60
|
+
th = Throttler(cfg)
|
|
61
|
+
for obj in _iter_stdin():
|
|
62
|
+
tool = obj.get("tool", "")
|
|
63
|
+
try:
|
|
64
|
+
d = th.acquire(tool)
|
|
65
|
+
print(json.dumps({"tool": d.tool, "allowed": True,
|
|
66
|
+
"wait_seconds": round(d.wait_seconds, 3),
|
|
67
|
+
"reason": d.reason, "pattern": d.pattern}),
|
|
68
|
+
flush=True)
|
|
69
|
+
except ThrottledError as e:
|
|
70
|
+
print(json.dumps({"tool": tool, "allowed": False,
|
|
71
|
+
"wait_seconds": round(e.retry_after, 3),
|
|
72
|
+
"reason": e.reason, "pattern": None}),
|
|
73
|
+
flush=True)
|
|
74
|
+
return 0
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def cmd_demo(_args) -> int:
|
|
78
|
+
print("== tool-throttle demo: an agent looping web_search ==")
|
|
79
|
+
cfg = ThrottlerConfig(
|
|
80
|
+
default=Quota(calls=60, per_seconds=60),
|
|
81
|
+
per_tool={"web_search": Quota(calls=5, per_seconds=60)},
|
|
82
|
+
on_exceed=REJECT,
|
|
83
|
+
)
|
|
84
|
+
th = Throttler(cfg)
|
|
85
|
+
allowed, rejected = 0, 0
|
|
86
|
+
for _ in range(12):
|
|
87
|
+
try:
|
|
88
|
+
th.acquire("web_search")
|
|
89
|
+
allowed += 1
|
|
90
|
+
except ThrottledError:
|
|
91
|
+
rejected += 1
|
|
92
|
+
print(f"web_search: {allowed} allowed, {rejected} rejected "
|
|
93
|
+
f"(quota 5/min) -- the loop is fused, not looping forever")
|
|
94
|
+
print(f"stats: {json.dumps(th.stats(), indent=2)}")
|
|
95
|
+
return 0
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def main(argv=None) -> int:
|
|
99
|
+
p = argparse.ArgumentParser(prog="tool-throttle",
|
|
100
|
+
description="MCP tool-call throttler")
|
|
101
|
+
p.add_argument("--version", action="version", version=__version__)
|
|
102
|
+
sub = p.add_subparsers(dest="cmd", required=True)
|
|
103
|
+
for name, fn, help_ in (
|
|
104
|
+
("check", cmd_check, "dry-run decisions for stdin JSON lines"),
|
|
105
|
+
("acquire", cmd_acquire, "consume tokens for stdin JSON lines"),
|
|
106
|
+
("demo", cmd_demo, "show throttling on a looping agent")):
|
|
107
|
+
s = sub.add_parser(name, help=help_)
|
|
108
|
+
if name != "demo":
|
|
109
|
+
s.add_argument("--config", default=None,
|
|
110
|
+
help="path to a throttle config JSON file")
|
|
111
|
+
s.set_defaults(func=fn)
|
|
112
|
+
args = p.parse_args(argv)
|
|
113
|
+
return args.func(args)
|
|
114
|
+
|
|
115
|
+
|
|
116
|
+
if __name__ == "__main__":
|
|
117
|
+
sys.exit(main())
|
|
@@ -0,0 +1,102 @@
|
|
|
1
|
+
"""Configuration for tool-throttle: per-tool quotas, JSON round-trip."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import fnmatch
|
|
6
|
+
import json
|
|
7
|
+
from dataclasses import dataclass, field
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
@dataclass
|
|
11
|
+
class Quota:
|
|
12
|
+
"""Rate quota: ``calls`` tokens per ``per_seconds`` seconds.
|
|
13
|
+
|
|
14
|
+
``burst`` caps the burst size (defaults to ``calls``).
|
|
15
|
+
"""
|
|
16
|
+
calls: int
|
|
17
|
+
per_seconds: float = 60.0
|
|
18
|
+
burst: int | None = None
|
|
19
|
+
|
|
20
|
+
def __post_init__(self) -> None:
|
|
21
|
+
if self.calls <= 0:
|
|
22
|
+
raise ValueError("calls must be > 0")
|
|
23
|
+
if self.per_seconds <= 0:
|
|
24
|
+
raise ValueError("per_seconds must be > 0")
|
|
25
|
+
if self.burst is not None and self.burst <= 0:
|
|
26
|
+
raise ValueError("burst must be > 0")
|
|
27
|
+
|
|
28
|
+
@property
|
|
29
|
+
def capacity(self) -> float:
|
|
30
|
+
return float(self.burst if self.burst is not None else self.calls)
|
|
31
|
+
|
|
32
|
+
@property
|
|
33
|
+
def refill_per_second(self) -> float:
|
|
34
|
+
return float(self.calls) / float(self.per_seconds)
|
|
35
|
+
|
|
36
|
+
def to_dict(self) -> dict:
|
|
37
|
+
return {"calls": self.calls, "per_seconds": self.per_seconds,
|
|
38
|
+
"burst": self.burst}
|
|
39
|
+
|
|
40
|
+
@classmethod
|
|
41
|
+
def from_dict(cls, d: dict) -> "Quota":
|
|
42
|
+
return cls(calls=d["calls"], per_seconds=d.get("per_seconds", 60.0),
|
|
43
|
+
burst=d.get("burst"))
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
#: What to do when a call exceeds its quota.
|
|
47
|
+
QUEUE = "queue" #: block until a token is available (bounded by max_queue_wait)
|
|
48
|
+
REJECT = "reject" #: fail fast with ThrottledError
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
@dataclass
|
|
52
|
+
class ThrottlerConfig:
|
|
53
|
+
"""Full throttle configuration.
|
|
54
|
+
|
|
55
|
+
``per_tool`` maps glob patterns (e.g. ``"web_*"``) to quotas.
|
|
56
|
+
Patterns are matched with :func:`fnmatch.fnmatchcase` in insertion
|
|
57
|
+
order; the first match wins. Tools matching nothing use ``default``.
|
|
58
|
+
"""
|
|
59
|
+
default: Quota = field(default_factory=lambda: Quota(calls=60, per_seconds=60))
|
|
60
|
+
per_tool: dict[str, Quota] = field(default_factory=dict)
|
|
61
|
+
on_exceed: str = QUEUE
|
|
62
|
+
max_queue_wait: float = 30.0
|
|
63
|
+
|
|
64
|
+
def __post_init__(self) -> None:
|
|
65
|
+
if self.on_exceed not in (QUEUE, REJECT):
|
|
66
|
+
raise ValueError("on_exceed must be 'queue' or 'reject'")
|
|
67
|
+
if self.max_queue_wait < 0:
|
|
68
|
+
raise ValueError("max_queue_wait must be >= 0")
|
|
69
|
+
|
|
70
|
+
def quota_for(self, tool: str) -> tuple[Quota, str | None]:
|
|
71
|
+
"""Return (quota, matched_pattern) for a tool name."""
|
|
72
|
+
for pattern, quota in self.per_tool.items():
|
|
73
|
+
if fnmatch.fnmatchcase(tool, pattern):
|
|
74
|
+
return quota, pattern
|
|
75
|
+
return self.default, None
|
|
76
|
+
|
|
77
|
+
def to_dict(self) -> dict:
|
|
78
|
+
return {
|
|
79
|
+
"default": self.default.to_dict(),
|
|
80
|
+
"per_tool": {p: q.to_dict() for p, q in self.per_tool.items()},
|
|
81
|
+
"on_exceed": self.on_exceed,
|
|
82
|
+
"max_queue_wait": self.max_queue_wait,
|
|
83
|
+
}
|
|
84
|
+
|
|
85
|
+
@classmethod
|
|
86
|
+
def from_dict(cls, d: dict) -> "ThrottlerConfig":
|
|
87
|
+
return cls(
|
|
88
|
+
default=Quota.from_dict(d.get("default", {"calls": 60})),
|
|
89
|
+
per_tool={p: Quota.from_dict(q)
|
|
90
|
+
for p, q in d.get("per_tool", {}).items()},
|
|
91
|
+
on_exceed=d.get("on_exceed", QUEUE),
|
|
92
|
+
max_queue_wait=d.get("max_queue_wait", 30.0),
|
|
93
|
+
)
|
|
94
|
+
|
|
95
|
+
def save(self, path: str) -> None:
|
|
96
|
+
with open(path, "w") as f:
|
|
97
|
+
json.dump(self.to_dict(), f, indent=2)
|
|
98
|
+
|
|
99
|
+
@classmethod
|
|
100
|
+
def load(cls, path: str) -> "ThrottlerConfig":
|
|
101
|
+
with open(path) as f:
|
|
102
|
+
return cls.from_dict(json.load(f))
|
|
@@ -0,0 +1,162 @@
|
|
|
1
|
+
"""The MCP tool-call throttle: intercept tool calls, apply token-bucket
|
|
2
|
+
quotas per tool name, and either queue or reject over-quota calls.
|
|
3
|
+
|
|
4
|
+
This is the semantic-layer "fuse" for agents: it does not care about HTTP
|
|
5
|
+
or API keys, it cares about *tool call intent* -- e.g. an agent stuck in a
|
|
6
|
+
loop hammering ``web_search`` 400 times a minute.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import threading
|
|
12
|
+
import time
|
|
13
|
+
from dataclasses import dataclass
|
|
14
|
+
|
|
15
|
+
from .bucket import TokenBucket
|
|
16
|
+
from .config import QUEUE, REJECT, Quota, ThrottlerConfig
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class ThrottledError(Exception):
|
|
20
|
+
"""Raised when a tool call exceeds its quota.
|
|
21
|
+
|
|
22
|
+
``retry_after`` is seconds until the call would be allowed (0 if the
|
|
23
|
+
caller should not retry, e.g. queue-wait exhausted).
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
def __init__(self, tool: str, retry_after: float, reason: str) -> None:
|
|
27
|
+
super().__init__(f"tool call throttled: {tool} ({reason})")
|
|
28
|
+
self.tool = tool
|
|
29
|
+
self.retry_after = float(retry_after)
|
|
30
|
+
self.reason = reason
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
@dataclass
|
|
34
|
+
class Decision:
|
|
35
|
+
"""The outcome of a throttling check for one tool call."""
|
|
36
|
+
tool: str
|
|
37
|
+
allowed: bool
|
|
38
|
+
wait_seconds: float # 0 if allowed immediately
|
|
39
|
+
reason: str # "ok" | "quota_exceeded"
|
|
40
|
+
pattern: str | None = None # which per_tool pattern matched (if any)
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class Throttler:
|
|
44
|
+
"""Intercept and throttle tool calls by tool name.
|
|
45
|
+
|
|
46
|
+
Each distinct tool name gets its own :class:`TokenBucket` sized by the
|
|
47
|
+
first matching quota in ``config.per_tool`` (or ``config.default``).
|
|
48
|
+
"""
|
|
49
|
+
|
|
50
|
+
def __init__(self, config: ThrottlerConfig | None = None,
|
|
51
|
+
clock=time.monotonic) -> None:
|
|
52
|
+
self.config = config or ThrottlerConfig()
|
|
53
|
+
self._clock = clock
|
|
54
|
+
self._lock = threading.Lock()
|
|
55
|
+
self._cond = threading.Condition(self._lock)
|
|
56
|
+
self._buckets: dict[str, tuple[TokenBucket, Quota, str | None]] = {}
|
|
57
|
+
self._stats: dict[str, dict[str, float]] = {}
|
|
58
|
+
|
|
59
|
+
# -- internals ------------------------------------------------------
|
|
60
|
+
|
|
61
|
+
def _entry(self, tool: str) -> tuple[TokenBucket, Quota, str | None]:
|
|
62
|
+
entry = self._buckets.get(tool)
|
|
63
|
+
if entry is None:
|
|
64
|
+
quota, pattern = self.config.quota_for(tool)
|
|
65
|
+
entry = (TokenBucket(quota.capacity, quota.refill_per_second,
|
|
66
|
+
clock=self._clock), quota, pattern)
|
|
67
|
+
self._buckets[tool] = entry
|
|
68
|
+
self._stats[tool] = {"allowed": 0, "rejected": 0, "queued": 0,
|
|
69
|
+
"waited_seconds": 0.0,
|
|
70
|
+
"quota": f"{quota.calls}/{quota.per_seconds}s",
|
|
71
|
+
"pattern": pattern or "default"}
|
|
72
|
+
return entry
|
|
73
|
+
|
|
74
|
+
# -- public API -----------------------------------------------------
|
|
75
|
+
|
|
76
|
+
def decide(self, tool: str, n: float = 1.0) -> Decision:
|
|
77
|
+
"""Check whether a call would be allowed *without* consuming tokens."""
|
|
78
|
+
with self._lock:
|
|
79
|
+
bucket, _, pattern = self._entry(tool)
|
|
80
|
+
wait = bucket.deficit_wait(n)
|
|
81
|
+
if wait <= 0:
|
|
82
|
+
return Decision(tool, True, 0.0, "ok", pattern)
|
|
83
|
+
return Decision(tool, False, wait, "quota_exceeded", pattern)
|
|
84
|
+
|
|
85
|
+
def acquire(self, tool: str, n: float = 1.0) -> Decision:
|
|
86
|
+
"""Consume a token for ``tool``.
|
|
87
|
+
|
|
88
|
+
With ``on_exceed="queue"`` this blocks (up to ``max_queue_wait``
|
|
89
|
+
seconds) until a token is available, then proceeds. With
|
|
90
|
+
``on_exceed="reject"`` -- or when the wait would exceed the bound --
|
|
91
|
+
raises :class:`ThrottledError`.
|
|
92
|
+
"""
|
|
93
|
+
with self._lock:
|
|
94
|
+
bucket, quota, pattern = self._entry(tool)
|
|
95
|
+
stats = self._stats[tool]
|
|
96
|
+
if bucket.try_consume(n):
|
|
97
|
+
with self._lock:
|
|
98
|
+
stats["allowed"] += 1
|
|
99
|
+
return Decision(tool, True, 0.0, "ok", pattern)
|
|
100
|
+
|
|
101
|
+
wait_needed = bucket.deficit_wait(n)
|
|
102
|
+
remaining = self.config.max_queue_wait
|
|
103
|
+
if self.config.on_exceed == REJECT or wait_needed > remaining:
|
|
104
|
+
with self._lock:
|
|
105
|
+
stats["rejected"] += 1
|
|
106
|
+
raise ThrottledError(tool, wait_needed, "quota_exceeded")
|
|
107
|
+
|
|
108
|
+
# queue: wait in small slices until a token frees up or we time out
|
|
109
|
+
queued = True
|
|
110
|
+
start = self._clock()
|
|
111
|
+
deadline = start + remaining
|
|
112
|
+
while True:
|
|
113
|
+
sleep_for = min(wait_needed, 0.05)
|
|
114
|
+
with self._cond:
|
|
115
|
+
self._cond.wait(timeout=sleep_for)
|
|
116
|
+
if bucket.try_consume(n):
|
|
117
|
+
waited = max(0.0, self._clock() - start)
|
|
118
|
+
with self._lock:
|
|
119
|
+
stats["allowed"] += 1
|
|
120
|
+
if queued:
|
|
121
|
+
stats["queued"] += 1
|
|
122
|
+
stats["waited_seconds"] += waited
|
|
123
|
+
return Decision(tool, True, waited, "queued", pattern)
|
|
124
|
+
wait_needed = bucket.deficit_wait(n)
|
|
125
|
+
remaining = deadline - self._clock()
|
|
126
|
+
if remaining <= 0:
|
|
127
|
+
with self._lock:
|
|
128
|
+
stats["rejected"] += 1
|
|
129
|
+
raise ThrottledError(tool, wait_needed, "queue_wait_exceeded")
|
|
130
|
+
|
|
131
|
+
def wrap(self, tool: str, fn):
|
|
132
|
+
"""Return a wrapper that throttles calls to ``fn`` as ``tool``."""
|
|
133
|
+
def _wrapped(*args, **kwargs):
|
|
134
|
+
self.acquire(tool)
|
|
135
|
+
return fn(*args, **kwargs)
|
|
136
|
+
_wrapped.__name__ = getattr(fn, "__name__", f"throttled_{tool}")
|
|
137
|
+
_wrapped.__doc__ = getattr(fn, "__doc__", None)
|
|
138
|
+
return _wrapped
|
|
139
|
+
|
|
140
|
+
def call(self, tool: str, fn, *args, **kwargs):
|
|
141
|
+
"""Call ``fn(*args, **kwargs)`` through the throttle as ``tool``."""
|
|
142
|
+
self.acquire(tool)
|
|
143
|
+
return fn(*args, **kwargs)
|
|
144
|
+
|
|
145
|
+
def stats(self) -> dict[str, dict[str, float]]:
|
|
146
|
+
"""Per-tool counters: allowed / rejected / queued / waited_seconds."""
|
|
147
|
+
with self._lock:
|
|
148
|
+
return {t: dict(s) for t, s in self._stats.items()}
|
|
149
|
+
|
|
150
|
+
def reset(self, tool: str | None = None) -> None:
|
|
151
|
+
"""Refill bucket(s) to full. ``tool=None`` resets everything."""
|
|
152
|
+
with self._lock:
|
|
153
|
+
targets = [tool] if tool else list(self._buckets)
|
|
154
|
+
for t in targets:
|
|
155
|
+
if t in self._buckets:
|
|
156
|
+
self._buckets[t][0].reset()
|
|
157
|
+
|
|
158
|
+
def clear(self) -> None:
|
|
159
|
+
"""Drop all buckets and stats (fresh state, same config)."""
|
|
160
|
+
with self._lock:
|
|
161
|
+
self._buckets.clear()
|
|
162
|
+
self._stats.clear()
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: tool-throttle
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: MCP tool-call throttling: token-bucket rate limits per tool name, a fuse against runaway agents
|
|
5
|
+
Author: hao li
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/hahahahahahahahah6/tool-throttle
|
|
8
|
+
Keywords: mcp,ai-agents,rate-limit,throttling,token-bucket
|
|
9
|
+
Requires-Python: >=3.9
|
|
10
|
+
Description-Content-Type: text/markdown
|
|
11
|
+
License-File: LICENSE
|
|
12
|
+
Dynamic: license-file
|
|
13
|
+
|
|
14
|
+
# tool-throttle
|
|
15
|
+
|
|
16
|
+
[](LICENSE)
|
|
17
|
+
[](pyproject.toml)
|
|
18
|
+
[](pyproject.toml)
|
|
19
|
+
[](https://pypi.org/project/tool-throttle/)
|
|
20
|
+
|
|
21
|
+
A token-bucket rate limiter for **MCP tool calls** — a fuse against runaway
|
|
22
|
+
agents. When an agent gets stuck in a loop hammering `web_search` 400 times a
|
|
23
|
+
minute, tool-throttle slows it down or stops it, per tool name, with per-tool
|
|
24
|
+
quotas you configure in JSON.
|
|
25
|
+
|
|
26
|
+
Zero dependencies. Standard library only. Thread-safe.
|
|
27
|
+
|
|
28
|
+
```bash
|
|
29
|
+
pip install tool-throttle
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
```python
|
|
33
|
+
from tool_throttle import Throttler, ThrottlerConfig, Quota, REJECT
|
|
34
|
+
|
|
35
|
+
th = Throttler(ThrottlerConfig(
|
|
36
|
+
default=Quota(calls=60, per_seconds=60), # 60 calls/min for everything
|
|
37
|
+
per_tool={
|
|
38
|
+
"web_search": Quota(calls=10, per_seconds=60), # tight leash
|
|
39
|
+
"expensive_*": Quota(calls=2, per_seconds=60), # glob patterns work
|
|
40
|
+
},
|
|
41
|
+
on_exceed=REJECT, # or "queue": wait for a token, up to max_queue_wait
|
|
42
|
+
))
|
|
43
|
+
|
|
44
|
+
# wrap any callable -- your MCP tool dispatcher, a client method, anything:
|
|
45
|
+
safe_search = th.wrap("web_search", mcp_client.call_tool)
|
|
46
|
+
safe_search("web_search", {"query": "..."}) # -> ThrottledError past quota
|
|
47
|
+
|
|
48
|
+
# or decide explicitly without consuming a token:
|
|
49
|
+
d = th.decide("web_search")
|
|
50
|
+
# Decision(tool='web_search', allowed=False, wait_seconds=42.1, ...)
|
|
51
|
+
|
|
52
|
+
th.stats() # per-tool allowed / rejected / queued / waited_seconds
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
Config round-trips through JSON, so quotas can live in a file:
|
|
56
|
+
|
|
57
|
+
```python
|
|
58
|
+
cfg = ThrottlerConfig(
|
|
59
|
+
default=Quota(calls=60, per_seconds=60),
|
|
60
|
+
per_tool={"web_*": Quota(calls=10, per_seconds=60, burst=3)},
|
|
61
|
+
on_exceed="queue",
|
|
62
|
+
max_queue_wait=30.0,
|
|
63
|
+
)
|
|
64
|
+
cfg.save("throttle.json")
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
## CLI
|
|
68
|
+
|
|
69
|
+
```bash
|
|
70
|
+
# dry-run: print a decision per {"tool": ...} JSON line on stdin
|
|
71
|
+
echo '{"tool": "web_search"}' | tool-throttle check --config throttle.json
|
|
72
|
+
|
|
73
|
+
# enforce: consume tokens (queues or rejects over-quota calls)
|
|
74
|
+
echo '{"tool": "web_search"}' | tool-throttle acquire --config throttle.json
|
|
75
|
+
|
|
76
|
+
# built-in demo of a looping agent being fused
|
|
77
|
+
tool-throttle demo
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
## How it works
|
|
81
|
+
|
|
82
|
+
Each distinct tool name gets its own `TokenBucket`, sized by the first
|
|
83
|
+
matching quota in `per_tool` (glob patterns, insertion order) or `default`.
|
|
84
|
+
Buckets refill steadily over time, so a `Quota(calls=10, per_seconds=60)`
|
|
85
|
+
allows 10 calls per minute sustained, with a burst of up to `burst`
|
|
86
|
+
(default: `calls`).
|
|
87
|
+
|
|
88
|
+
On quota exhaustion:
|
|
89
|
+
- `on_exceed="reject"` (default in the demo): raise `ThrottledError` with a
|
|
90
|
+
`retry_after` hint — fail fast.
|
|
91
|
+
- `on_exceed="queue"`: block until a token frees up, bounded by
|
|
92
|
+
`max_queue_wait` seconds; if the wait would exceed the bound, raise
|
|
93
|
+
`ThrottledError(reason="queue_wait_exceeded")`.
|
|
94
|
+
|
|
95
|
+
## Differentiation: why not just rate-limit at the API gateway?
|
|
96
|
+
|
|
97
|
+
API gateways throttle **HTTP requests**. That is the wrong layer for agent
|
|
98
|
+
safety:
|
|
99
|
+
|
|
100
|
+
- One HTTP request can fan out into many tool calls; one tool call pattern
|
|
101
|
+
can span many HTTP endpoints. Gateway counters never see *tool-call
|
|
102
|
+
semantics* — they see bytes and paths.
|
|
103
|
+
- Gateway limits are per API key / per IP. An agent looping on a *single
|
|
104
|
+
expensive tool* (say, a search tool that costs $0.02 per call) looks
|
|
105
|
+
perfectly normal at the HTTP layer while burning money at the tool layer.
|
|
106
|
+
- Quotas here are expressed in the language of the agent: *tool names and
|
|
107
|
+
glob patterns* (`expensive_*`), with per-tool bursts and queue-or-reject
|
|
108
|
+
policies — the thing an agent developer actually wants to tune.
|
|
109
|
+
|
|
110
|
+
tool-throttle is the **fuse at the MCP tool-call semantic layer**: it sits
|
|
111
|
+
between the agent's tool-call intent and execution, trips before the loop
|
|
112
|
+
does real damage, and reports per-tool stats so you can see which tool the
|
|
113
|
+
agent was abusing. It complements gateway rate limiting; it does not
|
|
114
|
+
replace it.
|
|
115
|
+
|
|
116
|
+
## API reference
|
|
117
|
+
|
|
118
|
+
| Symbol | Description |
|
|
119
|
+
|---|---|
|
|
120
|
+
| `Throttler(config)` | the interceptor; one per agent/process |
|
|
121
|
+
| `ThrottlerConfig(default, per_tool, on_exceed, max_queue_wait)` | quotas + policy |
|
|
122
|
+
| `Quota(calls, per_seconds=60, burst=None)` | rate quota for one pattern |
|
|
123
|
+
| `th.acquire(tool)` | consume a token; queues or raises `ThrottledError` |
|
|
124
|
+
| `th.decide(tool)` | dry-run check, does not consume |
|
|
125
|
+
| `th.wrap(tool, fn)` / `th.call(tool, fn, *a, **k)` | throttle a callable |
|
|
126
|
+
| `th.stats()` / `th.reset(tool?)` / `th.clear()` | observability and reset |
|
|
127
|
+
| `ThrottledError` | `.tool`, `.retry_after`, `.reason` |
|
|
128
|
+
| `TokenBucket(capacity, refill_per_second, clock=...)` | the primitive underneath |
|
|
129
|
+
|
|
130
|
+
## License
|
|
131
|
+
|
|
132
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
tests/test_smoke.py
|
|
5
|
+
tool_throttle/__init__.py
|
|
6
|
+
tool_throttle/bucket.py
|
|
7
|
+
tool_throttle/cli.py
|
|
8
|
+
tool_throttle/config.py
|
|
9
|
+
tool_throttle/throttler.py
|
|
10
|
+
tool_throttle.egg-info/PKG-INFO
|
|
11
|
+
tool_throttle.egg-info/SOURCES.txt
|
|
12
|
+
tool_throttle.egg-info/dependency_links.txt
|
|
13
|
+
tool_throttle.egg-info/entry_points.txt
|
|
14
|
+
tool_throttle.egg-info/top_level.txt
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
tool_throttle
|