diff --git a/README.md b/README.md index df5ef20e..839f2526 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ key-based auth, quotas, rate limits, failover, and a billing ledger. - **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 tokens the gateway fetches and refreshes itself (client-credentials or refresh-token grant) ([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 +- **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 (another account on an upstream 5xx or 401–403 refusal, 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/*`) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index d7d28085..eec7db55 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -639,8 +639,7 @@ pub struct TenantConf { /// Default per-model daily token caps, metered per key. #[serde(default)] pub model_quotas: HashMap, - /// Where an over-quota request degrades to instead of hard-failing; - /// None = pass through unmetered (the per-AK daily cap still backstops). + /// Where an over-quota request degrades to; None = the request is refused. #[serde(default)] pub fallback_model: Option, /// Env var holding this tenant's scoped admin token; empty = none. @@ -1238,8 +1237,7 @@ impl GatewayConfig { Ok(()) } - /// Config generation (a stable hash of the source document); mixed into - /// cache keys so a reload to a different config can't serve a stale entry. + /// Config generation: a stable hash of the source document. pub fn generation(&self) -> u64 { self.generation } diff --git a/crates/consts/src/lib.rs b/crates/consts/src/lib.rs index 53a6e162..b1cfb460 100644 --- a/crates/consts/src/lib.rs +++ b/crates/consts/src/lib.rs @@ -20,7 +20,6 @@ pub mod role { pub const MODEL: &str = "model"; pub const SYSTEM: &str = "system"; pub const TOOL: &str = "tool"; - pub const STORAGE: &str = "storage"; } /// Account tiers for PTU (provisioned) vs pay-as-you-go routing. diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 58056c83..5d94003b 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -858,7 +858,6 @@ fn usage_tokens(ctx: &DagContext) -> GResult { .response; let rate = gw_state::model_token_rate(&ctx.cfg, served_model(ctx.request.model_param_v2.as_ref())); - // saturating sums so a malformed usage subtree can't overflow the totals let mut tokens = match &resp.common_usage { Some(u) => { let ti = gw_models::TokenInput { diff --git a/crates/dag/src/token_estimate.rs b/crates/dag/src/token_estimate.rs index 23871acb..9f46cf0f 100644 --- a/crates/dag/src/token_estimate.rs +++ b/crates/dag/src/token_estimate.rs @@ -115,10 +115,6 @@ pub fn estimate_prompt_tokens( let mut num = 0usize; for msg in messages { - // storage-role messages are internal bookkeeping, not sent upstream. - if msg.role == gw_consts::role::STORAGE { - continue; - } num += per_msg; num += enc.encode_len(&message_text(msg)); num += enc.encode_len(&msg.role); diff --git a/crates/engines/src/claude_engine.rs b/crates/engines/src/claude_engine.rs index 20d5c2da..95cf14ca 100644 --- a/crates/engines/src/claude_engine.rs +++ b/crates/engines/src/claude_engine.rs @@ -661,17 +661,18 @@ pub fn anthropic_native_chunks( .filter(|usage| usage.is_object()) .cloned() .unwrap_or_else(|| json!({"input_tokens":response.prompt_tokens,"output_tokens":0})); - let mut events = vec![json!({ + let mut message_start = json!({ "type":"message_start", "message":{ "type":"message", "role":"assistant", "model":stream_model, "content":[], - "stop_reason":null, - "usage":start_usage + "stop_reason":null } - })]; + }); + message_start["message"]["usage"] = start_usage; + let mut events = vec![message_start]; if let Value::Array(blocks) = content { for (index, mut start) in blocks.into_iter().enumerate() { let mut deltas = Vec::new(); diff --git a/crates/engines/src/families.rs b/crates/engines/src/families.rs index 490e4fe9..c4b8534f 100644 --- a/crates/engines/src/families.rs +++ b/crates/engines/src/families.rs @@ -565,8 +565,7 @@ fn video_dialect(vendor: &str, wire: &str) -> VideoDialect { /// The vendor's async handle in a submit reply, whichever dialect answered; /// `None` for a synchronous reply that already carries the video. pub fn video_handle(v: &Value) -> Option<&str> { - // task ids before `request_id`: DashScope and Kling replies carry both, and - // their top-level request_id is the HTTP call's trace id, not the job + // task ids first: a DashScope or Kling request_id traces the HTTP call, not the job v["output"]["task_id"] .as_str() .or_else(|| v["data"]["task_id"].as_str()) @@ -738,14 +737,8 @@ fn video_outcome(model: &str, v: Value, status: u16) -> EngineOutcome { .or_else(|| v["video_url"].as_str()) .unwrap_or_default() .to_owned(); - let step = v["status"] - .as_str() - .or_else(|| v["output"]["task_status"].as_str()) - .unwrap_or_default() - .to_owned(); let units = whole_seconds(&v["video"]["duration"]).unwrap_or(0); let mut out = family_outcome(message, model.to_owned(), v, status); - out.response.step = step; out.response.billed_units = units; out } @@ -2070,7 +2063,6 @@ mod tests { ); let out = v.run().await.unwrap(); assert_eq!(out.response.message, "mock://videos/out.mp4"); - assert_eq!(out.response.step, "succeeded"); let mut s = SearchEngine::new( req( diff --git a/crates/engines/src/pump.rs b/crates/engines/src/pump.rs index 8e8b37e8..87467d98 100644 --- a/crates/engines/src/pump.rs +++ b/crates/engines/src/pump.rs @@ -128,9 +128,6 @@ where } } } - if out.aborted { - break; - } } out.streamed_live = sent_any; } diff --git a/crates/models/src/response.rs b/crates/models/src/response.rs index 7370ed46..e2ea4c95 100644 --- a/crates/models/src/response.rs +++ b/crates/models/src/response.rs @@ -47,9 +47,6 @@ pub struct GatewayResponse { /// seconds, rerank search units); priced by the model's `unit_price_micros`. #[serde(default)] pub billed_units: i64, - /// sora2 async step marker. - #[serde(default, skip_serializing_if = "String::is_empty")] - pub step: String, /// normalized usage view, filled by the CommonUsage post-processor. #[serde(default, skip_serializing_if = "Option::is_none")] @@ -152,13 +149,6 @@ mod tests { assert!(r.common_usage.is_none()); } - #[test] - fn empty_step_is_omitted() { - let r = GatewayResponse::default(); - let j = serde_json::to_value(&r).unwrap(); - assert!(j.get("step").is_none()); - } - #[test] fn detects_protected_anthropic_content() { let protected = GatewayResponse { diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 33a0c1bf..0d2aab72 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -5,6 +5,7 @@ use std::env; use std::sync::Arc; use std::time::Duration; +use axum::serve::ListenerExt as _; use gw_config::GatewayConfig; use gw_state::GatewayState; use gw_views::AppState; @@ -194,7 +195,11 @@ async fn main() -> anyhow::Result<()> { }), ); - let listener = tokio::net::TcpListener::bind(&addr).await?; + let listener = tokio::net::TcpListener::bind(&addr).await?.tap_io(|tcp| { + if let Err(e) = tcp.set_nodelay(true) { + tracing::debug!(error = %e, "TCP_NODELAY not set on an accepted connection"); + } + }); tracing::info!("gw listening on http://{addr}"); // connect-info so the audit trail can root the source IP at the TCP peer diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 87ac0286..1298ce8b 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -796,8 +796,7 @@ pub fn month_prefixes() -> [String; 2] { [month_prefix(month), month_prefix(previous_month(month))] } -/// Accrue usage under `over`, the override admission resolved (`None` resolves it now), settling -/// what `hold` held; a scope reaching its cap raises a `budget_exhausted` alert (deduped by the bus). +/// Accrue usage under `over` (`None` resolves it now) and settle `hold`; a scope at its cap alerts. async fn consume_budgets( state: &GatewayState, cfg: &GatewayConfig, diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 81de0b3b..21e3fcef 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -935,7 +935,6 @@ impl SharedConfig { pub async fn reload(&self, cfg: GatewayConfig) -> gw_models::GResult<()> { let _serialized = self.reload_lock.lock().await; let prev = self.inner.load(); - // cfg.generation keys the preserved cache so it can't serve a pre-reload entry let state = GatewayState::reload_from(&cfg, &prev.state).await?; self.inner.store(Arc::new(Snapshot { cfg: Arc::new(cfg), diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 53447e7a..0d69f0f7 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -34,9 +34,7 @@ const ROLLUP_BUCKET_SECS: i64 = 60; /// watermark trails it), so late rows and missed ticks self-heal. const ROLLUP_BACKFILL_SECS: i64 = 20 * 60; -/// A minute is rolled only once closed for this long, so no in-flight write or -/// trailing replica lands a row in a rolled minute; a spool replay that does adds -/// the row to the bucket in the same transaction, which keeps the max-upsert sound. +/// A minute rolls once closed this long, past in-flight writes; a late spool row joins its bucket on replay. const ROLLUP_SETTLE_SECS: i64 = ROLLUP_BUCKET_SECS; /// In-process rollup retention: the monthly-budget window, then a bucket cap. const ROLLUP_RETENTION_SECS: i64 = 62 * 86_400; @@ -291,10 +289,7 @@ impl UserUsageRow { self.billed_units = self.billed_units.saturating_add(r.billed_units); } - /// Keep the larger of each counter. A bucket's source rows only shrink via ledger - /// pruning, and a replayed late row is added to its bucket as it lands, so per-column - /// max always picks the more complete snapshot — a recompute from a partially pruned - /// ledger can never shrink a bucket. + /// Keep the larger of each counter: pruning only shrinks a bucket's sources, so a recompute never shrinks it. fn keep_max(&mut self, o: &UserUsageRow) { self.requests = self.requests.max(o.requests); self.prompt_tokens = self.prompt_tokens.max(o.prompt_tokens); @@ -614,8 +609,7 @@ impl Store for MemoryStore { .fetch_add(1, Ordering::Relaxed) .is_multiple_of(LEDGER_PRUNE_EVERY) { - // spare unrolled rows (the cap yields to billing integrity); a full scan, since - // completion order can wedge a young row ahead of prunable ones + // spare unrolled rows; a full scan, since completion order can wedge a young row first let low_water = self.ledger_max_rows - self.ledger_max_rows / 10; let mut excess = ledger.rows.len() - low_water; let MemoryLedger { rows, request_ids } = &mut *ledger; @@ -643,19 +637,7 @@ impl Store for MemoryStore { } ledger.rows.push(r.clone()); if r.created_at_epoch_secs < watermark { - rollup - .entry(( - bucket_floor(r.created_at_epoch_secs), - r.tenant.clone(), - r.user_id.clone(), - r.model.clone(), - )) - .or_insert_with(|| UserUsageRow { - user_id: r.user_id.clone(), - model: r.model.clone(), - ..Default::default() - }) - .add_record(r); + roll_into(&mut rollup, r); } } Ok(()) @@ -684,7 +666,6 @@ impl Store for MemoryStore { model: r.model.clone(), ..Default::default() }); - // saturating: a hostile record can carry i64::MAX counts (usage is floored, not capped) e.requests += 1; e.prompt_tokens = e.prompt_tokens.saturating_add(r.prompt_tokens); e.completion_tokens = e.completion_tokens.saturating_add(r.completion_tokens); @@ -784,19 +765,7 @@ impl Store for MemoryStore { .iter() .filter(|r| r.created_at_epoch_secs >= lo && r.created_at_epoch_secs < hi) { - fresh - .entry(( - bucket_floor(r.created_at_epoch_secs), - r.tenant.clone(), - r.user_id.clone(), - r.model.clone(), - )) - .or_insert_with(|| UserUsageRow { - user_id: r.user_id.clone(), - model: r.model.clone(), - ..Default::default() - }) - .add_record(r); + roll_into(&mut fresh, r); } } let written = fresh.len() as u64; @@ -2281,8 +2250,7 @@ impl PostgresStore { "ALTER TABLE batch_results ADD COLUMN IF NOT EXISTS user_id TEXT NOT NULL DEFAULT ''", "ALTER TABLE batch_results ADD COLUMN IF NOT EXISTS finish_reason TEXT NOT NULL DEFAULT ''", "ALTER TABLE batch_results ADD COLUMN IF NOT EXISTS tool_calls TEXT NOT NULL DEFAULT ''", - // pre-upgrade results predate user_id; backfill it from the items rows so erasure - // reaches history (idempotent) + // backfill pre-upgrade results' user_id from the items rows so erasure reaches history "UPDATE batch_results r SET user_id = i.user_id FROM batch_items i WHERE r.batch_id = i.batch_id AND r.idx = i.idx AND r.user_id = '' AND i.user_id <> ''", @@ -2522,8 +2490,7 @@ sql_store_impl!(PostgresStore, postgres, { } async fn batch_finalize(&self, id: &str, claim: i64) -> GResult> { - // lock the row, THEN aggregate: a single UPDATE's subquery reads the statement-start - // snapshot and would miss a result committing while it waits for the lock + // lock, then aggregate: an UPDATE's subquery snapshot misses results committed during the wait let mut tx = self .pool .begin() @@ -2811,6 +2778,21 @@ fn usage_of(r: &BillingRecord) -> UserUsageRow { } } +fn roll_into( + buckets: &mut BTreeMap<(i64, String, String, String), UserUsageRow>, + r: &BillingRecord, +) { + buckets + .entry(( + bucket_floor(r.created_at_epoch_secs), + r.tenant.clone(), + r.user_id.clone(), + r.model.clone(), + )) + .and_modify(|e| e.add_record(r)) + .or_insert_with(|| usage_of(r)); +} + /// Fold grouped usage rows into `map` by (user, model). fn fold_user_usage( map: &mut BTreeMap<(String, String), UserUsageRow>, diff --git a/crates/state/src/thinking_signature.rs b/crates/state/src/thinking_signature.rs index d22f44eb..c9556ae0 100644 --- a/crates/state/src/thinking_signature.rs +++ b/crates/state/src/thinking_signature.rs @@ -174,8 +174,7 @@ impl ThinkingSignatureAudit { } } - // thinking disabled: the client must strip the replayed blocks and the upstream - // arbitrates; enabled: a known anchor missing its blocks is still a Mismatch + // thinking off: the upstream arbitrates stripped blocks; on: a known anchor without them mismatches if sequence.blocks.is_empty() && !request.reasoning_engaged() { return (Some(context), ReviewVerdict::Miss); } diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 176adcd9..86e4bc40 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -1350,8 +1350,7 @@ fn user_hint(hint: Option, field: &Value) -> Option { hint.or_else(|| field.as_str().map(cap_user_hint)) } -/// Bound an attribution hint at `USER_HINT_MAX_BYTES` (it keys governance counters) -/// and drop NUL, which Postgres TEXT refuses in every row the user lands in. +/// Cap an attribution hint at `USER_HINT_MAX_BYTES` and drop NUL, which Postgres TEXT refuses. fn cap_user_hint(s: &str) -> String { let mut end = s.len().min(USER_HINT_MAX_BYTES); while !s.is_char_boundary(end) { @@ -3782,10 +3781,9 @@ fn messages_stream_response( BlockKind::Text => json!({"type":"text","text":""}), BlockKind::Thinking => json!({"type":"thinking","thinking":"","signature":""}), }; - self.queue.push_back(Self::ev( - "content_block_start", - json!({"type":"content_block_start","index":idx,"content_block":content_block}), - )); + let mut start = json!({"type":"content_block_start","index":idx}); + start["content_block"] = content_block; + self.queue.push_back(Self::ev("content_block_start", start)); idx } @@ -4259,8 +4257,7 @@ fn responses_stream_response( }) => { let err = *err; self.ensure_created(); - // carry both documented shapes: the docs put code/message at the top - // level, the vendor nests them, and a client models only one + // both shapes, since clients model one: the docs put code/message on top, the vendor nests them self.queue.push_back( Event::default().event("error").data( json!({ @@ -4314,14 +4311,14 @@ fn responses_stream_response( }; if !self.native { self.seq += 1; + let mut done = json!({"type":"response.completed","response":{ + "model": self.model, "status": self.status, + }}); + done["response"]["usage"] = responses_usage(pt, ct, tt, c.common_usage); self.queue.push_back( - Event::default().event("response.completed").data( - json!({"type":"response.completed","response":{ - "model": self.model, "status": self.status, - "usage": responses_usage(pt, ct, tt, c.common_usage), - }}) - .to_string(), - ), + Event::default() + .event("response.completed") + .data(done.to_string()), ); } self.queue.push_back(Event::default().data("[DONE]")); diff --git a/crates/views/src/mcp.rs b/crates/views/src/mcp.rs index f99d6e4b..7e898d22 100644 --- a/crates/views/src/mcp.rs +++ b/crates/views/src/mcp.rs @@ -522,7 +522,9 @@ fn jsonrpc_error(id: Value, code: i64, message: String) -> Response { } fn jsonrpc_error_value(id: Value, code: i64, message: &str) -> Value { - json!({"jsonrpc": "2.0", "id": id, "error": {"code": code, "message": message}}) + let mut error = json!({"jsonrpc": "2.0", "error": {"code": code, "message": message}}); + error["id"] = id; + error } #[cfg(test)] diff --git a/docs/performance.md b/docs/performance.md index 25dbb66c..2f8b95fc 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -55,6 +55,12 @@ speed sit orders of magnitude below that; the gateway is not the ceiling. `ak-bench`'s 10⁹ daily-token quota mid-run and every further request was a clean `service_quota_exceeded_exception` — reserve-then-settle admission does not overshoot at 40k requests/s. +- Streams on a pooled connection: until 2026-10-18 the listener left Nagle on, + so on Linux every streamed response on a keep-alive connection ended ~44 ms + late, waiting on the client's delayed ACK (one stream at a time: 22 + requests/s). With `TCP_NODELAY` on every accepted connection the same probe + runs ~7,000 streams/s at p50 0.12 ms, and 64 concurrent streams ~53k/s; the + macOS table above never showed the stall. - Resources: the gateway used ~8 cores at 256 concurrency. Memory grows only in the in-memory store, which keeps every ledger row (~1 KB per request); point `storage` at SQLite/Postgres for anything longer than a benchmark.