From 356e381d4ca50e1b0db256d997a33a8e75248f8f Mon Sep 17 00:00:00 2001 From: Vishu Bhatnagar Date: Tue, 13 Oct 2026 17:31:25 +0100 Subject: [PATCH] fix(rate-limiter): exclude prompts from per-server tool quotas Signed-off-by: Vishu Bhatnagar --- Cargo.lock | 2 +- .../python-package/rate_limiter/Cargo.toml | 2 +- .../python-package/rate_limiter/README.md | 14 ++- .../cpex_rate_limiter/plugin-manifest.yaml | 2 +- .../python-package/rate_limiter/src/engine.rs | 99 +++++++++++++++++++ .../python-package/rate_limiter/src/plugin.rs | 19 ++-- .../tests/rate_limiter/test_integration.py | 39 ++++++-- .../rate_limiter/test_redis_integration.py | 68 +++++++++++++ 8 files changed, 218 insertions(+), 27 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 0ef7e5b..957e05a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1545,7 +1545,7 @@ checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" [[package]] name = "rate_limiter" -version = "0.1.11" +version = "0.1.12" dependencies = [ "cpex_framework_bridge", "criterion", diff --git a/plugins/rust/python-package/rate_limiter/Cargo.toml b/plugins/rust/python-package/rate_limiter/Cargo.toml index f402081..d804147 100644 --- a/plugins/rust/python-package/rate_limiter/Cargo.toml +++ b/plugins/rust/python-package/rate_limiter/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "rate_limiter" -version = "0.1.11" +version = "0.1.12" edition.workspace = true authors.workspace = true license.workspace = true diff --git a/plugins/rust/python-package/rate_limiter/README.md b/plugins/rust/python-package/rate_limiter/README.md index 90b5089..38063f1 100644 --- a/plugins/rust/python-package/rate_limiter/README.md +++ b/plugins/rust/python-package/rate_limiter/README.md @@ -2,7 +2,7 @@ > Author: ContextForge Contributors -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). +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). ## Runtime Requirements @@ -13,7 +13,7 @@ This plugin depends on `cpex>=0.1.0,<0.2` and imports hook models from `cpex.fra | Hook | When it runs | |---|---| | `tool_pre_invoke` | Before every tool call — checks `by_user`, `by_user_per_server`, `by_tenant`, `by_tool` | -| `prompt_pre_fetch` | Before every prompt fetch — checks `by_user`, `by_user_per_server`, `by_tenant`, `by_tool` | +| `prompt_pre_fetch` | Before every prompt fetch — checks `by_user`, `by_tenant`, `by_tool` using the prompt name | 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). @@ -28,7 +28,7 @@ If any configured dimension is exceeded, the plugin returns a violation with HTT mode: enforce # enforce | permissive | disabled config: by_user: "300/m" # per-user limit across all servers - by_user_per_server: "60/m" # per-user limit for each MCP server + by_user_per_server: "60/m" # per-user tool limit for each upstream MCP catalog server by_tenant: "300/m" # shared limit across all users in a tenant by_tool: # per-tool overrides (applied on top of by_user) search: "10/m" @@ -56,7 +56,7 @@ If any configured dimension is exceeded, the plugin returns a violation with HTT | Field | Type | Default | Description | |---|---|---|---| | `by_user` | string | `null` | Per-user rate limit, e.g. `"60/m"` | -| `by_user_per_server` | string | `null` | Per-user rate limit independently enforced for each MCP server, e.g. `"60/m"` | +| `by_user_per_server` | string | `null` | Per-user tool rate limit independently enforced for each upstream MCP catalog server, e.g. `"60/m"` | | `by_tenant` | string | `null` | Per-tenant rate limit, e.g. `"600/m"` | | `by_tool` | dict | `{}` | Per-tool overrides, e.g. `{"search": "10/m"}` | | `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 **Omitting a dimension** (e.g. no `by_tenant`) means that dimension is unlimited — no counter is tracked for it. -`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. +`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. + +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. ## Response headers @@ -273,6 +275,8 @@ result.metadata["rate_limiter"] = { ## Migration Note +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. + 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. Version `0.1.7` is a **breaking change** for any existing consumer reading rate-limit metadata: diff --git a/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/plugin-manifest.yaml b/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/plugin-manifest.yaml index c7209de..24be2b8 100644 --- a/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/plugin-manifest.yaml +++ b/plugins/rust/python-package/rate_limiter/cpex_rate_limiter/plugin-manifest.yaml @@ -1,6 +1,6 @@ description: "Rate limiting by user, user/server, tenant, and tool — memory or Redis backends" author: "ContextForge Contributors" -version: "0.1.11" +version: "0.1.12" kind: "cpex_rate_limiter.rate_limiter.RateLimiterPlugin" available_hooks: - "prompt_pre_fetch" diff --git a/plugins/rust/python-package/rate_limiter/src/engine.rs b/plugins/rust/python-package/rate_limiter/src/engine.rs index 9f47ba5..41f9bca 100644 --- a/plugins/rust/python-package/rate_limiter/src/engine.rs +++ b/plugins/rust/python-package/rate_limiter/src/engine.rs @@ -75,6 +75,14 @@ impl RateLimiterEngine { } } + /// Prompt policy shares all counters and backend resources with tools, + /// but does not require or consume an upstream catalog quota. + pub(crate) fn for_prompts(&self) -> Self { + let mut engine = self.clone(); + engine.config.by_user_per_server = None; + engine + } + pub fn uses_async_backend(&self) -> bool { matches!(self.backend, EngineBackend::Redis(_)) } @@ -813,6 +821,97 @@ mod tests { ); } + #[test] + fn prompts_do_not_consume_server_quota() { + init_python(); + let (clock, _handle) = FakeClock::new(1_000_000); + let cfg = EngineConfig { + by_user: None, + by_user_per_server: Some(crate::config::parse_rate("1/h").unwrap()), + by_tenant: None, + by_tool: HashMap::new(), + algorithm: Algorithm::FixedWindow, + }; + let tools = RateLimiterEngine::new_with_clock(cfg, Arc::new(clock)); + let prompts = tools.for_prompts(); + for identity in [ + None, + Some("virtual-a"), + Some("virtual-b"), + Some("unsafe:id"), + ] { + let server = prompts.validated_server_id(identity).unwrap(); + let checks = prompts.build_checks("alice", None, "summary", None, server); + assert!(checks.is_empty()); + assert!(prompts.evaluate_many(checks, 1_000_000).unwrap().allowed); + } + let checks = tools.build_checks("alice", None, "search", None, Some("virtual-a")); + assert!( + tools + .evaluate_many(checks.clone(), 1_000_000) + .unwrap() + .allowed + ); + assert!(!tools.evaluate_many(checks, 1_000_000).unwrap().allowed); + assert!(tools.validated_server_id(None).is_err()); + } + + #[test] + fn prompt_limits_share_counters_with_tools_for_every_algorithm() { + init_python(); + for algorithm in [ + Algorithm::FixedWindow, + Algorithm::SlidingWindow, + Algorithm::TokenBucket, + ] { + for dimension in ["user", "tenant", "name"] { + let (clock, _handle) = FakeClock::new(1_000_000); + let limit = crate::config::parse_rate("2/h").unwrap(); + let cfg = EngineConfig { + by_user: (dimension == "user").then_some(limit.clone()), + by_user_per_server: Some(crate::config::parse_rate("10/h").unwrap()), + by_tenant: (dimension == "tenant").then_some(limit.clone()), + by_tool: if dimension == "name" { + HashMap::from([("summary".to_string(), limit)]) + } else { + HashMap::new() + }, + algorithm, + }; + let tools = RateLimiterEngine::new_with_clock(cfg, Arc::new(clock)); + let prompts = tools.for_prompts(); + let prompt_checks = + prompts.build_checks("alice", Some("acme"), "summary", Some("acme"), None); + let tool_checks = tools.build_checks( + "alice", + Some("acme"), + "summary", + Some("acme"), + Some("catalog-a"), + ); + assert!( + prompts + .evaluate_many(prompt_checks.clone(), 1_000_000) + .unwrap() + .allowed + ); + assert!( + tools + .evaluate_many(tool_checks.clone(), 1_000_000) + .unwrap() + .allowed + ); + assert!( + !prompts + .evaluate_many(prompt_checks, 1_000_000) + .unwrap() + .allowed + ); + assert!(!tools.evaluate_many(tool_checks, 1_000_000).unwrap().allowed); + } + } + } + #[test] fn per_server_user_limit_requires_safe_server_id() { init_python(); diff --git a/plugins/rust/python-package/rate_limiter/src/plugin.rs b/plugins/rust/python-package/rate_limiter/src/plugin.rs index 2b0fadb..aceb6fe 100644 --- a/plugins/rust/python-package/rate_limiter/src/plugin.rs +++ b/plugins/rust/python-package/rate_limiter/src/plugin.rs @@ -40,6 +40,7 @@ fn ensure_crypto_provider() { #[pyclass] pub struct RateLimiterPluginCore { engine: Arc, + prompt_engine: Arc, use_async: bool, /// When true, backend failures produce a BACKEND_UNAVAILABLE violation /// (fail-closed) instead of the default allow result (fail-open). Read @@ -57,10 +58,12 @@ impl RateLimiterPluginCore { // first TLS handshake. See `ensure_crypto_provider` doc above. ensure_crypto_provider(); let engine = Arc::new(RateLimiterEngine::new(config)?); + let prompt_engine = Arc::new(engine.for_prompts()); let fail_closed = parse_fail_mode(config)?; Ok(Self { use_async: engine.uses_async_backend(), engine, + prompt_engine, fail_closed, }) } @@ -87,25 +90,19 @@ impl RateLimiterPluginCore { .extract::()? .trim() .to_ascii_lowercase(); - let (user, tenant, server_id) = extract_request_context(context)?; - if let Err(message) = self.engine.validated_server_id(server_id.as_deref()) { - warn!("rate limiter: refusing prompt request: {message}"); - return Ok( - rate_limit_context_error_result(py, "PromptPrehookResult", message)?.into_bound(py), - ); - } + let (user, tenant, _) = extract_request_context(context)?; // Use tenant_id as the context prefix so that each team's rate limit // counters are isolated in Redis. Without this, all teams share keys. let context_prefix = tenant.as_deref(); let fail_closed = self.fail_closed; if !self.use_async { return match evaluate_sync_request( - &self.engine, + &self.prompt_engine, &user, tenant.as_deref(), &prompt, context_prefix, - server_id.as_deref(), + None, ) { Ok((allowed, headers, meta)) => Ok(build_prehook_result( py, @@ -127,7 +124,7 @@ impl RateLimiterPluginCore { }; } - let engine = Arc::clone(&self.engine); + let engine = Arc::clone(&self.prompt_engine); let context_prefix_owned = context_prefix.map(|s| s.to_string()); let trace_id_owned = trace_id.clone(); future_into_py(py, async move { @@ -137,7 +134,7 @@ impl RateLimiterPluginCore { tenant.as_deref(), &prompt, context_prefix_owned.as_deref(), - server_id.as_deref(), + None, ) .await { diff --git a/plugins/tests/rate_limiter/test_integration.py b/plugins/tests/rate_limiter/test_integration.py index 832f134..952fc18 100644 --- a/plugins/tests/rate_limiter/test_integration.py +++ b/plugins/tests/rate_limiter/test_integration.py @@ -323,19 +323,42 @@ async def test_blocked_over_limit(self, plugin): assert result.violation is not None assert result.violation.http_status_code == 429 - async def test_per_server_prompt_limits_are_independent(self): + @pytest.mark.parametrize("server_id", [None, "virtual-a", "virtual-b", "unsafe:id"]) + async def test_prompts_do_not_consume_per_server_tool_quota(self, server_id): plugin = RateLimiterPlugin( - _make_config(by_user=None, by_user_per_server="1/s") + _make_config(by_user=None, by_user_per_server="1/h") ) payload = PromptPrehookPayload(prompt_id="my-prompt") - server_a = _make_context(user="alice", server_id="server-a") - server_b = _make_context(user="alice", server_id="server-b") - - assert (await plugin.prompt_pre_fetch(payload, server_a)).continue_processing - blocked = await plugin.prompt_pre_fetch(payload, server_a) + context = _make_context(user="alice", server_id=server_id) + for _ in range(2): + result = await plugin.prompt_pre_fetch(payload, context) + assert result.continue_processing + assert result.violation is None + + tool = ToolPreInvokePayload(name="search") + catalog = _make_context(user="alice", server_id="virtual-a") + assert (await plugin.tool_pre_invoke(tool, catalog)).continue_processing + blocked = await plugin.tool_pre_invoke(tool, catalog) assert not blocked.continue_processing assert blocked.violation.code == "RATE_LIMIT" - assert (await plugin.prompt_pre_fetch(payload, server_b)).continue_processing + + async def test_prompt_and_tool_share_global_limit_with_per_server_enabled(self): + plugin = RateLimiterPlugin( + _make_config(by_user="2/h", by_user_per_server="5/h") + ) + prompt = PromptPrehookPayload(prompt_id="summary") + tool = ToolPreInvokePayload(name="search") + ordinary = _make_context(user="alice") + catalog = _make_context(user="alice", server_id="catalog-a") + assert (await plugin.prompt_pre_fetch(prompt, ordinary)).continue_processing + assert (await plugin.tool_pre_invoke(tool, catalog)).continue_processing + for hook, payload, context in [ + (plugin.prompt_pre_fetch, prompt, ordinary), + (plugin.tool_pre_invoke, tool, catalog), + ]: + blocked = await hook(payload, context) + assert not blocked.continue_processing + assert blocked.violation.code == "RATE_LIMIT" # --------------------------------------------------------------------------- diff --git a/plugins/tests/rate_limiter/test_redis_integration.py b/plugins/tests/rate_limiter/test_redis_integration.py index 79686c3..927e7fc 100644 --- a/plugins/tests/rate_limiter/test_redis_integration.py +++ b/plugins/tests/rate_limiter/test_redis_integration.py @@ -1161,6 +1161,74 @@ async def test_redis_per_server_user_counters_are_shared_and_isolated(self, redi assert keys_a == ["rl:user:alice:server:server-a:1"] assert keys_b == ["rl:user:alice:server:server-b:1"] + @pytest.mark.asyncio + async def test_redis_prompts_leave_catalog_quota_untouched(self, redis_url_for_integration): + """Prompts never create server keys or consume the tools' catalog quota.""" + await _flush_redis(redis_url_for_integration) + plugins = [ + _make_redis_plugin(redis_url_for_integration, limit=None, by_user_per_server="1/h") + for _ in range(2) + ] + try: + prompt = PromptPrehookPayload(prompt_id="summary", args={}) + for server_id in [None, "catalog-a", "virtual-b", "unsafe:id"]: + context = PluginContext( + global_context=GlobalContext(request_id="prompt", user="alice", server_id=server_id) + ) + for plugin in plugins: + result = await plugin.prompt_pre_fetch(prompt, context) + assert result.continue_processing + assert result.violation is None + assert await _keys_in_redis(redis_url_for_integration, "rl:*:server:*") == [] + + tool = ToolPreInvokePayload(name="search", arguments={}) + catalog = PluginContext( + global_context=GlobalContext(request_id="tool", user="alice", server_id="catalog-a") + ) + assert (await plugins[0].tool_pre_invoke(tool, catalog)).continue_processing + blocked = await plugins[1].tool_pre_invoke(tool, catalog) + assert not blocked.continue_processing + assert blocked.violation.code == "RATE_LIMIT" + finally: + for plugin in plugins: + await plugin.shutdown() + + @pytest.mark.asyncio + @pytest.mark.parametrize("dimension", ["user", "tenant", "prompt_name"]) + async def test_redis_prompt_limits_share_counters_with_tools_when_per_server_enabled( + self, redis_url_for_integration, dimension + ): + """The prompt policy preserves existing limits across hooks and replicas.""" + await _flush_redis(redis_url_for_integration) + config = { + "by_user": "2/h" if dimension == "user" else None, + "by_tenant": "2/h" if dimension == "tenant" else None, + "by_tool": {"summary": "2/h"} if dimension == "prompt_name" else {}, + "by_user_per_server": "10/h", + } + plugins = [_make_redis_plugin(redis_url_for_integration, **config) for _ in range(2)] + try: + prompt = PromptPrehookPayload(prompt_id="summary", args={}) + tool = ToolPreInvokePayload(name="summary", arguments={}) + ordinary = PluginContext( + global_context=GlobalContext(request_id="prompt", user="alice", tenant_id="acme") + ) + catalog = PluginContext( + global_context=GlobalContext(request_id="tool", user="alice", tenant_id="acme", server_id="catalog-a") + ) + assert (await plugins[0].prompt_pre_fetch(prompt, ordinary)).continue_processing + assert (await plugins[1].tool_pre_invoke(tool, catalog)).continue_processing + for hook, payload, context in [ + (plugins[1].prompt_pre_fetch, prompt, ordinary), + (plugins[0].tool_pre_invoke, tool, catalog), + ]: + blocked = await hook(payload, context) + assert not blocked.continue_processing + assert blocked.violation.code == "RATE_LIMIT" + finally: + for plugin in plugins: + await plugin.shutdown() + @pytest.mark.asyncio async def test_redis_window_resets_after_ttl(self, redis_url_for_integration): """After the rate window expires, Redis TTL resets counters and requests are allowed again."""