security: harden the admin API and the MCP proxy

README.md

@@ -10,14 +10,14 @@ key-based auth, quotas, rate limits, failover, and a billing ledger.
 
 - **OpenAI + Anthropic compatible surface** — `/v1/chat/completions`, `/v1/completions`, `/v1/responses`, `/v1/messages`, `/v1/embeddings`, `/v1/images/{generations,edits}`, `/v1/videos/generations` + `/v1/videos/{id}`, `/v1/audio/{speech,transcriptions,translations}`, `/v1/moderations`, `/v1/search`, `/v1/rerank`, `/v1/batches` + `/v1/files`, `/v1/models`, `/v1/realtime` (WebSocket) — streaming and non-streaming
 - **Cross-protocol conversion** — serve Anthropic-style `/v1/messages` on OpenAI-protocol models and vice versa, including streaming event mapping
-- **MCP gateway** — `/mcp/{server}` proxies Model Context Protocol servers (Streamable HTTP) behind the same access keys: per-key server entitlement and tool allowlists, `tools/list` filtered to the allowlist, tool results reviewed by the tenant's moderator (mask or block), every tool call, denial and intervention audited, server credentials kept in the gateway's environment as static bearers or OAuth 2.0 client credentials the gateway refreshes itself
+- **MCP gateway** — `/mcp/{server}` proxies Model Context Protocol servers (Streamable HTTP) behind the same access keys: per-key server entitlement and tool allowlists, `tools/list` filtered to the allowlist, tool results reviewed by the tenant's moderator (mask or block), every tool call, denial and intervention audited, sessions bound to the key that opened them, server credentials kept in the gateway's environment as static bearers or OAuth 2.0 client credentials the gateway refreshes itself ([Security model](docs/security.md))
 - **Coding agents drop in** — Claude Code, Codex CLI, VS Code chat (Copilot), Cursor and opencode work with a base URL and an access key; the `anthropic-beta` header rides through to Anthropic-wire upstreams and each client's captured wire shape replays in the live matrix ([Examples](docs/examples.md#coding-agents))
 - **Reasoning on every surface** — `reasoning_effort` / `reasoning{}` on `/v1/chat/completions` maps to each family's thinking dialect (Anthropic budget or adaptive by model generation, OpenAI effort, compatible vendors verbatim); reasoning comes back as `reasoning_content` + signed `reasoning_details` and replays into tool loops; `/v1/responses` forwards reasoning items and its native event stream verbatim; `/v1/messages` preserves signed `thinking`/`redacted_thinking` blocks end-to-end, pins reasoning traffic to its requested model, and audits tool-loop continuations against what was actually served (fail-open; tampering is a local 400)
 - **Staged request pipeline** — a 4-layer DAG per request: model resolve / quota / cache lookup → account selection (priority, PTU-first, round-robin or latency-ranked within a tier, failover) → rate limits + engine call (retry on upstream 5xx, then the model's `fallback_models` chain before any byte is sent) → usage extraction, billing, cache store
 - **Governance built in** — access-key auth, daily token quotas, QPS / QPM / TPM limits at key, product, and model level, request-level TTL cache, account cooldown and recovery, DLP redaction and blocklist plugins. Admission reserves then settles, so concurrent requests can't overshoot a quota
 - **Multi-tenant** — keys carry a tenant; tenants get a pooled QPS bucket, a model entitlement allowlist, per-(key, model) quota defaults with an optional fallback-model degrade, key lifecycle (expiry/ban), and tenant-scoped admin tokens. Billing records charged cost and (optionally) vendor cost per row, so margin is queryable per tenant × model
 - **Per-user billing & enterprise audit** — every ledger row attributes to an effective end user (the key's `owner`, else the request's `x-gw-user` / `user` hint) with a `request_id`, so cost rolls up per user (`/admin/usage/users`) and soft per-user daily budgets apply on every surface; daily and calendar-month cost budgets (charged micro-dollars) cap a tenant pool, each key and each end user, with optional month-to-month rollover of the unspent remainder, raising a webhook alert when reached. Per-tenant content policy adds blocklist action tiers (block / flag / shadow), regex recognizers, secret masking, and an external-moderation seam with an AWS Bedrock Guardrails backend (deny on blocked policies, mask anonymized PII); every hit is recorded without prompt text. An admin-operation trail (key CRUD / config / reload, with source IP) and optional at-rest content retention complete the audit surfaces (`/admin/audit/*`)
-- **Fleet-ready** — run N instances behind a load balancer: Postgres shares config (versioned + a change feed), the access-key table, the ledger/files/batches store, and a distributed batch queue any instance drains; Redis shares rate/quota/TPM counters, account health, and optionally the response cache. Single-node stays zero-dependency
+- **Fleet-ready** — run N instances behind a load balancer: Postgres shares config (versioned + a change feed), the access-key table, the ledger/files/batches store, and a distributed batch queue any instance drains; Redis shares rate/quota/TPM counters, monthly cost counters, account health, and optionally the response cache. Single-node stays zero-dependency
 - **Providers behind traits** — engines talk to upstreams through a `Transport` seam; accounts with a real endpoint go over HTTP (reqwest + rustls), accounts without one are served by a deterministic in-process mock; AWS Bedrock (Claude, Llama, Cohere natively; every model through Converse) via SigV4 or API key with EventStream streaming
 - **Fast** — the whole pipeline (auth, admission, DLP, engine, billing) costs ~25 µs per request in-process; over HTTP one node serves ~90k requests/s at p99 under 10 ms on small bodies and ~40k/s at 256 concurrency on 52 KB / 13k-token prompts, mock upstream ([numbers and method](docs/performance.md))
 - **Observability built in** — Prometheus `/metrics` (per-route request/status counters, per-pipeline-stage latency, token counters), structured access logs, and one OTLP span per request (route, model, tenant, user, tokens, routing decisions; W3C `traceparent` joins the caller's trace) as soon as `OTEL_EXPORTER_OTLP_ENDPOINT` names a collector
@@ -47,7 +47,7 @@ GW_CONFIG=conf/gateway.yaml cargo run -p gw-server
 # GW_TRANSPORT=mock forces zero egress; GW_TRANSPORT=http disables the mock.
 ```
 
-Guides: [Examples](docs/examples.md) · [API](docs/api.md) · [Providers](docs/providers.md) · [Governance](docs/governance.md) · [Observability](docs/observability.md) · [Deployment](docs/deployment.md) · [Configuration](docs/configuration.md) · [Architecture](docs/architecture.md) · [Development](docs/development.md) · [Performance](docs/performance.md) · [Roadmap](https://github.com/cocoonstack/gateway/issues/1)
+Guides: [Examples](docs/examples.md) · [API](docs/api.md) · [Providers](docs/providers.md) · [Governance](docs/governance.md) · [Observability](docs/observability.md) · [Deployment](docs/deployment.md) · [Configuration](docs/configuration.md) · [Architecture](docs/architecture.md) · [Development](docs/development.md) · [Performance](docs/performance.md) · [Security](docs/security.md) · [Roadmap](https://github.com/cocoonstack/gateway/issues/1)
 
 ## Docker
 

crates/config/src/lib.rs

@@ -130,6 +130,9 @@ pub struct McpServerConf {
     pub oauth: Option<McpOAuthConf>,
     #[serde(default = "default_mcp_timeout")]
     pub timeout_seconds: u64,
+    /// Largest reply the proxy buffers for filtering or review; a larger one is refused.
+    #[serde(default = "default_mcp_max_reply_bytes")]
+    pub max_reply_bytes: usize,
 }
 
 impl McpServerConf {
@@ -743,6 +746,9 @@ pub struct GatewayConfig {
     /// MCP servers reachable through `/mcp/{server}` by entitled keys.
     #[serde(default)]
     pub mcp_servers: Vec<McpServerConf>,
+    /// Concurrent realtime sessions plus MCP listen streams one key may hold; 0 = unlimited.
+    #[serde(default = "default_max_live_streams")]
+    pub max_live_streams_per_key: usize,
     /// Trust `x-real-ip` / `x-forwarded-for` for the audit source IP. Off by
     /// default: the audit records the real TCP peer, which a client can't forge.
     /// Enable only when a trusted proxy fronts the gateway and sets those headers.
@@ -1100,6 +1106,12 @@ impl GatewayConfig {
             self.mcp_servers.iter().map(|m| m.name.as_str()),
         )?;
         for m in &self.mcp_servers {
+            if m.max_reply_bytes == 0 {
+                return Err(ConfigError::BadMcpServer {
+                    server: m.name.clone(),
+                    reason: "max_reply_bytes must be greater than 0",
+                });
+            }
             let Some(o) = &m.oauth else {
                 continue;
             };
@@ -1436,6 +1448,14 @@ fn default_mcp_timeout() -> u64 {
     60
 }
 
+fn default_mcp_max_reply_bytes() -> usize {
+    16 * 1024 * 1024
+}
+
+fn default_max_live_streams() -> usize {
+    64
+}
+
 fn default_alert_dedup_seconds() -> u64 {
     300
 }

crates/engines/src/http_transport.rs

@@ -166,11 +166,16 @@ impl Transport for HttpTransport {
                     tokio::time::sleep(RETRY_BACKOFF * attempt).await;
                 }
                 Err(e) => {
-                    return Err(GatewayError::new(
-                        upstream_fault_code(e.is_timeout()),
-                        502,
-                        format!("upstream request failed: {e}"),
-                    ));
+                    let what = if e.is_timeout() {
+                        "upstream request timed out"
+                    } else {
+                        "upstream request failed"
+                    };
+                    tracing::warn!(account = %req.account, error = %e, what);
+                    return Err(
+                        GatewayError::new(upstream_fault_code(e.is_timeout()), 502, what)
+                            .with_source(e),
+                    );
                 }
             }
         };

crates/handler/src/lib.rs

@@ -276,9 +276,11 @@ impl OnlineHandler {
                     ctx.quota_at,
                 )
                 .await;
-            if let Some((i, next)) = fallback_after(&e)
-                .then(|| next_fallback(&snap.cfg, &ctx, tried))
-                .flatten()
+            // signed thinking replays only against the model that produced it
+            if let Some((i, next)) = (is_upstream_fault(&e)
+                && !ctx.request.replays_reasoning_output())
+            .then(|| next_fallback(&snap.cfg, &ctx, tried))
+            .flatten()
             {
                 tried = i + 1;
                 switch_model(&mut ctx, next, &e.message);
@@ -557,10 +559,13 @@ async fn note_abuse(ctx: &DagContext) {
         .emit("abuse_suspend", ctx.ak.ak.clone(), summary);
 }
 
-/// Whether a pipeline error is the upstream's fault: a 5xx (vendor or
-/// connection failure) or a vendor 429; gateway-side denials never fall back.
-fn fallback_after(e: &GatewayError) -> bool {
-    e.http_status >= 500 || e.original_status() == Some(429)
+/// Whether a pipeline error came from upstream: a vendor 5xx or 429, or a
+/// 502/503 the gateway raised for a connection failure or an exhausted pool.
+fn is_upstream_fault(e: &GatewayError) -> bool {
+    match e.original_status() {
+        Some(status) => status >= 500 || status == 429,
+        None => e.http_status >= 502,
+    }
 }
 
 /// The next entry of the requested model's fallback chain past `tried` that the caller's tenant may use.
@@ -1021,7 +1026,9 @@ mod tests {
 
     async fn vendor_by_model() -> (String, Arc<std::sync::atomic::AtomicU32>) {
         use std::sync::atomic::{AtomicU32, Ordering};
+
         use tokio::io::{AsyncReadExt, AsyncWriteExt};
+
         let hits = Arc::new(AtomicU32::new(0));
         let seen = Arc::clone(&hits);
         let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();

crates/models/src/request.rs

@@ -105,7 +105,22 @@ impl GatewayRequest {
     /// replays only against the model that produced it, so the request that
     /// engages reasoning and the continuation carrying its output both pin.
     pub fn pins_reasoning_route(&self) -> bool {
-        self.reasoning_engaged()
+        self.reasoning_engaged() || self.replays_reasoning_output()
+    }
+
+    /// Whether this request carries reasoning output that must replay against
+    /// the model that produced it: prior signed thinking, an Anthropic
+    /// protected block, or a Responses `input` reasoning item. A bare reasoning
+    /// request that has produced nothing yet does not, so it may still fall back.
+    pub fn replays_reasoning_output(&self) -> bool {
+        let responses_input_replay = self
+            .model_param_v2
+            .as_ref()
+            .filter(|p| p.protocol == gw_consts::Protocol::Responses)
+            .and_then(|p| p.raw.get("input"))
+            .and_then(serde_json::Value::as_array)
+            .is_some_and(|items| items.iter().any(|item| item["type"] == "reasoning"));
+        responses_input_replay
             || self.message.iter().any(|m| {
                 m.reasoning_details.is_some()
                     || m.parts

crates/state/src/admission.rs

@@ -77,7 +77,7 @@ enum BudgetScope {
 }
 
 impl BudgetScope {
-    fn per_user(self) -> bool {
+    fn is_per_user(self) -> bool {
         matches!(self, Self::UserTokens | Self::UserCost)
     }
 
@@ -91,14 +91,7 @@ impl BudgetScope {
 
     /// The governance counter: user scopes carry the tenant, month counters their calendar month.
     fn key(self, month: Option<(i64, u32)>, ak: &AkInfo, user: &str) -> String {
-        let tag;
-        let prefix = match month {
-            Some((y, m)) => {
-                tag = format!("m:{y}{m:02}:");
-                tag.as_str()
-            }
-            None => "",
-        };
+        let prefix = month.map_or(String::new(), month_prefix);
         match self {
             Self::UserTokens => format!("{prefix}ub:{}:{user}", ak.tenant),
             Self::TenantCost => format!("{prefix}cb:tenant:{}", ak.tenant),
@@ -404,7 +397,7 @@ pub async fn check_tenant_rate(
 /// Per-AK QPS.
 pub async fn check_ak_rate(gov: &dyn Governance, ak: &AkInfo) -> Result<(), String> {
     admit(gov.rate_allow(&ak.ak, ak.qps).await, || {
-        format!("rate limit exceeded for ak {} (qps {})", ak.ak, ak.qps)
+        format!("rate limit exceeded for key {} (qps {})", ak.ak_id, ak.qps)
     })
 }
 
@@ -450,7 +443,7 @@ pub async fn reserve_daily(
     admit(
         gov.quota_reserve(&ak.ak, amount, ak.daily_token_quota, at)
             .await,
-        || format!("daily token quota exhausted for ak {}", ak.ak),
+        || format!("daily token quota exhausted for key {}", ak.ak_id),
     )
 }
 
@@ -470,8 +463,8 @@ pub async fn reserve_tpm(
         Ok(Some(amount))
     } else {
         Err(format!(
-            "token-per-minute limit exceeded for ak {} (tpm {tpm})",
-            ak.ak
+            "token-per-minute limit exceeded for key {} (tpm {tpm})",
+            ak.ak_id
         ))
     }
 }
@@ -583,7 +576,7 @@ async fn budgets(
         let Some(mut limit) = limit else {
             continue;
         };
-        if scope.per_user() && user.is_empty() {
+        if scope.is_per_user() && user.is_empty() {
             continue;
         }
         let key = match window {
@@ -622,6 +615,16 @@ fn previous_month((y, m): (i64, u32)) -> (i64, u32) {
     if m == 1 { (y - 1, 12) } else { (y, m - 1) }
 }
 
+fn month_prefix((y, m): (i64, u32)) -> String {
+    format!("m:{y}{m:02}:")
+}
+
+/// The counter prefixes of the current and previous month: everything the rollover still reads.
+pub fn month_prefixes() -> [String; 2] {
+    let month = civil_month(crate::epoch_secs());
+    [month_prefix(month), month_prefix(previous_month(month))]
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

crates/state/src/governance.rs

@@ -39,6 +39,8 @@ pub trait Governance: Send + Sync + std::fmt::Debug {
     async fn counter_get(&self, key: &str) -> i64;
     /// Add to a calendar-window counter, arming `ttl` on first use; returns the new total.
     async fn counter_add(&self, key: &str, amount: i64, ttl: Duration) -> i64;
+    /// Drop in-process counters outside the windows `prefixes` name; a TTL backend has nothing to do.
+    async fn counter_retain(&self, prefixes: &[String]);
 
     /// Fixed-window request limit (QPM): take one permit.
     async fn window_allow(&self, key: &str, limit: i64, window: Duration) -> bool;
@@ -113,6 +115,9 @@ impl Governance for MemoryGovernance {
     async fn counter_add(&self, key: &str, amount: i64, _ttl: Duration) -> i64 {
         self.counters.consume(key, amount)
     }
+    async fn counter_retain(&self, prefixes: &[String]) {
+        self.counters.retain_prefixed(prefixes);
+    }
     async fn window_allow(&self, key: &str, limit: i64, window: Duration) -> bool {
         self.qpm.reserve(key, 1, limit, window)
     }
@@ -296,6 +301,7 @@ impl Governance for RedisGovernance {
     async fn counter_add(&self, key: &str, amount: i64, ttl: Duration) -> i64 {
         self.incr_window(&counter_key(key), amount, ttl).await
     }
+    async fn counter_retain(&self, _prefixes: &[String]) {}
     async fn window_allow(&self, key: &str, limit: i64, window: Duration) -> bool {
         self.incr_window(&format!("gw:qpm:{key}"), 1, window).await <= limit
     }

crates/state/src/lib.rs

@@ -24,6 +24,7 @@ pub mod health;
 pub mod keystore;
 pub mod latency;
 pub mod store;
+pub mod streams;
 pub mod thinking_signature;
 
 pub use alerts::{AlertBus, AlertEvent};
@@ -391,6 +392,12 @@ impl QuotaStore {
     pub fn reset_all(&self) {
         self.used.clear();
     }
+
+    /// Drop every counter whose key starts with none of `prefixes`.
+    pub fn retain_prefixed(&self, prefixes: &[String]) {
+        self.used
+            .retain(|k, _| prefixes.iter().any(|p| k.starts_with(p.as_str())));
+    }
 }
 
 /// Account pool: the highest-priority slot serving a model type, round-robin
@@ -761,6 +768,8 @@ pub struct GatewayState {
     pub thinking_signatures: ThinkingSignatureAudit,
     /// Per-account call latency for `stability.latency_routing`; per instance.
     pub latency: latency::Latency,
+    /// Open long-lived streams per key, for `max_live_streams_per_key`.
+    pub streams: Arc<streams::LiveStreams>,
 }
 
 impl Default for GatewayState {
@@ -778,6 +787,7 @@ impl Default for GatewayState {
             alerts: Arc::new(alerts::AlertBus::default()),
             thinking_signatures: ThinkingSignatureAudit::new(),
             latency: latency::Latency::default(),
+            streams: Arc::default(),
         }
     }
 }
@@ -896,6 +906,7 @@ impl GatewayState {
             alerts: prev.alerts.clone(),
             thinking_signatures: prev.thinking_signatures.clone(),
             latency: prev.latency.clone(),
+            streams: prev.streams.clone(),
         })
     }
 }