Skip to content

Commit 1b4577f

Browse files
authored
fix(rate-limiter): exclude prompts from per-server tool quotas (#203)
Signed-off-by: Vishu Bhatnagar <vishu.bhatnagar@ibm.com>
1 parent 57fc9f0 commit 1b4577f

8 files changed

Lines changed: 218 additions & 27 deletions

File tree

‎Cargo.lock‎

Lines changed: 1 addition & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎plugins/rust/python-package/rate_limiter/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name = "rate_limiter"
3-
version = "0.1.11"
3+
version = "0.1.12"
44
edition.workspace = true
55
authors.workspace = true
66
license.workspace = true

‎plugins/rust/python-package/rate_limiter/README.md‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
> Author: ContextForge Contributors
44
5-
Enforces global per-user, per-user/per-server, tenant, and tool rate limits across `tool_pre_invoke` and `prompt_pre_fetch` hooks. Supports pluggable counting algorithms (fixed window, sliding window, token bucket), an in-process memory backend (single-instance), and a Redis backend (shared across all gateway instances).
5+
Enforces global per-user, tenant, and tool/prompt-name rate limits across `tool_pre_invoke` and `prompt_pre_fetch` hooks, with an additional per-user/per-server quota for tool invocations. Supports pluggable counting algorithms (fixed window, sliding window, token bucket), an in-process memory backend (single-instance), and a Redis backend (shared across all gateway instances).
66

77
## Runtime Requirements
88

@@ -13,7 +13,7 @@ This plugin depends on `cpex>=0.1.0,<0.2` and imports hook models from `cpex.fra
1313
| Hook | When it runs |
1414
|---|---|
1515
| `tool_pre_invoke` | Before every tool call — checks `by_user`, `by_user_per_server`, `by_tenant`, `by_tool` |
16-
| `prompt_pre_fetch` | Before every prompt fetch — checks `by_user`, `by_user_per_server`, `by_tenant`, `by_tool` |
16+
| `prompt_pre_fetch` | Before every prompt fetch — checks `by_user`, `by_tenant`, `by_tool` using the prompt name |
1717

1818
If any configured dimension is exceeded, the plugin returns a violation with HTTP 429. All requests include `X-RateLimit-*` headers. The most restrictive active dimension is surfaced (e.g. if both user and tenant limits are active, the one closest to exhaustion is reported).
1919

@@ -28,7 +28,7 @@ If any configured dimension is exceeded, the plugin returns a violation with HTT
2828
mode: enforce # enforce | permissive | disabled
2929
config:
3030
by_user: "300/m" # per-user limit across all servers
31-
by_user_per_server: "60/m" # per-user limit for each MCP server
31+
by_user_per_server: "60/m" # per-user tool limit for each upstream MCP catalog server
3232
by_tenant: "300/m" # shared limit across all users in a tenant
3333
by_tool: # per-tool overrides (applied on top of by_user)
3434
search: "10/m"
@@ -56,7 +56,7 @@ If any configured dimension is exceeded, the plugin returns a violation with HTT
5656
| Field | Type | Default | Description |
5757
|---|---|---|---|
5858
| `by_user` | string | `null` | Per-user rate limit, e.g. `"60/m"` |
59-
| `by_user_per_server` | string | `null` | Per-user rate limit independently enforced for each MCP server, e.g. `"60/m"` |
59+
| `by_user_per_server` | string | `null` | Per-user tool rate limit independently enforced for each upstream MCP catalog server, e.g. `"60/m"` |
6060
| `by_tenant` | string | `null` | Per-tenant rate limit, e.g. `"600/m"` |
6161
| `by_tool` | dict | `{}` | Per-tool overrides, e.g. `{"search": "10/m"}` |
6262
| `algorithm` | string | `"fixed_window"` | Counting algorithm: `"fixed_window"`, `"sliding_window"`, or `"token_bucket"` |
@@ -77,7 +77,9 @@ If any configured dimension is exceeded, the plugin returns a violation with HTT
7777

7878
**Omitting a dimension** (e.g. no `by_tenant`) means that dimension is unlimited — no counter is tracked for it.
7979

80-
`by_user` and `by_user_per_server` are independent and may be enabled together. For example, `by_user: "300/m"` plus `by_user_per_server: "60/m"` gives each user at most 60 requests per server and 300 requests across all servers. This prevents one user from multiplying their effective quota without relying on the shared `by_tenant` ceiling. When `by_user_per_server` is configured, requests without a safe non-empty `server_id` are blocked with `RATE_LIMIT_CONTEXT_MISSING` (HTTP 503) rather than bypassing the limit.
80+
`by_user` and `by_user_per_server` are independent and may be enabled together. For example, `by_user: "300/m"` plus `by_user_per_server: "60/m"` gives each user at most 60 tool calls per catalog server and 300 combined tool calls and prompt fetches within the same tenant context. This prevents one user from multiplying their effective quota without relying on the shared `by_tenant` ceiling. When `by_user_per_server` is configured, tool calls without a safe non-empty `server_id` are blocked with `RATE_LIMIT_CONTEXT_MISSING` (HTTP 503) rather than bypassing the limit.
81+
82+
Prompt fetches neither require a catalog `server_id` nor consume `by_user_per_server` counters, even when a virtual-server ID is present. They still enforce `by_user`, `by_tenant`, and prompt-name overrides in `by_tool`. These counters remain shared with tool invocations within the same tenant context.
8183

8284
## Response headers
8385

@@ -273,6 +275,8 @@ result.metadata["rate_limiter"] = {
273275

274276
## Migration Note
275277

278+
Version `0.1.12` fixes prompt retrieval when `by_user_per_server` is enabled. The per-server quota now applies only to tool invocations; prompt fetches retain existing global and prompt-name limits. Counter keys and tool identity validation are unchanged. Upgrade every replica to `0.1.12` or later before enabling per-server tool quotas with prompt hooks registered.
279+
276280
Version `0.1.11` adds the optional `by_user_per_server` dimension. Existing configurations and global `by_user` Redis keys remain unchanged. Enabling the new field creates fresh user/server counters; deploy `0.1.11` to every replica before enabling it so all replicas enforce the same dimensions. Use the stable catalog UUID carried in `GlobalContext.server_id`; display names or slugs can reset quotas when renamed.
277281

278282
Version `0.1.7` is a **breaking change** for any existing consumer reading rate-limit metadata:

‎plugins/rust/python-package/rate_limiter/cpex_rate_limiter/plugin-manifest.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
description: "Rate limiting by user, user/server, tenant, and tool — memory or Redis backends"
22
author: "ContextForge Contributors"
3-
version: "0.1.11"
3+
version: "0.1.12"
44
kind: "cpex_rate_limiter.rate_limiter.RateLimiterPlugin"
55
available_hooks:
66
- "prompt_pre_fetch"

‎plugins/rust/python-package/rate_limiter/src/engine.rs‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,14 @@ impl RateLimiterEngine {
7575
}
7676
}
7777

78+
/// Prompt policy shares all counters and backend resources with tools,
79+
/// but does not require or consume an upstream catalog quota.
80+
pub(crate) fn for_prompts(&self) -> Self {
81+
let mut engine = self.clone();
82+
engine.config.by_user_per_server = None;
83+
engine
84+
}
85+
7886
pub fn uses_async_backend(&self) -> bool {
7987
matches!(self.backend, EngineBackend::Redis(_))
8088
}
@@ -813,6 +821,97 @@ mod tests {
813821
);
814822
}
815823

824+
#[test]
825+
fn prompts_do_not_consume_server_quota() {
826+
init_python();
827+
let (clock, _handle) = FakeClock::new(1_000_000);
828+
let cfg = EngineConfig {
829+
by_user: None,
830+
by_user_per_server: Some(crate::config::parse_rate("1/h").unwrap()),
831+
by_tenant: None,
832+
by_tool: HashMap::new(),
833+
algorithm: Algorithm::FixedWindow,
834+
};
835+
let tools = RateLimiterEngine::new_with_clock(cfg, Arc::new(clock));
836+
let prompts = tools.for_prompts();
837+
for identity in [
838+
None,
839+
Some("virtual-a"),
840+
Some("virtual-b"),
841+
Some("unsafe:id"),
842+
] {
843+
let server = prompts.validated_server_id(identity).unwrap();
844+
let checks = prompts.build_checks("alice", None, "summary", None, server);
845+
assert!(checks.is_empty());
846+
assert!(prompts.evaluate_many(checks, 1_000_000).unwrap().allowed);
847+
}
848+
let checks = tools.build_checks("alice", None, "search", None, Some("virtual-a"));
849+
assert!(
850+
tools
851+
.evaluate_many(checks.clone(), 1_000_000)
852+
.unwrap()
853+
.allowed
854+
);
855+
assert!(!tools.evaluate_many(checks, 1_000_000).unwrap().allowed);
856+
assert!(tools.validated_server_id(None).is_err());
857+
}
858+
859+
#[test]
860+
fn prompt_limits_share_counters_with_tools_for_every_algorithm() {
861+
init_python();
862+
for algorithm in [
863+
Algorithm::FixedWindow,
864+
Algorithm::SlidingWindow,
865+
Algorithm::TokenBucket,
866+
] {
867+
for dimension in ["user", "tenant", "name"] {
868+
let (clock, _handle) = FakeClock::new(1_000_000);
869+
let limit = crate::config::parse_rate("2/h").unwrap();
870+
let cfg = EngineConfig {
871+
by_user: (dimension == "user").then_some(limit.clone()),
872+
by_user_per_server: Some(crate::config::parse_rate("10/h").unwrap()),
873+
by_tenant: (dimension == "tenant").then_some(limit.clone()),
874+
by_tool: if dimension == "name" {
875+
HashMap::from([("summary".to_string(), limit)])
876+
} else {
877+
HashMap::new()
878+
},
879+
algorithm,
880+
};
881+
let tools = RateLimiterEngine::new_with_clock(cfg, Arc::new(clock));
882+
let prompts = tools.for_prompts();
883+
let prompt_checks =
884+
prompts.build_checks("alice", Some("acme"), "summary", Some("acme"), None);
885+
let tool_checks = tools.build_checks(
886+
"alice",
887+
Some("acme"),
888+
"summary",
889+
Some("acme"),
890+
Some("catalog-a"),
891+
);
892+
assert!(
893+
prompts
894+
.evaluate_many(prompt_checks.clone(), 1_000_000)
895+
.unwrap()
896+
.allowed
897+
);
898+
assert!(
899+
tools
900+
.evaluate_many(tool_checks.clone(), 1_000_000)
901+
.unwrap()
902+
.allowed
903+
);
904+
assert!(
905+
!prompts
906+
.evaluate_many(prompt_checks, 1_000_000)
907+
.unwrap()
908+
.allowed
909+
);
910+
assert!(!tools.evaluate_many(tool_checks, 1_000_000).unwrap().allowed);
911+
}
912+
}
913+
}
914+
816915
#[test]
817916
fn per_server_user_limit_requires_safe_server_id() {
818917
init_python();

‎plugins/rust/python-package/rate_limiter/src/plugin.rs‎

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ fn ensure_crypto_provider() {
4040
#[pyclass]
4141
pub struct RateLimiterPluginCore {
4242
engine: Arc<RateLimiterEngine>,
43+
prompt_engine: Arc<RateLimiterEngine>,
4344
use_async: bool,
4445
/// When true, backend failures produce a BACKEND_UNAVAILABLE violation
4546
/// (fail-closed) instead of the default allow result (fail-open). Read
@@ -57,10 +58,12 @@ impl RateLimiterPluginCore {
5758
// first TLS handshake. See `ensure_crypto_provider` doc above.
5859
ensure_crypto_provider();
5960
let engine = Arc::new(RateLimiterEngine::new(config)?);
61+
let prompt_engine = Arc::new(engine.for_prompts());
6062
let fail_closed = parse_fail_mode(config)?;
6163
Ok(Self {
6264
use_async: engine.uses_async_backend(),
6365
engine,
66+
prompt_engine,
6467
fail_closed,
6568
})
6669
}
@@ -87,25 +90,19 @@ impl RateLimiterPluginCore {
8790
.extract::<String>()?
8891
.trim()
8992
.to_ascii_lowercase();
90-
let (user, tenant, server_id) = extract_request_context(context)?;
91-
if let Err(message) = self.engine.validated_server_id(server_id.as_deref()) {
92-
warn!("rate limiter: refusing prompt request: {message}");
93-
return Ok(
94-
rate_limit_context_error_result(py, "PromptPrehookResult", message)?.into_bound(py),
95-
);
96-
}
93+
let (user, tenant, _) = extract_request_context(context)?;
9794
// Use tenant_id as the context prefix so that each team's rate limit
9895
// counters are isolated in Redis. Without this, all teams share keys.
9996
let context_prefix = tenant.as_deref();
10097
let fail_closed = self.fail_closed;
10198
if !self.use_async {
10299
return match evaluate_sync_request(
103-
&self.engine,
100+
&self.prompt_engine,
104101
&user,
105102
tenant.as_deref(),
106103
&prompt,
107104
context_prefix,
108-
server_id.as_deref(),
105+
None,
109106
) {
110107
Ok((allowed, headers, meta)) => Ok(build_prehook_result(
111108
py,
@@ -127,7 +124,7 @@ impl RateLimiterPluginCore {
127124
};
128125
}
129126

130-
let engine = Arc::clone(&self.engine);
127+
let engine = Arc::clone(&self.prompt_engine);
131128
let context_prefix_owned = context_prefix.map(|s| s.to_string());
132129
let trace_id_owned = trace_id.clone();
133130
future_into_py(py, async move {
@@ -137,7 +134,7 @@ impl RateLimiterPluginCore {
137134
tenant.as_deref(),
138135
&prompt,
139136
context_prefix_owned.as_deref(),
140-
server_id.as_deref(),
137+
None,
141138
)
142139
.await
143140
{

‎plugins/tests/rate_limiter/test_integration.py‎

Lines changed: 31 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -323,19 +323,42 @@ async def test_blocked_over_limit(self, plugin):
323323
assert result.violation is not None
324324
assert result.violation.http_status_code == 429
325325

326-
async def test_per_server_prompt_limits_are_independent(self):
326+
@pytest.mark.parametrize("server_id", [None, "virtual-a", "virtual-b", "unsafe:id"])
327+
async def test_prompts_do_not_consume_per_server_tool_quota(self, server_id):
327328
plugin = RateLimiterPlugin(
328-
_make_config(by_user=None, by_user_per_server="1/s")
329+
_make_config(by_user=None, by_user_per_server="1/h")
329330
)
330331
payload = PromptPrehookPayload(prompt_id="my-prompt")
331-
server_a = _make_context(user="alice", server_id="server-a")
332-
server_b = _make_context(user="alice", server_id="server-b")
333-
334-
assert (await plugin.prompt_pre_fetch(payload, server_a)).continue_processing
335-
blocked = await plugin.prompt_pre_fetch(payload, server_a)
332+
context = _make_context(user="alice", server_id=server_id)
333+
for _ in range(2):
334+
result = await plugin.prompt_pre_fetch(payload, context)
335+
assert result.continue_processing
336+
assert result.violation is None
337+
338+
tool = ToolPreInvokePayload(name="search")
339+
catalog = _make_context(user="alice", server_id="virtual-a")
340+
assert (await plugin.tool_pre_invoke(tool, catalog)).continue_processing
341+
blocked = await plugin.tool_pre_invoke(tool, catalog)
336342
assert not blocked.continue_processing
337343
assert blocked.violation.code == "RATE_LIMIT"
338-
assert (await plugin.prompt_pre_fetch(payload, server_b)).continue_processing
344+
345+
async def test_prompt_and_tool_share_global_limit_with_per_server_enabled(self):
346+
plugin = RateLimiterPlugin(
347+
_make_config(by_user="2/h", by_user_per_server="5/h")
348+
)
349+
prompt = PromptPrehookPayload(prompt_id="summary")
350+
tool = ToolPreInvokePayload(name="search")
351+
ordinary = _make_context(user="alice")
352+
catalog = _make_context(user="alice", server_id="catalog-a")
353+
assert (await plugin.prompt_pre_fetch(prompt, ordinary)).continue_processing
354+
assert (await plugin.tool_pre_invoke(tool, catalog)).continue_processing
355+
for hook, payload, context in [
356+
(plugin.prompt_pre_fetch, prompt, ordinary),
357+
(plugin.tool_pre_invoke, tool, catalog),
358+
]:
359+
blocked = await hook(payload, context)
360+
assert not blocked.continue_processing
361+
assert blocked.violation.code == "RATE_LIMIT"
339362

340363

341364
# ---------------------------------------------------------------------------

‎plugins/tests/rate_limiter/test_redis_integration.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1161,6 +1161,74 @@ async def test_redis_per_server_user_counters_are_shared_and_isolated(self, redi
11611161
assert keys_a == ["rl:user:alice:server:server-a:1"]
11621162
assert keys_b == ["rl:user:alice:server:server-b:1"]
11631163

1164+
@pytest.mark.asyncio
1165+
async def test_redis_prompts_leave_catalog_quota_untouched(self, redis_url_for_integration):
1166+
"""Prompts never create server keys or consume the tools' catalog quota."""
1167+
await _flush_redis(redis_url_for_integration)
1168+
plugins = [
1169+
_make_redis_plugin(redis_url_for_integration, limit=None, by_user_per_server="1/h")
1170+
for _ in range(2)
1171+
]
1172+
try:
1173+
prompt = PromptPrehookPayload(prompt_id="summary", args={})
1174+
for server_id in [None, "catalog-a", "virtual-b", "unsafe:id"]:
1175+
context = PluginContext(
1176+
global_context=GlobalContext(request_id="prompt", user="alice", server_id=server_id)
1177+
)
1178+
for plugin in plugins:
1179+
result = await plugin.prompt_pre_fetch(prompt, context)
1180+
assert result.continue_processing
1181+
assert result.violation is None
1182+
assert await _keys_in_redis(redis_url_for_integration, "rl:*:server:*") == []
1183+
1184+
tool = ToolPreInvokePayload(name="search", arguments={})
1185+
catalog = PluginContext(
1186+
global_context=GlobalContext(request_id="tool", user="alice", server_id="catalog-a")
1187+
)
1188+
assert (await plugins[0].tool_pre_invoke(tool, catalog)).continue_processing
1189+
blocked = await plugins[1].tool_pre_invoke(tool, catalog)
1190+
assert not blocked.continue_processing
1191+
assert blocked.violation.code == "RATE_LIMIT"
1192+
finally:
1193+
for plugin in plugins:
1194+
await plugin.shutdown()
1195+
1196+
@pytest.mark.asyncio
1197+
@pytest.mark.parametrize("dimension", ["user", "tenant", "prompt_name"])
1198+
async def test_redis_prompt_limits_share_counters_with_tools_when_per_server_enabled(
1199+
self, redis_url_for_integration, dimension
1200+
):
1201+
"""The prompt policy preserves existing limits across hooks and replicas."""
1202+
await _flush_redis(redis_url_for_integration)
1203+
config = {
1204+
"by_user": "2/h" if dimension == "user" else None,
1205+
"by_tenant": "2/h" if dimension == "tenant" else None,
1206+
"by_tool": {"summary": "2/h"} if dimension == "prompt_name" else {},
1207+
"by_user_per_server": "10/h",
1208+
}
1209+
plugins = [_make_redis_plugin(redis_url_for_integration, **config) for _ in range(2)]
1210+
try:
1211+
prompt = PromptPrehookPayload(prompt_id="summary", args={})
1212+
tool = ToolPreInvokePayload(name="summary", arguments={})
1213+
ordinary = PluginContext(
1214+
global_context=GlobalContext(request_id="prompt", user="alice", tenant_id="acme")
1215+
)
1216+
catalog = PluginContext(
1217+
global_context=GlobalContext(request_id="tool", user="alice", tenant_id="acme", server_id="catalog-a")
1218+
)
1219+
assert (await plugins[0].prompt_pre_fetch(prompt, ordinary)).continue_processing
1220+
assert (await plugins[1].tool_pre_invoke(tool, catalog)).continue_processing
1221+
for hook, payload, context in [
1222+
(plugins[1].prompt_pre_fetch, prompt, ordinary),
1223+
(plugins[0].tool_pre_invoke, tool, catalog),
1224+
]:
1225+
blocked = await hook(payload, context)
1226+
assert not blocked.continue_processing
1227+
assert blocked.violation.code == "RATE_LIMIT"
1228+
finally:
1229+
for plugin in plugins:
1230+
await plugin.shutdown()
1231+
11641232
@pytest.mark.asyncio
11651233
async def test_redis_window_resets_after_ttl(self, redis_url_for_integration):
11661234
"""After the rate window expires, Redis TTL resets counters and requests are allowed again."""

0 commit comments

Comments
 (0)