Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion plugins/rust/python-package/rate_limiter/Cargo.toml
Original file line number Diff line number Diff line change
@@ -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
Expand Down
14 changes: 9 additions & 5 deletions plugins/rust/python-package/rate_limiter/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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).

Expand All @@ -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"
Expand Down Expand Up @@ -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"` |
Expand All @@ -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

Expand Down Expand Up @@ -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:
Expand Down
Original file line number Diff line number Diff line change
@@ -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"
Expand Down
99 changes: 99 additions & 0 deletions plugins/rust/python-package/rate_limiter/src/engine.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(_))
}
Expand Down Expand Up @@ -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();
Expand Down
19 changes: 8 additions & 11 deletions plugins/rust/python-package/rate_limiter/src/plugin.rs
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@ fn ensure_crypto_provider() {
#[pyclass]
pub struct RateLimiterPluginCore {
engine: Arc<RateLimiterEngine>,
prompt_engine: Arc<RateLimiterEngine>,
use_async: bool,
/// When true, backend failures produce a BACKEND_UNAVAILABLE violation
/// (fail-closed) instead of the default allow result (fail-open). Read
Expand All @@ -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,
})
}
Expand All @@ -87,25 +90,19 @@ impl RateLimiterPluginCore {
.extract::<String>()?
.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,
Expand All @@ -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 {
Expand All @@ -137,7 +134,7 @@ impl RateLimiterPluginCore {
tenant.as_deref(),
&prompt,
context_prefix_owned.as_deref(),
server_id.as_deref(),
None,
)
.await
{
Expand Down
39 changes: 31 additions & 8 deletions plugins/tests/rate_limiter/test_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"


# ---------------------------------------------------------------------------
Expand Down
68 changes: 68 additions & 0 deletions plugins/tests/rate_limiter/test_redis_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down
Loading