From 40140e7812d019b843c19a2c5f8b02bc5e2e8289 Mon Sep 17 00:00:00 2001 From: nexus Date: Tue, 30 Jun 2026 21:42:00 +0800 Subject: [PATCH 1/8] chore(repo): sync docs, tooling, and CI config from internal main (sync 9) Bring repository-level docs, schema, and lint/config assets up to date with internal main (TLS-bump streaming-compliance + provider-adapter cycle). - Architecture docs: SSE streaming-compliance, provider-adapter, normalization, hook, cost-estimation, audit-pipeline - API + feature docs: control-plane providers OpenAPI; CP-UI ai-gateway-routing + overview - Tooling: scripts/doc-lockstep.config.mjs - Schema: tools/db-migrate providers.prisma + traffic.prisma - README, CHANGELOG --- CHANGELOG.md | 63 ++++++ README.md | 50 +---- .../audit-pipeline-architecture.md | 4 +- .../sse-streaming-compliance-architecture.md | 198 ++++++++++++++---- .../cost-estimation-architecture.md | 1 + .../services/ai-gateway/hook-architecture.md | 4 + .../ai-gateway/normalization-architecture.md | 2 + .../provider-adapter-architecture.md | 45 +++- .../api/openapi/control-plane/providers.yaml | 29 +++ .../features/cp-ui/ai-gateway-routing.md | 2 + docs/users/features/cp-ui/overview.md | 2 +- scripts/doc-lockstep.config.mjs | 11 + tools/db-migrate/schema/providers.prisma | 9 + tools/db-migrate/schema/traffic.prisma | 6 + 14 files changed, 334 insertions(+), 92 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index be548ed3..1d9b014e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,29 @@ All notable changes to this project are documented here. The format follows ## [Unreleased] +### Fixed — Request/Response hook timing + +- **Streamed responses now record response-hook timing, exactly once per hook.** + The streaming response pipeline runs the response stage at every checkpoint, so + the live audit-only path previously recorded nothing (`response_hooks_ms` NULL) + while the chunked_async path recorded the same hook once per checkpoint (N + duplicate rows, an N×-inflated aggregate — observed as a "RESPONSE PIPELINE (63)" + list of identical rows). The trace is now folded to one record per hook (summed + latency, latest decision) across the ai-gateway live + Model A paths and the + shared compliance-proxy/agent path. The audit drawer also collapses any residual + duplicates (historical rows) into a single `×N` card. + +### Added — microsecond-precision hook timing (additive, backward compatible) + +- Per-hook latency is now measured in **microseconds** (`latencyUs`) alongside the + existing truncated-millisecond `latencyMs`, with new aggregate columns + `request_hooks_us` / `response_hooks_us` beside the unchanged `_ms` columns. + Hooks run at microsecond scale, so the millisecond aggregates floored a + sub-millisecond hook to `0`; the µs fields carry the real value, surfaced + precisely per hook in the control-plane audit drawer. The `_ms` columns / wire + ids / values are unchanged. The new binwire field ids are forward-incompatible, + so the deploy order is **schema → Hub → producers**. + ### Changed (BREAKING — major version bump) - **Hook `onMatch` collapses to a single `action` (approve | redact | block).** The orthogonal `onMatch.inflightAction` (approve / block-hard / block-soft / @@ -40,6 +63,46 @@ All notable changes to this project are documented here. The format follows as a 403 reject. The Agent signals a block by dropping the connection (no rich error body); the proxies return an attributed 403 whose response-stage reason carries rule-ID labels only, never the upstream value. +### Fixed — co-firing redact + soft-block no longer drops the redaction (security) + +- **A redact hook co-firing with a soft-block hook now masks-and-delivers instead of + leaking or failing closed.** When a redact hook (`Modify` + masked content/spans) and + a soft-block hook fired on the same request or response, the pipeline aggregator + promoted the reported `Decision` to `BlockSoft` (the strictest) but DROPPED the redact + hook's replacement content, leaving spans without content. Downstream this produced a + no-op rewrite that, depending on the path, either failed closed (canonical response) + or replayed/forwarded the ORIGINAL unredacted body — a PII leak on the shared buffer + pipeline (compliance-proxy appliance included), the agent Model A wire, and both + request stages. `mergeResults` now carries the redact's `ModifiedContent` + unconditionally, and every redaction consumer gates on the new + `decision.CompliancePipelineResult.CarriesRedaction()` predicate (Modify OR a + BlockSoft masking a co-firing redact) rather than `Decision==Modify`, so the masked + body is applied and delivered on all paths. The audit row stamps the disposition + `action=redact` even when the (soft-block) `Decision` ceiling is `BlockSoft`. No + config or schema change; behavior is compliance-safe (a hard `block`/`RejectHard` + still rejects; a standalone soft-block still delivers-with-warning). The no-redactor + buffer degrade is now posture-aware (appliance fail-closed, agent fail-open). + +### Changed — three-end streaming-compliance parity via a shared Model A engine + +- **The Model A streaming-compliance algorithm is now a single shared engine driving + three ends.** The prescan-gated real-time streaming path (bounded tail-hold + + union prescan + confirm + escalate-to-buffer redaction) for a redact-scope + `chunked_async` stream is extracted into a substrate-agnostic engine + (`shared/transport/streaming/modela`). The AI Gateway drives it with a canonical + substrate (fail-closed) and the transparent proxy used by the Agent + + Compliance Proxy drives it with a raw-SSE-wire substrate (fail-open, NE + host-packet safety) — so hooks/compliance behave identically across all three + ends while each keeps its own ingress and delivery. The transparent-proxy live + path becomes **audit-only** (real-time write-through, observe-only checkpoints, + never blocks/rewrites): scope-derived routing sends a `block` scope to buffer and + a `redact` scope to Model A (or buffer), so only non-enforcing traffic reaches + live. The adoption also closed two latent PII-leak paths in the shipped AI Gateway + Model A (a redact masked behind a co-firing soft-block; a memory-pressure eviction + of an incomplete content unit). No config or contract change; behavior is + compliance-safe (a sub-window value is never delivered raw; storage never persists + a raw prefix on an enforcing outcome). + ### Changed — normalized projection is now fully view-time (no migration required) - **The normalized traffic projection is no longer written on the hot path; it is diff --git a/README.md b/README.md index 3cfd403f..10eb28c4 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # Nexus Gateway -[![CI](https://github.com/AlphaBitCore/nexus-gateway/actions/workflows/ci.yml/badge.svg?branch=main)](.github/workflows/ci.yml) -[![Go CI](https://github.com/AlphaBitCore/nexus-gateway/actions/workflows/go-ci.yml/badge.svg?branch=main)](.github/workflows/go-ci.yml) +[![CI](https://github.com/itechchoice/abc-nexus-gateway/actions/workflows/ci.yml/badge.svg?branch=main)](.github/workflows/ci.yml) +[![Go CI](https://github.com/itechchoice/abc-nexus-gateway/actions/workflows/go-ci.yml/badge.svg?branch=main)](.github/workflows/go-ci.yml) [![Coverage gate](https://img.shields.io/badge/coverage-%E2%89%A595%25%20per%20package-brightgreen)](./scripts/check-go-coverage.sh) [![Status: 1.1.0](https://img.shields.io/badge/status-1.1.0-brightgreen)](./CHANGELOG.md) [![License: Apache 2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](./LICENSE) @@ -102,52 +102,6 @@ Nexus is designed so that the compliance layer adds as little latency as possibl **Quota accounting is write-behind.** Per-request quota costs accumulate in-process and flush to Redis on a 250ms interval, removing the synchronous Redis round-trip from the hot path. Configurable back to synchronous mode (`NEXUS_QUOTA_WRITE_BEHIND=0`) for strict accounting requirements. -### Benchmarking & load-testing toolkit - -The numbers above are produced and re-verified by a three-repo toolkit, each maintained as its own standalone repository. The load generator was previously the in-tree `tools/loadtest` and was extracted to `nexus-loadtest` (see [CHANGELOG](./CHANGELOG.md)). - -| Repo | Role | Key docs | -|---|---|---| -| **[llm-gateway-benchmark](https://github.com/AlphaBitCore/llm-gateway-benchmark)** | On-demand AWS rig (CloudFormation + Ansible) that benchmarks Nexus head-to-head against 5 other gateways (Bifrost, LiteLLM, Kong, Portkey, TensorZero) — each isolated on its own box, all hitting one shared mock upstream. Deploy to run, `delete-stack` to tear down. | [ARCHITECTURE](https://github.com/AlphaBitCore/llm-gateway-benchmark/blob/main/ARCHITECTURE.md) · [LOADTEST-RUNBOOK](https://github.com/AlphaBitCore/llm-gateway-benchmark/blob/main/docs/LOADTEST-RUNBOOK.md) · [CONTROL-BOX-RUNBOOK](https://github.com/AlphaBitCore/llm-gateway-benchmark/blob/main/docs/CONTROL-BOX-RUNBOOK.md) | -| **[nexus-mock-provider](https://github.com/AlphaBitCore/nexus-mock-provider)** | High-performance mock upstream that speaks the real OpenAI / Gemini / Anthropic wire formats (streaming + non-streaming) and echoes requests back with plausible token usage. Removes the real, paid, rate-limited provider from the measurement so you benchmark the gateway, not the model. Listens on `:3062`. | [README](https://github.com/AlphaBitCore/nexus-mock-provider/blob/main/README.md) · [CONFIGURE-NEXUS](https://github.com/AlphaBitCore/nexus-mock-provider/blob/main/CONFIGURE-NEXUS.md) | -| **[nexus-loadtest](https://github.com/AlphaBitCore/nexus-loadtest)** | Scenario-driven load generator for any OpenAI- or Anthropic-compatible endpoint. Simulates realistic, weighted, multi-turn traffic and reports TTFT, inter-token latency, and token throughput. Scales to tens of thousands of concurrent virtual users from one host. | [README](https://github.com/AlphaBitCore/nexus-loadtest/blob/main/README.md) · [DESIGN](https://github.com/AlphaBitCore/nexus-loadtest/blob/main/DESIGN.md) | - -#### Run a quick local benchmark (single gateway) - -Measure your local AI Gateway (`:3050`) against the mock upstream — no real provider, no cost: - -1. **Start the mock upstream** (in a `nexus-mock-provider` checkout): - ```bash - make run # serves :3062, all three specs (OpenAI / Gemini / Anthropic) - ``` -2. **Point Nexus at the mock.** Add a provider credential / routing rule whose upstream base URL is `http://localhost:3062` (the mock accepts any API key). See [`CONFIGURE-NEXUS.md`](https://github.com/AlphaBitCore/nexus-mock-provider/blob/main/CONFIGURE-NEXUS.md) for the exact routing setup. -3. **Run the load generator** (in a `nexus-loadtest` checkout) against the gateway, authenticating with a Nexus virtual key: - ```bash - go run ./cmd/loadtest -config profiles/realistic.json \ - -target http://localhost:3050 -vk -out runs/ - ``` -4. **Read the report** at `runs//` — TTFT, inter-token latency, throughput, and per-tier breakdown. Use `-compare` against an earlier `summary.json` to gate regressions. - -#### Run the full head-to-head matrix (AWS) - -To compare Nexus against the other gateways on isolated boxes, use the `llm-gateway-benchmark` rig (binaries ship prebuilt in `artifacts/` — nothing compiles on-box): - -```bash -# 1. Deploy infra -aws cloudformation deploy --stack-name nexus-perf-matrix \ - --template-file cloudformation/perf-matrix-stack.yaml --capabilities CAPABILITY_IAM \ - --parameter-overrides KeyName= AdminCidr=/32 -# 2. Provision every box (host-native) -scripts/gen-inventory.sh nexus-perf-matrix ~/.ssh/.pem -cd ansible && ansible-playbook -i inventory.ini site.yml -# 3. Run the benchmark (one gateway, all tiers → a report each) -GATEWAY=nexus scripts/bench/run-tiers.sh -# 4. Tear down when idle (on-demand, cost control) -aws cloudformation delete-stack --stack-name nexus-perf-matrix -``` - -Compare gateways by **TTFT delta** — the shared mock's latency cancels out, so the difference is the gateway's own overhead. Full procedure and the report-validity gate are in the rig's [LOADTEST-RUNBOOK](https://github.com/AlphaBitCore/llm-gateway-benchmark/blob/main/docs/LOADTEST-RUNBOOK.md). - --- ## Architecture in one minute diff --git a/docs/developers/architecture/cross-cutting/observability/audit-pipeline-architecture.md b/docs/developers/architecture/cross-cutting/observability/audit-pipeline-architecture.md index 5452004b..8dc70fbf 100644 --- a/docs/developers/architecture/cross-cutting/observability/audit-pipeline-architecture.md +++ b/docs/developers/architecture/cross-cutting/observability/audit-pipeline-architecture.md @@ -64,6 +64,8 @@ A second consumer group, `hub-alerting`, attaches to the same traffic queues plu - The defer reads the upstream `PhaseSink` populated by the singleton tracing transport (`UpstreamTtfbMs`, `UpstreamTotalMs`), snapshots the per-request `PhaseTimer` (`auth_ms`, `quota_ms`, `routing_ms`, `cache_lookup_ms`, `req_adapter_ms`, `resp_adapter_ms`, plus the request-hooks sub-phases `hook_extract_ms`, `hook_build_ms`, `hook_pipeline_ms`, `hook_rewrite_ms`), computes `upstream_body_ms` from the TTFB / total gap, computes `audit_emit_ms` for the time spent in the defer up to the hand-off, finalizes `LatencyMs` with **ceiling-millisecond rounding** (so a sub-millisecond cache hit reports as `1` instead of `0`, which the wire format would have treated as "field absent"), and only then calls `h.deps.AuditWriter.Enqueue(rec)`. - **`RequestHooksMs` vs the `hook_*_ms` sub-phases.** `RequestHooksMs` aggregates only each hook's own self-timed `Execute` duration. The hooks *stage* also does framing work — content extraction (`hook_extract_ms`), pipeline build (`hook_build_ms`), the whole-`Execute` wall-clock (`hook_pipeline_ms`, ≥ the per-hook sum; the gap is in-pipeline framework overhead), and Modify body rewrite (`hook_rewrite_ms`) — that `RequestHooksMs` never counted — and that, before these keys, was unattributed in `latency_breakdown` entirely (the hooks stage was the only stage with no phase key). Because `RequestHooksMs` covers only per-hook `Execute`, a hook whose cost scales with body size (e.g. a full-body PII scan over a large prompt) *would* register near-zero in `request_hooks_ms` while its real cost lands in the framing segments; these four keys make each segment attributable so that case is diagnosable. (Measured cost is workload-dependent — for small request bodies these segments are sub-millisecond.) `RequestHooksMs` semantics are unchanged (additive observability, no contract break). +- **Microsecond precision (`request_hooks_us` / `response_hooks_us`).** Hooks run at microsecond scale (rule-pack scans are tens of µs), so the millisecond aggregates floor a sub-millisecond hook to `0`. Each hook is therefore self-timed in **both** units — `latencyMs` (the truncated integer-ms floor, kept for backward compatibility, never clamped because it is *summed* downstream) and `latencyUs` (the precise microsecond value) — carried per hook in the `*_hooks_pipeline` JSONB and aggregated into the additive `request_hooks_us` / `response_hooks_us` columns alongside the unchanged `_ms` columns. The control-plane audit drawer renders the precise per-hook µs; the latency waterfall stays on the `_ms` aggregate (sub-millisecond hooks are a visual sliver there). New binwire field ids `103` / `104` carry the µs aggregates; the wire is forward-incompatible (an old Hub rejects an unknown field id), so the deploy order is **schema → Hub → producers**. +- **Streaming response hooks fold to one row per hook.** The streaming response pipeline runs the response stage once per checkpoint (a byte-window cadence plus a mandatory EOF checkpoint), so the same hook is scanned many times in one stream. The trace is folded to **one record per hook** — latency summed (the real scan CPU across the stream), the last checkpoint's decision authoritative — before it lands in `*_hooks_pipeline` and the `_ms` / `_us` aggregates: the ai-gateway live + Model A paths via the `responseHookAccumulator`, and the shared chunked_async path (compliance-proxy + agent) via `foldHookResults`. Without the fold a streamed response recorded the hook N times (N duplicate rows, N×-inflated aggregate); with it the drawer shows one card. The control-plane drawer additionally collapses any residual duplicate rows (historical data) into a single `×N` card defensively. - **`cache_lookup_ms` is a superset that spans the hooks stage.** The per-stage cursor advances via the next stage's `Mark()`, and the request-hooks stage records its sub-phases with `MarkBetween` (which does not advance the cursor). Consequently the cache stage's `Mark(cache_lookup_ms)` measures from the end of quota through the hooks stage *and* the cache lookup — i.e. `cache_lookup_ms` already contains the hooks-stage wall-clock (`hook_extract_ms + hook_build_ms + hook_pipeline_ms + hook_rewrite_ms`). This is pre-existing behaviour; the hook sub-phase keys are additive and do not change it. Analytics must therefore **not** sum all `*_ms` keys as if disjoint — the hook sub-phases are a drill-down *within* `cache_lookup_ms`, not siblings of it. - Latency rounding lives in `proxy.go`'s defer **and** in `finalize()` because a handful of failure paths return before the defer can finish snapshot computation; both sites round identically. - A `LatencyDetail` operator flag widens the snapshot to surface sub-ms phases as `1`, used during perf investigations. The default keeps prod rows compact. @@ -188,7 +190,7 @@ Producer-side back-pressure: - The in-memory buffer absorbs short bursts (10000 records). - **Overflow is governed by `LossMode`** (`AI_GATEWAY_AUDIT_LOSS_MODE`); the config default is **`spillblock`** (zero-loss). On a full in-heap buffer the record is handed to the durable on-disk spool off the request path — identical to `spill` in the normal regime, so the request path is not blocked there — and only if the spool channel *itself* saturates does Enqueue back-pressure the request goroutine (bounded by `backpressureMaxWait`) instead of taking `spill`'s last-resort bounded drop. So durable audit stays a true zero-loss promise at `spill`'s throughput. `block` (the empty/unknown fallback) hard-back-pressures from the first full buffer; `spill`/`drop` are the explicit lossy opt-outs for non-compliance callers (a saturated `spill`, or any `drop`, counts on `nexus_audit_mq_dropped_total`). The resolved mode is logged at boot (`audit overflow policy resolved configured=… effective=…`). -- `Close()` retries draining for 15 seconds before counting the remainder as dropped, so a graceful rollout window cleanly flushes pending audit; a kill-9 path drops whatever is still in memory. +- `Close()` retries draining for 15 seconds before counting the remainder as dropped, so a graceful rollout window cleanly flushes pending audit; a kill-9 path drops whatever is still in memory. On stop each bounded-queue consumer worker drains its remaining queued records and publishes them (bounded at `batchMaxCount` per publish, `drainOnStop`) before exiting, so a clean shutdown never leaves buffered audit unpublished. Consumer-side back-pressure: diff --git a/docs/developers/architecture/cross-cutting/safety/sse-streaming-compliance-architecture.md b/docs/developers/architecture/cross-cutting/safety/sse-streaming-compliance-architecture.md index f9c09b08..a7ee0fa4 100644 --- a/docs/developers/architecture/cross-cutting/safety/sse-streaming-compliance-architecture.md +++ b/docs/developers/architecture/cross-cutting/safety/sse-streaming-compliance-architecture.md @@ -66,13 +66,17 @@ changes on the ai-gateway path would be lossy for in-flight requests. The admin's `streamingMode` policy (`agent_settings.streamingMode` / per-domain override on `interception_domain`) resolves to one of -three modes via `shared/transport/streaming/policy.Resolve`: +three modes via `shared/transport/streaming/policy.Resolve`. A fourth +row — `modela` — is NOT a `policy.Resolve` output; it is a +scope-routing-derived overlay applied on top of the admin mode (see +"Scope-derived routing" below): | Mode | Behavior | Implemented by | |---|---|---| | `passthrough` | Bytes copied through without parsing; no hook executor (no compliance inspection). Used for non-AI SSE / cert-pinned clients where compliance can't introspect. | `shared/transport/tlsbump/sse.go::handleSSEResponse case "passthrough"` (agent + compliance-proxy); `ai-gateway/internal/ingress/proxy/proxy_cache_passthrough.go::runPassthroughStream` (ai-gateway) — three-service consistent. | -| `chunked_async` (live) | Bytes copied through immediately; LivePipeline runs hook executor at every checkpoint (cumulative bytes hit `FirstInspectChars`, then every `ReinspectStepChars`). PreHook callback stamps `ci.Normalized` BEFORE each hook run. Low-latency + full compliance. | **Two implementations on purpose** — see "Two LivePipeline implementations" below. `shared/transport/streaming/live.go` (tlsbump-driven: agent + compliance-proxy transparent-forwarder shape); `ai-gateway/internal/platform/streaming/live.go` (cache-replay path with format transform + hold-back + Modify-rewrite + OpenAI-DONE toggle). | -| `buffer_full_block` (buffer) | Read full body into memory; run hook executor ONCE on complete content; replay to client only on Approve / Abstain; on a `block` (RejectHard) decision write a single error event. PreHook callback stamps `ci.Normalized` between Phase 1 (read) and Phase 2 (run hooks). Strongest enforcement; highest latency. | `shared/transport/streaming/buffer.go` — tlsbump callers (agent + compliance-proxy) + ai-gateway (`proxy_cache_buffer.go::runBufferStream`). Three-service consistent. **Limitation**: Modify (`redact`) decisions are not supported by this pipeline (no rewrite arm in Phase 3); ai-gateway logs WARN + treats as Approve. | +| `chunked_async` (live) | Bytes copied through immediately; the LivePipeline runs the hook executor at every checkpoint **observe-only** — it records the audit outcome but NEVER blocks, holds back, or rewrites the wire. An enforcing scope (block/redact) is routed away from this path by scope-derived routing (see below), so live carries only non-enforcing (audit) traffic and always returns Approve. Low-latency real-time write-through. | **Two implementations on purpose** — see "Two LivePipeline implementations" below. `shared/transport/streaming/live.go` (tlsbump-driven: agent + compliance-proxy, **audit-only**); `ai-gateway/internal/platform/streaming/live.go` (cache-replay path with format transform + OpenAI-DONE toggle). | +| `modela` (redact-scope chunked_async) | Prescan-gated REAL-TIME streaming with a bounded held tail: forwards the response live except a trailing window, runs a cheap union prescan after each unit, and on a confirmed redact/block hit ESCALATES to buffer-the-remainder redaction so the held tail (and the rest) are delivered redacted, never raw. A sub-window value is always fully held when its completing bytes arrive; a value larger than the window leaks a bounded fragment (the disclosed best-effort-wire surface). | **One shared engine, two substrate adapters** (see "The shared Model A engine" below): `shared/transport/streaming/modela/engine.go` drives the algorithm; `ai-gateway/internal/ingress/proxy/proxy_cache_modela_substrate.go` (canonical chunks, fail-CLOSED) and `shared/transport/tlsbump/sse_modela.go` (raw SSE frames, fail-OPEN) supply the transport seams. | +| `buffer_full_block` (buffer) | Read full body into memory; run hook executor ONCE on complete content; replay to client only on Approve / Abstain; on a `block` (RejectHard) decision write a single error event; on a `redact` (Modify) decision the body is redacted before delivery (the masked copy is what the client receives — so it is also the stored copy). The two services redact via DIFFERENT mechanisms (see "Implemented by"). PreHook callback stamps `ci.Normalized` between Phase 1 (read) and Phase 2 (run hooks). Strongest enforcement; highest latency. | tlsbump (agent + compliance-proxy): `shared/transport/streaming/buffer.go` Phase-3 `Modify` arm splices via the per-host `sse_frame_redactor.go` (`SpliceTextFrames`), which is **fail-OPEN** — on an unreconstructable mask it relays the original + stamps `REDACT_INFLIGHT_UNSUPPORTED` (correct for the agent NE host-packet path; making the compliance-proxy appliance fail-CLOSED here is tracked work). ai-gateway: `proxy_cache_buffer.go::runCanonicalBufferStream` redacts on the canonical waist (`redactCanonicalBuffer`) and re-encodes — fail-CLOSED; it does NOT use the wire FrameRedactor. | The two adjacent fast-paths: @@ -90,6 +94,82 @@ The two adjacent fast-paths: as upstream responses, with the cached canonical chunks re-encoded into the current ingress's wire shape by a `StreamTranscoder`. +## Scope-derived routing (enforcing scopes never reach the audit-only live path) + +The admin streaming mode is a UX/latency preference; an **enforcing response scope +overrides it**. Because the live path is audit-only and cannot enforce, a scope that +may block or redact must never land there. The router builds the response pipeline +once and reads two content-INDEPENDENT predicates — `Pipeline.MayBlock()` / +`MayRedact()` (`shared/policy/pipeline/enforcement.go`), derived from the bound +hooks' declared `onMatch` action, an over-approximation that errs toward buffer: + +| Resolved scope | Routed mode | +|---|---| +| `MayBlock` | `buffer` (zero-leak full-buffer hard block) | +| `MayRedact` + admin `chunked_async` + a per-host adapter | `modela` (prescan-gated stream + escalate) | +| `MayRedact` + admin `chunked_async`, no adapter (no per-host text extraction) | `buffer` | +| `MayRedact` + admin `passthrough` (raw forwarding cannot redact) | `buffer` | +| non-enforcing | admin mode unchanged (`live` stays audit-only) | + +ai-gateway implements this in `stream_shape.go` §B2 (`s.modelAArmed`); tlsbump mirrors +it in `sse.go::scopeRouteSSEMode` / `overrideStreamingModeForScope` (`sse_modela.go`). +The same probe pipeline is reused as the Model A prescan/confirm executor. + +> **Predicate soundness caveat.** `MayBlock`/`MayRedact` read the *declared* +> `onMatch` action. A hook whose runtime decision can EXCEED its declared ceiling +> (e.g. a `webhook-forward` hook with `onMatch.action: approve` whose remote verdict +> reconciles via `StrictestDecision` to a reject/redact), and a hook with a malformed +> `onMatch`, are not over-routed to buffer today — closing that gap is tracked work +> the audit-only-live safety guarantee depends on. + +## The shared Model A engine + +`shared/transport/streaming/modela/engine.go` is the substrate-agnostic Model A +algorithm — `Run[U any](ctx, Substrate[U], Config)`. It owns the policy-neutral +control flow (bounded tail-hold window accounting, the windowed union prescan over +still-held content, one full confirm on a prescan hit, the escalate-to-buffer +decision keyed on the enforcing **action**) and nothing transport-specific. The +`Substrate[U]` interface supplies the seams a concrete transport fills: + +| Seam | ai-gateway (canonical) | tlsbump (wire) | +|---|---|---| +| unit `U` | `provcore.Chunk` | one parsed `*SSEEvent` (cached with its extracted text) | +| `Next` / text extraction | `ChunkSubscription` + canonical fields | `SSEParser` + the matched adapter's `ExtractStreamChunk` | +| deliver / re-encode | `canonicalbridge.StreamTranscoder` | `WriteSSEEvent` | +| `Confirm` / `Prescan` | the relay hook runner + `MayMatchRawContent` | `bo.policyResolver` pipeline + `MayMatchRawContent` | +| `Escalate` redaction | `redactCanonicalBuffer` (canonical waist) | `SpliceTextFrames` / `FrameRedactor` over the held+drained remainder | +| **fail posture** | fail-CLOSED (appliance) | fail-OPEN (agent NE host-packet safety) | + +The engine is posture-neutral; fail-open vs fail-closed is injected entirely by the +substrate. The escalate gate keys on `ActionFromDecision(decision)` (not the raw +decision enum) so a redact whose spans are masked behind a co-firing soft-block is +still escalated. A `MaxBufferBytes` eviction of an incomplete content unit, and the +escalation drain, both bound memory and escalate/block rather than deliver raw. + +**Aggregate-count rules are streaming-best-effort on Model A.** The confirm runs over +the still-held window only (`scanBuf[deliveredScanLen:]`), not the whole transcript, so +a rule that triggers on a *cumulative count* across the full response ("block if more +than N occurrences of pattern X") sees only the occurrences inside the held tail at any +one confirm — occurrences already delivered live are no longer in the window. Such a +rule may therefore under-count on Model A and fail to fire. This is the same disclosed +best-effort family as the over-window single-value surface above: strong compliance for +aggregate-count rules is the buffer or block modes, not Model A. + +**Long-pattern coverage and the config-time signal.** The flush-before-deliver lookahead +is sized per rule set from the longest CONTIGUOUS enforceable pattern (`pipeline.MaxPatternBound`), +floored at `modela.DefaultMaxPatternBytes`. Two pattern shapes remain best-effort, by design: +(1) **unbounded** patterns (`*` / `+` / `{n,}`) — their match length has no finite upper +bound, so no finite window can guarantee they are fully held; this is the normal, expected +state (token / PEM / JWT rules commonly carry unbounded quantifiers) and the engine package +header discloses it. (2) A **bounded** pattern whose length meets or exceeds the tail window +(`modela.DefaultTailWindowBytes`, ~8 KB) — the engine clamps the lookahead below the window, +so a value that long straddling unit boundaries may leak a bounded fragment. Case (2) is +rare (realistic patterns sit well under 8 KB); when it occurs the substrate emits a one-time +operator warning (`modela.WarnStreamingCoverageGap`) at stream setup. For full coverage of +either shape, route the affected policy through buffered streaming mode (or narrow the rule) +— the tail window is not an admin knob. The signal is observability-only; it changes no +enforcement. + ## The `PreHookCallback` contract The canonical type lives in `shared/policy/hooks/core/types.go`: @@ -166,6 +246,17 @@ LivePipeline fires PreHook at **every** checkpoint with the full body and `ci.Normalized` reflects the latest claim, which is what the per-checkpoint hook executor needs. +Because the executor runs once per checkpoint, the same hook is scanned many +times in one stream. The aggregated `finalResult.HookResults` therefore folds the +per-checkpoint results to **one record per hook** before it is emitted to the +audit pipeline — latency summed (the real scan CPU across the stream), the last +checkpoint's decision authoritative — via `foldHookResults` in +`shared/transport/streaming/live.go` (the agent + compliance-proxy path) and the +`responseHookAccumulator` on the ai-gateway side. Without the fold a streamed +response recorded the hook N times (N duplicate rows, an N×-inflated +`response_hooks_ms` / `response_hooks_us`); the fold keeps one row carrying the +true total. See `audit-pipeline-architecture.md` for the aggregate columns. + ### Service-side closures Neither ingress side wires an `OnPayload` closure — the PreHook's @@ -205,9 +296,11 @@ dormant complexity. - `Decision` enum (`Approve` / `Abstain` / `BlockSoft` / `RejectHard` / `Modify`) — single source. - `responseprehook.Build` — single builder, both impls call it. -- `nexus_streaming_modify_degraded_total{reason="buffer_mode"}` counter — - emitted from `shared.BufferPipeline.Process`, fires for all three - services when a Modify decision lands under buffer mode. +- `nexus_streaming_modify_degraded_total` counter — emitted from + `shared.BufferPipeline.Process` (`reason="buffer_mode"`, only the rare no-redactor + degrade) and from `sse_frame_redactor.go` (`tlsbump_splice_divergence` / + `tlsbump_tool_arg_undeliverable`, each tlsbump fail-open). It is a tlsbump signal: + ai-gateway's canonical buffer path does not flow through `shared.BufferPipeline`. - `nexus_normalize_panic_total{location="registry"|"on_payload"}` and `nexus_prehook_normalize_drop_total{adapter}` counters — both emitted from `shared/transport/normalize/responseprehook`. Single source of @@ -217,6 +310,19 @@ dormant complexity. planes route their pre-hook through `responseprehook.Build`. - `streampolicy.Store` + `BootStore` — single Store shape, single boot helper, three-service aligned. +- **The Model A streaming-compliance algorithm** (`shared/transport/streaming/modela/engine.go`) + — unlike the two transport-specific LivePipelines, the prescan-gated tail-hold / + confirm / escalate algorithm IS genuinely unified: one shared engine driven by the + canonical substrate (ai-gateway) and the wire substrate (tlsbump), so hooks / + compliance behave identically across all three ends while entry + delivery stay + per-service. See "The shared Model A engine" above. +- The per-host wire redactor (`shared/transport/streaming/frame_redactor_splice.go` + `SpliceTextFrames` + `FrameRedactor`) — single splice impl for buffer mode and the + Model A wire escalation. + +Note: tlsbump's shared `LivePipeline` is now **audit-only** (observe-only checkpoints, +real-time write-through, never blocks/rewrites) — enforcing scopes are routed to +`buffer` / `modela` upstream, so the live path carries only non-enforcing traffic. - Default fall-back arm: unknown enum → passthrough (matches tlsbump's `resolveStreamingMode`; pinned by `TestDispatchStreamMode_UnknownEnumFallsBackToPassthrough`). - `passthrough` mode relay — single `shared/transport/streaming.Passthrough`, @@ -242,40 +348,52 @@ and `proxy_cache_dispatch_test.go`. ## Asymmetries (intentional + visible) -### Buffer mode: Modify decisions degrade to Approve - -ai-gateway honors `buffer_full_block` by routing the -SSE handler through `shared.BufferPipeline` when -`StreamingPolicy.Get().Mode == ModeBufferFullBlock`. One residual -asymmetry remains by architecture: - -`shared.BufferPipeline.Process` Phase 3 handles `RejectHard` (the -`block` action — a single error event) and the default (Approve / -Abstain replay) but has no `Modify` arm — buffer mode replays the -buffered events verbatim, so a hook that returns `Modify` (`redact`) -with `ModifiedContent` cannot rewrite the body the way `LivePipeline`'s -held-back deltas can be edited mid-stream before the first flush. (A -`BlockSoft` decision, if one is produced internally, folds to the -`block` action and takes the same RejectHard arm — there is no soft -path.) - -**Three-service unified degradation signal**: the -degradation is detected inside `shared.BufferPipeline.Process` -itself — when `result.Decision == Modify` arrives from the executor, -the pipeline: - -1. Emits a `WARN` log line with the `requestId` and rejection reason. -2. Bumps the Prometheus counter - `nexus_streaming_modify_degraded_total{reason="buffer_mode"}`. - -Because all three data planes (`ai-gateway`, `compliance-proxy`, -`agent`) buffer through the same shared pipeline, the metric and log -fire from a single source of truth — Prometheus scrape job/instance -labels distinguish which data plane saw the degradation. The -ai-gateway `bufferModeExecutor` adapter (struct in -`proxy_cache_buffer.go`) is now a pure type bridge from -`StreamHookRunner` (func) to `PipelineExecutor` (interface); it owns -no log or metric. +### Buffer mode: Modify is redacted (degradation is now narrow) + +Buffer mode HONORS a `Modify` (`redact`) decision — it no longer degrades to +Approve. The two data planes redact via different mechanisms: + +- **tlsbump** (agent + compliance-proxy): `shared.BufferPipeline.Process` Phase 3 has + a `Modify` arm that runs the per-host `sseFrameRedactor` (`SpliceTextFrames`) over + the buffered timeline — splicing the masked text into the text frames, passing + non-text frames byte-verbatim. The masked stream is both delivered and stored. The + redactor is **fail-OPEN**: on an unreconstructable mask (tool-arg masking + undeliverable on the wire, or a wire/normalized divergence) it relays the original + + stamps `REDACT_INFLIGHT_UNSUPPORTED` (correct for the agent NE host-packet path; + the compliance-proxy appliance fail-CLOSED variant is tracked work). +- **ai-gateway**: `runCanonicalBufferStream` does NOT touch the wire frames — it + redacts on the canonical waist (`redactCanonicalBuffer`) and re-encodes to the + ingress wire, **fail-CLOSED** (a redact the policy demanded but cannot apply is + rejected, never forwarded unredacted). + +A `BlockSoft` decision folds to the `block` action and takes the RejectHard arm — +there is no soft delivery path. + +**Degradation signal**: `nexus_streaming_modify_degraded_total` fires only on the +NARROW residual, not on every Modify. `sse_frame_redactor.go` bumps it with reason +`tlsbump_splice_divergence` / `tlsbump_tool_arg_undeliverable` on each tlsbump +fail-open; `shared.BufferPipeline` bumps `buffer_mode` only when a `Modify` lands with +no FrameRedactor wired (a no-adapter degrade, not a production path — production +always resolves a per-host adapter). ai-gateway's canonical buffer path does NOT flow +through `shared.BufferPipeline`, so it never emits this counter. Prometheus scrape +labels distinguish which data plane saw a tlsbump degradation. + +**Enforcement re-emit matches the ingress wire shape**: when buffer +mode (or the Model-A cache re-emit) re-encodes the redacted canonical +body with no cross-format transcoder selected (a same-shape route), +the fallback encoder is chosen by ingress, not hardcoded to chat: +`fallbackStreamEncoder(ingress, model)` in `proxy_cache_buffer.go` +returns `NewResponsesStreamEncoder` for a `/v1/responses` +(`FormatOpenAIResponses`) ingress and `NewChatCompletionsStreamEncoder` +otherwise. So an enforced `/v1/responses` stream still emits +`response.*` events with a terminal `response.completed`, never +`chat.completion.chunk`. Enforcement takes precedence over the native +Responses passthrough (see *Responses egress: two signals* in +[provider-adapter-architecture.md](../../services/ai-gateway/provider-adapter-architecture.md)); +because the redacted body is rebuilt from the canonical waist, the +upstream's built-in-tool / audio events are forfeited under an +enforcing response scope — the accepted trade for a redacted or blocked +stream. **Admin-visible warning surface**: the Control Plane streaming-compliance settings endpoint diff --git a/docs/developers/architecture/services/ai-gateway/cost-estimation-architecture.md b/docs/developers/architecture/services/ai-gateway/cost-estimation-architecture.md index 8fb40232..fcafd950 100644 --- a/docs/developers/architecture/services/ai-gateway/cost-estimation-architecture.md +++ b/docs/developers/architecture/services/ai-gateway/cost-estimation-architecture.md @@ -81,6 +81,7 @@ The usage extractor delegates to the shared normalize codecs (`packages/shared/t - **Anthropic SSE** is folded by the codec's stream state machine: input-side usage (including `cache_read_input_tokens` / `cache_creation_input_tokens`) comes from `message_start`, output-side from `message_delta`, merged into the canonical convention (`PromptTokens` = uncached + cache read + cache creation; `CompletionTokens` = output). Tool-use-only and thinking-only streams carry full usage like text streams. - **Reasoning tokens**: a wire-explicit count (`output_tokens_details.thinking_tokens` on Anthropic; `completion_tokens_details.reasoning_tokens` on OpenAI-compatible wires) always wins; the character-based derivation is only a fallback when the wire omits the count. The OpenAI-compatible `reasoning` field is an accepted wire alias of `reasoning_content` and feeds the same reasoning text accounting. +- **`/v1/responses` native passthrough**: when a genuine Responses upstream streams verbatim on the non-enforced lane, the egress copier forwards each SSE frame byte-for-byte yet a usage tee still decodes the canonical `Usage` (Responses `input_tokens` / `output_tokens`) onto the same chunk, so cost and token stamping land unchanged — the raw-byte forwarding bypasses the re-encode, not the usage accounting. The two-signal egress model is described in [`provider-adapter-architecture.md`](./provider-adapter-architecture.md) (*Responses egress: two signals*). See [`sse-streaming-compliance-architecture.md`](../../cross-cutting/safety/sse-streaming-compliance-architecture.md) for the streaming dispatch contract. ### Normalized projection is not on the cost path — cost/field neutral diff --git a/docs/developers/architecture/services/ai-gateway/hook-architecture.md b/docs/developers/architecture/services/ai-gateway/hook-architecture.md index d972e3e9..3ed6ee68 100644 --- a/docs/developers/architecture/services/ai-gateway/hook-architecture.md +++ b/docs/developers/architecture/services/ai-gateway/hook-architecture.md @@ -57,6 +57,8 @@ The same stream-entry probe drives the streaming-mode routing. The resolved resp A pipeline runs its hooks under a total timeout with a per-hook timeout — the AI Gateway sets these to 15 and 5 seconds (the per-hook value overridable per config via `timeoutMs`), and the framework falls back to 30 and 5 seconds when a caller leaves them unset. Every `Execute` call is wrapped so a panicking hook becomes an error rather than crashing the data plane. On an error or timeout the hook's `failBehavior` decides the outcome — `fail-closed` yields `RejectHard`, the default `fail-open` yields `Approve`. A nil result is treated as `Abstain`. +Each hook is self-timed from a single captured `elapsed` and stamped in **both** units: `LatencyMs` (the truncated integer-millisecond floor, kept for backward compatibility and **never clamped** because it is summed downstream) and `LatencyUs` (the precise microsecond value). Hooks run at microsecond scale, so `LatencyMs` floors a sub-millisecond hook to `0`; `LatencyUs` carries the real value into the per-hook trace and the additive `request_hooks_us` / `response_hooks_us` aggregates. See `audit-pipeline-architecture.md` for the aggregation and the streaming one-row-per-hook fold. + The runner has two modes: - **Sequential** (the AI Gateway) — hooks run in priority order, short-circuiting on the first `RejectHard`. When a hook returns `Modify`, its transform spans are applied to the normalized payload before the next hook runs, so later hooks see the redacted content; emitted tags accumulate across hooks. @@ -64,6 +66,8 @@ The runner has two modes: `mergeResults` aggregates by priority order: the first `RejectHard` wins outright; otherwise a `BlockSoft` (e.g. a soft AI-Guard verdict folded into the merge); otherwise a `Modify`; otherwise `Approve`. Tags are unioned, and the strictest action across hooks (`StrictestAction`) is carried onto the result. A merged `BlockSoft` carries no distinct client-facing response — `ActionFromDecision` folds it to the `block` action, so dispatch treats it identically to `RejectHard` (reject; see §3). The AI Gateway enables two flags on its pipeline: `allowModify` (MODIFY passes through instead of being downgraded to APPROVE) and `clearSoftOnApprove` (a later APPROVE clears a pending soft block). +When a redacting hook (`Modify`) co-fires with a soft-block hook the aggregate is promoted to `BlockSoft` while the redact's payload still rides along, so consumers must apply the mask rather than forward raw. They gate on `CompliancePipelineResult.CarriesRedaction()`, not on `Decision == Modify`. For a `BlockSoft` aggregate that predicate consults `RedactionApplicable`, a flag `mergeResults` stamps when an **enforcing** per-hook decision (`Modify` or `BlockSoft`) contributed an **applicable** artifact — `ModifiedContent` (carried only from `Modify` hooks) or a span whose `ContentAddress` is not an audit-only sentinel (`normalize.IsAuditOnlySentinelAddress`, currently the `webhook.flat` address a webhook-forward hook emits for advisory `redactions[]` it cannot apply in-flight). Keying on raw span/content presence instead would let an approve-webhook's (or an approve-ceiling-reconciled `Modify`-webhook's) audit-only spans masked behind a soft-block report a redaction the aggregate cannot apply, over-blocking a stream that should soft-deliver; keying on `Decision == Modify` alone would drop a real co-firing redaction. The sentinel check is a deliberate denylist (an unknown address counts as applicable) so the failure direction is over-block, never leak. Advisory spans remain unioned into the result for the audit record regardless. + ## 6. Config flow `HookConfigCache` is the bridge from stored config to the resolver: a loader reads the `HookConfig` rows and `Swap`s them into the `PolicyResolver`. On the server-side data planes it reloads when the Hub pushes a config change (via the thing-client `OnConfigChanged` callback) with a TTL backstop; the Agent has no direct database access, so it is push-only. Before the swap, `rulepack.Enrich` binds each installed rule pack into the relevant hook's config under `_rulePackInstalls`, so the `rulepack-engine` hook evaluates packs without holding a database handle inside `Execute`. diff --git a/docs/developers/architecture/services/ai-gateway/normalization-architecture.md b/docs/developers/architecture/services/ai-gateway/normalization-architecture.md index 9bd66e30..a987fd73 100644 --- a/docs/developers/architecture/services/ai-gateway/normalization-architecture.md +++ b/docs/developers/architecture/services/ai-gateway/normalization-architecture.md @@ -15,6 +15,8 @@ type Normalizer interface { `Meta` carries the call context: `AdapterType` (the wire key), `Model`, `ContentType`, `Direction` (`DirectionRequest` / `DirectionResponse`), `EndpointPath`, and `Stream`. `NormalizedPayload` is the canonical output — `Kind` (`ai-chat` / `ai-embedding` / `http-json` / …), `Protocol`, `Model`, `Stream`, `Messages[]`, `Tools[]`, `Params`, `Usage`, `FinishReason`, `Inputs[]` (embeddings), `Confidence`, and `DetectedSpec`. A normalizer that does not recognize the bytes returns `ErrUnsupported`, which the coordinator uses to fall through to the next candidate. +A `TransformSpan` describes one byte-level modification against a `NormalizedPayload` (a hook redaction, AI-Guard suggestion, cache-normaliser strip, or cache_control inject); `ApplySpans` reconstructs the wire body from the payload + its spans. A span's `ContentAddress` indexes the modified content (e.g. `messages..content.`, `http.bodyView`). The reserved address `AddressAuditOnlySentinel` (`"webhook.flat"`) marks an **audit-only** span: it records *what* a subsystem flagged but addresses a flat projection the gateway never reconstructs, so `ApplySpans` drops it — it lands in the audit record only and never mutates delivered or stored bytes. `IsAuditOnlySentinelAddress` is the single recogniser; compliance consumers use its negation as a denylist to tell an applicable redaction from an advisory audit-only span (see `hook-architecture.md` §5, `CarriesRedaction`). + ## 2. The tiered dispatch model `core.Registry` is the coordinator. `BuildRegistry` (`packages/shared/transport/normalize/buildregistry.go`) assembles it once per service and freezes it. `Registry.Normalize` dispatches in tiers: diff --git a/docs/developers/architecture/services/ai-gateway/provider-adapter-architecture.md b/docs/developers/architecture/services/ai-gateway/provider-adapter-architecture.md index 1a34f523..52a403f8 100644 --- a/docs/developers/architecture/services/ai-gateway/provider-adapter-architecture.md +++ b/docs/developers/architecture/services/ai-gateway/provider-adapter-architecture.md @@ -82,11 +82,52 @@ On an upstream error response the adapter's `ErrorNormalizer.Normalize` produces The caller's wire shape is preserved end-to-end: whatever ingress a client calls — `/v1/chat/completions`, `/v1/messages`, gemini `:generateContent`, `/v1/responses`, `/v1/embeddings`, plus the Azure and GLM native ingresses — receives a response in that same shape. The upstream target wire is an internal concern resolved at the call site, not the caller's: -- **Request.** The ingress body is canonicalized once (`IngressChatToCanonical` for chat-kind, `IngressEmbeddingsToCanonical` for embeddings), then `TargetExecutor` sets the call-time `WireShape` from the *target* format — `ChatWireShapeForTarget` for chat-kind, `EmbeddingsWireShapeForTarget` for embeddings — so `Transport.BuildURL` and `SchemaCodec.EncodeRequest` target the correct wire for the primary target and every failover target. The per-request `Ingress.WireShape` is not mutated to the target shape; the `/v1/responses` → chat-completions downgrade is the one exception, because Responses canonicalizes to chat before dispatch. +- **Request.** The ingress body is canonicalized once (`IngressChatToCanonical` for chat-kind, `IngressEmbeddingsToCanonical` for embeddings), then `TargetExecutor` sets the call-time `WireShape` from the *target* format — `ChatWireShapeForTarget` for chat-kind, `EmbeddingsWireShapeForTarget` for embeddings — so `Transport.BuildURL` and `SchemaCodec.EncodeRequest` target the correct wire for the primary target and every failover target. The per-request `Ingress.WireShape` is not mutated to the target shape; the `/v1/responses` → chat-completions downgrade is the one exception, and it is **capability-driven** — applied only when the resolved target does *not* serve the Responses API, in which case Responses canonicalizes to chat before dispatch (see *Responses egress: two signals*, below). - **Embeddings endpoint selection (single vs batch).** A codec may serve two upstream endpoints from one `WireShape`. Gemini's `WireShapeGeminiEmbedContent` covers both `:embedContent` (single string `input`) and `:batchEmbedContents` (array `input`); the choice is encoded only by the embeddings codec, which inspects the canonical `input` cardinality and returns an `EncodeResult.URLOverride` of `:embedContent` or `:batchEmbedContents`. Because embeddings always skip the gateway cache, the cross-format request is translated by `IngressEmbeddingsToWire`, which now **surfaces that override** alongside the wire body; `TargetExecutor` threads it into `Adapter.ExecuteWithBody`, where `applyURLOverride` swaps the action suffix on the `Transport.BuildURL` result. Dropping the override sends the batch body (`{"requests":[…]}`) to the single-embed URL and Gemini rejects it with `Unknown name "requests": Cannot find field` (regression guard: `TestIngressEmbeddingsToWire_GeminiEndpointSelection`, `TestExecute_EmbeddingsBridgeURLOverride_ReachesAdapter`). No Gemini-native embeddings *ingress* exists, so an embeddings request to a Gemini target is always cross-format and always flows through this path. - **Response.** The upstream wire body is decoded to canonical, then reshaped back to the caller's format with `ResponseCanonicalToIngress` (chat) / `ResponseCanonicalToIngressEmbeddings` (embeddings), keyed on the ingress read from the request context (not the mutable per-request copy). The reshape fires when the ingress format differs from the target and is an identity no-op for same-format native routes. -The cross-format decision is driven by `typology.KindFromWireShape` (chat / embeddings) plus the responses-native rule rather than a hardcoded ingress list, so a new chat or embeddings ingress is covered without changing the dispatch gates. +The cross-format decision is driven by `typology.KindFromWireShape` (chat / embeddings) plus the per-provider Responses capability (see *Responses egress: two signals*, below) rather than a hardcoded ingress list or a Format-level guess, so a new chat or embeddings ingress is covered without changing the dispatch gates. + +### Responses egress: two signals (capability + content) + +The `/v1/responses` ingress is not a codec on top of chat — it is a co-equal OpenAI standard with strictly greater expressive power (typed `output[]` items, reasoning items, built-in tools, audio streaming, stateful `previous_response_id`). Routing it correctly uses **two independent signals**, never a single Format guess. The **request side** decides what wire to send upstream from a per-provider capability; the **response side** decides how to decode/encode from the actual bytes that came back. The two are deliberately separate: the request-side capability can be wrong (a provider may mis-declare or change behaviour), and the content signal is what makes the egress bulletproof regardless. + +**Request side — capability-driven.** Whether a target receives a Responses-shape body or the downgraded chat-completions body is the per-provider capability `servesResponsesAPI`, resolved by `canonicalbridge.Bridge.ServesResponses(target, override)`: + +``` +/v1/responses ingress, resolved target T: + ServesResponses(T)? + yes → send Responses-shape body to T's /v1/responses (built-in tools preserved) + no → responses → canonical(chat) → T's native chat wire (the downgrade) +``` + +Resolution order: the per-provider override (`Provider.serves_responses_api`, a nullable Boolean column) is **downgrade-only** — it can turn the capability *off* for a chat-only OpenAI-compatible endpoint, but cannot claim a capability the adapter lacks; with no override (the common case) the default comes from the adapter's `RequestShapes` (`RequestShapes ⊇ WireShapeOpenAIResponses`; today only `FormatOpenAI`). A lockstep test pins the Format default against the OpenAI adapter's declared `RequestShapes`. The capability is resolved per-target from the hydrated routing snapshot inside the failover loop — never a per-request DB read — and rides the already-threaded `CallTarget`. The request-side wire decision consults it at the same value everywhere: the executor's `nativeResponses` decision, `stage_routing.go`, `stage_cache_body.go`, and `bridge.IngressChatToWire`. This keeps real OpenAI / Azure working out of the box (default true) while letting a mock or chat-only endpoint opt out, so the gateway never POSTs a Responses body to an endpoint that would 404 it. + +**Response side — content-driven (authoritative).** The decode/encode decision is driven by the **actual upstream bytes**, not by Format and not by the request-side capability. `specs/openai/responses/classify.go` performs exactly one classification: + +``` +Non-stream (ClassifyNonStreamBody — top-level "object"): + "response" → verbatim passthrough (no re-encode; built-in tools preserved) + "chat.completion" → canonical → EncodeResponsesResponse + else → fail closed (502 — never verbatim) + +Stream (ClassifyFirstSSEFrame — first decoded SSE frame): + event: response.* / data {"type":"response.*"} → copier mode (verbatim frames) + data {"object":"chat.completion.chunk"} → chat mode → responsesStreamEncoder + else → fail closed (canonical chat lane — never verbatim) +``` + +The streaming classification happens **exactly once**, lazily on the first decoded frame at the raw-byte boundary (`specs/openai/stream/stream_responses_egress.go`), reusing the shared `SSEScanner` buffer — one per-stream hold, zero per-chunk allocation. The resolved wire shape is carried forward on the stream; the proxy layer does not sniff a second time. Trusting the bytes (not the declared Format) is what keeps a chat-shaped reply from being forwarded to a `/v1/responses` client: even a provider that mis-declares its capability cannot leak `chat.completion.chunk` frames, because the encoder follows the sniffed shape. Content authority is **wire-shape only**, never a compliance signal; the sniffed shape is cross-checked against the resolved capability. + +**Raw-byte copier (non-enforced native path).** On the verbatim path the copier forwards each upstream SSE frame byte-for-byte (`Chunk.Verbatim` + `RawBytes`) so built-in-tool / audio events reach the client unparsed, while a usage tee decodes the canonical `Delta` / tool-call / reasoning / usage fields onto the **same** chunk so token and cost accounting survive. "Preserved" here means no re-encode through the canonical waist — not zero-cost. + +**Precedence: enforcement > passthrough.** An enforcing response scope (redact / hard-block) forces canonical **buffer** mode, which rewrites the canonical body and therefore cannot also forward verbatim frames. Verbatim passthrough is allowed **only** on the non-enforced `/v1/responses` live lane (`allowVerbatim` in `stream_shape.go` requires `FormatOpenAIResponses` ingress AND no enforcing block/redact AND not the chat-ingress auto-upgrade). When an enforcing scope applies, built-in-tool / audio fidelity is **forfeited** — an accepted, documented blind spot. The enforcement and Model-A re-emit fallback is ingress-shape-aware: for `FormatOpenAIResponses` ingress it builds `NewResponsesStreamEncoder` (never `chat.completion.chunk`), so an enforced `/v1/responses` stream still emits `event: response.*` with a terminal `response.completed`. A second blind spot: built-in-tool content forwarded on the raw-byte path is absent from normalized text and compliance scanning. + +**Fail closed.** An unclassifiable / empty / keep-alive-only first frame is **never** forwarded verbatim — the non-stream path returns 502 and the stream path falls back to the canonical chat (or enforced) lane. SSE comments and keep-alives are skipped via the shared `SSEScanner` before classification. + +**Cache-HIT scope.** Content-peek runs only on the LIVE / cache-MISS lane. Cache-HIT replay chunks carry no `RawBytes`, so the origin-tag override (`StreamHitOrigin`) stays the authoritative wire-shape selector on a hit. + +**Parity holds (§3a Rule 6).** Streaming and non-streaming `/v1/responses` egress stay at parity: both start `response.created`, end `response.completed`, carry a terminal event, and report the same `finish_reason`. ### Round-trip equivalence standard (the shape-conversion test of record) diff --git a/docs/users/api/openapi/control-plane/providers.yaml b/docs/users/api/openapi/control-plane/providers.yaml index 44b0f3b2..eb96b91c 100644 --- a/docs/users/api/openapi/control-plane/providers.yaml +++ b/docs/users/api/openapi/control-plane/providers.yaml @@ -131,6 +131,18 @@ paths: - boolean - "null" description: Whether the provider is active. Defaults to true when omitted. + servesResponsesApi: + type: + - boolean + - "null" + description: >- + Whether this provider's endpoint serves the OpenAI Responses + API (/v1/responses). null (or omitted) uses the adapter + type's default; false forces a downgrade to + /v1/chat/completions for OpenAI-compatible endpoints that only + implement chat completions. Downgrade-only: a true value + cannot exceed the wire adapter's actual capabilities. + example: false headers: description: Optional arbitrary JSON of extra upstream headers. models: @@ -320,6 +332,15 @@ paths: - boolean - "null" description: New enabled state. + servesResponsesApi: + type: + - boolean + - "null" + description: >- + New OpenAI Responses API (/v1/responses) capability override. + null clears the override back to the adapter type's default; + false forces a downgrade to /v1/chat/completions. Downgrade-only: + a true value cannot exceed the wire adapter's actual capabilities. example: displayName: OpenAI (Prod) enabled: false @@ -1103,6 +1124,14 @@ components: - "null" enabled: type: boolean + servesResponsesApi: + type: + - boolean + - "null" + description: >- + Whether this provider's endpoint serves the OpenAI Responses API + (/v1/responses). null means "use the adapter type's default"; false + downgrades to /v1/chat/completions. Downgrade-only. headers: description: arbitrary JSON value createdAt: diff --git a/docs/users/features/cp-ui/ai-gateway-routing.md b/docs/users/features/cp-ui/ai-gateway-routing.md index d4426a97..257e00b9 100644 --- a/docs/users/features/cp-ui/ai-gateway-routing.md +++ b/docs/users/features/cp-ui/ai-gateway-routing.md @@ -14,6 +14,8 @@ These five pages form a setup chain: register a **Provider**, attach a **Credent **Key concepts.** `adapterType` identifies the provider's wire spec (which in-tree codec talks to it). A provider with no enabled credential cannot serve traffic. +**Serves OpenAI Responses API.** The provider info form carries a "Serves OpenAI Responses API" setting with three choices: **Use adapter default** (the default — the gateway infers the capability from the provider's adapter type, so a real OpenAI provider serves `/v1/responses` out of the box), **Enabled**, and **Disabled (chat completions only)**. It controls whether a `/v1/responses` request routed to this provider is sent upstream in the Responses wire shape or downgraded to chat completions first. Leave it on the default unless the provider is an OpenAI-compatible endpoint that only implements `/v1/chat/completions` — set it to **Disabled** so the gateway never sends a Responses-shape body the endpoint would reject. The setting can only narrow the adapter default (it can disable a capability, not invent one); regardless of how it is set, a `/v1/responses` caller always receives a Responses-shaped reply. + **Fetching models from the provider (OpenAI-compatible only).** On the Models step of the create-provider wizard, a "Fetch from /v1/models" button is shown for custom providers. Clicking it calls `POST /api/admin/providers/discover-models` with the base URL, adapter type, and API key entered in the earlier wizard steps. The response pre-fills the model table with the ids returned by the upstream provider's model-listing endpoint. Each row also receives a suggested model type (chat, embedding, audio, or image) derived by a best-effort heuristic from the model id; the admin can change the type before saving. The button is disabled until both a base URL and an API key (or "skip credential") are provided. This feature is limited to providers whose adapter type speaks the OpenAI wire format — `openai`, `deepseek`, and all OpenAI-compatible adapters (Groq, Mistral, Fireworks, Together, xAI, Perplexity, HuggingFace, Moonshot). If the selected adapter does not support the standard `/v1/models` endpoint, the wizard shows an inline message that model fetch is not available for that adapter and prompts the admin to add models manually. Pricing fields are always left blank after a fetch — the `/v1/models` endpoint carries no pricing data — and must be filled manually before saving. diff --git a/docs/users/features/cp-ui/overview.md b/docs/users/features/cp-ui/overview.md index 03d82bcd..e7abb90d 100644 --- a/docs/users/features/cp-ui/overview.md +++ b/docs/users/features/cp-ui/overview.md @@ -39,7 +39,7 @@ Four of the five pages — Dashboard, Analytics & Metrics, Quota Usage, and Cach The **Provider** and **Model** filters — and the Analytics group-by axes — attribute by the model/provider that actually **served** the request (the routed target), so a filter on "Provider X" returns everything X handled regardless of what the client originally asked for. The **requested** provider/model is shown only as the separate "requested model" column and is blank when the client did not pin one (for example `model="auto"` or an OpenAI-style request). -**Key actions.** Switch source tab. Open the filter panel for time range and advanced filters; apply, clear, or refresh. Click any row to open the event drawer (full request and response payload, hook trace, downstream timings). In the drawer's normalized payload view, content that a compliance hook redacted is marked inline as a badge over the replacement text, with a tooltip naming the rule, source, action, and reason. When the storage policy kept no readable copy at all, the payload view shows a notice instead of the content, and the notice distinguishes why: content dropped because the operator's policy says to drop it; or — when the policy was to redact but the redaction could not be safely applied to the stored copy — a separate notice explaining that the copy was dropped as the safe fallback, with the reason in plain words (the machine token shown alongside), the parts of the payload that could not be resolved, and the rules that matched. Events recorded before this distinction existed show a neutral "content not stored per the storage policy" notice that does not guess between the two. Paginate. Deep-link to a single event via `?thingId=`. +**Key actions.** Switch source tab. Open the filter panel for time range and advanced filters; apply, clear, or refresh. Click any row to open the event drawer (full request and response payload, hook trace, downstream timings). The hook trace lists each hook that ran with its decision and execution time shown to microsecond precision (hooks routinely complete in well under a millisecond, so a millisecond figure would read as zero); a hook that a streamed response evaluated repeatedly is shown once with a ×N count and its total time, rather than one row per evaluation. In the drawer's normalized payload view, content that a compliance hook redacted is marked inline as a badge over the replacement text, with a tooltip naming the rule, source, action, and reason. When the storage policy kept no readable copy at all, the payload view shows a notice instead of the content, and the notice distinguishes why: content dropped because the operator's policy says to drop it; or — when the policy was to redact but the redaction could not be safely applied to the stored copy — a separate notice explaining that the copy was dropped as the safe fallback, with the reason in plain words (the machine token shown alongside), the parts of the payload that could not be resolved, and the rules that matched. Events recorded before this distinction existed show a neutral "content not stored per the storage policy" notice that does not guess between the two. Paginate. Deep-link to a single event via `?thingId=`. **Normalized view.** The drawer's payload tab renders each direction as a typed, readable projection rather than raw bytes. A provenance badge says which decoder produced the view: **Tier 1** (exact protocol decoder, matched by key, host, or content sniff), **Tier 2** (pattern probe for consumer web surfaces) — both shown with a confidence score meaning the fraction of the input the decoder recognized — or the neutral **Structural** badge for rows where no AI protocol was identified and the body is shown as a typed projection of the raw HTTP content (JSON tree, text, form fields, binary digest, or an event-stream frame list). Structural rows deliberately carry no confidence numeral — the projection is faithful by construction, but it makes no claim about AI semantics. The normalized view is a derived projection: it is recomputed automatically when decoders improve, so a historical row's rendering (and its confidence) can get better over time — the raw bytes in the Raw tab are the immutable audit record and never change. An unrecognized event stream renders as a frame list — one row per frame with its event-name chip and the frame data pretty-printed when it is JSON — collapsed beyond the first 50 frames behind a "show all" control; very long streams note that the frame view is truncated while the full stream remains available in the Raw tab. Chat-style rows render role bubbles; a tool call's multi-line inputs (for example a shell command an agent ran) display with real line breaks, and the usage row includes reasoning tokens when the provider reported them. diff --git a/scripts/doc-lockstep.config.mjs b/scripts/doc-lockstep.config.mjs index 2ef493c4..db6b8a77 100644 --- a/scripts/doc-lockstep.config.mjs +++ b/scripts/doc-lockstep.config.mjs @@ -274,6 +274,17 @@ export default [ 'packages/shared/transport/streaming/policy/**', 'packages/shared/transport/tlsbump/sse.go', 'packages/shared/transport/tlsbump/bump.go', + // Substrate-agnostic Model A engine + the per-host wire redactor, the + // scope-routing predicate, and both substrate adapters (canonical + wire) + // that drive the engine — three-end hooks/compliance parity. + 'packages/shared/transport/streaming/modela/**', + 'packages/shared/transport/streaming/frame_redactor.go', + 'packages/shared/transport/streaming/frame_redactor_splice.go', + 'packages/shared/transport/tlsbump/sse_modela.go', + 'packages/shared/transport/tlsbump/sse_frame_redactor.go', + 'packages/shared/policy/pipeline/enforcement.go', + 'packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela.go', + 'packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_substrate.go', // ai-gateway streaming format + ingress dispatch (R1/R3 fixes) 'packages/ai-gateway/internal/platform/streaming/live.go', 'packages/ai-gateway/internal/platform/streaming/format/**', diff --git a/tools/db-migrate/schema/providers.prisma b/tools/db-migrate/schema/providers.prisma index ce742005..ac0397cc 100644 --- a/tools/db-migrate/schema/providers.prisma +++ b/tools/db-migrate/schema/providers.prisma @@ -35,6 +35,15 @@ model Provider { captureRequestBody Boolean? @map("capture_request_body") captureResponseBody Boolean? @map("capture_response_body") rawBodySpillEnabled Boolean? @map("raw_body_spill_enabled") + /// Whether this provider's upstream natively serves the OpenAI + /// /v1/responses API (Responses-shape request body + response bytes). + /// NULL = adapter RequestShapes default (FormatOpenAI → true; every + /// other adapter → false). An explicit value only DOWNGRADES: false + /// forces /v1/responses ingress through the canonical(chat) codec even + /// for an OpenAI-family adapter (e.g. an OpenAI-compatible endpoint that + /// only implements /v1/chat/completions); true cannot grant a capability + /// the adapter lacks. + servesResponsesApi Boolean? @map("serves_responses_api") createdAt DateTime @default(now()) @db.Timestamptz(3) updatedAt DateTime @updatedAt @db.Timestamptz(3) diff --git a/tools/db-migrate/schema/traffic.prisma b/tools/db-migrate/schema/traffic.prisma index 4a93414e..adca17aa 100644 --- a/tools/db-migrate/schema/traffic.prisma +++ b/tools/db-migrate/schema/traffic.prisma @@ -59,6 +59,12 @@ model traffic_event { request_hooks_ms Int? @map("request_hooks_ms") /// Aggregate of per-hook latency on the response side. response_hooks_ms Int? @map("response_hooks_ms") + /// Microsecond-precision hook aggregates (SUM(*_hooks_pipeline[*].latencyUs)). + /// Hooks run at microsecond scale, so the _ms columns truncate sub-millisecond + /// hooks to 0; these carry the real value. Additive — the _ms columns are + /// unchanged. NULL when no hook ran on that stage. + request_hooks_us Int? @map("request_hooks_us") + response_hooks_us Int? @map("response_hooks_us") /// Long-tail per-source phase durations (ms). Closed key set per `source`: /// ai-gateway → auth_ms / quota_ms / routing_ms / cache_lookup_ms / /// req_adapter_ms / resp_adapter_ms From c4e7a0cb5d094fc7517294e5ed03cfd31ca731d1 Mon Sep 17 00:00:00 2001 From: nexus Date: Tue, 30 Jun 2026 21:42:27 +0800 Subject: [PATCH 2/8] feat(services): sync backend services and shared libs from internal main (sync 9) Backend service implementations and shared Go libraries (TLS-bump streaming-compliance + provider-adapter cycle): - packages/ai-gateway: provider adapters, normalization, SSE streaming compliance, cost estimation - packages/shared: transport (normalize/codecs), policy, audit - packages/control-plane: provider store wiring - packages/nexus-hub: hub service - packages/compliance-proxy: pipeline - go.work: workspace member set Runtime artifacts (logs/, dev-certs/) intentionally excluded. --- .../cmd/ai-gateway/wiring/providers.go | 13 +- .../ai-gateway/wiring/providers_layer_test.go | 4 +- .../cmd/ai-gateway/wiring/wiring_init_test.go | 10 +- .../internal/cache/layer/helpers_test.go | 2 +- .../internal/cache/layer/layer_test.go | 8 +- .../internal/cache/layer/loaders.go | 4 +- .../internal/cache/layer/lookups_test.go | 4 +- .../internal/cache/layer/metrics_test.go | 2 +- .../internal/execution/canonicalbridge/api.go | 7 +- .../execution/canonicalbridge/bridge.go | 103 +- .../canonicalbridge/bridge_responses_test.go | 82 +- .../execution/canonicalbridge/matrix_test.go | 35 +- .../canonicalbridge/stream_encoders.go | 211 ++-- .../stream_encoders_golden_test.go | 177 ++++ .../stream_encoders_ingress_test.go | 62 ++ .../stream_encoders_openai_types.go | 70 ++ .../internal/execution/executor/executor.go | 2 +- .../executor/executor_internals_test.go | 5 +- .../ingress/debug/debug_branches_test.go | 2 +- .../ingress/proxy/audit_findings_test.go | 8 +- .../ingress/proxy/chunk_stream_reader_test.go | 12 +- .../ingress/proxy/hook_extraction_test.go | 2 +- .../internal/ingress/proxy/ingress.go | 16 +- .../internal/ingress/proxy/ingress_test.go | 8 +- .../ingress/proxy/proxy_cache_buffer.go | 23 +- .../ingress/proxy/proxy_cache_buffer_test.go | 33 + .../proxy/proxy_cache_crossingress_test.go | 122 +++ .../ingress/proxy/proxy_cache_modela.go | 215 ++-- .../proxy/proxy_cache_modela_helpers.go | 29 +- .../proxy/proxy_cache_modela_substrate.go | 194 ++++ .../ingress/proxy/proxy_cache_modela_test.go | 251 ++++- .../ingress/proxy/proxy_canonical_redact.go | 37 +- .../ingress/proxy/proxy_egress_shape_test.go | 59 +- .../ingress/proxy/proxy_fakeexec_test.go | 5 +- .../ingress/proxy/proxy_helpers_test.go | 10 +- .../internal/ingress/proxy/proxy_hooks.go | 62 ++ .../ingress/proxy/proxy_modify_test.go | 145 +++ .../ingress/proxy/proxy_residuals_test.go | 6 +- .../internal/ingress/proxy/proxy_responses.go | 39 +- .../internal/ingress/proxy/proxy_upstream.go | 18 +- .../proxy/response_hook_accumulator_test.go | 106 ++ .../proxy/responses_egress_relay_test.go | 123 +++ .../ingress/proxy/stage_cache_body.go | 2 +- .../internal/ingress/proxy/stage_context.go | 6 + .../internal/ingress/proxy/stage_hooks.go | 17 +- .../ingress/proxy/stage_hooks_test.go | 46 + .../internal/ingress/proxy/stage_routing.go | 5 +- .../internal/ingress/proxy/stream_context.go | 6 + .../ingress/proxy/stream_pipeline_test.go | 2 +- .../internal/ingress/proxy/stream_relay.go | 24 +- .../internal/ingress/proxy/stream_shape.go | 21 +- .../internal/platform/audit/enums.go | 7 +- .../internal/platform/audit/message.go | 2 + .../platform/audit/phase_g_batch_test.go | 153 +++ .../internal/platform/audit/record.go | 13 +- .../platform/audit/record_message_helpers.go | 31 +- .../audit/record_message_helpers_test.go | 37 + .../internal/platform/audit/writer_publish.go | 47 +- .../internal/platform/store/provider.go | 10 +- .../internal/platform/store/provider_test.go | 4 +- .../internal/providers/core/types.go | 16 + .../specs/openai/responses/classify.go | 58 ++ .../specs/openai/responses/classify_test.go | 41 + .../providers/specs/openai/stream/stream.go | 17 +- .../specs/openai/stream/stream_responses.go | 25 +- .../openai/stream/stream_responses_egress.go | 179 ++++ .../stream/stream_responses_egress_test.go | 254 +++++ .../internal/providers/target/resolver.go | 22 +- .../ai-gateway/internal/routing/core/types.go | 6 + .../ai-gateway/internal/routing/resolver.go | 19 +- .../internal/audit/event_message.go | 2 + .../ai/providers/handler/helpers_test.go | 8 +- .../handler/propagation_fail_test.go | 4 +- .../ai/providers/handler/providers.go | 140 +-- .../ai/providers/handler/providers_test.go | 162 ++- .../ai/providers/providerstore/provider.go | 103 +- .../providerstore/provider_pgxmock_test.go | 132 ++- .../fleet/handler/agent/helpers_test.go | 8 +- .../handler/traffic/traffic_cov_test.go | 20 +- .../handler/traffic/traffic_handler_test.go | 16 +- .../store/trafficstore/traffic_queries.go | 14 +- .../trafficstore/trafficstore_pgxmock_test.go | 4 +- .../observability/consumer/alert_view.go | 3 +- .../observability/consumer/binwire_decode.go | 4 + .../consumer/consumer_behaviors_test.go | 4 +- .../consumer/flush_pgxmock_test.go | 40 +- .../observability/consumer/message.go | 11 +- .../observability/consumer/traffic_copy.go | 3 + .../consumer/traffic_copy_test.go | 4 +- .../observability/consumer/traffic_inserts.go | 11 +- packages/shared/audit/event_types.go | 11 +- packages/shared/policy/decision/types.go | 51 +- packages/shared/policy/decision/types_test.go | 26 + .../hooks/core/capability_profiles_test.go | 11 + packages/shared/policy/hooks/core/noop.go | 4 +- packages/shared/policy/hooks/core/onmatch.go | 9 +- packages/shared/policy/hooks/core/types.go | 37 +- .../policy/hooks/matcher/match_bound.go | 162 +++ .../policy/hooks/matcher/match_bound_test.go | 138 +++ .../shared/policy/hooks/webhook/webhook.go | 23 +- .../hooks/webhook/webhook_escalatable_test.go | 37 + .../shared/policy/pipeline/audit_emitter.go | 36 +- .../pipeline/audit_emitter_helpers_test.go | 44 +- .../shared/policy/pipeline/enforcement.go | 117 ++- .../pipeline/maxpatternbound_cache_test.go | 175 ++++ packages/shared/policy/pipeline/merge.go | 179 ++++ packages/shared/policy/pipeline/pipeline.go | 142 +-- .../pipeline/pipeline_maxpattern_test.go | 85 ++ .../pipeline/pipeline_mayredact_test.go | 110 +++ .../shared/policy/pipeline/pipeline_test.go | 217 +++++ packages/shared/policy/pipeline/policy.go | 150 ++- .../policy/pipeline/resolver_prewarm.go | 73 ++ .../shared/policy/pipeline/unionprescan.go | 54 +- .../policy/pipeline/unionprescan_test.go | 16 +- packages/shared/transport/mq/binwire.go | 9 + packages/shared/transport/mq/messages.go | 13 +- .../shared/transport/normalize/core/spans.go | 29 + .../transport/normalize/core/spans_test.go | 27 + packages/shared/transport/streaming/buffer.go | 77 +- .../shared/transport/streaming/buffer_test.go | 93 ++ .../streaming/fold_hook_results_test.go | 61 ++ .../transport/streaming/frame_redactor.go | 7 + .../streaming/frame_redactor_test.go | 114 ++- packages/shared/transport/streaming/live.go | 284 ++---- .../shared/transport/streaming/live_test.go | 140 +-- .../shared/transport/streaming/metrics.go | 32 + .../streaming/modela/coverage_warn.go | 47 + .../streaming/modela/coverage_warn_test.go | 58 ++ .../transport/streaming/modela/engine.go | 413 ++++++++ .../transport/streaming/modela/engine_test.go | 919 ++++++++++++++++++ .../transport/streaming/pipelines_test.go | 57 +- .../tlsbump/forward_handler_pins_test.go | 173 ++++ .../tlsbump/forward_request_phase.go | 67 +- packages/shared/transport/tlsbump/sse.go | 29 +- .../tlsbump/sse_buffer_redact_med_test.go | 14 +- .../transport/tlsbump/sse_frame_redactor.go | 94 +- .../tlsbump/sse_frame_redactor_test.go | 137 ++- .../shared/transport/tlsbump/sse_modela.go | 488 ++++++++++ .../transport/tlsbump/sse_modela_test.go | 527 ++++++++++ .../transport/typology/defaults_test.go | 5 + .../shared/transport/typology/endpointkind.go | 8 + .../transport/typology/endpointkind_test.go | 2 + 142 files changed, 8468 insertions(+), 1452 deletions(-) create mode 100644 packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_golden_test.go create mode 100644 packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_ingress_test.go create mode 100644 packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_openai_types.go create mode 100644 packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_substrate.go create mode 100644 packages/ai-gateway/internal/ingress/proxy/response_hook_accumulator_test.go create mode 100644 packages/ai-gateway/internal/ingress/proxy/responses_egress_relay_test.go create mode 100644 packages/ai-gateway/internal/platform/audit/record_message_helpers_test.go create mode 100644 packages/ai-gateway/internal/providers/specs/openai/responses/classify.go create mode 100644 packages/ai-gateway/internal/providers/specs/openai/responses/classify_test.go create mode 100644 packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses_egress.go create mode 100644 packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses_egress_test.go create mode 100644 packages/shared/policy/hooks/matcher/match_bound.go create mode 100644 packages/shared/policy/hooks/matcher/match_bound_test.go create mode 100644 packages/shared/policy/hooks/webhook/webhook_escalatable_test.go create mode 100644 packages/shared/policy/pipeline/maxpatternbound_cache_test.go create mode 100644 packages/shared/policy/pipeline/merge.go create mode 100644 packages/shared/policy/pipeline/pipeline_maxpattern_test.go create mode 100644 packages/shared/policy/pipeline/resolver_prewarm.go create mode 100644 packages/shared/transport/normalize/core/spans_test.go create mode 100644 packages/shared/transport/streaming/fold_hook_results_test.go create mode 100644 packages/shared/transport/streaming/modela/coverage_warn.go create mode 100644 packages/shared/transport/streaming/modela/coverage_warn_test.go create mode 100644 packages/shared/transport/streaming/modela/engine.go create mode 100644 packages/shared/transport/streaming/modela/engine_test.go create mode 100644 packages/shared/transport/tlsbump/sse_modela.go create mode 100644 packages/shared/transport/tlsbump/sse_modela_test.go diff --git a/packages/ai-gateway/cmd/ai-gateway/wiring/providers.go b/packages/ai-gateway/cmd/ai-gateway/wiring/providers.go index 04159ef6..db1a7446 100644 --- a/packages/ai-gateway/cmd/ai-gateway/wiring/providers.go +++ b/packages/ai-gateway/cmd/ai-gateway/wiring/providers.go @@ -66,12 +66,13 @@ func (p *providerStoreAdapter) GetProviderByID(ctx context.Context, providerID s extras["provider.pathPrefix"] = prov.PathPrefix } return provtarget.ProviderRow{ - ID: prov.ID, - Name: prov.Name, - AdapterType: prov.AdapterType, - BaseURL: prov.BaseURL, - Extras: extras, - Disabled: !prov.Enabled, + ID: prov.ID, + Name: prov.Name, + AdapterType: prov.AdapterType, + BaseURL: prov.BaseURL, + Extras: extras, + Disabled: !prov.Enabled, + ServesResponsesAPI: prov.ServesResponsesAPI, }, nil } diff --git a/packages/ai-gateway/cmd/ai-gateway/wiring/providers_layer_test.go b/packages/ai-gateway/cmd/ai-gateway/wiring/providers_layer_test.go index f9bc5e37..c2ed9262 100644 --- a/packages/ai-gateway/cmd/ai-gateway/wiring/providers_layer_test.go +++ b/packages/ai-gateway/cmd/ai-gateway/wiring/providers_layer_test.go @@ -17,7 +17,7 @@ import ( // for provider queries. var providerColsForLayer = []string{ "id", "name", "displayName", "adapter_type", "baseUrl", - "pathPrefix", "apiVersion", "region", "enabled", + "pathPrefix", "apiVersion", "region", "enabled", "serves_responses_api", } // modelColsForLayer mirrors the SELECT list for model queries (23 columns). @@ -95,6 +95,7 @@ func TestProviderStoreAdapter_GetProviderByID_extras(t *testing.T) { &apiVerStr, // apiVersion (*string) ®ionStr, // region (*string) true, // enabled + (*bool)(nil), // serves_responses_api ) mock.ExpectQuery(`FROM "Provider"`).WillReturnRows(provRows) @@ -140,6 +141,7 @@ func TestProviderStoreAdapter_GetProviderByID_noExtras(t *testing.T) { (*string)(nil), // apiVersion (NULL) (*string)(nil), // region (NULL) true, // enabled + (*bool)(nil), // serves_responses_api ) mock.ExpectQuery(`FROM "Provider"`).WillReturnRows(provRows) diff --git a/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_init_test.go b/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_init_test.go index 19daf4e7..1c8f8128 100644 --- a/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_init_test.go +++ b/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_init_test.go @@ -110,7 +110,7 @@ func TestInitGeminiCacheMgrSet_listGeminiProviders_withLayer(t *testing.T) { // Seed the layer with a gemini provider. provRows := pgxmock.NewRows(providerColsForLayer). - AddRow("gemini-prov-1", "gemini", nil, "gemini", "https://generativelanguage.googleapis.com", "", (*string)(nil), (*string)(nil), true) + AddRow("gemini-prov-1", "gemini", nil, "gemini", "https://generativelanguage.googleapis.com", "", (*string)(nil), (*string)(nil), true, (*bool)(nil)) mock.ExpectQuery(`FROM "Provider"`).WillReturnRows(provRows) if err := l.ReloadProviders(context.Background()); err != nil { t.Fatalf("ReloadProviders: %v", err) @@ -271,8 +271,8 @@ func newTestCacheLayerWithAnthropicProvider(t *testing.T) (*cachelayer.Layer, pg t.Helper() mock, l := newLayerWithMock(t) provRows := pgxmock.NewRows(providerColsForLayer). - AddRow("prov-anthropic", "anthropic", nil, "anthropic", "https://api.anthropic.com", "", (*string)(nil), (*string)(nil), true). - AddRow("prov-openai", "openai", nil, "openai", "https://api.openai.com", "", (*string)(nil), (*string)(nil), true) + AddRow("prov-anthropic", "anthropic", nil, "anthropic", "https://api.anthropic.com", "", (*string)(nil), (*string)(nil), true, (*bool)(nil)). + AddRow("prov-openai", "openai", nil, "openai", "https://api.openai.com", "", (*string)(nil), (*string)(nil), true, (*bool)(nil)) mock.ExpectQuery(`FROM "Provider"`).WillReturnRows(provRows) if err := l.ReloadProviders(context.Background()); err != nil { t.Fatalf("ReloadProviders: %v", err) @@ -305,7 +305,7 @@ func TestProjectCacheBlobToNormaliserConfig_withAnthropicProvider(t *testing.T) func TestProjectCacheBlobToNormaliserConfig_withBedrockProvider(t *testing.T) { mock, l := newLayerWithMock(t) provRows := pgxmock.NewRows(providerColsForLayer). - AddRow("prov-bedrock", "bedrock", nil, "bedrock", "https://bedrock.aws.com", "", (*string)(nil), (*string)(nil), true) + AddRow("prov-bedrock", "bedrock", nil, "bedrock", "https://bedrock.aws.com", "", (*string)(nil), (*string)(nil), true, (*bool)(nil)) mock.ExpectQuery(`FROM "Provider"`).WillReturnRows(provRows) if err := l.ReloadProviders(context.Background()); err != nil { t.Fatalf("ReloadProviders: %v", err) @@ -1036,7 +1036,7 @@ func TestInitIntrospectRegistry_snapshotWithPopulatedLayer(t *testing.T) { // Seed a provider. provRows := pgxmock.NewRows(providerColsForLayer).AddRow( "prov-snap-1", "openai", nil, "openai", "https://api.openai.com", - "", (*string)(nil), (*string)(nil), true, + "", (*string)(nil), (*string)(nil), true, (*bool)(nil), ) mock.ExpectQuery(`FROM "Provider"`).WillReturnRows(provRows) if err := l.ReloadProviders(context.Background()); err != nil { diff --git a/packages/ai-gateway/internal/cache/layer/helpers_test.go b/packages/ai-gateway/internal/cache/layer/helpers_test.go index cab75b8a..a68c371e 100644 --- a/packages/ai-gateway/internal/cache/layer/helpers_test.go +++ b/packages/ai-gateway/internal/cache/layer/helpers_test.go @@ -42,7 +42,7 @@ func discardLogger() *slog.Logger { var ( providerCols = []string{ "id", "name", "displayName", "adapter_type", "baseUrl", - "pathPrefix", "apiVersion", "region", "enabled", + "pathPrefix", "apiVersion", "region", "enabled", "serves_responses_api", } modelCols = []string{ "id", "code", "name", "providerId", "p_name", "p_adapter_type", diff --git a/packages/ai-gateway/internal/cache/layer/layer_test.go b/packages/ai-gateway/internal/cache/layer/layer_test.go index ed18205d..becba99e 100644 --- a/packages/ai-gateway/internal/cache/layer/layer_test.go +++ b/packages/ai-gateway/internal/cache/layer/layer_test.go @@ -103,7 +103,7 @@ func TestStart_HappyPath_PopulatesEverySnapshot(t *testing.T) { mock.ExpectQuery(`FROM "Provider"`). WillReturnRows(pgxmock.NewRows(providerCols). AddRow("p1", "openai", strPtr("OpenAI"), "openai", - "https://api.openai.com", "/v1", strPtr("2024-01"), strPtr("us-east-1"), true)) + "https://api.openai.com", "/v1", strPtr("2024-01"), strPtr("us-east-1"), true, (*bool)(nil))) mock.ExpectQuery(`FROM "Model" m`). WillReturnRows(pgxmock.NewRows(modelCols). AddRow(makeModelRow("m1", "gpt-4o", "p1", true)...)) @@ -165,7 +165,7 @@ func TestReloadSnapshots_RoundTripsAndFiresMetricsHook(t *testing.T) { mock.ExpectQuery(`FROM "Provider"`). WillReturnRows(pgxmock.NewRows(providerCols). - AddRow("p1", "openai", nil, "openai", "https://x", "/v1", nil, nil, true)) + AddRow("p1", "openai", nil, "openai", "https://x", "/v1", nil, nil, true, (*bool)(nil))) if err := l.ReloadProviders(context.Background()); err != nil { t.Fatalf("ReloadProviders: %v", err) } @@ -297,8 +297,8 @@ func TestStats_ReportsLiveValues(t *testing.T) { // Two providers, one model, one cred. mock.ExpectQuery(`FROM "Provider"`). WillReturnRows(pgxmock.NewRows(providerCols). - AddRow("p1", "openai", nil, "openai", "https://x", "/v1", nil, nil, true). - AddRow("p2", "anthropic", nil, "anthropic", "https://y", "/v1", nil, nil, true)) + AddRow("p1", "openai", nil, "openai", "https://x", "/v1", nil, nil, true, (*bool)(nil)). + AddRow("p2", "anthropic", nil, "anthropic", "https://y", "/v1", nil, nil, true, (*bool)(nil))) mock.ExpectQuery(`FROM "Model" m`). WillReturnRows(pgxmock.NewRows(modelCols). AddRow(makeModelRow("m1", "gpt-4o", "p1", true)...)) diff --git a/packages/ai-gateway/internal/cache/layer/loaders.go b/packages/ai-gateway/internal/cache/layer/loaders.go index c3e0439c..14533496 100644 --- a/packages/ai-gateway/internal/cache/layer/loaders.go +++ b/packages/ai-gateway/internal/cache/layer/loaders.go @@ -12,7 +12,7 @@ import ( // loadProviders reads every Provider row. func (l *Layer) loadProviders(ctx context.Context) (map[string]store.Provider, error) { rows, err := l.pool.Query(ctx, ` - SELECT id, name, "displayName", adapter_type, "baseUrl", "pathPrefix", "apiVersion", region, enabled + SELECT id, name, "displayName", adapter_type, "baseUrl", "pathPrefix", "apiVersion", region, enabled, serves_responses_api FROM "Provider" `) if err != nil { @@ -24,7 +24,7 @@ func (l *Layer) loadProviders(ctx context.Context) (map[string]store.Provider, e for rows.Next() { var p store.Provider if err := rows.Scan(&p.ID, &p.Name, &p.DisplayName, &p.AdapterType, &p.BaseURL, - &p.PathPrefix, &p.APIVersion, &p.Region, &p.Enabled); err != nil { + &p.PathPrefix, &p.APIVersion, &p.Region, &p.Enabled, &p.ServesResponsesAPI); err != nil { return nil, fmt.Errorf("cachelayer: scan provider: %w", err) } out[p.ID] = p diff --git a/packages/ai-gateway/internal/cache/layer/lookups_test.go b/packages/ai-gateway/internal/cache/layer/lookups_test.go index 183fde41..39043aab 100644 --- a/packages/ai-gateway/internal/cache/layer/lookups_test.go +++ b/packages/ai-gateway/internal/cache/layer/lookups_test.go @@ -19,9 +19,9 @@ func primeSnapshots(t *testing.T, mock pgxmock.PgxPoolIface, l *Layer) { mock.ExpectQuery(`FROM "Provider"`). WillReturnRows(pgxmock.NewRows(providerCols). AddRow("p1", "openai", strPtr("OpenAI"), "openai", - "https://api.openai.com", "/v1", nil, nil, true). + "https://api.openai.com", "/v1", nil, nil, true, (*bool)(nil)). AddRow("p2", "anthropic", strPtr("Anthropic"), "anthropic", - "https://api.anthropic.com", "/v1", nil, nil, true)) + "https://api.anthropic.com", "/v1", nil, nil, true, (*bool)(nil))) mock.ExpectQuery(`FROM "Model" m`). WillReturnRows(pgxmock.NewRows(modelCols). AddRow(makeModelRow("m1", "gpt-4o", "p1", true)...). diff --git a/packages/ai-gateway/internal/cache/layer/metrics_test.go b/packages/ai-gateway/internal/cache/layer/metrics_test.go index e0809d2b..5670559b 100644 --- a/packages/ai-gateway/internal/cache/layer/metrics_test.go +++ b/packages/ai-gateway/internal/cache/layer/metrics_test.go @@ -120,7 +120,7 @@ func TestMetrics_AllHooksInvokeUnderlyingInstruments(t *testing.T) { mock.MatchExpectationsInOrder(false) mock.ExpectQuery(`FROM "Provider"`). WillReturnRows(pgxmock.NewRows(providerCols). - AddRow("p1", "openai", nil, "openai", "https://x", "/v1", nil, nil, true)) + AddRow("p1", "openai", nil, "openai", "https://x", "/v1", nil, nil, true, (*bool)(nil))) if err := l.ReloadProviders(context.Background()); err != nil { t.Fatalf("ReloadProviders: %v", err) } diff --git a/packages/ai-gateway/internal/execution/canonicalbridge/api.go b/packages/ai-gateway/internal/execution/canonicalbridge/api.go index 509178d5..dc6fd1fb 100644 --- a/packages/ai-gateway/internal/execution/canonicalbridge/api.go +++ b/packages/ai-gateway/internal/execution/canonicalbridge/api.go @@ -22,9 +22,10 @@ type API interface { // embeddings or model listing keep the legacy OpenAI-only // translation rule. EndpointRoutable(ep typology.WireShape, ingress, target provcore.Format) bool - // TargetNativelyServesResponsesAPI reports whether a target - // provider wire format natively serves /v1/responses. - TargetNativelyServesResponsesAPI(target provcore.Format) bool + // ServesResponses reports whether the resolved target serves the + // OpenAI /v1/responses wire end-to-end. override is the per-provider + // downgrade-only signal (nil = adapter RequestShapes default). + ServesResponses(target provcore.Format, override *bool) bool // IngressChatToCanonical converts the client ingress JSON to // canonical OpenAI chat.completions request JSON. IngressChatToCanonical(ingress provcore.Format, body []byte, ct provcore.CallTarget) ([]byte, error) diff --git a/packages/ai-gateway/internal/execution/canonicalbridge/bridge.go b/packages/ai-gateway/internal/execution/canonicalbridge/bridge.go index faf3f8dc..ef20cfc3 100644 --- a/packages/ai-gateway/internal/execution/canonicalbridge/bridge.go +++ b/packages/ai-gateway/internal/execution/canonicalbridge/bridge.go @@ -15,19 +15,20 @@ import ( "github.com/tidwall/sjson" ) -// formatsNativelyServingResponsesAPI is the set of provider wire formats whose -// adapter declares AdapterSpec.RequestShapes ⊇ "responses-api". The bridge -// consults this set on the /v1/responses ingress path — when the target's -// Format is in this set, the body is forwarded verbatim and no -// Responses↔canonical codec runs. Kept in lockstep with each adapter's -// spec.go RequestShapes declaration (canonical truth: -// providers/spec_openai/spec.go). +// formatDefaultServesResponses reports the adapter-RequestShapes default for +// whether a target Format natively serves the OpenAI /v1/responses wire +// (RequestShapes ⊇ WireShapeOpenAIResponses). Today only FormatOpenAI's +// adapter declares responses-api; every other adapter serves chat-completions +// only, so /v1/responses ingress to them goes through the canonical(chat) +// codec. The lockstep test pins this against the openai adapter's RequestShapes +// so the default cannot drift from the adapter that actually talks the wire. // // Adding a sibling adapter: per provider-adapter-architecture.md §3a Rule 7 // (binding), confirm via a captured 200 from that provider's real /v1/responses -// endpoint before adding to this set AND the adapter's RequestShapes. -var formatsNativelyServingResponsesAPI = map[provcore.Format]bool{ - provcore.FormatOpenAI: true, +// endpoint before declaring responses-api in its RequestShapes AND extending +// this default. +func formatDefaultServesResponses(target provcore.Format) bool { + return target == provcore.FormatOpenAI } // Bridge performs ingress ↔ canonical ↔ target wire conversions for chat. @@ -100,10 +101,24 @@ func (b *Bridge) ChatRoutable(ingress, target provcore.Format) bool { return ingressSupportsHubCanonicalChat(ingress) && formatSupportsChat(target) } -// TargetNativelyServesResponsesAPI reports whether a target provider wire format -// natively serves /v1/responses (adapter declares RequestShapes ⊇ "responses-api"). -func (b *Bridge) TargetNativelyServesResponsesAPI(target provcore.Format) bool { - return formatsNativelyServingResponsesAPI[target] +// ServesResponses reports whether the resolved target serves the OpenAI +// /v1/responses wire end-to-end — the single signal that decides, on a +// /v1/responses ingress, whether the request body is sent Responses-shape +// (true) or canonicalized to chat-completions (false). +// +// Resolution order (downgrade-only override): +// - override == nil → adapter RequestShapes default (FormatOpenAI → true). +// - override == false → false, always (a chat-only OpenAI-compatible endpoint +// opts out so the gateway never POSTs a Responses body it would 404). +// - override == true → the adapter default; a true override cannot grant a +// capability the adapter lacks (it is ignored unless RequestShapes already +// includes WireShapeOpenAIResponses). +func (b *Bridge) ServesResponses(target provcore.Format, override *bool) bool { + base := formatDefaultServesResponses(target) + if override != nil && !*override { + return false + } + return base } // ResponsesRoutable reports whether /v1/responses ingress traffic @@ -120,7 +135,7 @@ func (b *Bridge) TargetNativelyServesResponsesAPI(target provcore.Format) bool { // Bedrock uses AWS binary event-stream framing on streams (no SSE, // no chat-completions wire shape on responses). func (b *Bridge) ResponsesRoutable(target provcore.Format) bool { - if formatsNativelyServingResponsesAPI[target] { + if formatDefaultServesResponses(target) { return true } // Cross-format path requires the target to have a chat-completions @@ -269,32 +284,31 @@ func (b *Bridge) NewStreamTranscoder(ingress, target provcore.Format, model stri if openAILike(ingress) && openAILike(target) { return nil } - // /v1/responses ingress + native responses-api target (today only - // spec_openai): the upstream's Responses SSE bytes flow through - // unchanged. No transcoder needed. - if ingress == provcore.FormatOpenAIResponses && formatsNativelyServingResponsesAPI[target] { + // Unknown ingress (not OpenAI-like, not a known native streaming ingress): + // no transcoder — passthrough, and let the executor surface any wire-format + // mismatch as an upstream error. Every recognised ingress re-encodes canonical + // chunks into its native wire via the single source of truth. For /v1/responses + // this returns the Responses encoder (the decode session forwards genuine + // Responses frames verbatim via provcore.Chunk.Verbatim; anything it could not + // classify as Responses is re-encoded here) — keeping a chat-shaped reply from + // ever leaking to a non-chat client. + if !openAILike(ingress) && !isNativeStreamIngress(ingress) { return nil } - switch { - case ingress == provcore.FormatOpenAIResponses: - // Cross-format: target is chat-completions wire. The reverse encoder - // re-shapes canonical chunks into Responses SSE event grammar before - // forwarding to the /v1/responses client. - return newResponsesStreamEncoder(model) - case openAILike(ingress): - return newOpenAIStreamEncoder(model) - case ingress == provcore.FormatAnthropic: - return newAnthropicStreamEncoder() - case ingress == provcore.FormatGemini, ingress == provcore.FormatVertex: - return &geminiStreamEncoder{} - case ingress == provcore.FormatCohere: - return &cohereStreamEncoder{} - case ingress == provcore.FormatReplicate: - return &replicateStreamEncoder{} + return IngressStreamEncoder(ingress, model) +} + +// isNativeStreamIngress reports whether ingress is a non-OpenAI ingress with a +// dedicated stream encoder (the formats IngressStreamEncoder re-encodes beyond +// the OpenAI-family chat default). +func isNativeStreamIngress(ingress provcore.Format) bool { + switch ingress { + case provcore.FormatOpenAIResponses, provcore.FormatAnthropic, + provcore.FormatGemini, provcore.FormatVertex, + provcore.FormatCohere, provcore.FormatReplicate: + return true default: - // Unknown ingress: fall back to passthrough and let the executor - // surface any wire-format mismatch as an upstream error. - return nil + return false } } @@ -368,12 +382,13 @@ func (b *Bridge) IngressChatToWire(ingress, target provcore.Format, body []byte, if ingress == target { return body, nil } - // /v1/responses ingress + a target whose adapter natively serves - // responses-api (today: spec_openai): forward the body verbatim. - // This is the capability-driven same-shape passthrough — adding a - // new sibling to formatsNativelyServingResponsesAPI activates the - // fast path for them without further code changes here. - if ingress == provcore.FormatOpenAIResponses && formatsNativelyServingResponsesAPI[target] { + // /v1/responses ingress + a target that serves the Responses wire + // (per-provider capability, downgrade-only override on ct): forward the + // body verbatim so built-in tools / stateful fields survive. Otherwise + // fall through to canonical(chat). The executor resolves the same + // predicate before deciding whether to call this at all, so the two sites + // agree on the wire shape sent upstream. + if ingress == provcore.FormatOpenAIResponses && b.ServesResponses(target, ct.ServesResponsesAPI) { return body, nil } canon, err := b.IngressChatToCanonical(ingress, body, ct) diff --git a/packages/ai-gateway/internal/execution/canonicalbridge/bridge_responses_test.go b/packages/ai-gateway/internal/execution/canonicalbridge/bridge_responses_test.go index 1fa8ee3f..694af4f5 100644 --- a/packages/ai-gateway/internal/execution/canonicalbridge/bridge_responses_test.go +++ b/packages/ai-gateway/internal/execution/canonicalbridge/bridge_responses_test.go @@ -4,6 +4,7 @@ import ( "testing" provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/specs/openai" "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/typology" "github.com/tidwall/gjson" ) @@ -150,13 +151,55 @@ func TestBridge_ResponseCanonicalToIngress_Responses(t *testing.T) { } } +// TestBridge_IngressChatToWire_ResponsesCapability proves the request-side wire +// shape follows the resolved capability: a /v1/responses body to an OpenAI +// target is forwarded verbatim by default, but a downgrade override (false) +// forces canonicalization to chat-completions so the gateway never POSTs a +// Responses body to an endpoint that only serves /v1/chat/completions. +func TestBridge_IngressChatToWire_ResponsesCapability(t *testing.T) { + b := testBridge(t) + body := []byte(`{"model":"gpt-4o","input":"hi"}`) + + // Default (override nil) → native passthrough: body forwarded verbatim. + ct := dummyCallTarget(provcore.FormatOpenAI) + out, err := b.IngressChatToWire(provcore.FormatOpenAIResponses, provcore.FormatOpenAI, body, ct, false) + if err != nil { + t.Fatalf("IngressChatToWire (default): %v", err) + } + if string(out) != string(body) { + t.Fatalf("default capability must forward the Responses body verbatim; got %s", out) + } + + // override=false → canonicalize to chat-completions (messages[], no verbatim input). + no := false + ctOff := dummyCallTarget(provcore.FormatOpenAI) + ctOff.ServesResponsesAPI = &no + out, err = b.IngressChatToWire(provcore.FormatOpenAIResponses, provcore.FormatOpenAI, body, ctOff, false) + if err != nil { + t.Fatalf("IngressChatToWire (override=false): %v", err) + } + if string(out) == string(body) { + t.Fatal("override=false must canonicalize the Responses body, not forward it verbatim") + } + if !gjson.GetBytes(out, "messages").Exists() { + t.Fatalf("override=false must produce a chat-completions body with messages[]; got %s", out) + } +} + // TestBridge_NewStreamTranscoder_Responses pins the per-direction // transcoder choice for Responses ingress. func TestBridge_NewStreamTranscoder_Responses(t *testing.T) { b := testBridge(t) - // Same-shape passthrough = nil transcoder. - if tr := b.NewStreamTranscoder(provcore.FormatOpenAIResponses, provcore.FormatOpenAI, "gpt-5.2"); tr != nil { - t.Error("Responses → OpenAI same-shape passthrough should yield nil transcoder") + // Responses ingress always yields the Responses encoder: the egress wire + // shape is decided from the actual upstream bytes, not the target Format. + // A genuine Responses upstream is forwarded verbatim by the relay's + // Verbatim path (not via a nil transcoder); an upstream that returns + // chat.completion is decoded to canonical and re-encoded here, so a + // chat-shaped reply can never leak to a /v1/responses client. + if tr := b.NewStreamTranscoder(provcore.FormatOpenAIResponses, provcore.FormatOpenAI, "gpt-5.2"); tr == nil { + t.Error("Responses → OpenAI must yield the Responses encoder, not a nil passthrough") + } else if _, ok := tr.(*responsesStreamEncoder); !ok { + t.Errorf("Responses → OpenAI transcoder should be *responsesStreamEncoder, got %T", tr) } // Cross-format target = responsesStreamEncoder (re-encode canonical → Responses SSE). tr := b.NewStreamTranscoder(provcore.FormatOpenAIResponses, provcore.FormatAnthropic, "claude-sonnet-4-6") @@ -298,19 +341,26 @@ func TestBridge_ResponseAcrossFormats_UnknownFromFormat(t *testing.T) { } } -// TestBridge_LockstepCheck pins the formatsNativelyServingResponsesAPI -// lockstep with each adapter's RequestShapes declaration. If a sibling -// adapter starts declaring responses-api support but the bridge's -// lockstep map is not updated, this test surfaces the drift. +// TestBridge_LockstepCheck pins the formatDefaultServesResponses default +// against the openai adapter's actual RequestShapes declaration, so the +// bridge default cannot drift from the adapter that talks the wire. If a +// sibling adapter starts declaring responses-api support, its Format must be +// added to formatDefaultServesResponses (and this test extended). func TestBridge_LockstepCheck(t *testing.T) { - // Today only spec_openai is in the lockstep map; no other adapter - // declares responses-api. Verify both sides. - got := len(formatsNativelyServingResponsesAPI) - want := 1 - if got != want { - t.Errorf("formatsNativelyServingResponsesAPI has %d entries, want %d. If you added an adapter declaring responses-api in its RequestShapes, also add its Format here.", got, want) - } - if !formatsNativelyServingResponsesAPI[provcore.FormatOpenAI] { - t.Error("FormatOpenAI must be in formatsNativelyServingResponsesAPI") + // Both sides of the lockstep for the one format that serves Responses. + if !openai.NewSpec(nil).SupportsShape(typology.WireShapeOpenAIResponses) { + t.Fatal("spec_openai must declare WireShapeOpenAIResponses in RequestShapes") + } + if !formatDefaultServesResponses(provcore.FormatOpenAI) { + t.Error("formatDefaultServesResponses(FormatOpenAI) must be true") + } + // No other registered format may default to serving Responses. + for _, f := range provcore.AllFormats() { + if f == provcore.FormatOpenAI { + continue + } + if formatDefaultServesResponses(f) { + t.Errorf("format %q must NOT default to serving /v1/responses (no adapter declares it)", f) + } } } diff --git a/packages/ai-gateway/internal/execution/canonicalbridge/matrix_test.go b/packages/ai-gateway/internal/execution/canonicalbridge/matrix_test.go index ad8f584c..7db1a154 100644 --- a/packages/ai-gateway/internal/execution/canonicalbridge/matrix_test.go +++ b/packages/ai-gateway/internal/execution/canonicalbridge/matrix_test.go @@ -341,21 +341,28 @@ func TestSubsetFields_CoversCanonicalContract(t *testing.T) { } } -func TestTargetNativelyServesResponsesAPI(t *testing.T) { - // The native-passthrough set governs which TARGET formats can - // receive /v1/responses traffic verbatim (no Responses↔canonical - // codec). The set is intentionally narrow — only providers whose - // /v1/responses endpoint has been verified with a captured 200. - // Drift here is the failure mode the cross-format guard - // catches; this test pins the current sanctioned set. +func TestServesResponses(t *testing.T) { + // ServesResponses governs whether a /v1/responses request is sent + // upstream Responses-shape (true) or canonicalized to chat (false). + // The default is the adapter RequestShapes default; the per-provider + // override is downgrade-only. b := testBridge(t) - if !b.TargetNativelyServesResponsesAPI(provcore.FormatOpenAI) { - t.Errorf("FormatOpenAI must natively serve /v1/responses (lockstep with spec_openai RequestShapes)") - } - // Bedrock must NOT be in the native-passthrough set (uses AWS - // event-stream framing — no SSE on streams). - if b.TargetNativelyServesResponsesAPI(provcore.FormatBedrock) { - t.Errorf("FormatBedrock must NOT serve /v1/responses natively") + if !b.ServesResponses(provcore.FormatOpenAI, nil) { + t.Errorf("FormatOpenAI must serve /v1/responses by default (lockstep with spec_openai RequestShapes)") + } + // Bedrock has no Responses codec (AWS event-stream framing). + if b.ServesResponses(provcore.FormatBedrock, nil) { + t.Errorf("FormatBedrock must NOT serve /v1/responses") + } + // Override is downgrade-only: false wins for an OpenAI target. + no := false + if b.ServesResponses(provcore.FormatOpenAI, &no) { + t.Errorf("override=false must force canonical(chat) even for FormatOpenAI") + } + // Override=true cannot grant a capability the adapter lacks. + yes := true + if b.ServesResponses(provcore.FormatAnthropic, &yes) { + t.Errorf("override=true must NOT make a non-Responses adapter serve /v1/responses") } } diff --git a/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders.go b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders.go index 79a61abe..7c00a6aa 100644 --- a/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders.go +++ b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders.go @@ -26,6 +26,12 @@ type openAIStreamEncoder struct { created int64 model string headerSent bool + // scratch is reused across Write calls; each Write truncates it to zero + // length before appending its frames. The slice returned by Write aliases + // scratch and may be overwritten by the next Write — callers MUST NOT retain + // it (the io.Writer "must not retain p" contract; every call site writes it + // out synchronously before the next Write). + scratch []byte } // NewChatCompletionsStreamEncoder returns an encoder that converts canonical @@ -49,6 +55,36 @@ func NewResponsesStreamEncoder(model string) StreamTranscoder { return newResponsesStreamEncoder(model) } +// IngressStreamEncoder returns the encoder that re-encodes canonical +// provider.Chunk values into the CALLER's ingress-native SSE wire shape. It is +// the single source of truth for "given canonical chunks, which wire does this +// ingress speak" — consumed both by [Bridge.NewStreamTranscoder] (the +// cross-format live path) and by the buffer / Model-A re-emit path +// (proxy.fallbackStreamEncoder). Keeping one switch eliminates the drift class +// where a second, partial copy defaulted every non-OpenAI ingress to the +// chat-completions encoder — which leaked chat.completion.chunk frames to a +// Gemini / Anthropic / Responses client whenever the enforcing buffer path ran. +// +// It NEVER returns nil: the buffer/re-emit caller must always have a concrete +// encoder. OpenAI-family ingresses (and any unrecognised ingress) get the +// chat-completions encoder, since canonical IS the chat-completions shape. +func IngressStreamEncoder(ingress provcore.Format, model string) StreamTranscoder { + switch ingress { + case provcore.FormatOpenAIResponses: + return newResponsesStreamEncoder(model) + case provcore.FormatAnthropic: + return newAnthropicStreamEncoder() + case provcore.FormatGemini, provcore.FormatVertex: + return &geminiStreamEncoder{} + case provcore.FormatCohere: + return &cohereStreamEncoder{} + case provcore.FormatReplicate: + return &replicateStreamEncoder{} + default: + return newOpenAIStreamEncoder(model) + } +} + func newOpenAIStreamEncoder(model string) *openAIStreamEncoder { b := make([]byte, 12) _, _ = rand.Read(b) @@ -59,110 +95,82 @@ func newOpenAIStreamEncoder(model string) *openAIStreamEncoder { } } -// oaiEnvelope wraps a choices slice in the full OpenAI chat.completion.chunk envelope. -func (e *openAIStreamEncoder) oaiEnvelope(choices []any) []byte { - return oaiDeltaSSE(map[string]any{ - "id": e.id, - "object": "chat.completion.chunk", - "created": e.created, - "model": e.model, - "choices": choices, +// emit marshals one envelope (single choice + optional usage) and appends the +// SSE frame to the reused scratch buffer. The payload is ALWAYS produced by +// json.Marshal (preserving go-json's HTML-safe escape); only the `data: ` / +// `\n\n` framing is hand-assembled — no user/upstream data is hand-serialised. +func (e *openAIStreamEncoder) emit(choice oaiStreamChoice, usage *oaiStreamUsage) { + data, _ := json.Marshal(oaiStreamEnvelope{ + Choices: []oaiStreamChoice{choice}, + Created: e.created, + ID: e.id, + Model: e.model, + Object: "chat.completion.chunk", + Usage: usage, }) + e.scratch = append(e.scratch, "data: "...) + e.scratch = append(e.scratch, data...) + e.scratch = append(e.scratch, '\n', '\n') } func (e *openAIStreamEncoder) Write(_ context.Context, chunk provcore.Chunk) ([]byte, error) { - var buf bytes.Buffer + e.scratch = e.scratch[:0] // Emit role-assignment chunk before any content. if !e.headerSent { e.headerSent = true - buf.Write(e.oaiEnvelope([]any{ - map[string]any{ - "index": 0, - "delta": map[string]any{"role": "assistant", "content": ""}, - "finish_reason": nil, - }, - })) + empty := "" + e.emit(oaiStreamChoice{Delta: oaiStreamDelta{Content: &empty, Role: "assistant"}}, nil) } // Check content before Done: providers like Gemini 2.5 combine text, // finishReason, and usageMetadata into a single SSE frame, so chunk.Delta // can be non-empty even when chunk.Done is also true. if chunk.Delta != "" { - buf.Write(e.oaiEnvelope([]any{ - map[string]any{ - "index": 0, - "delta": map[string]any{"content": chunk.Delta}, - "finish_reason": nil, - }, - })) + d := chunk.Delta + e.emit(oaiStreamChoice{Delta: oaiStreamDelta{Content: &d}}, nil) } if len(chunk.ToolCallDeltas) > 0 { - buf.Write(e.oaiToolCallEnvelope(chunk.ToolCallDeltas)) + e.emit(oaiStreamChoice{Delta: oaiStreamDelta{ToolCalls: buildOAIToolCalls(chunk.ToolCallDeltas)}}, nil) } if chunk.ReasoningDelta != "" { - buf.Write(e.oaiEnvelope([]any{ - map[string]any{ - "index": 0, - "delta": map[string]any{"reasoning_content": chunk.ReasoningDelta}, - "finish_reason": nil, - }, - })) + e.emit(oaiStreamChoice{Delta: oaiStreamDelta{ReasoningContent: chunk.ReasoningDelta}}, nil) } if chunk.Done { - // Emit finish_reason=stop chunk before [DONE]. - finishChunk := map[string]any{ - "id": e.id, - "object": "chat.completion.chunk", - "created": e.created, - "model": e.model, - "choices": []any{ - map[string]any{ - "index": 0, - "delta": map[string]any{}, - "finish_reason": finishReasonOrStop(chunk.FinishReason), - }, - }, - } - if chunk.Usage != nil { - u := map[string]any{} - if chunk.Usage.PromptTokens != nil { - u["prompt_tokens"] = *chunk.Usage.PromptTokens - } - if chunk.Usage.CompletionTokens != nil { - u["completion_tokens"] = *chunk.Usage.CompletionTokens - } - if chunk.Usage.TotalTokens != nil { - u["total_tokens"] = *chunk.Usage.TotalTokens - } - // Project the same detail sub-blocks the non-stream projector - // emits (openai_projection.go projectUsage), so stream clients - // can read cache-read + reasoning splits the same way as on a - // non-stream response. Without these, gemini-2.5-pro / o1 / - // kimi-k2.6 stream responses look reasoning_tokens=0 even when - // thoughtsTokenCount / - // completion_tokens_details.reasoning_tokens / etc. were - // populated on the canonical chunk. - if chunk.Usage.CacheReadTokens != nil && *chunk.Usage.CacheReadTokens > 0 { - u["prompt_tokens_details"] = map[string]any{ - "cached_tokens": *chunk.Usage.CacheReadTokens, - } - } - if chunk.Usage.ReasoningTokens != nil && *chunk.Usage.ReasoningTokens > 0 { - u["completion_tokens_details"] = map[string]any{ - "reasoning_tokens": *chunk.Usage.ReasoningTokens, - } - } - finishChunk["usage"] = u - } - buf.Write(oaiDeltaSSE(finishChunk)) - // [DONE] appended by LivePipeline.EmitOpenAIDone — return nil after. + // Emit finish_reason chunk before [DONE] (appended by + // LivePipeline.EmitOpenAIDone). Detail sub-blocks (cache-read + + // reasoning) mirror the non-stream projector so stream clients read the + // same splits. + fr := finishReasonOrStop(chunk.FinishReason) + e.emit(oaiStreamChoice{Delta: oaiStreamDelta{}, FinishReason: &fr}, buildOAIStreamUsage(chunk.Usage)) } - if buf.Len() == 0 { + if len(e.scratch) == 0 { return nil, nil } - return buf.Bytes(), nil + return e.scratch, nil +} + +// buildOAIStreamUsage maps a canonical usage into the wire usage block: a token +// field is set whenever its source pointer is non-nil (so a non-nil 0 still +// renders), and a detail sub-block appears only when its source is non-nil AND +// > 0. Returns nil when no usage was reported (the usage key is omitted). +func buildOAIStreamUsage(u *provcore.Usage) *oaiStreamUsage { + if u == nil { + return nil + } + out := &oaiStreamUsage{ + PromptTokens: u.PromptTokens, + CompletionTokens: u.CompletionTokens, + TotalTokens: u.TotalTokens, + } + if u.CacheReadTokens != nil && *u.CacheReadTokens > 0 { + out.PromptTokensDetails = &oaiPromptTokensDetails{CachedTokens: *u.CacheReadTokens} + } + if u.ReasoningTokens != nil && *u.ReasoningTokens > 0 { + out.CompletionTokensDetails = &oaiCompletionTokensDetails{ReasoningTokens: *u.ReasoningTokens} + } + return out } // anthropicStreamEncoder converts canonical chunks to the Anthropic Messages @@ -580,11 +588,6 @@ func buildGeminiUsage(u *provcore.Usage) map[string]any { return out } -func oaiDeltaSSE(payload map[string]any) []byte { - data, _ := json.Marshal(payload) - return fmt.Appendf(nil, "data: %s\n\n", data) -} - // finishReasonOrStop returns the canonical OpenAI finish_reason, defaulting to // "stop" when the upstream stream never reported one. An empty FinishReason on // the terminal chunk is the case for the live cross-format path (whose Done @@ -645,39 +648,13 @@ func canonicalFinishToCohere(fr string) string { } } -func (e *openAIStreamEncoder) oaiToolCallEnvelope(deltas []provcore.ToolCallDelta) []byte { - return oaiDeltaSSE(map[string]any{ - "id": e.id, - "object": "chat.completion.chunk", - "created": e.created, - "model": e.model, - "choices": []any{ - map[string]any{ - "index": 0, - "delta": map[string]any{"tool_calls": buildToolCallDeltas(deltas)}, - "finish_reason": nil, - }, - }, - }) -} - -func buildToolCallDeltas(deltas []provcore.ToolCallDelta) []any { - type fnCall struct { - Name string `json:"name,omitempty"` - Arguments string `json:"arguments,omitempty"` - } - type tcDelta struct { - Index int `json:"index"` - ID string `json:"id,omitempty"` - Type string `json:"type,omitempty"` - Function fnCall `json:"function"` - } - calls := make([]any, 0, len(deltas)) +// buildOAIToolCalls converts canonical tool-call deltas into the OpenAI +// streaming tool_calls shape. Type=function is set only when an id starts a new +// call; continuation deltas carry just index + arguments. +func buildOAIToolCalls(deltas []provcore.ToolCallDelta) []oaiToolCall { + calls := make([]oaiToolCall, 0, len(deltas)) for _, d := range deltas { - tc := tcDelta{ - Index: d.Index, - Function: fnCall{Name: d.Name, Arguments: d.Arguments}, - } + tc := oaiToolCall{Index: d.Index, Function: oaiToolFunc{Name: d.Name, Arguments: d.Arguments}} if d.ID != "" { tc.ID = d.ID tc.Type = "function" diff --git a/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_golden_test.go b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_golden_test.go new file mode 100644 index 00000000..86137e46 --- /dev/null +++ b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_golden_test.go @@ -0,0 +1,177 @@ +package canonicalbridge + +import ( + "bytes" + "context" + "strings" + "testing" + + provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" +) + +// goldenEnc returns an openAIStreamEncoder with deterministic id/created so the +// golden byte output is stable (the production constructor seeds id from +// crypto/rand and created from time.Now). These golden frames PIN the exact wire +// bytes the encoder emits; the map[string]any→struct refactor must keep them +// byte-for-byte identical (zero wire change — see plan §12.1). +func goldenEnc() *openAIStreamEncoder { + return &openAIStreamEncoder{id: "chatcmpl-GOLD", created: 1700000000, model: "test-model"} +} + +func gpi(i int) *int { return &i } + +// gFrame wraps a choice JSON in the full envelope (alphabetical key order, no usage). +func gFrame(choice string) string { + return `data: {"choices":[` + choice + `],"created":1700000000,"id":"chatcmpl-GOLD","model":"test-model","object":"chat.completion.chunk"}` + "\n\n" +} + +// gFrameU wraps a choice + usage block (usage sorts after object, alphabetically). +func gFrameU(choice, usage string) string { + return `data: {"choices":[` + choice + `],"created":1700000000,"id":"chatcmpl-GOLD","model":"test-model","object":"chat.completion.chunk","usage":` + usage + `}` + "\n\n" +} + +type goldenFrame struct { + name string + headerSent bool + chunk provcore.Chunk + want string +} + +func goldenFrames() []goldenFrame { + return []goldenFrame{ + { + name: "role_content", headerSent: false, chunk: provcore.Chunk{Delta: "hello"}, + want: gFrame(`{"delta":{"content":"","role":"assistant"},"finish_reason":null,"index":0}`) + + gFrame(`{"delta":{"content":"hello"},"finish_reason":null,"index":0}`), + }, + { + name: "content", headerSent: true, chunk: provcore.Chunk{Delta: "more"}, + want: gFrame(`{"delta":{"content":"more"},"finish_reason":null,"index":0}`), + }, + { + // escape vectors: < > & " \ \n \r \t + unicode/emoji + name: "content_escape", headerSent: true, chunk: provcore.Chunk{Delta: "a&\"c\\d\n\r\t中😀"}, + want: gFrame(`{"delta":{"content":"` + "a\\u003cb\\u003e\\u0026\\\"c\\\\d\\n\\r\\t中😀" + `"},"finish_reason":null,"index":0}`), + }, + { + // SSE frame-injection vector: \n\n must stay inside the JSON string, never split the frame + name: "content_sse_inject", headerSent: true, chunk: provcore.Chunk{Delta: "data: evil\n\ndata: {\"x\":1}"}, + want: gFrame(`{"delta":{"content":"data: evil\n\ndata: {\"x\":1}"},"finish_reason":null,"index":0}`), + }, + { + // JSON breakout vector: quotes escaped, envelope not broken + name: "content_breakout", headerSent: true, chunk: provcore.Chunk{Delta: "\",\"injected\":\"x"}, + want: gFrame(`{"delta":{"content":"\",\"injected\":\"x"},"finish_reason":null,"index":0}`), + }, + { + // XSS vector: angle brackets HTML-safe escaped + name: "content_xss", headerSent: true, chunk: provcore.Chunk{Delta: ""}, + want: gFrame(`{"delta":{"content":"` + "\\u003c/script\\u003e\\u003csvg onload=alert(1)\\u003e" + `"},"finish_reason":null,"index":0}`), + }, + { + // control chars escaped as \u00XX; build want via concatenation to avoid + // embedding real NUL bytes in source. + name: "content_ctrl", headerSent: true, chunk: provcore.Chunk{Delta: "\x00\x01\x1f"}, + want: gFrame(`{"delta":{"content":"` + "\\u0000\\u0001\\u001f" + `"},"finish_reason":null,"index":0}`), + }, + { + name: "toolcall", headerSent: true, chunk: provcore.Chunk{ToolCallDeltas: []provcore.ToolCallDelta{ + {Index: 0, ID: "call_1", Name: "get_weather", Arguments: `{"city":`}, + {Index: 1, Arguments: `"NYC"}`}, + }}, + want: gFrame(`{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":"}},{"index":1,"function":{"arguments":"\"NYC\"}"}}]},"finish_reason":null,"index":0}`), + }, + { + // redacted tool args delivered as a placeholder must survive intact + name: "toolcall_redact", headerSent: true, chunk: provcore.Chunk{ToolCallDeltas: []provcore.ToolCallDelta{ + {Index: 0, ID: "call_1", Name: "f", Arguments: "[REDACTED]"}, + }}, + want: gFrame(`{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"f","arguments":"[REDACTED]"}}]},"finish_reason":null,"index":0}`), + }, + { + name: "reasoning", headerSent: true, chunk: provcore.Chunk{ReasoningDelta: "think<>&"}, + want: gFrame(`{"delta":{"reasoning_content":"` + "think\\u003c\\u003e\\u0026" + `"},"finish_reason":null,"index":0}`), + }, + { + name: "done_no_usage", headerSent: true, chunk: provcore.Chunk{Done: true}, + want: gFrame(`{"delta":{},"finish_reason":"stop","index":0}`), + }, + { + name: "done_usage", headerSent: true, chunk: provcore.Chunk{Done: true, FinishReason: "length", Usage: &provcore.Usage{ + PromptTokens: gpi(10), CompletionTokens: gpi(5), TotalTokens: gpi(15), + }}, + want: gFrameU(`{"delta":{},"finish_reason":"length","index":0}`, `{"completion_tokens":5,"prompt_tokens":10,"total_tokens":15}`), + }, + { + name: "done_usage_details", headerSent: true, chunk: provcore.Chunk{Done: true, Usage: &provcore.Usage{ + PromptTokens: gpi(10), CompletionTokens: gpi(5), TotalTokens: gpi(15), + CacheReadTokens: gpi(4), ReasoningTokens: gpi(3), + }}, + want: gFrameU(`{"delta":{},"finish_reason":"stop","index":0}`, `{"completion_tokens":5,"completion_tokens_details":{"reasoning_tokens":3},"prompt_tokens":10,"prompt_tokens_details":{"cached_tokens":4},"total_tokens":15}`), + }, + { + name: "content_done", headerSent: true, chunk: provcore.Chunk{Delta: "tail", Done: true, FinishReason: "stop", Usage: &provcore.Usage{ + PromptTokens: gpi(1), CompletionTokens: gpi(2), TotalTokens: gpi(3), + }}, + want: gFrame(`{"delta":{"content":"tail"},"finish_reason":null,"index":0}`) + + gFrameU(`{"delta":{},"finish_reason":"stop","index":0}`, `{"completion_tokens":2,"prompt_tokens":1,"total_tokens":3}`), + }, + } +} + +// TestOAIStreamEncoderGolden pins the exact wire bytes for every frame type and +// escape vector. The map→struct refactor must keep these byte-for-byte identical. +func TestOAIStreamEncoderGolden(t *testing.T) { + ctx := context.Background() + for _, f := range goldenFrames() { + enc := goldenEnc() + enc.headerSent = f.headerSent + b, err := enc.Write(ctx, f.chunk) + if err != nil { + t.Fatalf("%s: Write: %v", f.name, err) + } + if string(b) != f.want { + t.Errorf("%s byte mismatch:\n got: %q\nwant: %q", f.name, string(b), f.want) + } + } +} + +// TestOAIStreamEncoderNoResidue proves that reusing one encoder across chunks +// (the scratch-buffer reuse introduced by the refactor) never leaks bytes from a +// prior frame into a later one. A long delta followed by a short delta must +// yield a frame whose content is exactly the short delta — no tail residue. +func TestOAIStreamEncoderNoResidue(t *testing.T) { + ctx := context.Background() + enc := goldenEnc() + enc.headerSent = true + + long, _ := enc.Write(ctx, provcore.Chunk{Delta: strings.Repeat("LONGSECRET", 50)}) + if !bytes.Contains(long, []byte("LONGSECRET")) { + t.Fatalf("first frame should contain the long delta; got %q", long) + } + short, _ := enc.Write(ctx, provcore.Chunk{Delta: "x"}) + if want := gFrame(`{"delta":{"content":"x"},"finish_reason":null,"index":0}`); string(short) != want { + t.Errorf("second frame leaked prior-frame residue:\n got: %q\nwant: %q", string(short), want) + } + if bytes.Contains(short, []byte("LONGSECRET")) { + t.Errorf("second frame must not contain any byte of the first frame; got %q", short) + } +} + +// TestOAIStreamEncoderCrossRequestIsolation proves two independent encoder +// instances (two requests) never share state — a fresh encoder's first frame is +// the role header, carrying nothing from a previously-used encoder. +func TestOAIStreamEncoderCrossRequestIsolation(t *testing.T) { + ctx := context.Background() + encA := goldenEnc() + _, _ = encA.Write(ctx, provcore.Chunk{Delta: "AAAAA-request-A-secret"}) + + encB := goldenEnc() + b, _ := encB.Write(ctx, provcore.Chunk{Delta: "B"}) + if bytes.Contains(b, []byte("request-A-secret")) { + t.Errorf("request B leaked request A bytes; got %q", b) + } + if !bytes.Contains(b, []byte(`"role":"assistant"`)) { + t.Errorf("fresh encoder must emit role header first; got %q", b) + } +} diff --git a/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_ingress_test.go b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_ingress_test.go new file mode 100644 index 00000000..86fbc6be --- /dev/null +++ b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_ingress_test.go @@ -0,0 +1,62 @@ +package canonicalbridge + +import ( + "context" + "strings" + "testing" + + provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" +) + +// TestIngressStreamEncoder_NativeWireShape is the regression test for the +// cross-format streaming egress bug: the buffer / Model-A re-emit path +// (proxy.fallbackStreamEncoder) and the cross-format live transcoder must encode +// canonical chunks into the CALLER's ingress wire shape. A Gemini / Anthropic / +// Responses client must NEVER receive chat.completion.chunk frames — which is +// exactly what leaked when a non-OpenAI ingress streamed under an enforcing +// (redact/block) response scope. +func TestIngressStreamEncoder_NativeWireShape(t *testing.T) { + cases := []struct { + name string + ingress provcore.Format + mustHave string // a marker proving the native wire shape + mustNotHave string // the wrong wire shape that must never leak + }{ + {"openai_chat", provcore.FormatOpenAI, "chat.completion.chunk", "response.output_text"}, + {"responses", provcore.FormatOpenAIResponses, "response.", "chat.completion.chunk"}, + {"anthropic", provcore.FormatAnthropic, "content_block", "chat.completion.chunk"}, + {"gemini", provcore.FormatGemini, "\"text\":\"hi\"", "chat.completion.chunk"}, + {"vertex", provcore.FormatVertex, "\"text\":\"hi\"", "chat.completion.chunk"}, + {"cohere", provcore.FormatCohere, "hi", "chat.completion.chunk"}, + {"replicate", provcore.FormatReplicate, "hi", "chat.completion.chunk"}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + enc := IngressStreamEncoder(c.ingress, "test-model") + if enc == nil { + t.Fatalf("IngressStreamEncoder(%s) returned nil — the buffer re-emit path must always have a concrete encoder", c.ingress) + } + // Drive a realistic short stream: a text delta then a terminal chunk, + // concatenating all emitted bytes (encoders are stateful — preamble on + // the first Write, terminal on the last). + var out strings.Builder + for _, chunk := range []provcore.Chunk{ + {Delta: "hi"}, + {Done: true, FinishReason: "stop"}, + } { + b, err := enc.Write(context.Background(), chunk) + if err != nil { + t.Fatalf("%s: Write: %v", c.name, err) + } + out.Write(b) + } + s := out.String() + if c.mustHave != "" && !strings.Contains(s, c.mustHave) { + t.Fatalf("%s: emitted wire missing native marker %q; got:\n%s", c.name, c.mustHave, s) + } + if c.mustNotHave != "" && strings.Contains(s, c.mustNotHave) { + t.Fatalf("%s: emitted wire LEAKED foreign shape %q to the client; got:\n%s", c.name, c.mustNotHave, s) + } + }) + } +} diff --git a/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_openai_types.go b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_openai_types.go new file mode 100644 index 00000000..d273677b --- /dev/null +++ b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_openai_types.go @@ -0,0 +1,70 @@ +package canonicalbridge + +// stream_encoders_openai_types.go — wire-shape typed structs for the OpenAI +// chat.completion.chunk stream encoder, split from stream_encoders.go to keep +// that file under the size ratchet. +// +// Field order is deliberately ALPHABETICAL to mirror go-json's map-key sorting, +// so the struct encoder emits byte-identical output to the prior map[string]any +// path (zero wire change). Field order is NOT an external contract — it is +// pinned only to keep the wire bytes stable for clients, response caches, and +// provider prefix-caches. go-json applies the same HTML-safe string escape to +// struct fields as it did to map values. + +type oaiStreamEnvelope struct { + Choices []oaiStreamChoice `json:"choices"` + Created int64 `json:"created"` + ID string `json:"id"` + Model string `json:"model"` + Object string `json:"object"` + Usage *oaiStreamUsage `json:"usage,omitempty"` +} + +type oaiStreamChoice struct { + Delta oaiStreamDelta `json:"delta"` + // FinishReason renders `null` on delta frames (nil pointer) and the value on + // the terminal frame; NO omitempty so the explicit null is preserved. + FinishReason *string `json:"finish_reason"` + Index int `json:"index"` +} + +// oaiStreamDelta covers every delta shape. Content is *string so a nil pointer +// is OMITTED (tool / reasoning / done frames) while a non-nil empty pointer +// renders the explicit `"content":""` the role-header frame requires. +type oaiStreamDelta struct { + Content *string `json:"content,omitempty"` + ReasoningContent string `json:"reasoning_content,omitempty"` + Role string `json:"role,omitempty"` + ToolCalls []oaiToolCall `json:"tool_calls,omitempty"` +} + +type oaiToolCall struct { + Index int `json:"index"` + ID string `json:"id,omitempty"` + Type string `json:"type,omitempty"` + Function oaiToolFunc `json:"function"` +} + +type oaiToolFunc struct { + Name string `json:"name,omitempty"` + Arguments string `json:"arguments,omitempty"` +} + +// oaiStreamUsage uses *int per token field so a non-nil zero still renders +// (matching the prior `if ptr != nil` map logic); detail sub-blocks are pointers +// with omitempty so they appear only when their source is set AND > 0. +type oaiStreamUsage struct { + CompletionTokens *int `json:"completion_tokens,omitempty"` + CompletionTokensDetails *oaiCompletionTokensDetails `json:"completion_tokens_details,omitempty"` + PromptTokens *int `json:"prompt_tokens,omitempty"` + PromptTokensDetails *oaiPromptTokensDetails `json:"prompt_tokens_details,omitempty"` + TotalTokens *int `json:"total_tokens,omitempty"` +} + +type oaiPromptTokensDetails struct { + CachedTokens int `json:"cached_tokens"` +} + +type oaiCompletionTokensDetails struct { + ReasoningTokens int `json:"reasoning_tokens"` +} diff --git a/packages/ai-gateway/internal/execution/executor/executor.go b/packages/ai-gateway/internal/execution/executor/executor.go index 2f7ba641..7896ef00 100644 --- a/packages/ai-gateway/internal/execution/executor/executor.go +++ b/packages/ai-gateway/internal/execution/executor/executor.go @@ -287,7 +287,7 @@ func (e *TargetExecutor) executeInner( // proxy-level needsCanonicalization=false rule and the egress // native-passthrough skip — all three sites must agree. nativeResponses := base.WireShape == typology.WireShapeOpenAIResponses && - e.bridge != nil && e.bridge.TargetNativelyServesResponsesAPI(callTarget.Format) + e.bridge != nil && e.bridge.ServesResponses(callTarget.Format, callTarget.ServesResponsesAPI) if e.bridge != nil && !nativeResponses { switch ingressKind { case typology.EndpointKindChat: diff --git a/packages/ai-gateway/internal/execution/executor/executor_internals_test.go b/packages/ai-gateway/internal/execution/executor/executor_internals_test.go index 87bc5923..a4f0b9f1 100644 --- a/packages/ai-gateway/internal/execution/executor/executor_internals_test.go +++ b/packages/ai-gateway/internal/execution/executor/executor_internals_test.go @@ -828,8 +828,9 @@ func TestExecute_ResponsesNative_KeepsResponsesWireShape(t *testing.T) { }} reg := newRegistry(t, adapter) res := okResolver() - // Real bridge: TargetNativelyServesResponsesAPI(FormatOpenAI)==true (package - // map, codec-independent), so the native-responses passthrough branch fires. + // Real bridge: ServesResponses(FormatOpenAI, nil)==true (adapter + // RequestShapes default, codec-independent), so the native-responses + // passthrough branch fires. bridge := canonicalbridge.New(map[provcore.Format]provcore.SchemaCodec{}) exec := New(reg, res, nil, bridge) diff --git a/packages/ai-gateway/internal/ingress/debug/debug_branches_test.go b/packages/ai-gateway/internal/ingress/debug/debug_branches_test.go index 19610c6b..64dedf1f 100644 --- a/packages/ai-gateway/internal/ingress/debug/debug_branches_test.go +++ b/packages/ai-gateway/internal/ingress/debug/debug_branches_test.go @@ -225,7 +225,7 @@ func (s stubBridgeForDebug) EndpointRoutable(ep typology.WireShape, ingress, pro return s.routable } -func (s stubBridgeForDebug) TargetNativelyServesResponsesAPI(target provcore.Format) bool { +func (s stubBridgeForDebug) ServesResponses(target provcore.Format, override *bool) bool { return false } diff --git a/packages/ai-gateway/internal/ingress/proxy/audit_findings_test.go b/packages/ai-gateway/internal/ingress/proxy/audit_findings_test.go index 439b6266..5a3902a1 100644 --- a/packages/ai-gateway/internal/ingress/proxy/audit_findings_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/audit_findings_test.go @@ -178,7 +178,7 @@ func TestStampUnpricedCost(t *testing.T) { func TestChunkSSEReader_TerminalError_UpstreamError(t *testing.T) { sub := &queuedChunkSub{entries: []queuedChunkEntry{{err: errors.New("boom")}}} - rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI) + rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI, false) rd.usageSink = &chunkUsageHolder{} // Drain until EOF so the terminal error is published. @@ -201,7 +201,7 @@ func TestChunkSSEReader_TerminalError_ClientAbort(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() // client gone sub := &queuedChunkSub{entries: []queuedChunkEntry{{err: context.Canceled}}} - rd := newChunkSSEReaderFromSubscription(ctx, sub, nil, provcore.FormatOpenAI) + rd := newChunkSSEReaderFromSubscription(ctx, sub, nil, provcore.FormatOpenAI, false) rd.usageSink = &chunkUsageHolder{} if _, err := rd.Read(make([]byte, 64)); err == nil { @@ -217,7 +217,7 @@ func TestChunkSSEReader_TerminalError_CleanEOF(t *testing.T) { sub := &queuedChunkSub{entries: []queuedChunkEntry{ {chunk: provcore.Chunk{Done: true, RawBytes: []byte("data: [DONE]\n\n")}}, }} - rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI) + rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI, false) rd.usageSink = &chunkUsageHolder{} buf := make([]byte, 64) @@ -244,7 +244,7 @@ func (errTranscoder) Write(_ context.Context, chunk provcore.Chunk) ([]byte, err func TestChunkSSEReader_TerminalError_TranscoderError(t *testing.T) { sub := &queuedChunkSub{entries: []queuedChunkEntry{{chunk: provcore.Chunk{Delta: "x"}}}} - rd := newChunkSSEReaderFromSubscription(context.Background(), sub, errTranscoder{}, provcore.FormatOpenAI) + rd := newChunkSSEReaderFromSubscription(context.Background(), sub, errTranscoder{}, provcore.FormatOpenAI, false) rd.usageSink = &chunkUsageHolder{} if _, err := rd.Read(make([]byte, 64)); err == nil { diff --git a/packages/ai-gateway/internal/ingress/proxy/chunk_stream_reader_test.go b/packages/ai-gateway/internal/ingress/proxy/chunk_stream_reader_test.go index ee32ccb8..d331e308 100644 --- a/packages/ai-gateway/internal/ingress/proxy/chunk_stream_reader_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/chunk_stream_reader_test.go @@ -173,7 +173,7 @@ func TestChunkUsageHolder_NilReceiverAndNilUsage(t *testing.T) { func TestChunkSSEReader_NilSubReturnsEOF(t *testing.T) { // nil sub triggers the "closed = true; return 0, io.EOF" arm. - rd := newChunkSSEReaderFromSubscription(context.Background(), nil, nil, provcore.FormatOpenAI) + rd := newChunkSSEReaderFromSubscription(context.Background(), nil, nil, provcore.FormatOpenAI, false) buf := make([]byte, 16) n, err := rd.Read(buf) if n != 0 || err == nil { @@ -241,7 +241,7 @@ func (f *queuedChunkSub) Close() error { return nil } func TestChunkSSEReader_ProviderErrorEmitsSSEFrame(t *testing.T) { pErr := errors.New("upstream disconnected") sub := &queuedChunkSub{entries: []queuedChunkEntry{{err: pErr}}} - rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI) + rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI, false) rd.usageSink = &chunkUsageHolder{} buf := make([]byte, 4096) @@ -263,7 +263,7 @@ func TestChunkSSEReader_DoneChunkWithRawBytes(t *testing.T) { sub := &queuedChunkSub{entries: []queuedChunkEntry{ {chunk: provcore.Chunk{Done: true, RawBytes: terminal}}, }} - rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI) + rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI, false) rd.usageSink = &chunkUsageHolder{} buf := make([]byte, 64) @@ -307,7 +307,7 @@ func TestChunkSSEReader_PassthroughMultiChunk_ByteLossless(t *testing.T) { // Read with a 1-byte buffer to force maximal partial reads across the sliding // window + backing-array reuse — the harshest aliasing exposure. sub := &queuedChunkSub{entries: append([]queuedChunkEntry(nil), entries...)} - rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI) + rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI, false) rd.usageSink = &chunkUsageHolder{} var got bytes.Buffer one := make([]byte, 1) @@ -388,7 +388,7 @@ func TestChunkSSEReader_DeltaSynthesisedOpenAIEnvelope(t *testing.T) { sub := &queuedChunkSub{entries: []queuedChunkEntry{ {chunk: provcore.Chunk{Delta: "hello "}}, }} - rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI) + rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI, false) rd.usageSink = &chunkUsageHolder{} buf := make([]byte, 256) @@ -408,7 +408,7 @@ func TestChunkSSEReader_DefaultArmYieldsZero(t *testing.T) { sub := &queuedChunkSub{entries: []queuedChunkEntry{ {chunk: provcore.Chunk{}}, }} - rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI) + rd := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI, false) rd.usageSink = &chunkUsageHolder{} buf := make([]byte, 16) diff --git a/packages/ai-gateway/internal/ingress/proxy/hook_extraction_test.go b/packages/ai-gateway/internal/ingress/proxy/hook_extraction_test.go index 827e6544..897a7e37 100644 --- a/packages/ai-gateway/internal/ingress/proxy/hook_extraction_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/hook_extraction_test.go @@ -485,7 +485,7 @@ func TestChunkSSEReader_PartialReadAcrossChunks(t *testing.T) { {RawBytes: []byte("data: 2\n\n")}, {Done: true, RawBytes: []byte("data: [DONE]\n\n")}, }} - r := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI) + r := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI, false) r.usageSink = &chunkUsageHolder{} // Read in 4-byte chunks to force the internal buf re-fill loop. buf := make([]byte, 4) diff --git a/packages/ai-gateway/internal/ingress/proxy/ingress.go b/packages/ai-gateway/internal/ingress/proxy/ingress.go index 10d592e7..1ffe08f7 100644 --- a/packages/ai-gateway/internal/ingress/proxy/ingress.go +++ b/packages/ai-gateway/internal/ingress/proxy/ingress.go @@ -174,11 +174,11 @@ func stickyKeyFromCtx(ctx context.Context) string { // dispatch in canonicalbridge.NewStreamTranscoder (per-Format-family // grammar selection — see api.go for the asymmetry rationale). // -// Returns (Format, ok=false) for wire shapes that are not yet routed -// through Format-keyed SSE transcoder construction (Gemini, Vertex, -// Bedrock, Cohere, Voyage). Callers MUST check ok and skip the dependent -// operation when the mapping is undefined — silently returning the empty -// Format would drive downstream dispatch to an unconfigured codec. +// Returns (Format, ok=false) for wire shapes with no Format-keyed SSE +// transcoder construction (Bedrock event-stream framing, Voyage embeddings). +// Callers MUST check ok and skip the dependent operation when the mapping is +// undefined — silently returning the empty Format would drive downstream +// dispatch to an unconfigured codec. func WireShapeToBodyFormat(w typology.WireShape) (provcore.Format, bool) { switch w { case typology.WireShapeOpenAIChat, @@ -189,6 +189,12 @@ func WireShapeToBodyFormat(w typology.WireShape) (provcore.Format, bool) { return provcore.FormatOpenAIResponses, true case typology.WireShapeAnthropicMessages: return provcore.FormatAnthropic, true + case typology.WireShapeGeminiGenerateContent: + return provcore.FormatGemini, true + case typology.WireShapeVertexGenerateContent: + return provcore.FormatVertex, true + case typology.WireShapeCohereChat: + return provcore.FormatCohere, true } return provcore.Format(""), false } diff --git a/packages/ai-gateway/internal/ingress/proxy/ingress_test.go b/packages/ai-gateway/internal/ingress/proxy/ingress_test.go index 3c3fc3d7..9065c0a6 100644 --- a/packages/ai-gateway/internal/ingress/proxy/ingress_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/ingress_test.go @@ -103,10 +103,12 @@ func TestWireShapeToBodyFormat(t *testing.T) { {"openai-embeddings", typology.WireShapeOpenAIEmbeddings, provcore.FormatOpenAI, true}, {"openai-responses", typology.WireShapeOpenAIResponses, provcore.FormatOpenAIResponses, true}, {"anthropic-messages", typology.WireShapeAnthropicMessages, provcore.FormatAnthropic, true}, - {"unmapped-gemini", typology.WireShapeGeminiGenerateContent, provcore.Format(""), false}, - {"unmapped-vertex", typology.WireShapeVertexEmbedContent, provcore.Format(""), false}, + {"gemini-generate", typology.WireShapeGeminiGenerateContent, provcore.FormatGemini, true}, + {"vertex-generate", typology.WireShapeVertexGenerateContent, provcore.FormatVertex, true}, + {"cohere-chat", typology.WireShapeCohereChat, provcore.FormatCohere, true}, + {"unmapped-vertex-embed", typology.WireShapeVertexEmbedContent, provcore.Format(""), false}, {"unmapped-bedrock", typology.WireShapeBedrockConverse, provcore.Format(""), false}, - {"unmapped-cohere", typology.WireShapeCohereEmbed, provcore.Format(""), false}, + {"unmapped-cohere-embed", typology.WireShapeCohereEmbed, provcore.Format(""), false}, {"unmapped-voyage", typology.WireShapeVoyageEmbeddings, provcore.Format(""), false}, {"sentinel-none", typology.WireShapeNone, provcore.Format(""), false}, } diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_buffer.go b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_buffer.go index 8a447490..4d9e8abc 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_buffer.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_buffer.go @@ -43,10 +43,10 @@ import ( // the final aggregated usage, so per-chunk Done/Usage are stripped from the body // chunks. The OpenAI encoder does not emit `data: [DONE]` (the live pipeline // normally appends it), so this helper appends it when emitDone is set. -func emitCanonicalStream(ctx context.Context, w io.Writer, transcoder canonicalbridge.StreamTranscoder, model string, emitDone bool, chunks []provcore.Chunk, finalUsage *provcore.Usage, finishReason string) error { +func emitCanonicalStream(ctx context.Context, w io.Writer, transcoder canonicalbridge.StreamTranscoder, ingress provcore.Format, model string, emitDone bool, chunks []provcore.Chunk, finalUsage *provcore.Usage, finishReason string) error { enc := transcoder if enc == nil { - enc = canonicalbridge.NewChatCompletionsStreamEncoder(model) + enc = fallbackStreamEncoder(ingress, model) } for _, c := range chunks { c.Done = false @@ -82,6 +82,23 @@ func emitCanonicalStream(ctx context.Context, w io.Writer, transcoder canonicalb return nil } +// fallbackStreamEncoder builds the re-emit encoder used by an enforcement path +// (buffer / Model A) when no cross-format transcoder was selected (same-shape +// passthrough). Enforcement re-encodes the redacted canonical body to the +// ingress wire, so the encoder MUST match the ingress shape: a /v1/responses +// stream emits response.* events, NEVER chat.completion.chunk. Enforcement +// takes precedence over native passthrough — built-in-tool / audio fidelity is +// forfeited under an enforcing response scope (the canonical waist cannot carry +// it), which is the accepted trade for a redacted/blocked stream. +func fallbackStreamEncoder(ingress provcore.Format, model string) canonicalbridge.StreamTranscoder { + // Re-encode the redacted/buffered canonical chunks into the CALLER's ingress + // wire shape. Delegating to the shared IngressStreamEncoder (the same source + // the live transcoder uses) covers every ingress — Gemini, Anthropic, Cohere, + // Replicate, Responses, OpenAI — so an enforcing buffer / Model-A stream never + // leaks chat.completion.chunk frames to a non-chat client. + return canonicalbridge.IngressStreamEncoder(ingress, model) +} + // runCanonicalBufferStream is the streaming BUFFER-mode handler. It consumes the // canonical ChunkSubscription pre-transcode, accumulates the full canonical // response, redacts on the canonical waist via redactCanonicalBuffer, and @@ -206,7 +223,7 @@ func (h *Handler) runCanonicalBufferStream(ctx context.Context, s *streamState, // unchanged to the ingress wire. emit = collected } - if err := emitCanonicalStream(ctx, tee, s.transcoder, s.target.ModelCode, s.emitDone, emit, finalUsage, acc.finishReason); err != nil { + if err := emitCanonicalStream(ctx, tee, s.transcoder, s.ingressFormat, s.target.ModelCode, s.emitDone, emit, finalUsage, acc.finishReason); err != nil { term.termErr.Store(&streamTerminalError{code: streamErrCodeUpstream, err: err}) } return term diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_buffer_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_buffer_test.go index 6a40692d..2c65952b 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_buffer_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_buffer_test.go @@ -276,6 +276,39 @@ func TestCanonicalBuffer_HardBlock_ZeroContent(t *testing.T) { } } +// TestCanonicalBuffer_BlockSoftMaskedRedact_RedactDelivers guards the shared redaction +// LOCUS (redactCanonicalBuffer) on the canonical BUFFER path: when a co-firing soft-block +// masks a redact, mergeResults now carries the redact's ModifiedContent, so the LOCUS +// (keyed on CarriesRedaction(), not Decision==Modify) APPLIES the redaction and delivers +// the masked body — the original PII never reaches the client and the stream is not +// blocked. Before the #13 fix the dropped ModifiedContent produced a no-op rewrite that +// failed closed on this canonical path. +func TestCanonicalBuffer_BlockSoftMaskedRedact_RedactDelivers(t *testing.T) { + cache := newBlockSoftPlusRedactResponseHookCache(t) + chunks := []provcore.Chunk{ + {Delta: "contact "}, + {Delta: "alice@example.com"}, + {Delta: " now", Done: true}, + } + s, w, usage := bufferTestState(t, cache, openAIChatIngress, nil, chunks) + + s.h.runCanonicalBufferStream(context.Background(), s, w, usage) + + out := w.String() + if strings.Contains(out, "alice@example.com") { + t.Fatalf("BlockSoft-masked redact leaked the original PII in buffer mode: %q", out) + } + if !strings.Contains(out, "[REDACTED_EMAIL]") { + t.Fatalf("BlockSoft-masked redact must DELIVER the masked content, got %q", out) + } + if strings.Contains(out, `"error"`) { + t.Fatalf("a co-firing redact must redact-deliver, not fail closed, got %q", out) + } + if !s.rec.ResponseHookRewritten || s.rec.ResponseAction != goHooks.ActionRedact { + t.Errorf("redact-deliver must stamp ResponseHookRewritten + ResponseAction=redact; got rewritten=%v action=%q", s.rec.ResponseHookRewritten, s.rec.ResponseAction) + } +} + // TestCanonicalBuffer_Approve_ForwardsContent proves the common case: no PII → // Approve → the buffered content is delivered to the client unchanged. func TestCanonicalBuffer_Approve_ForwardsContent(t *testing.T) { diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_crossingress_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_crossingress_test.go index 000ea777..92205730 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_crossingress_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_crossingress_test.go @@ -300,6 +300,128 @@ func TestCacheHit_Stream_OriginDiffers_StampsContext(t *testing.T) { } } +// TestCacheHit_Stream_CrossIngress_WireGrammar is the Amendment E +// streaming proof, both directions: a stream cache entry written under +// one OpenAI wire shape and replayed to a client on the OTHER shape must +// reach the client re-encoded into the READER's SSE grammar — never the +// writer's. Replay ChunkRecords carry no RawBytes/Verbatim, so the +// response-side content-peek classifier cannot fire on a HIT; the +// OriginWireShape tag stamped on the request context (StreamHitOrigin) is +// the ONLY signal that selects the transcoder in streamShapeStage. This +// pins that the origin override survived the Format-backdoor deletion: +// the target-Format reassignment now drives NewStreamTranscoder to the +// reader's encoder in both directions (chat→responses and responses→chat). +func TestCacheHit_Stream_CrossIngress_WireGrammar(t *testing.T) { + const relayToken = "xrelaymarker" + cases := []struct { + name string + originShape typology.WireShape + ingressShape typology.WireShape + ingressFormat provcore.Format + path string + reqBody []byte + keyWireShape typology.WireShape + wantContains []string + wantAbsent []string + }{ + { + name: "chat_written_responses_replay", + originShape: typology.WireShapeOpenAIChat, + ingressShape: typology.WireShapeOpenAIResponses, + ingressFormat: provcore.FormatOpenAIResponses, + path: "/v1/responses", + reqBody: []byte(`{"model":"gpt-4o","input":"cross stream chat to responses","stream":true}`), + keyWireShape: typology.WireShapeOpenAIResponses, + // The Responses encoder opens with response.created and closes + // with response.completed; the cached delta rides an + // output_text.delta event in between. + wantContains: []string{"event: response.created", "event: response.completed", relayToken}, + wantAbsent: []string{"chat.completion.chunk"}, + }, + { + name: "responses_written_chat_replay", + originShape: typology.WireShapeOpenAIResponses, + ingressShape: typology.WireShapeOpenAIChat, + ingressFormat: provcore.FormatOpenAI, + path: "/v1/chat/completions", + reqBody: []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"cross stream responses to chat"}],"stream":true}`), + keyWireShape: typology.WireShapeOpenAIChat, + wantContains: []string{"chat.completion.chunk", relayToken}, + wantAbsent: []string{"event: response.created", "event: response.completed"}, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + cacheOpt, cleanup := withCache(t) + defer cleanup() + + deps := makeOpenAIDeps(t, "", emptyHookCache(t), cacheOpt) + streamEntry := &cache.StreamEntry{ + Provider: "openai", + Model: "gpt-4o", + Chunks: []cache.ChunkRecord{ + {Delta: relayToken}, + {Done: true, Usage: &provcore.Usage{ + PromptTokens: iPtr(2), CompletionTokens: iPtr(3), TotalTokens: iPtr(5), + }}, + }, + CachedAt: time.Now().UTC(), + OriginWireShape: tc.originShape, + } + // Derive the cache key exactly as the handler does for this + // reader: PrepareBody at the reader's wire shape, then BuildKey. + adapter, ok := deps.ProviderReg.Get(provcore.FormatOpenAI) + if !ok { + t.Fatal("openai adapter missing") + } + prepReq := provcore.Request{ + WireShape: tc.keyWireShape, + Body: tc.reqBody, + BodyFormat: tc.ingressFormat, + Stream: true, + } + prepReq.Target.ProviderModelID = "gpt-4o" + finalBody, _, _, err := adapter.PrepareBody(prepReq) + if err != nil { + t.Fatalf("PrepareBody: %v", err) + } + cacheKey := deps.Cache.BuildKey("openai", "gpt-4o", finalBody, "") + if _, err := deps.Cache.StoreStream(context.Background(), cacheKey, streamEntry); err != nil { + t.Fatalf("StoreStream: %v", err) + } + + h := NewHandler(deps).ServeProxy(Ingress{ + WireShape: tc.ingressShape, + BodyFormat: tc.ingressFormat, + Stream: true, + }) + req := httptest.NewRequest(http.MethodPost, tc.path, strings.NewReader(string(tc.reqBody))) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer vk") + w := httptest.NewRecorder() + h(w, req) + + if w.Code != http.StatusOK { + t.Fatalf("status=%d body=%s", w.Code, w.Body.String()) + } + if !strings.EqualFold(w.Header().Get("X-Nexus-Cache"), "hit") { + t.Fatalf("x-nexus-cache=%q want HIT (cross-ingress stream replay did not hit)", w.Header().Get("X-Nexus-Cache")) + } + out := w.Body.String() + for _, want := range tc.wantContains { + if !strings.Contains(out, want) { + t.Errorf("replay missing %q — origin override did not re-encode to the reader's grammar; body=%s", want, out) + } + } + for _, absent := range tc.wantAbsent { + if strings.Contains(out, absent) { + t.Errorf("replay leaked the writer's grammar %q to the reader; body=%s", absent, out) + } + } + }) + } +} + // TestCacheHit_Stream_ResponsesIngress_OriginOverride exercises the // B2 stream-HIT override branch in handleStreamWithSubscription: // when origin override is set AND the standard NewStreamTranscoder diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela.go b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela.go index d896f05a..d0e2db2b 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela.go @@ -48,6 +48,7 @@ import ( "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/platform/middleware" provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" hookcore "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/streaming/modela" ) // modelATailWindowBytes bounds the trailing canonical content held undelivered. @@ -63,8 +64,24 @@ const modelATailWindowBytes = 8 * 1024 // HIT is confirmed, at which point it escalates to canonical-buffer redaction for // the remainder. prescan is the cheap union prefilter (closing over the response // probe's MayMatchRawContent); a nil prescan fails safe to "always confirm". -func (h *Handler) runModelAStream(ctx context.Context, s *streamState, tee http.ResponseWriter, usage *chunkUsageHolder, prescan func([]byte) bool) *chunkSSEReader { +func (h *Handler) runModelAStream(ctx context.Context, s *streamState, tee http.ResponseWriter, usage *chunkUsageHolder, prescan func([]byte) bool, maxPattern int) *chunkSSEReader { term := &chunkSSEReader{ctx: ctx, ingressFormat: s.ingressFormat} + + // respAcc folds every false-positive confirm's per-hook latency into one record + // per hook; appended once at stream end. authoritativeAppended is set true only + // after escalation's redactCanonicalBuffer has written the authoritative + // full-buffer scan — so on escalation success the accumulator is discarded (no + // double count), while on a normal end OR an escalate-then-abort (overflow / + // client-abort / upstream-error before the redact append) the defer appends the + // accumulated confirms so the trace is never lost. + var respAcc responseHookAccumulator + authoritativeAppended := false + defer func() { + if s.rec != nil && !authoritativeAppended { + s.rec.HooksPipeline = appendHookTrace(s.rec.HooksPipeline, "response", respAcc.finalize()) + } + }() + // Defensive nil-guard, symmetric with the canonical buffer handler. if s.sub == nil || tee == nil { return term @@ -79,7 +96,7 @@ func (h *Handler) runModelAStream(ctx context.Context, s *streamState, tee http. enc := s.transcoder if enc == nil { - enc = canonicalbridge.NewChatCompletionsStreamEncoder(s.target.ModelCode) + enc = fallbackStreamEncoder(s.ingressFormat, s.target.ModelCode) } maxBuf := s.streamMaxBufferBytes @@ -87,166 +104,35 @@ func (h *Handler) runModelAStream(ctx context.Context, s *streamState, tee http. maxBuf = defaultCanonicalBufferMaxBytes } - // held holds the trailing bounded-tail chunks that have NOT been delivered to - // the client. scanBuf accumulates the redactable canonical content the cheap - // union prescan reads: assistant deltas + tool-call arguments AND tool-call - // name/id. The enforcer (extractChatResponse) scans the whole tool_calls - // payload (function name + id + arguments), so the gate must cover the same - // channels — otherwise a value present only in a tool-call name/id would be - // invisible to Model A and delivered raw (reasoning is omitted, matching the - // canonical redaction's coverage). scannedLen high-waters the bytes already - // prefilter-scanned so the prescan is WINDOWED (new content + a bounded - // lookbehind) instead of re-scanning the whole buffer every chunk — O(N), not - // O(N²). contentBytes mirrors heldBytes but counts only scanned-content bytes - // so reasoning/non-content chunks don't evict content from the tail window - // early. confirmedLen is the scan length already cleared by a confirm, so a - // benign false-positive trigger is not re-confirmed until NEW content arrives. - var ( - held []provcore.Chunk - heldBytes int - contentBytes int - scanBuf []byte - scannedLen int - confirmedLen int - finishReason string - anyConfirm bool - ) - for { - chunk, err := s.sub.Next(ctx) - if err != nil { - if errors.Is(err, io.EOF) { - break - } - if ctx.Err() != nil { - term.termErr.Store(&streamTerminalError{code: streamErrCodeClientAbort, err: err}) - return term - } - var pe *provcore.ProviderError - if !errors.As(err, &pe) { - pe = &provcore.ProviderError{Status: http.StatusBadGateway, Code: provcore.CodeUpstreamError, Message: err.Error()} - } - _, _ = tee.Write(envelope.SynthesizeSSEErrorFrame(s.ingressFormat, pe)) - term.termErr.Store(&streamTerminalError{code: streamErrCodeUpstream, err: err}) - return term - } - if chunk.Usage != nil { - usage.record(chunk.Usage) - } - if chunk.FinishReason != "" { - finishReason = chunk.FinishReason - } - - // Accumulate the redactable canonical content for the prescan, then HOLD - // the chunk. Tool-call name/id are scanned alongside arguments — the - // enforcer extracts all three; fields are newline-separated so a pattern - // cannot span two fields (over-matching a prefilter only wastes a confirm, - // never misses a match). - contentStart := len(scanBuf) - scanBuf = append(scanBuf, chunk.Delta...) - for _, d := range chunk.ToolCallDeltas { - if d.Arguments != "" { - scanBuf = append(scanBuf, d.Arguments...) - scanBuf = append(scanBuf, '\n') - } - if d.Name != "" { - scanBuf = append(scanBuf, d.Name...) - scanBuf = append(scanBuf, '\n') - } - if d.ID != "" { - scanBuf = append(scanBuf, d.ID...) - scanBuf = append(scanBuf, '\n') - } - } - held = append(held, chunk) - heldBytes += canonicalChunkSize(chunk) - contentBytes += len(scanBuf) - contentStart - - // Prescan GATE: a HIT over content not yet confirmed pays for ONE full - // confirm. The prescan is WINDOWED — it scans only the content not yet - // prefilter-scanned plus a bounded lookbehind (the tail window), so a value - // spanning chunk boundaries within the still-held tail is detected, without - // re-scanning the whole accumulated buffer every chunk (O(N²)) or copying - // it ([]byte(scan.String()) was a full per-chunk copy). A prefilter hit on - // already-delivered content (scrolled past the lookbehind) is irrelevant — - // that content can no longer be redacted. confirmedLen gates "new content - // since the last confirm". The prescan is checked BEFORE the tail is - // released below, so a chunk that completes a confirmed match is always - // still held (undelivered) when the escalation redacts it. - if len(scanBuf) > confirmedLen { - winStart := scannedLen - modelATailWindowBytes - if winStart < 0 { - winStart = 0 - } - scannedLen = len(scanBuf) - if prescan(scanBuf[winStart:]) { - confirmedLen = len(scanBuf) - res := h.modelAConfirm(ctx, s, string(scanBuf), usage) - if res != nil { - anyConfirm = true - if res.Decision == hookcore.RejectHard || res.Decision == hookcore.Modify { - // Confirmed hit → escalate to buffer-to-END: deliver only the - // redacted REMAINDER (held + drained rest) through the SAME - // encoder. The already-delivered prefix is never re-delivered. - return h.escalateModelA(ctx, s, tee, enc, usage, held, finishReason, term) - } - // False positive (Approve / non-enforcing): stamp the response - // audit outcome so a SIEM sees "hooks evaluated" + any tags - // (parity with the wire/buffer paths), then resume real-time. - stampModelAResponseHook(s.rec, res) - } - } - } - - // Release the tail: deliver any held chunk once the held CONTENT exceeds the - // bounded window — measured in scanned-content bytes (Delta + tool args / - // name / id) so reasoning and other non-content chunks don't evict - // redactable content from the window early. A secondary cap on total held - // bytes bounds memory if a long non-content (reasoning) phase holds many - // chunks. Sound to deliver here — either the prescan MISSed (no hook can - // match) or the latest confirm Approved the accumulated content. - for (contentBytes > modelATailWindowBytes || heldBytes > maxBuf) && len(held) > 0 { - front := held[0] - if werr := writeEncodedChunk(ctx, enc, tee, front); werr != nil { - term.termErr.Store(&streamTerminalError{code: streamErrCodeUpstream, err: werr}) - return term - } - heldBytes -= canonicalChunkSize(front) - contentBytes -= modelAContentSize(front) - held = held[1:] - } - - if chunk.Done { - break - } - } - - // EOF without escalation → every prescan was a MISS or a false-positive - // Approve, so the held tail is provably benign. If NO confirm ever ran the - // sound prescan cleared the whole stream, so the response is APPROVE — stamp - // it so the audit row distinguishes "response hooks evaluated, approved" from - // "no response hook configured" (a false-positive confirm already stamped its - // own outcome above). Deliver the held tail raw, then the terminal frame. - if !anyConfirm { - s.rec.ResponseHookDecision = string(hookcore.Approve) - s.rec.ResponseAction = hookcore.ActionApprove - } - for _, c := range held { - if werr := writeEncodedChunk(ctx, enc, tee, c); werr != nil { - term.termErr.Store(&streamTerminalError{code: streamErrCodeUpstream, err: werr}) - return term - } - } - snap := usage.snapshot() - var finalUsage *provcore.Usage - if usageHasAny(snap) { - finalUsage = &snap - } - if werr := writeEncodedTerminal(ctx, enc, tee, s.emitDone, finalUsage, finishReason); werr != nil { - term.termErr.Store(&streamTerminalError{code: streamErrCodeUpstream, err: werr}) + // The shared engine owns the prescan-gated tail-hold / confirm / escalate + // algorithm; modelACanonicalSubstrate supplies the canonical seams (chunk + // subscription, redactable-channel extraction, single-encoder delivery, the + // per-checkpoint confirm, and the canonical-buffer escalation). The substrate + // records all terminal state on `term`, so the engine's returned error — already + // reflected there — is not consulted here. + sub := &modelACanonicalSubstrate{ + h: h, + s: s, + tee: tee, + enc: enc, + usage: usage, + prescan: prescan, + term: term, + respAcc: &respAcc, + authoritativeAppended: &authoritativeAppended, } + // Config-time operator signal (#16): warn once if the derived contiguous-pattern bound + // meets/exceeds the tail window the engine clamps the lookahead below. Off the per-byte + // path; maxPattern is already derived (buildResponsePrescan), not recomputed here. + modela.WarnStreamingCoverageGap(s.logger, maxPattern, modelATailWindowBytes) + _ = modela.Run(ctx, sub, modela.Config{TailWindowBytes: modelATailWindowBytes, MaxBufferBytes: maxBuf, MaxPatternBytes: maxPattern}) return term } +// modelACanonicalSubstrate satisfies the shared engine's Substrate over canonical +// provider chunks. +var _ modela.Substrate[provcore.Chunk] = (*modelACanonicalSubstrate)(nil) + // modelAConfirm runs ONE full response-stage confirm over the accumulated // canonical content via the relay's per-checkpoint hookRunner (the same pipeline // the live path runs). It returns the decision only; the actual redaction on a @@ -285,7 +171,7 @@ func (h *Handler) modelAConfirm(ctx context.Context, s *streamState, content str // fails closed (in-band error frame) on overflow; // - fail closed: a RejectHard / unproducible-rewrite is surfaced as an error // frame inside redactCanonicalBuffer — the original remainder is never sent. -func (h *Handler) escalateModelA(ctx context.Context, s *streamState, tee http.ResponseWriter, enc canonicalbridge.StreamTranscoder, usage *chunkUsageHolder, held []provcore.Chunk, finishReason string, term *chunkSSEReader) *chunkSSEReader { +func (h *Handler) escalateModelA(ctx context.Context, s *streamState, tee http.ResponseWriter, enc canonicalbridge.StreamTranscoder, usage *chunkUsageHolder, held []provcore.Chunk, finishReason string, term *chunkSSEReader, authoritativeAppended *bool) *chunkSSEReader { maxBuf := s.streamMaxBufferBytes if maxBuf <= 0 { maxBuf = defaultCanonicalBufferMaxBytes @@ -354,6 +240,17 @@ func (h *Handler) escalateModelA(ctx context.Context, s *streamState, tee http.R tokenTotal := int64(usageInt(snap.TotalTokens)) outcome := h.redactCanonicalBuffer(s.r, s.rec, ingress, s.target, acc.canonicalBody(), tokenTotal, s.requestID, s.logger) + // Mark the authoritative append done ONLY when redactCanonicalBuffer actually + // appended (outcome.appended). On its non-appending early returns — BypassHooks, + // pipeline build error, nil pipeline — the flag stays false so the caller's defer + // falls back to appending the accumulated confirm traces (no silent loss). On the + // normal executed path it appended the authoritative full-buffer scan, so the + // defer must NOT also append (which would duplicate the response hook rows). Every + // early return ABOVE this line also leaves the flag false → fallback still covers + // escalate-then-abort. + if authoritativeAppended != nil { + *authoritativeAppended = outcome.appended + } if outcome.failClosed { // Hard block / unproducible rewrite: deliver ONLY an in-band SSE error // frame — zero remainder content. ResponseHookRewritten stays false so the diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_helpers.go b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_helpers.go index 9ef5ae98..f64f8302 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_helpers.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_helpers.go @@ -14,6 +14,7 @@ import ( "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/platform/audit" provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" hookcore "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/streaming/modela" "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/typology" ) @@ -35,7 +36,7 @@ func modelAContentSize(c provcore.Chunk) int { // Model-A stream that approved or hit a false positive would leave // ResponseHookDecision empty — indistinguishable to a SIEM from "no response // hook configured" — and would lose any tags the confirm carried. -func stampModelAResponseHook(rec *audit.Record, res *hookcore.CompliancePipelineResult) { +func stampModelAResponseHook(rec *audit.Record, acc *responseHookAccumulator, res *hookcore.CompliancePipelineResult) { if rec == nil || res == nil { return } @@ -43,7 +44,14 @@ func stampModelAResponseHook(rec *audit.Record, res *hookcore.CompliancePipeline rec.ResponseHookReason = res.Reason rec.ResponseHookReasonCode = res.ReasonCode rec.ComplianceTags = mergeTagSets(rec.ComplianceTags, res.Tags) - rec.HooksPipeline = appendHookTrace(rec.HooksPipeline, "response", res.HookResults) + // Fold into the accumulator instead of appending here: a stream can run + // several false-positive confirms, and appending each would write the same + // hook N times and N×-inflate the response hook aggregates. runModelAStream + // appends the folded trace once at stream end (unless escalation already wrote + // the authoritative full-buffer scan via redactCanonicalBuffer). + if acc != nil { + acc.add(res.HookResults) + } rec.ResponseAction = hookcore.ActionFromDecision(res.Decision) } @@ -54,10 +62,15 @@ func stampModelAResponseHook(rec *audit.Record, res *hookcore.CompliancePipeline // probe cannot be built (no cache, build error, or no response rules), it returns // an "always confirm" prescan so the gate fails safe to a full confirm on every // checkpoint rather than silently skipping enforcement. -func (h *Handler) buildResponsePrescan(ctx context.Context, s *streamState) func([]byte) bool { +// buildResponsePrescan returns the cheap union prefilter for the Model-A relay AND the +// flush-before-deliver lookahead (modela.Config.MaxPatternBytes) DERIVED from the +// resolved response pipeline's longest contiguous enforceable match — floored at +// modela.DefaultMaxPatternBytes so it never drops below the proven-safe baseline. The +// lookahead is the package default whenever no pipeline resolves (always-confirm path). +func (h *Handler) buildResponsePrescan(ctx context.Context, s *streamState) (func([]byte) bool, int) { alwaysConfirm := func([]byte) bool { return true } if h.deps == nil || h.deps.HookConfigCache == nil { - return alwaysConfirm + return alwaysConfirm, modela.DefaultMaxPatternBytes } var epType hookcore.EndpointType if ingress, ok := IngressFromContext(s.r.Context()); ok { @@ -70,9 +83,13 @@ func (h *Handler) buildResponsePrescan(ctx context.Context, s *streamState) func s.logger, ) if err != nil || probe == nil { - return alwaysConfirm + return alwaysConfirm, modela.DefaultMaxPatternBytes + } + maxPattern := modela.DefaultMaxPatternBytes + if bound, _ := probe.MaxPatternBound(); bound > maxPattern { + maxPattern = bound // a longer contiguous enforceable pattern needs a wider lookahead } - return probe.MayMatchRawContent + return probe.MayMatchRawContent, maxPattern } // writeEncodedChunk forward-encodes one canonical BODY chunk through the shared diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_substrate.go b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_substrate.go new file mode 100644 index 00000000..bfa0800a --- /dev/null +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_substrate.go @@ -0,0 +1,194 @@ +package proxy + +// proxy_cache_modela_substrate.go — the canonical-subscription substrate that +// drives the shared streaming Model-A engine (shared/transport/streaming/modela). +// The engine owns the prescan-gated tail-hold algorithm; this adapter supplies the +// canonical seams: pulling provider chunks off the ChunkSubscription, extracting +// the redactable channels, forward-encoding through the single long-lived stream +// encoder, confirming via the relay's per-checkpoint hook runner, and escalating to +// the canonical-buffer redaction LOCUS. Its fail posture is fail-CLOSED (the +// gateway is not in the host packet path) — surfaced inside escalateModelA / +// redactCanonicalBuffer as an in-band SSE error frame, never the original body. + +import ( + "context" + "errors" + "net/http" + + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/execution/canonicalbridge" + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/ingress/envelope" + provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" + hookcore "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" + sharedstreaming "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/streaming" +) + +// modelACanonicalSubstrate adapts a canonical ChunkSubscription stream to the +// shared Model-A engine. It carries the per-stream delivery context (encoder, tee, +// usage accumulator, terminal-error carrier) and the cheap union prescan; the +// finishReason is observed as chunks arrive so the terminal frame and any +// escalation carry it. +type modelACanonicalSubstrate struct { + h *Handler + s *streamState + tee http.ResponseWriter + enc canonicalbridge.StreamTranscoder + usage *chunkUsageHolder + prescan func([]byte) bool + term *chunkSSEReader + + // respAcc folds every false-positive confirm's per-hook latency into one record + // per hook (set from the owning runModelAStream local); authoritativeAppended is + // flipped true once escalation's redactCanonicalBuffer has written the + // authoritative full-buffer scan, so the owner's defer does not double-append. + respAcc *responseHookAccumulator + authoritativeAppended *bool + + finishReason string +} + +// Next pulls the next canonical chunk, recording usage and the latest finish reason +// as it passes (the engine holds the chunk thereafter and does not surface it back +// except via Deliver / Escalate). io.EOF ends the stream; any other error is +// terminal and routed to OnError by the engine. +func (m *modelACanonicalSubstrate) Next(ctx context.Context) (provcore.Chunk, error) { + chunk, err := m.s.sub.Next(ctx) + if err != nil { + return provcore.Chunk{}, err + } + if chunk.Usage != nil { + m.usage.record(chunk.Usage) + } + if chunk.FinishReason != "" { + m.finishReason = chunk.FinishReason + } + return chunk, nil +} + +// AppendRedactableText appends the chunk's scannable channels — assistant delta plus +// tool-call arguments / name / id — onto dst, newline-separating the tool-call +// fields so a pattern cannot span two unrelated fields. Reasoning is omitted to +// match the canonical redaction's coverage. Appending onto the engine's buffer keeps +// the hot miss path allocation-free. +func (m *modelACanonicalSubstrate) AppendRedactableText(dst []byte, chunk provcore.Chunk) []byte { + dst = append(dst, chunk.Delta...) + for _, d := range chunk.ToolCallDeltas { + if d.Arguments != "" { + dst = append(dst, d.Arguments...) + dst = append(dst, '\n') + } + if d.Name != "" { + dst = append(dst, d.Name...) + dst = append(dst, '\n') + } + if d.ID != "" { + dst = append(dst, d.ID...) + dst = append(dst, '\n') + } + } + return dst +} + +// UnitBytes is the chunk's total canonical size (the held-bytes ceiling budget). +func (m *modelACanonicalSubstrate) UnitBytes(chunk provcore.Chunk) int { + return canonicalChunkSize(chunk) +} + +// ContentBytes is the chunk's redactable-content size without field separators (the +// tail-window budget): assistant delta + tool-call arguments / name / id, never +// reasoning, so a non-content reasoning chunk does not evict content from the window. +func (m *modelACanonicalSubstrate) ContentBytes(chunk provcore.Chunk) int { + return modelAContentSize(chunk) +} + +// IsDone reports the canonical terminal chunk. +func (m *modelACanonicalSubstrate) IsDone(chunk provcore.Chunk) bool { return chunk.Done } + +// Deliver forward-encodes one held chunk's body through the single long-lived +// encoder so the real-time → escalation switch preserves SSE frame continuity. A +// write error is terminal; it is recorded on the carrier and returned. +func (m *modelACanonicalSubstrate) Deliver(ctx context.Context, chunk provcore.Chunk) error { + if err := writeEncodedChunk(ctx, m.enc, m.tee, chunk); err != nil { + m.term.termErr.Store(&streamTerminalError{code: streamErrCodeUpstream, err: err}) + return err + } + return nil +} + +// DeliverTerminal emits the one terminal frame (finish reason + aggregated usage) +// plus the OpenAI sentinel when the ingress expects it. +func (m *modelACanonicalSubstrate) DeliverTerminal(ctx context.Context) error { + snap := m.usage.snapshot() + var finalUsage *provcore.Usage + if usageHasAny(snap) { + finalUsage = &snap + } + if err := writeEncodedTerminal(ctx, m.enc, m.tee, m.s.emitDone, finalUsage, m.finishReason); err != nil { + m.term.termErr.Store(&streamTerminalError{code: streamErrCodeUpstream, err: err}) + return err + } + return nil +} + +// Prescan is the cheap union prefilter (closing over the response probe's +// MayMatchRawContent); a nil prescan was already replaced with "always confirm" at +// construction, so this is never nil. +func (m *modelACanonicalSubstrate) Prescan(content []byte) bool { return m.prescan(content) } + +// Confirm runs ONE full response-stage confirm over the accumulated canonical +// content via the relay's per-checkpoint hook runner. modelAConfirm ALWAYS returns +// a non-nil result — RejectHard on a missing/failed runner, otherwise the pipeline's +// non-nil aggregate (Execute never returns nil) — so the engine's nil-as-approve +// branch and OnConfirmApproved(nil) are never exercised on the canonical path, and a +// missing enforcer never streams unredacted. +func (m *modelACanonicalSubstrate) Confirm(ctx context.Context, content string) *hookcore.CompliancePipelineResult { + return m.h.modelAConfirm(ctx, m.s, content, m.usage) +} + +// Escalate hands off to the canonical buffer-to-end redaction. trigger is the +// triggering decision on a confirmed hit, or nil on a memory-pressure escalation; +// escalateModelA RE-EVALUATES the buffered remainder on the canonical waist (it does +// not trust trigger's spans and stamps the authoritative response-hook outcome itself), +// so trigger is advisory only and the same hand-off serves both triggers. The +// confirmed-vs-pressure distinction is recorded as a metric dimension only (NOT a +// persisted field, and kept off ResponseHookReasonCode which redactCanonicalBuffer +// overwrites). +func (m *modelACanonicalSubstrate) Escalate(ctx context.Context, held []provcore.Chunk, trigger *hookcore.CompliancePipelineResult) error { + cause := sharedstreaming.ModelAEscalationConfirmed + if trigger == nil { + cause = sharedstreaming.ModelAEscalationMemoryPressure + } + sharedstreaming.RecordModelAEscalation(cause) + m.h.escalateModelA(ctx, m.s, m.tee, m.enc, m.usage, held, m.finishReason, m.term, m.authoritativeAppended) + return nil +} + +// OnConfirmApproved stamps a non-enforcing confirm outcome (false positive or +// nil-as-approve) onto the audit record. stampModelAResponseHook no-ops on a nil +// result. +func (m *modelACanonicalSubstrate) OnConfirmApproved(res *hookcore.CompliancePipelineResult) { + stampModelAResponseHook(m.s.rec, m.respAcc, res) +} + +// OnApproveEOF stamps the response as evaluated-approved when the stream reached EOF +// without any confirm — the sound prescan cleared the whole stream. +func (m *modelACanonicalSubstrate) OnApproveEOF() { + m.s.rec.ResponseHookDecision = string(hookcore.Approve) + m.s.rec.ResponseAction = hookcore.ActionApprove +} + +// OnError emits the substrate's terminal error framing for a non-EOF Next error: a +// client abort is recorded without a wire frame; an upstream error synthesizes an +// in-band SSE error frame on the ingress wire. Fail-closed: never the original body. +func (m *modelACanonicalSubstrate) OnError(ctx context.Context, err error) error { + if ctx.Err() != nil { + m.term.termErr.Store(&streamTerminalError{code: streamErrCodeClientAbort, err: err}) + return err + } + var pe *provcore.ProviderError + if !errors.As(err, &pe) { + pe = &provcore.ProviderError{Status: http.StatusBadGateway, Code: provcore.CodeUpstreamError, Message: err.Error()} + } + _, _ = m.tee.Write(envelope.SynthesizeSSEErrorFrame(m.s.ingressFormat, pe)) + m.term.termErr.Store(&streamTerminalError{code: streamErrCodeUpstream, err: err}) + return err +} diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_test.go index c0cbf7f4..2382536c 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_modela_test.go @@ -14,9 +14,11 @@ import ( "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/execution/canonicalbridge" provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/builtins" goHooks "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/payloadcapture" compliance "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/pipeline" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/streaming/modela" streampolicy "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/streaming/policy" "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/typology" ) @@ -58,6 +60,94 @@ func containsPrescan(marker string) func([]byte) bool { return func(b []byte) bool { return bytes.Contains(b, []byte(marker)) } } +// softBlockHook always returns a BlockSoft decision with no spans — a content-safety +// "soft flag" that, when co-firing with a redact hook, makes the pipeline aggregate +// report BlockSoft (it outranks Modify) while still carrying the redact's spans. +type softBlockHook struct{ goHooks.AnyEndpointAnyModality } + +func (softBlockHook) Execute(_ context.Context, _ *goHooks.HookInput) (*goHooks.HookResult, error) { + return &goHooks.HookResult{Decision: goHooks.BlockSoft, Reason: "content flagged (soft)", ReasonCode: "SOFT_FLAG"}, nil +} + +// newBlockSoftPlusRedactResponseHookCache wires TWO response-stage hooks: a +// PII-redact (email → [REDACTED_EMAIL], Modify) and an always-on soft-block. On a +// response carrying an email the aggregate Decision is BlockSoft (soft-block +// outranks Modify) while TransformSpans still carries the email redaction — the +// masked-redact case redactCanonicalBuffer must still apply. +func newBlockSoftPlusRedactResponseHookCache(t *testing.T) *compliance.HookConfigCache { + t.Helper() + reg := builtins.Registry.Clone() + reg.Register("softblock-resp-hook", func(_ *goHooks.HookConfig) (goHooks.Hook, error) { + return softBlockHook{}, nil + }) + reg.Freeze() + loader := func(_ context.Context) ([]goHooks.HookConfig, error) { + return []goHooks.HookConfig{ + { + ID: "pii-resp-1", ImplementationID: "pii-detector", Name: "pii-detect-response", + Priority: 10, Enabled: true, Stage: "response", FailBehavior: "fail-closed", TimeoutMs: 1000, + ApplicableIngress: []string{"ALL"}, + Config: map[string]any{ + "onMatch": map[string]any{"inflightAction": "redact", "storageAction": "redact"}, + "patternDefinitions": []any{ + map[string]any{"id": "email", "regex": `[a-z0-9._%+-]+@[a-z0-9.-]+\.[a-z]{2,}`, "flags": "i", "replacement": "[REDACTED_EMAIL]"}, + }, + }, + }, + { + ID: "softblock-resp-1", ImplementationID: "softblock-resp-hook", Name: "softblock-response", + Priority: 5, Enabled: true, Stage: "response", FailBehavior: "fail-closed", TimeoutMs: 1000, + ApplicableIngress: []string{"ALL"}, Config: map[string]any{}, + }, + }, nil + } + cache := compliance.NewHookConfigCache(loader, reg, 0, noopLogger()) + if err := cache.Start(context.Background()); err != nil { + t.Fatalf("cache.Start: %v", err) + } + time.Sleep(50 * time.Millisecond) + return cache +} + +// TestModelAStream_BlockSoftMaskedRedact_RedactDelivers is the end-to-end guard for the +// #13 security fix: when a soft-block hook co-fires with a redact hook, the aggregate +// Decision is BlockSoft but mergeResults now carries the redact's ModifiedContent (and +// spans). The engine escalates on the enforcing action, and redactCanonicalBuffer keys +// on CarriesRedaction() (not Decision==Modify), so the redaction is APPLIED and the +// masked body ([REDACTED_EMAIL]) is delivered — the complete email never reaches the +// wire AND the stream is not blocked. Before the #13 fix mergeResults dropped the +// ModifiedContent, producing a no-op rewrite that (on this canonical path) failed closed +// and (on the wire/buffer paths) leaked the original raw. +func TestModelAStream_BlockSoftMaskedRedact_RedactDelivers(t *testing.T) { + cache := newBlockSoftPlusRedactResponseHookCache(t) + chunks := []provcore.Chunk{ + {Delta: "reach me at "}, + {Delta: "alice@example.com", Done: true}, + } + s, w, usage := bufferTestState(t, cache, openAIChatIngress, nil, chunks) + runner, _ := modelACountingRunner(t, cache, openAIChatIngress) + s.hookRunner = runner + + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) + + out := w.String() + if strings.Contains(out, "alice@example.com") { + t.Fatalf("BlockSoft-masked redact leaked the complete PII raw: %q", out) + } + if !strings.Contains(out, "[REDACTED_EMAIL]") { + t.Fatalf("BlockSoft-masked redact must DELIVER the masked content, got %q", out) + } + if strings.Contains(out, `"error"`) { + t.Fatalf("a co-firing redact must redact-deliver, not fail closed, got %q", out) + } + if !s.rec.ResponseHookRewritten { + t.Error("a redact-delivered stream must stamp ResponseHookRewritten") + } + if s.rec.ResponseAction != goHooks.ActionRedact { + t.Errorf("rec.ResponseAction = %q, want redact (Decision=BlockSoft but disposition is redact-deliver)", s.rec.ResponseAction) + } +} + // TestModelAStream_PrescanMiss_RealTime_ZeroConfirm pins the common case: a // prescan that never matches streams the full body in real time and pays ZERO // full confirms. @@ -71,7 +161,7 @@ func TestModelAStream_PrescanMiss_RealTime_ZeroConfirm(t *testing.T) { runner, confirms := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("NEVER")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("NEVER"), 0) out := w.String() if !strings.Contains(out, "hello") || !strings.Contains(out, "world") { @@ -102,7 +192,7 @@ func TestModelAStream_FalsePositive_OneConfirm_FullBody(t *testing.T) { s.hookRunner = runner // Prescan HITs on "TRIGGER" (a false positive: no email → the hook Approves). - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("TRIGGER")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("TRIGGER"), 0) out := w.String() if !strings.Contains(out, "benign TRIGGER text") || !strings.Contains(out, "totally") { @@ -129,7 +219,7 @@ func TestModelAStream_ConfirmedHit_Escalates_Redacts(t *testing.T) { runner, confirms := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) out := w.String() if strings.Contains(out, "alice@example.com") { @@ -163,7 +253,7 @@ func TestModelAStream_NoHit_StampsApprove(t *testing.T) { runner, confirms := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("NEVER")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("NEVER"), 0) if *confirms != 0 { t.Fatalf("clean stream must not confirm, ran %d", *confirms) @@ -186,7 +276,7 @@ func TestModelAStream_FalsePositive_StampsDecision(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("TRIGGER")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("TRIGGER"), 0) if s.rec.ResponseHookDecision != string(goHooks.Approve) { t.Errorf("false-positive confirm must stamp ResponseHookDecision=%q, got %q", string(goHooks.Approve), s.rec.ResponseHookDecision) @@ -207,7 +297,7 @@ func TestModelAStream_ToolCallName_Scanned(t *testing.T) { runner, confirms := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) if *confirms < 1 { t.Fatalf("a value in a tool-call name must be scanned and confirmed (FIX-2), ran %d confirms", *confirms) @@ -228,7 +318,7 @@ func TestModelAStream_EscalationPreservesWireToolIndex(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) out := w.String() if !strings.Contains(out, `"index":1`) { @@ -255,7 +345,7 @@ func TestModelAStream_BoundedFragment_CompleteValueNeverDelivered(t *testing.T) runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) out := w.String() if !strings.Contains(out, "benign filler") { @@ -292,7 +382,7 @@ func TestModelAStream_StorageCopyRedactedWithinTailWindow(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) // w.String() == tee.captured() == the bytes the relay persists as the audit // storage copy (rec.ResponseBody / rec.ResponseBodyRedacted on a redact rewrite). @@ -320,7 +410,7 @@ func TestModelAStream_HardBlock_ZeroContentAfterEscalation(t *testing.T) { runner, confirms := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("secret")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("secret"), 0) out := w.String() if strings.Contains(out, "hunter2") { @@ -356,7 +446,7 @@ func TestModelAStream_EscalationBufferCap_FailsClosed(t *testing.T) { s.hookRunner = runner s.streamMaxBufferBytes = 256 // smaller than the drained remainder - term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) out := w.String() if !strings.Contains(out, `"error"`) || !strings.Contains(out, "maximum buffer size") { @@ -367,6 +457,42 @@ func TestModelAStream_EscalationBufferCap_FailsClosed(t *testing.T) { } } +// TestModelAStream_MemoryPressure_HoldsIncompleteContent pins the memory-pressure +// safety guard: when a reasoning burst trips the held-bytes ceiling while a content +// unit carrying an INCOMPLETE sensitive value is still at the front of the window, +// the engine escalates rather than force-delivering that content raw. The prescan +// never hits (the "@" only arrives in the final, still-held chunk), so this is a +// pure memory-pressure escalation (zero confirms) — and the incomplete prefix +// "alice" must never reach the wire (the old force-deliver-raw behavior would have +// leaked it before the value completed). +func TestModelAStream_MemoryPressure_HoldsIncompleteContent(t *testing.T) { + cache := newPiiRedactResponseHookCache(t) + chunks := []provcore.Chunk{ + {Delta: "alice"}, // incomplete email (no "@" yet) — prescan misses + {ReasoningDelta: strings.Repeat("y", 400)}, // reasoning burst: 0 content bytes, ~400 held bytes + {Delta: "@example.com", Done: true}, // completes the email, still held at escalation + } + s, w, usage := bufferTestState(t, cache, openAIChatIngress, nil, chunks) + runner, confirms := modelACountingRunner(t, cache, openAIChatIngress) + s.hookRunner = runner + s.streamMaxBufferBytes = 200 // tripped by the reasoning burst while content window (8KB) is not + + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) + + out := w.String() + if strings.Contains(out, "alice") { + t.Fatalf("memory-pressure eviction leaked the incomplete content raw: %q", out) + } + if *confirms != 0 { + t.Errorf("memory-pressure escalation must not run a prescan confirm, ran %d", *confirms) + } + // The held tail already exceeds the cap, so the escalation drain fails closed + // with an in-band error frame — the safe outcome (block, never leak). + if !strings.Contains(out, `"error"`) { + t.Errorf("expected a fail-closed error frame under memory pressure, got %q", out) + } +} + // TestModelAStream_UsagePreserved pins that usage observed on the canonical // chunks lands on the usageHolder and is re-emitted on the wire (real-time path). func TestModelAStream_UsagePreserved(t *testing.T) { @@ -379,7 +505,7 @@ func TestModelAStream_UsagePreserved(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("NEVER")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("NEVER"), 0) snap := usage.snapshot() if snap.TotalTokens == nil || *snap.TotalTokens != 12 { @@ -399,7 +525,7 @@ func TestModelAStream_ProviderError_SynthesizesFrame(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) if !strings.Contains(w.String(), `"error"`) { t.Errorf("expected SSE error frame on provider fault, got %q", w.String()) @@ -420,7 +546,7 @@ func TestModelAStream_ClientAbort_NoFrame(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() - term := s.h.runModelAStream(ctx, s, w, usage, containsPrescan("@")) + term := s.h.runModelAStream(ctx, s, w, usage, containsPrescan("@"), 0) if w.Len() != 0 { t.Errorf("client abort must not write to the gone peer, got %q", w.String()) @@ -435,7 +561,7 @@ func TestModelAStream_NilSubOrTee_NoOp(t *testing.T) { cache := newPiiRedactResponseHookCache(t) s, w, usage := bufferTestState(t, cache, openAIChatIngress, nil, nil) s.sub = nil - if term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")); term == nil { + if term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0); term == nil { t.Fatal("expected a non-nil terminal carrier even on nil sub") } if w.Len() != 0 { @@ -465,7 +591,7 @@ func TestModelAStream_NilPrescan_FailsSafeToConfirm(t *testing.T) { runner, confirms := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, nil) + s.h.runModelAStream(context.Background(), s, w, usage, nil, 0) if *confirms == 0 { t.Error("nil prescan must fail safe to always-confirm (≥1 confirm)") @@ -483,7 +609,7 @@ func TestBuildResponsePrescan_NoCache_AlwaysConfirms(t *testing.T) { req = req.WithContext(WithIngress(req.Context(), openAIChatIngress)) s := &streamState{h: h, r: req, logger: noopLogger()} - prescan := h.buildResponsePrescan(context.Background(), s) + prescan, _ := h.buildResponsePrescan(context.Background(), s) if !prescan([]byte("anything")) { t.Error("no-cache prescan must fail safe to always-confirm") } @@ -499,7 +625,7 @@ func TestBuildResponsePrescan_RealCache_GatesOnContent(t *testing.T) { req = req.WithContext(WithIngress(req.Context(), openAIChatIngress)) s := &streamState{h: h, r: req, logger: noopLogger()} - prescan := h.buildResponsePrescan(context.Background(), s) + prescan, _ := h.buildResponsePrescan(context.Background(), s) if !prescan([]byte("write to alice@example.com")) { t.Error("email content must trip the union prefilter") } @@ -508,6 +634,28 @@ func TestBuildResponsePrescan_RealCache_GatesOnContent(t *testing.T) { } } +// TestBuildResponsePrescan_DerivesFlooredLookahead pins the substrate glue that wires the +// derived Model-A flush-before-deliver lookahead: typical PII patterns are far under the +// 4096 floor, so the derivation returns exactly DefaultMaxPatternBytes; and the no-cache +// fail-safe path also returns the default rather than 0. +func TestBuildResponsePrescan_DerivesFlooredLookahead(t *testing.T) { + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil) + req = req.WithContext(WithIngress(req.Context(), openAIChatIngress)) + + cache := newPiiRedactResponseHookCache(t) + h := &Handler{deps: &Deps{HookConfigCache: cache, Logger: noopLogger()}} + s := &streamState{h: h, r: req, logger: noopLogger()} + if _, maxPattern := h.buildResponsePrescan(context.Background(), s); maxPattern != modela.DefaultMaxPatternBytes { + t.Fatalf("typical PII patterns must floor the lookahead at DefaultMaxPatternBytes, got %d", maxPattern) + } + + h2 := &Handler{deps: &Deps{Logger: noopLogger()}} // no hook cache → always-confirm fail-safe + s2 := &streamState{h: h2, r: req, logger: noopLogger()} + if _, maxPattern := h2.buildResponsePrescan(context.Background(), s2); maxPattern != modela.DefaultMaxPatternBytes { + t.Fatalf("no-cache path must return DefaultMaxPatternBytes (not 0), got %d", maxPattern) + } +} + // TestStreamRelayStage_ModelAArmed_RedactsViaCanonical pins the dispatch: a // MayRedact chunked_async stream (modelAArmed) is routed to runModelAStream — the // canonical Model-A path — so the delivered stream is redacted, distinct from the @@ -559,7 +707,7 @@ func TestModelAStream_Concurrent_PerCallIsolation(t *testing.T) { s, w, usage := bufferTestState(t, cache, openAIChatIngress, nil, chunks) runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) out := w.String() if strings.Contains(out, "carol@example.com") { t.Errorf("per-call leak under concurrency: %q", out) @@ -584,7 +732,7 @@ func TestModelAStream_EscalationProviderError_Classified(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) if !strings.Contains(w.String(), `"error"`) { t.Errorf("expected synthesized error frame on drain fault, got %q", w.String()) @@ -607,7 +755,7 @@ func TestModelAStream_ToolArgs_ConfirmedRedact(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) out := w.String() if strings.Contains(out, "alice@example.com") { @@ -631,7 +779,7 @@ func TestModelAStream_FinishReason_Preserved_RealTime(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("NEVER")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("NEVER"), 0) if !strings.Contains(w.String(), `"finish_reason":"length"`) { t.Errorf("real-time path must preserve finish_reason=length, got %q", w.String()) @@ -647,7 +795,7 @@ func TestModelAStream_EOFWithoutDone_DeliversTail(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("NEVER")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("NEVER"), 0) out := w.String() if !strings.Contains(out, "tail without done") { @@ -667,7 +815,7 @@ func TestModelAStream_GenericUpstreamFault_RealTime(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) if !strings.Contains(w.String(), `"error"`) { t.Errorf("expected synthesized error frame, got %q", w.String()) @@ -688,7 +836,7 @@ func TestModelAStream_RealTimeFlushWriteError_Classified(t *testing.T) { s.hookRunner = runner fw := &failingWriter{header: http.Header{}} - term := s.h.runModelAStream(context.Background(), s, fw, usage, containsPrescan("NEVER")) + term := s.h.runModelAStream(context.Background(), s, fw, usage, containsPrescan("NEVER"), 0) if te := term.terminalError(); te == nil || te.code != streamErrCodeUpstream { t.Errorf("expected %q on flush write error, got %+v", streamErrCodeUpstream, te) @@ -705,13 +853,52 @@ func TestModelAStream_EOFTailWriteError_Classified(t *testing.T) { s.hookRunner = runner fw := &failingWriter{header: http.Header{}} - term := s.h.runModelAStream(context.Background(), s, fw, usage, containsPrescan("NEVER")) + term := s.h.runModelAStream(context.Background(), s, fw, usage, containsPrescan("NEVER"), 0) if te := term.terminalError(); te == nil || te.code != streamErrCodeUpstream { t.Errorf("expected %q on EOF tail write error, got %+v", streamErrCodeUpstream, te) } } +// failOnMarkerWriter writes through until a payload contains marker, which it +// fails — letting a test fail one specific late frame (the terminal sentinel) +// while every earlier body write succeeds. +type failOnMarkerWriter struct { + *bytes.Buffer + header http.Header + marker string +} + +func (w *failOnMarkerWriter) Header() http.Header { return w.header } +func (w *failOnMarkerWriter) WriteHeader(int) {} +func (w *failOnMarkerWriter) Write(p []byte) (int, error) { + if bytes.Contains(p, []byte(w.marker)) { + return 0, io.ErrClosedPipe + } + return w.Buffer.Write(p) +} + +// TestModelAStream_TerminalWriteError_Classified pins a write failure on the +// terminal frame specifically: the body and finish frame deliver, but the [DONE] +// sentinel write fails → the terminal carrier is classified upstream. +func TestModelAStream_TerminalWriteError_Classified(t *testing.T) { + cache := newPiiRedactResponseHookCache(t) + chunks := []provcore.Chunk{{Delta: "clean body", Done: true}} + s, _, usage := bufferTestState(t, cache, openAIChatIngress, nil, chunks) + runner, _ := modelACountingRunner(t, cache, openAIChatIngress) + s.hookRunner = runner + w := &failOnMarkerWriter{Buffer: &bytes.Buffer{}, header: http.Header{}, marker: "[DONE]"} + + term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("NEVER"), 0) + + if !strings.Contains(w.String(), "clean body") { + t.Fatalf("body must deliver before the terminal write fails, got %q", w.String()) + } + if te := term.terminalError(); te == nil || te.code != streamErrCodeUpstream { + t.Errorf("terminal write failure must classify upstream, got %+v", te) + } +} + // TestModelAStream_EscalationSynthWriteError_Classified pins a write failure // delivering the redacted remainder on escalation → upstream classification. func TestModelAStream_EscalationSynthWriteError_Classified(t *testing.T) { @@ -722,7 +909,7 @@ func TestModelAStream_EscalationSynthWriteError_Classified(t *testing.T) { s.hookRunner = runner fw := &failingWriter{header: http.Header{}} - term := s.h.runModelAStream(context.Background(), s, fw, usage, containsPrescan("@")) + term := s.h.runModelAStream(context.Background(), s, fw, usage, containsPrescan("@"), 0) if te := term.terminalError(); te == nil || te.code != streamErrCodeUpstream { t.Errorf("expected %q on escalation synth write error, got %+v", streamErrCodeUpstream, te) @@ -743,7 +930,7 @@ func TestModelAStream_EscalationDrain_FinishReason(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) out := w.String() if strings.Contains(out, "alice@example.com") { @@ -767,7 +954,7 @@ func TestModelAStream_EscalationDrain_EOFNoDone(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) out := w.String() if strings.Contains(out, "alice@example.com") { @@ -789,7 +976,7 @@ func TestModelAStream_EscalationDrain_ClientAbort(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - term := s.h.runModelAStream(ctx, s, w, usage, containsPrescan("@")) + term := s.h.runModelAStream(ctx, s, w, usage, containsPrescan("@"), 0) if te := term.terminalError(); te == nil || te.code != streamErrCodeClientAbort { t.Errorf("expected %q on drain client abort, got %+v", streamErrCodeClientAbort, te) @@ -856,7 +1043,7 @@ func TestModelAStream_EscalationDrain_RecordsUsage(t *testing.T) { runner, _ := modelACountingRunner(t, cache, openAIChatIngress) s.hookRunner = runner - s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) if snap := usage.snapshot(); snap.TotalTokens == nil || *snap.TotalTokens != 7 { t.Fatalf("usage drained during escalation must be recorded, got %v", usage.snapshot().TotalTokens) @@ -877,7 +1064,7 @@ func TestModelAStream_EscalationTerminalWriteError_Classified(t *testing.T) { s.hookRunner = runner w := &kthFailWriter{header: http.Header{}, failOn: 2} // synth ok, terminal fails - term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@")) + term := s.h.runModelAStream(context.Background(), s, w, usage, containsPrescan("@"), 0) if te := term.terminalError(); te == nil || te.code != streamErrCodeUpstream { t.Errorf("expected %q on escalation terminal write error, got %+v", streamErrCodeUpstream, te) diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_canonical_redact.go b/packages/ai-gateway/internal/ingress/proxy/proxy_canonical_redact.go index b508c548..7a03b84e 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_canonical_redact.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_canonical_redact.go @@ -51,6 +51,13 @@ type canonicalRedactOutcome struct { // result is the pipeline result (decision / tags / blocking rule), set // whenever the pipeline executed. nil on bypass / no-pipeline. result *hookcore.CompliancePipelineResult + // appended is true only when this call actually appended a response-stage + // trace to rec.HooksPipeline (the pipeline executed). It is false on the + // early returns — BypassHooks, pipeline build error, nil pipeline — that + // return without appending. The Model-A escalation hand-off uses it to decide + // whether its accumulator fallback still needs to append (so a bypassed + // escalation does not silently drop the accumulated confirm traces). + appended bool } // redactCanonicalBuffer runs the response-stage compliance pipeline on the @@ -148,6 +155,7 @@ func (h *Handler) redactCanonicalBuffer( rec.ResponseHookReasonCode = hookResult.ReasonCode rec.ComplianceTags = mergeTagSets(rec.ComplianceTags, hookResult.Tags) rec.HooksPipeline = appendHookTrace(rec.HooksPipeline, "response", hookResult.HookResults) + out.appended = true if br := mapBlockingRule(hookResult.BlockingRule); br != nil { rec.BlockingRule = br } @@ -162,15 +170,20 @@ func (h *Handler) redactCanonicalBuffer( out.errMsg = hookResult.Reason return out } - if hookResult.Decision == hookcore.Modify { - // A Modify decision MUST produce an applied rewrite. The canonical locus - // is fail-closed (the ai-gateway is not in the host packet path → strong + // Apply the redaction rewrite whenever the aggregate carries redaction work: a + // Modify decision, OR a BlockSoft that masks a co-firing redact. The shared + // predicate is the single source of truth used by every consumer (request + + // response, stream + non-stream) so keying on Decision==Modify alone cannot + // silently drop a real redaction. A standalone soft-block (no spans, no + // ModifiedContent) is allow-with-warning and still delivers the original. + if hookResult.CarriesRedaction() { + // The redaction MUST produce an applied rewrite. The canonical locus is + // fail-closed (the ai-gateway is not in the host packet path → strong // compliance): a redaction the policy demanded but could not apply must // NEVER deliver the original. Attempt the rewrite UNCONDITIONALLY — a - // Modify carrying only TransformSpans (tool-arg masking) with empty - // ModifiedContent still has real work to apply, so the old - // len(ModifiedContent)>0 gate silently dropped it and returned the - // original unredacted body (fail-OPEN). + // rewrite carrying only TransformSpans (tool-arg masking) with empty + // ModifiedContent still has real work to apply, so a len(ModifiedContent)>0 + // gate would silently drop it and return the original unredacted body. redacted, n, rErr := extractor.RewriteResponseBody(r.Context(), canonicalBody, canonicalPath, rewriteContentWithToolArgs(hookResult.ModifiedContent, respContent, hookResult.TransformSpans)) if rErr != nil { // On canonical the rewrite is always supported: the sniff above only @@ -185,11 +198,11 @@ func (h *Handler) redactCanonicalBuffer( return out } if n == 0 { - // The Modify produced NO applied change: empty ModifiedContent and no + // The rewrite produced NO applied change: empty ModifiedContent and no // applicable tool-arg spans. Mirror the rewrite-error arm — fail // closed so the unredacted original is never delivered. out.rewritten // stays false, so the relay's redacted-copy guard keeps storage NULL. - logger.Error("canonical response Modify produced no applied rewrite — failing closed") + logger.Error("canonical response redaction produced no applied rewrite — failing closed") out.failClosed = true out.errStatus = http.StatusInternalServerError out.errMsg = "response rewrite produced no change" @@ -197,6 +210,12 @@ func (h *Handler) redactCanonicalBuffer( } rec.ResponseHookRewriteCount = n rec.ResponseHookRewritten = true + // Stamp the DISPOSITION action: an applied rewrite is a redact-deliver, even + // when the aggregate Decision is BlockSoft (a soft-block masking a co-firing + // redact). Without this the audit row would read Action=block + masked body, + // which misrepresents the outcome — the disposition is redact, the Decision is + // the (soft-block) ceiling. + rec.ResponseAction = hookcore.ActionRedact out.body = redacted out.rewritten = true return out diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_egress_shape_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_egress_shape_test.go index 85fdbee4..7e1fb27b 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_egress_shape_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_egress_shape_test.go @@ -36,47 +36,60 @@ func egressTarget(adapter string) routingcore.RoutingTarget { // TestEgressReshapeNonStream_RoutesByIngressShape pins that the reshape is // driven by the ingress shape A, identity for OpenAI-family, with the // /v1/responses native-passthrough exception — independent of the target B. +// responses-shape upstream body ({"object":"response"}) — what a genuine +// /v1/responses upstream returns; the content classifier forwards it verbatim. +const egressResponsesBody = `{"object":"response","id":"resp_x","output":[]}` + func TestEgressReshapeNonStream_RoutesByIngressShape(t *testing.T) { cases := []struct { name string wire typology.WireShape body provcore.Format + inBody string // upstream body; "" → egressCanonicalBody (chat.completion) target routingcore.RoutingTarget wantReshaped bool // expect ResponseCanonicalToIngress invoked wantFormat provcore.Format + wantVerbatim string // when not reshaped, the exact body the caller must receive }{ // OpenAI-family ingress: canonical already IS the caller's shape → identity, no call. - {"openai-chat identity", typology.WireShapeOpenAIChat, provcore.FormatOpenAI, egressTarget("anthropic"), false, ""}, + {"openai-chat identity", typology.WireShapeOpenAIChat, provcore.FormatOpenAI, "", egressTarget("anthropic"), false, "", egressCanonicalBody}, // Anthropic /v1/messages → MUST reshape to anthropic, for BOTH native and cross targets. - {"anthropic native", typology.WireShapeAnthropicMessages, provcore.FormatAnthropic, egressTarget("anthropic"), true, provcore.FormatAnthropic}, - {"anthropic cross→openai", typology.WireShapeAnthropicMessages, provcore.FormatAnthropic, egressTarget("openai"), true, provcore.FormatAnthropic}, + {"anthropic native", typology.WireShapeAnthropicMessages, provcore.FormatAnthropic, "", egressTarget("anthropic"), true, provcore.FormatAnthropic, ""}, + {"anthropic cross→openai", typology.WireShapeAnthropicMessages, provcore.FormatAnthropic, "", egressTarget("openai"), true, provcore.FormatAnthropic, ""}, // Gemini /v1beta → MUST reshape to gemini, native and cross. - {"gemini native", typology.WireShapeGeminiGenerateContent, provcore.FormatGemini, egressTarget("gemini"), true, provcore.FormatGemini}, - {"gemini cross→openai", typology.WireShapeGeminiGenerateContent, provcore.FormatGemini, egressTarget("openai"), true, provcore.FormatGemini}, - // /v1/responses: native passthrough (openai target serves responses) → no reshape; - // cross-format → reshape to Responses output[]. - {"responses native passthrough", typology.WireShapeOpenAIResponses, provcore.FormatOpenAIResponses, egressTarget("openai"), false, ""}, - {"responses cross→anthropic", typology.WireShapeOpenAIResponses, provcore.FormatOpenAIResponses, egressTarget("anthropic"), true, provcore.FormatOpenAIResponses}, + {"gemini native", typology.WireShapeGeminiGenerateContent, provcore.FormatGemini, "", egressTarget("gemini"), true, provcore.FormatGemini, ""}, + {"gemini cross→openai", typology.WireShapeGeminiGenerateContent, provcore.FormatGemini, "", egressTarget("openai"), true, provcore.FormatGemini, ""}, + // /v1/responses content-driven: a Responses-shape body is verbatim (zero-loss + // for built-in tools), independent of the target Format. + {"responses body verbatim (openai target)", typology.WireShapeOpenAIResponses, provcore.FormatOpenAIResponses, egressResponsesBody, egressTarget("openai"), false, "", egressResponsesBody}, + // A chat.completion body on a /v1/responses ingress is re-encoded to the + // Responses output[] grammar — even when the target is the "native" openai + // adapter (bulletproof: a chat-shaped reply never leaks to a Responses client). + {"responses chat-upstream re-encode (openai target)", typology.WireShapeOpenAIResponses, provcore.FormatOpenAIResponses, "", egressTarget("openai"), true, provcore.FormatOpenAIResponses, ""}, + {"responses cross→anthropic", typology.WireShapeOpenAIResponses, provcore.FormatOpenAIResponses, "", egressTarget("anthropic"), true, provcore.FormatOpenAIResponses, ""}, // Non-OpenAI embeddings ingress → embeddings reshape path (chat spy NOT // invoked; the embeddings re-encoder handles canonical→ingress). - {"gemini embeddings", typology.WireShapeGeminiEmbedContent, provcore.FormatGemini, egressTarget("openai"), false, ""}, + {"gemini embeddings", typology.WireShapeGeminiEmbedContent, provcore.FormatGemini, "", egressTarget("openai"), false, "", egressCanonicalBody}, } for _, tc := range cases { t.Run(tc.name, func(t *testing.T) { var called bool var gotFormat provcore.Format fb := &fakeBridge{ - targetNativelyServesResponsesAPI: func(target provcore.Format) bool { return target == provcore.FormatOpenAI }, responseCanonicalToIngress: func(ingress provcore.Format, canonical []byte) ([]byte, error) { called = true gotFormat = ingress return []byte(`{"reshaped_to":"` + string(ingress) + `"}`), nil }, } + inBody := tc.inBody + if inBody == "" { + inBody = egressCanonicalBody + } h := NewHandler(&Deps{CanonicalBridge: fb}) out, err := h.egressReshapeNonStream( Ingress{WireShape: tc.wire, BodyFormat: tc.body}, - tc.target, []byte(egressCanonicalBody)) + tc.target, []byte(inBody)) if err != nil { t.Fatalf("egressReshapeNonStream: %v", err) } @@ -90,13 +103,31 @@ func TestEgressReshapeNonStream_RoutesByIngressShape(t *testing.T) { if !bytes.Contains(out, []byte(`"reshaped_to":"`+string(tc.wantFormat)+`"`)) { t.Fatalf("reshaped output not returned to caller; got %s", out) } - } else if !bytes.Equal(out, []byte(egressCanonicalBody)) { - t.Fatalf("identity path must return the canonical body verbatim; got %s", out) + } else if !bytes.Equal(out, []byte(tc.wantVerbatim)) { + t.Fatalf("identity/verbatim path must return the body unchanged; got %s want %s", out, tc.wantVerbatim) } }) } } +// TestEgressReshapeNonStream_FailsClosedOnUnclassifiable proves the content +// classifier refuses to forward an unclassifiable /v1/responses upstream body +// verbatim — it returns an error so the handler can surface a 502 rather than +// leaking a garbage / wrong-shape payload to a Responses client. +func TestEgressReshapeNonStream_FailsClosedOnUnclassifiable(t *testing.T) { + fb := &fakeBridge{responseCanonicalToIngress: func(provcore.Format, []byte) ([]byte, error) { + t.Fatal("reshape must not be invoked on an unclassifiable body") + return nil, nil + }} + h := NewHandler(&Deps{CanonicalBridge: fb}) + out, err := h.egressReshapeNonStream( + Ingress{WireShape: typology.WireShapeOpenAIResponses, BodyFormat: provcore.FormatOpenAIResponses}, + egressTarget("openai"), []byte(`{"unexpected":"garbage"}`)) + if err == nil { + t.Fatalf("unclassifiable body must fail closed; got out=%s", out) + } +} + // TestEgressReshapeNonStream_Guards covers the early-return guards: an empty // body and a nil bridge both pass the body through untouched without invoking // any reshape. diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_fakeexec_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_fakeexec_test.go index 1748b54c..816c694d 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_fakeexec_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_fakeexec_test.go @@ -118,7 +118,10 @@ func (b *fakeBridge) EndpointRoutable(ep typology.WireShape, ingress, target pro return ingress.Valid() && target.Valid() } -func (b *fakeBridge) TargetNativelyServesResponsesAPI(target provcore.Format) bool { +func (b *fakeBridge) ServesResponses(target provcore.Format, override *bool) bool { + if override != nil && !*override { + return false + } if b.targetNativelyServesResponsesAPI != nil { return b.targetNativelyServesResponsesAPI(target) } diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_helpers_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_helpers_test.go index ad92e3e3..e55ca58a 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_helpers_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_helpers_test.go @@ -1007,7 +1007,7 @@ func TestProxyCache_ChunkSSEReader_RawBytesPassthrough(t *testing.T) { {RawBytes: []byte("data: {\"a\":1}\n\n")}, {Done: true, RawBytes: []byte("data: [DONE]\n\n")}, }} - r := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI) + r := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI, false) r.usageSink = &chunkUsageHolder{} got, err := io.ReadAll(r) if err != nil { @@ -1026,7 +1026,7 @@ func TestProxyCache_ChunkSSEReader_DeltaFallback(t *testing.T) { {Delta: "hello"}, {Done: true}, }} - r := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI) + r := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI, false) r.usageSink = &chunkUsageHolder{} got, _ := io.ReadAll(r) // Should synthesize {"choices":[{"delta":{"content":"hello"}}]} @@ -1036,7 +1036,7 @@ func TestProxyCache_ChunkSSEReader_DeltaFallback(t *testing.T) { } func TestProxyCache_ChunkSSEReader_NilSub(t *testing.T) { - r := newChunkSSEReaderFromSubscription(context.Background(), nil, nil, provcore.FormatOpenAI) + r := newChunkSSEReaderFromSubscription(context.Background(), nil, nil, provcore.FormatOpenAI, false) buf := make([]byte, 10) n, err := r.Read(buf) if n != 0 || !errors.Is(err, io.EOF) { @@ -1046,7 +1046,7 @@ func TestProxyCache_ChunkSSEReader_NilSub(t *testing.T) { func TestProxyCache_ChunkSSEReader_ProviderErrorSynthesizesFrame(t *testing.T) { sub := &fakeChunkSub{err: errors.New("upstream blew up")} - r := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI) + r := newChunkSSEReaderFromSubscription(context.Background(), sub, nil, provcore.FormatOpenAI, false) r.usageSink = &chunkUsageHolder{} got, _ := io.ReadAll(r) // Must contain an SSE frame with the error message @@ -1062,7 +1062,7 @@ func TestProxyCache_ChunkSSEReader_ContextCanceled(t *testing.T) { ctx, cancel := context.WithCancel(context.Background()) cancel() sub := &fakeChunkSub{err: context.Canceled} - r := newChunkSSEReaderFromSubscription(ctx, sub, nil, provcore.FormatOpenAI) + r := newChunkSSEReaderFromSubscription(ctx, sub, nil, provcore.FormatOpenAI, false) r.usageSink = &chunkUsageHolder{} buf := make([]byte, 10) _, err := r.Read(buf) diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_hooks.go b/packages/ai-gateway/internal/ingress/proxy/proxy_hooks.go index 2220d318..3914956a 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_hooks.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_hooks.go @@ -17,6 +17,67 @@ import ( // helpers + the audit tag-set + blocking-rule mappers split out of proxy.go (behavior // unchanged). The orchestrator runRequestHooks stays in proxy.go with the request flow. +// hookKey identifies a hook across repeated response-stage scans of one stream +// (live checkpoints / Model A confirms). Keyed by stable identity — NOT the +// per-scan slice index Order — so a hook is never split into two rows if the +// per-checkpoint pipeline rebuild ever shifts ordering. +type hookKey struct { + hookID string + implID string +} + +// responseHookAccumulator folds the per-hook results of multiple response-stage +// scans into ONE record per hook: LatencyMs/LatencyUs are SUMMED (the real scan +// CPU spent across the stream) and decision/reason/order are overwritten by the +// latest scan (the last checkpoint that ran is authoritative). The streaming +// response pipeline runs once per checkpoint/confirm, so appending each scan's +// results directly would write a hook N times and N×-inflate the response hook +// aggregates; this accumulator is the single-append guard. Single-writer (the +// live checkpoint runs on the relay's main goroutine; Model A on its one loop +// goroutine) so it needs no lock, and the map is lazily allocated so a stream +// that produces no response-hook results pays nothing. +type responseHookAccumulator struct { + order []hookKey + byKey map[hookKey]hookcore.HookResult +} + +func (a *responseHookAccumulator) add(results []hookcore.HookResult) { + if len(results) == 0 { + return + } + if a.byKey == nil { + a.byKey = make(map[hookKey]hookcore.HookResult, len(results)) + } + for _, r := range results { + k := hookKey{hookID: r.HookID, implID: r.ImplementationID} + if cur, ok := a.byKey[k]; ok { + // Latest scan is authoritative for ALL fields; only the two latency + // counters accumulate. Copying the whole struct (rather than a chosen + // subset) keeps any future-read field — Tags, Action, BlockingRule — in + // sync with the latest scan instead of serving a stale first-scan value. + r.LatencyMs += cur.LatencyMs + r.LatencyUs += cur.LatencyUs + a.byKey[k] = r + } else { + a.byKey[k] = r + a.order = append(a.order, k) + } + } +} + +// finalize returns the folded per-hook results in first-seen order, or nil when +// nothing was accumulated (so appendHookTrace is a no-op). +func (a *responseHookAccumulator) finalize() []hookcore.HookResult { + if len(a.order) == 0 { + return nil + } + out := make([]hookcore.HookResult, 0, len(a.order)) + for _, k := range a.order { + out = append(out, a.byKey[k]) + } + return out +} + func appendHookTrace(existing []audit.HookExecRecord, stage string, results []hookcore.HookResult) []audit.HookExecRecord { if len(results) == 0 { return existing @@ -32,6 +93,7 @@ func appendHookTrace(existing []audit.HookExecRecord, stage string, results []ho Reason: r.Reason, ReasonCode: r.ReasonCode, LatencyMs: r.LatencyMs, + LatencyUs: r.LatencyUs, Error: r.Error, }) } diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_modify_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_modify_test.go index 9ca8138d..03d1b5fa 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_modify_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_modify_test.go @@ -118,6 +118,151 @@ func TestRunRequestHooks_Modify_RewritesBody(t *testing.T) { } } +// newBlockSoftPlusRedactRequestHookCache is the REQUEST-stage co-firing analog of +// newBlockSoftPlusRedactResponseHookCache: a PII-redact hook (inflight+storage redact) +// AND an always-on soft-block hook, both at Stage "request". The pipeline aggregator +// ranks BlockSoft above Modify (StrictestDecision), so the resolved Decision is +// BLOCK_SOFT while the redact's ModifiedContent/TransformSpans ride along — exactly the +// co-firing shape that gated on Decision==Modify would have forwarded UNREDACTED upstream +// before #13. softBlockHook is the package-level impl shared with the response harness. +func newBlockSoftPlusRedactRequestHookCache(t *testing.T) *compliance.HookConfigCache { + t.Helper() + reg := builtins.Registry.Clone() + reg.Register("softblock-req-hook", func(_ *goHooks.HookConfig) (goHooks.Hook, error) { + return softBlockHook{}, nil + }) + reg.Freeze() + loader := func(_ context.Context) ([]goHooks.HookConfig, error) { + return []goHooks.HookConfig{ + { + ID: "pii-req-1", ImplementationID: "pii-detector", Name: "pii-detect-request", + Priority: 10, Enabled: true, Stage: "request", FailBehavior: "fail-closed", TimeoutMs: 1000, + ApplicableIngress: []string{"ALL"}, + Config: map[string]any{ + "onMatch": map[string]any{"inflightAction": "redact", "storageAction": "redact"}, + "patternDefinitions": []any{ + map[string]any{"id": "email", "regex": `[a-z0-9._%+-]+@[a-z0-9.-]+\.[a-z]{2,}`, "flags": "i", "replacement": "[REDACTED_EMAIL]"}, + }, + }, + }, + { + ID: "softblock-req-1", ImplementationID: "softblock-req-hook", Name: "softblock-request", + Priority: 5, Enabled: true, Stage: "request", FailBehavior: "fail-closed", TimeoutMs: 1000, + ApplicableIngress: []string{"ALL"}, Config: map[string]any{}, + }, + }, nil + } + cache := compliance.NewHookConfigCache(loader, reg, 0, slog.Default()) + if err := cache.Start(context.Background()); err != nil { + t.Fatalf("cache.Start: %v", err) + } + time.Sleep(50 * time.Millisecond) + return cache +} + +// TestRunRequestHooks_BlockSoftMaskedRedact_RedactsRequestBody pins #13 leak #4 (ai-gateway +// request side): when a redact hook co-fires with a soft-block hook, the aggregate Decision +// is BLOCK_SOFT, so the pre-#13 gate `Decision==Modify` skipped the rewrite and forwarded the +// ORIGINAL request body upstream. The gate now keys on CarriesRedaction(), so the body is +// redacted; the audit DISPOSITION is stamped Action=redact (not the BlockSoft ceiling) and the +// redacted wire copy reaches RequestBodyRedacted. This asserts the unit-level contract; the +// end-to-end "upstream never sees the original" evidence is in stage_hooks_test.go. +func TestRunRequestHooks_BlockSoftMaskedRedact_RedactsRequestBody(t *testing.T) { + cache := newBlockSoftPlusRedactRequestHookCache(t) + h := &Handler{deps: &Deps{ + HookConfigCache: cache, + TrafficAdapter: &openai.Adapter{}, + Logger: slog.Default(), + }} + + body := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"ping alice@example.com"}]}`) + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(string(body))) + rec := httptest.NewRecorder() + auditRec := &audit.Record{RequestID: "req-cofire"} + + rewritten, result, rejected := h.runRequestHooks(req, rec, auditRec, "req-cofire", body, routingcore.RoutingTarget{}, openAIIngress, nil, slog.Default()) + if rejected { + t.Fatalf("co-firing soft-block+redact must redact-and-forward, not reject; response=%s", rec.Body.String()) + } + if rewritten == nil { + t.Fatal("expected a redacted request body; got nil — the co-firing BlockSoft skipped the rewrite (leak)") + } + // The aggregate decision is the soft-block ceiling, but the redact still applied. + if result == nil || result.Decision != goHooks.BlockSoft { + t.Fatalf("aggregate Decision = %v, want BLOCK_SOFT (soft-block ranks above the redact)", result) + } + got := gjson.GetBytes(rewritten, "messages.0.content").String() + if got != "ping [REDACTED_EMAIL]" { + t.Errorf("forwarded body content = %q, want %q (email must be masked)", got, "ping [REDACTED_EMAIL]") + } + if strings.Contains(string(rewritten), "alice@example.com") { + t.Errorf("forwarded body still carries the original email: %s", rewritten) + } + if !auditRec.HookRewritten { + t.Error("audit.HookRewritten should be true on the applied co-firing redact") + } + // DISPOSITION action is redact even though the aggregate Decision is BLOCK_SOFT. + if auditRec.RequestAction != goHooks.ActionRedact { + t.Errorf("audit.RequestAction = %q, want %q (an applied rewrite is a redact, not the BlockSoft ceiling)", auditRec.RequestAction, goHooks.ActionRedact) + } + if auditRec.HookDecision != string(goHooks.BlockSoft) { + t.Errorf("audit.HookDecision = %q, want %q", auditRec.HookDecision, string(goHooks.BlockSoft)) + } + if string(auditRec.RequestBodyRedacted) != string(rewritten) { + t.Errorf("audit.RequestBodyRedacted = %q, want the redacted wire copy", auditRec.RequestBodyRedacted) + } +} + +// TestRunRequestHooks_StandaloneBlockSoft_Refuses pins the #15 fold-to-block fix on the +// ai-gateway request stage: a standalone soft-block (no co-firing redact) carries no +// applicable redaction, so ActionFromDecision folds it to the block action and the request +// is REFUSED (403, rejected=true), never forwarded. Before the fix the refuse arm keyed on +// Decision==RejectHard only, so a BlockSoft fell through and forwarded the original body. +func TestRunRequestHooks_StandaloneBlockSoft_Refuses(t *testing.T) { + reg := builtins.Registry.Clone() + reg.Register("softblock-req-hook", func(_ *goHooks.HookConfig) (goHooks.Hook, error) { + return softBlockHook{}, nil + }) + reg.Freeze() + loader := func(_ context.Context) ([]goHooks.HookConfig, error) { + return []goHooks.HookConfig{ + { + ID: "softblock-req-1", ImplementationID: "softblock-req-hook", Name: "softblock-request", + Priority: 5, Enabled: true, Stage: "request", FailBehavior: "fail-closed", TimeoutMs: 1000, + ApplicableIngress: []string{"ALL"}, Config: map[string]any{}, + }, + }, nil + } + cache := compliance.NewHookConfigCache(loader, reg, 0, slog.Default()) + if err := cache.Start(context.Background()); err != nil { + t.Fatalf("cache.Start: %v", err) + } + time.Sleep(50 * time.Millisecond) + + h := &Handler{deps: &Deps{HookConfigCache: cache, TrafficAdapter: &openai.Adapter{}, Logger: slog.Default()}} + body := []byte(`{"model":"gpt-4","messages":[{"role":"user","content":"flag me"}]}`) + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(string(body))) + rec := httptest.NewRecorder() + auditRec := &audit.Record{RequestID: "req-standalone-soft"} + + rewritten, result, rejected := h.runRequestHooks(req, rec, auditRec, "req-standalone-soft", body, routingcore.RoutingTarget{}, openAIIngress, nil, slog.Default()) + if !rejected { + t.Fatalf("standalone request BlockSoft must be refused (fold-to-block), got rejected=false") + } + if rec.Code != http.StatusForbidden { + t.Fatalf("status = %d, want 403", rec.Code) + } + if rewritten != nil { + t.Fatalf("a refused request must not forward a body, got %q", rewritten) + } + if result == nil || result.Decision != goHooks.BlockSoft { + t.Fatalf("aggregate Decision = %v, want BLOCK_SOFT", result) + } + if auditRec.RequestAction != goHooks.ActionBlock { + t.Errorf("audit.RequestAction = %q, want %q (BlockSoft folds to block)", auditRec.RequestAction, goHooks.ActionBlock) + } +} + func TestRunRequestHooks_NoHooks_ReturnsOriginalBody(t *testing.T) { loader := func(_ context.Context) ([]goHooks.HookConfig, error) { return nil, nil } cache := compliance.NewHookConfigCache(loader, builtins.Registry, 0, slog.Default()) diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_residuals_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_residuals_test.go index f27ff5f8..3a058e50 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_residuals_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_residuals_test.go @@ -211,7 +211,7 @@ func TestChunkSSEReader_TranscoderError_PropagatesError(t *testing.T) { sub := &fakeChunkSub{chunks: []provcore.Chunk{ {Delta: "hello"}, // non-Done chunk; transcoder fires }} - r := newChunkSSEReaderFromSubscription(context.Background(), sub, errorTranscoder{}, provcore.FormatOpenAI) + r := newChunkSSEReaderFromSubscription(context.Background(), sub, errorTranscoder{}, provcore.FormatOpenAI, false) r.usageSink = &chunkUsageHolder{} buf := make([]byte, 64) n, err := r.Read(buf) @@ -242,7 +242,7 @@ func TestChunkSSEReader_TranscoderDoneChunk_EmitsFrame(t *testing.T) { sub := &fakeChunkSub{chunks: []provcore.Chunk{ {Done: true, RawBytes: []byte("data: [DONE]\n\n")}, }} - r := newChunkSSEReaderFromSubscription(context.Background(), sub, terminalTranscoder{}, provcore.FormatAnthropic) + r := newChunkSSEReaderFromSubscription(context.Background(), sub, terminalTranscoder{}, provcore.FormatAnthropic, false) r.usageSink = &chunkUsageHolder{} got, _ := io.ReadAll(r) if !strings.Contains(string(got), "message_stop") { @@ -271,7 +271,7 @@ func TestChunkSSEReader_TranscoderSkipsChunk_ReturnsZeroBytes(t *testing.T) { {Delta: "skip-me"}, // transcoder skips this {Done: true}, }} - r := newChunkSSEReaderFromSubscription(context.Background(), sub, skipTranscoder{}, provcore.FormatAnthropic) + r := newChunkSSEReaderFromSubscription(context.Background(), sub, skipTranscoder{}, provcore.FormatAnthropic, false) r.usageSink = &chunkUsageHolder{} // First Read: transcoder skips the first chunk → returns 0, nil. buf := make([]byte, 64) diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_responses.go b/packages/ai-gateway/internal/ingress/proxy/proxy_responses.go index d40ff429..6d44cda4 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_responses.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_responses.go @@ -466,6 +466,11 @@ type chunkSSEReader struct { err error transcoder canonicalbridge.StreamTranscoder // non-nil for cross-format; nil for passthrough ingressFormat provcore.Format // ingress wire shape; drives SSE error-frame envelope (G4) + // allowVerbatim forwards a chunk's RawBytes byte-for-byte when the decode + // session marked it Verbatim (genuine Responses upstream on a non-enforced + // /v1/responses ingress), bypassing the transcoder so built-in-tool / audio + // events survive. Off everywhere else, so a Verbatim flag is ignored. + allowVerbatim bool // termErr publishes the reader's terminal failure (if any) for the // post-pump audit stamp. nil = the stream reached a clean EOF. It is // written once from the reader goroutine on the terminal Read and read @@ -498,8 +503,8 @@ func (r *chunkSSEReader) terminalError() *streamTerminalError { return r.termErr.Load() } -func newChunkSSEReaderFromSubscription(ctx context.Context, sub streamcache.ChunkSubscription, transcoder canonicalbridge.StreamTranscoder, ingressFormat provcore.Format) *chunkSSEReader { - return &chunkSSEReader{ctx: ctx, sub: sub, transcoder: transcoder, ingressFormat: ingressFormat} +func newChunkSSEReaderFromSubscription(ctx context.Context, sub streamcache.ChunkSubscription, transcoder canonicalbridge.StreamTranscoder, ingressFormat provcore.Format, allowVerbatim bool) *chunkSSEReader { + return &chunkSSEReader{ctx: ctx, sub: sub, transcoder: transcoder, ingressFormat: ingressFormat, allowVerbatim: allowVerbatim} } func (r *chunkSSEReader) Read(p []byte) (int, error) { @@ -558,23 +563,39 @@ func (r *chunkSSEReader) Read(p []byte) (int, error) { r.usageSink.record(chunk.Usage) } + // verbatim forwards the upstream frame byte-for-byte, bypassing the + // transcoder, when the decode session marked the chunk Verbatim and this + // lane allows it (non-enforced /v1/responses passthrough). It preserves + // built-in-tool / audio events the canonical waist cannot represent. + verbatim := r.allowVerbatim && chunk.Verbatim && len(chunk.RawBytes) > 0 + switch { case chunk.Done: - // Cross-format: transcoder synthesises the ingress-format terminal - // events (e.g. Anthropic message_stop, Gemini finishReason frame). - // Passthrough: forward the provider's raw terminal frame so that - // native ingress clients (Anthropic SDK, Gemini SDK, etc.) receive - // the typed terminator they expect. - if r.transcoder != nil { + // Verbatim: forward the provider's real terminal frame (response.completed + // with its full payload). Cross-format: transcoder synthesises the + // ingress-format terminal events (e.g. Anthropic message_stop, Gemini + // finishReason frame). Passthrough: forward the provider's raw terminal + // frame so native ingress clients receive the typed terminator they expect. + switch { + case verbatim: + r.scratch = append(r.scratch[:0], chunk.RawBytes...) + r.buf = r.scratch + case r.transcoder != nil: b, _ := r.transcoder.Write(r.ctx, chunk) if len(b) > 0 { r.buf = b } - } else if len(chunk.RawBytes) > 0 { + case len(chunk.RawBytes) > 0: r.scratch = append(r.scratch[:0], chunk.RawBytes...) r.buf = r.scratch } r.closed = true + case verbatim: + // Genuine Responses frame on the non-enforced passthrough lane: forward + // the original bytes (built-in-tool / audio events included) instead of + // re-encoding the decoded canonical fields. + r.scratch = append(r.scratch[:0], chunk.RawBytes...) + r.buf = r.scratch case r.transcoder != nil: // Cross-format: delegate all non-Done chunks to the transcoder so // provider-native RawBytes are never forwarded to the client. diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_upstream.go b/packages/ai-gateway/internal/ingress/proxy/proxy_upstream.go index 029946d9..d306a71e 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_upstream.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_upstream.go @@ -256,12 +256,22 @@ func (h *Handler) egressReshapeNonStream(ingress Ingress, target routingcore.Rou return body, nil } if ingress.WireShape == typology.WireShapeOpenAIResponses { - // Responses ingress: native passthrough already in shape; cross-format - // re-encodes canonical chat → Responses output[] via the bridge. - if h.deps.CanonicalBridge.TargetNativelyServesResponsesAPI(provcore.Format(target.AdapterType)) { + // Content-authoritative egress: the wire shape follows the ACTUAL + // response bytes, never the target Format. A Responses-shape body is + // already in the client's shape (verbatim — zero-loss for built-in + // tools); a chat.completion body is canonical and re-encodes to the + // Responses output[] grammar via EncodeResponsesResponse. An + // unclassifiable body must NEVER be forwarded verbatim to a + // /v1/responses client — fail closed with a 502 so a chat-shaped or + // garbage reply can't leak in the wrong wire shape. + switch openairesponses.ClassifyNonStreamBody(body) { + case openairesponses.ClassResponses: return body, nil + case openairesponses.ClassChat: + return h.deps.CanonicalBridge.ResponseCanonicalToIngress(ingress.BodyFormat, body) + default: + return nil, fmt.Errorf("egress: unclassifiable /v1/responses upstream body; refusing verbatim passthrough") } - return h.deps.CanonicalBridge.ResponseCanonicalToIngress(ingress.BodyFormat, body) } if ingress.BodyFormat.IsOpenAIFamily() { // Canonical == OpenAI shape == the caller's shape. Identity. diff --git a/packages/ai-gateway/internal/ingress/proxy/response_hook_accumulator_test.go b/packages/ai-gateway/internal/ingress/proxy/response_hook_accumulator_test.go new file mode 100644 index 00000000..e8ad81aa --- /dev/null +++ b/packages/ai-gateway/internal/ingress/proxy/response_hook_accumulator_test.go @@ -0,0 +1,106 @@ +package proxy + +import ( + "testing" + + hookcore "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// TestResponseHookAccumulator_FoldsRepeatedScans verifies that scanning the same +// hook across multiple checkpoints folds to ONE record with summed latency and the +// latest decision — the core guard against N× inflation of the response aggregates. +func TestResponseHookAccumulator_FoldsRepeatedScans(t *testing.T) { + var acc responseHookAccumulator + acc.add([]hookcore.HookResult{{Order: 0, HookID: "h1", HookName: "quality", Decision: hookcore.Approve, LatencyMs: 0, LatencyUs: 120}}) + acc.add([]hookcore.HookResult{{Order: 0, HookID: "h1", HookName: "quality", Decision: hookcore.Approve, LatencyMs: 0, LatencyUs: 130}}) + acc.add([]hookcore.HookResult{{Order: 0, HookID: "h1", HookName: "quality", Decision: hookcore.RejectHard, Reason: "blocked", LatencyMs: 1, LatencyUs: 1400}}) + + got := acc.finalize() + if len(got) != 1 { + t.Fatalf("expected 1 folded row, got %d", len(got)) + } + if got[0].LatencyUs != 120+130+1400 { + t.Fatalf("LatencyUs not summed: got %d want %d", got[0].LatencyUs, 120+130+1400) + } + if got[0].LatencyMs != 1 { + t.Fatalf("LatencyMs not summed: got %d want 1", got[0].LatencyMs) + } + if got[0].Decision != hookcore.RejectHard || got[0].Reason != "blocked" { + t.Fatalf("latest decision/reason not kept: %+v", got[0]) + } +} + +// TestResponseHookAccumulator_KeyIgnoresOrder proves the key uses stable identity, +// not the volatile per-scan Order — so a hook whose pipeline-rebuild index shifts +// still folds into one row (Order is carried as the latest output value). +func TestResponseHookAccumulator_KeyIgnoresOrder(t *testing.T) { + var acc responseHookAccumulator + acc.add([]hookcore.HookResult{{Order: 0, HookID: "h1", LatencyUs: 100}}) + acc.add([]hookcore.HookResult{{Order: 2, HookID: "h1", LatencyUs: 50}}) + + got := acc.finalize() + if len(got) != 1 { + t.Fatalf("order drift split the hook into %d rows", len(got)) + } + if got[0].LatencyUs != 150 { + t.Fatalf("LatencyUs=%d want 150", got[0].LatencyUs) + } + if got[0].Order != 2 { + t.Fatalf("Order=%d want 2 (latest scan wins)", got[0].Order) + } +} + +// TestResponseHookAccumulator_DistinctHooksKeepOrder verifies distinct hooks stay +// separate and emit in first-seen order, summed across scans. +func TestResponseHookAccumulator_DistinctHooksKeepOrder(t *testing.T) { + var acc responseHookAccumulator + acc.add([]hookcore.HookResult{{HookID: "a", LatencyUs: 10}, {HookID: "b", LatencyUs: 20}}) + acc.add([]hookcore.HookResult{{HookID: "a", LatencyUs: 5}, {HookID: "b", LatencyUs: 7}}) + + got := acc.finalize() + if len(got) != 2 { + t.Fatalf("expected 2 rows, got %d", len(got)) + } + if got[0].HookID != "a" || got[1].HookID != "b" { + t.Fatalf("first-seen order not preserved: %+v", got) + } + if got[0].LatencyUs != 15 || got[1].LatencyUs != 27 { + t.Fatalf("sums wrong: %+v", got) + } +} + +// TestResponseHookAccumulator_ImplementationIDDisambiguates ensures two bindings of +// the same ImplementationID with distinct HookIDs do not collapse together. +func TestResponseHookAccumulator_ImplementationIDDisambiguates(t *testing.T) { + var acc responseHookAccumulator + acc.add([]hookcore.HookResult{ + {HookID: "x", ImplementationID: "pii", LatencyUs: 1}, + {HookID: "y", ImplementationID: "pii", LatencyUs: 2}, + }) + if got := acc.finalize(); len(got) != 2 { + t.Fatalf("distinct hookIds sharing implId collapsed: %d rows", len(got)) + } +} + +// TestResponseHookAccumulator_EmptyAndSixtyThreeFold reproduces the observed +// "RESPONSE PIPELINE (63)" duplication: 63 identical checkpoint scans must fold to +// ONE appended row with summed microsecond latency — not 63 rows. +func TestResponseHookAccumulator_EmptyAndSixtyThreeFold(t *testing.T) { + var acc responseHookAccumulator + if acc.finalize() != nil { + t.Fatal("empty finalize should be nil") + } + for range 63 { + acc.add([]hookcore.HookResult{{HookID: "h1", HookName: "response-quality-signals", Decision: hookcore.Approve, LatencyUs: 100}}) + } + rows := appendHookTrace(nil, "response", acc.finalize()) + if len(rows) != 1 { + t.Fatalf("63 scans must fold to 1 row, got %d (N× inflation)", len(rows)) + } + if rows[0].LatencyUs != 6300 { + t.Fatalf("summed LatencyUs=%d want 6300", rows[0].LatencyUs) + } + if rows[0].Stage != "response" || rows[0].HookID != "h1" { + t.Fatalf("unexpected row: %+v", rows[0]) + } +} diff --git a/packages/ai-gateway/internal/ingress/proxy/responses_egress_relay_test.go b/packages/ai-gateway/internal/ingress/proxy/responses_egress_relay_test.go new file mode 100644 index 00000000..9c9ec1e0 --- /dev/null +++ b/packages/ai-gateway/internal/ingress/proxy/responses_egress_relay_test.go @@ -0,0 +1,123 @@ +// Package handler — responses_egress_relay_test.go covers the streaming egress +// pieces of the /v1/responses fix: the enforcement re-emit encoder must match +// the ingress shape (a /v1/responses stream emits response.* events, never +// chat.completion.chunk), and the live relay must forward a Verbatim chunk's +// RawBytes byte-for-byte only on the allowed passthrough lane. +package proxy + +import ( + "bytes" + "context" + "errors" + "io" + "testing" + + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/cache/stream" + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/execution/canonicalbridge" + provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" +) + +// TestFallbackStreamEncoder_IngressShape proves the enforcement re-emit encoder +// (buffer / Model A nil-transcoder fallback) is ingress-shape-aware: a +// /v1/responses stream emits Responses typed events and NEVER chat.completion +// chunks, while every other ingress keeps the chat-completions encoder. +func TestFallbackStreamEncoder_IngressShape(t *testing.T) { + ctx := context.Background() + render := func(ingress provcore.Format) string { + enc := fallbackStreamEncoder(ingress, "gpt-4o") + var out bytes.Buffer + b, _ := enc.Write(ctx, provcore.Chunk{Delta: "hi"}) + out.Write(b) + b, _ = enc.Write(ctx, provcore.Chunk{Done: true, FinishReason: "stop"}) + out.Write(b) + return out.String() + } + + resp := render(provcore.FormatOpenAIResponses) + if !bytes.Contains([]byte(resp), []byte("event: response.")) { + t.Fatalf("Responses ingress must emit response.* events; got:\n%s", resp) + } + if bytes.Contains([]byte(resp), []byte("chat.completion.chunk")) { + t.Fatalf("Responses ingress must NEVER emit chat.completion.chunk; got:\n%s", resp) + } + if !bytes.Contains([]byte(resp), []byte("response.completed")) { + t.Fatalf("Responses ingress enforced stream must carry a terminal response.completed; got:\n%s", resp) + } + + chat := render(provcore.FormatOpenAI) + if !bytes.Contains([]byte(chat), []byte("chat.completion.chunk")) { + t.Fatalf("OpenAI chat ingress must emit chat.completion.chunk; got:\n%s", chat) + } +} + +// verbatimSub is a minimal ChunkSubscription emitting a fixed chunk timeline. +type verbatimSub struct { + chunks []provcore.Chunk + i int +} + +func (s *verbatimSub) Next(_ context.Context) (provcore.Chunk, error) { + if s.i >= len(s.chunks) { + return provcore.Chunk{}, io.EOF + } + c := s.chunks[s.i] + s.i++ + return c, nil +} +func (s *verbatimSub) Close() error { return nil } + +var _ streamcache.ChunkSubscription = (*verbatimSub)(nil) + +func readAll(t *testing.T, r io.Reader) string { + t.Helper() + var out bytes.Buffer + buf := make([]byte, 256) + for { + n, err := r.Read(buf) + out.Write(buf[:n]) + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("read: %v", err) + } + } + return out.String() +} + +// TestChunkSSEReader_VerbatimForwarding proves a Verbatim chunk is forwarded +// byte-for-byte on the allowed passthrough lane (built-in-tool events preserved) +// and re-encoded through the transcoder when verbatim is not allowed. +func TestChunkSSEReader_VerbatimForwarding(t *testing.T) { + builtin := []byte("event: response.web_search_call.results\ndata: {\"type\":\"response.web_search_call.results\"}\n\n") + chunksOf := func() []provcore.Chunk { + return []provcore.Chunk{ + {Verbatim: true, RawBytes: builtin, NativeEvent: "response.web_search_call.results"}, + {Verbatim: true, RawBytes: []byte("event: response.completed\ndata: {}\n\n"), Done: true}, + } + } + + // A Responses encoder is the transcoder on both lanes; the only difference + // is whether Verbatim forwarding is allowed. The encoder, fed a built-in + // chunk with no canonical content, never reproduces the web_search bytes — + // so their presence proves the verbatim lane bypassed it. + enc := func() canonicalbridge.StreamTranscoder { return canonicalbridge.NewResponsesStreamEncoder("gpt-4o") } + + // allowVerbatim=true → original bytes forwarded, including the built-in event. + r := newChunkSSEReaderFromSubscription(context.Background(), &verbatimSub{chunks: chunksOf()}, enc(), provcore.FormatOpenAIResponses, true) + r.usageSink = &chunkUsageHolder{} + got := readAll(t, r) + if !bytes.Contains([]byte(got), []byte("response.web_search_call.results")) { + t.Fatalf("allowed verbatim lane must forward the built-in event byte-for-byte; got:\n%s", got) + } + + // allowVerbatim=false → Verbatim ignored; the transcoder governs and drops + // the built-in event (no canonical representation), proving the raw bytes + // were NOT forwarded. + r2 := newChunkSSEReaderFromSubscription(context.Background(), &verbatimSub{chunks: chunksOf()}, enc(), provcore.FormatOpenAIResponses, false) + r2.usageSink = &chunkUsageHolder{} + got2 := readAll(t, r2) + if bytes.Contains([]byte(got2), []byte("response.web_search_call.results")) { + t.Fatalf("verbatim must be ignored when not allowed; got:\n%s", got2) + } +} diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_cache_body.go b/packages/ai-gateway/internal/ingress/proxy/stage_cache_body.go index ecda2f1c..8c0cbc47 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stage_cache_body.go +++ b/packages/ai-gateway/internal/ingress/proxy/stage_cache_body.go @@ -124,7 +124,7 @@ func (st cacheStage) prepareUpstreamBody() (ok bool, prepared bool) { case s.resolved.WireShape == typology.WireShapeOpenAIResponses: // Responses is chat-kind but has its own native-passthrough // rule (only targets that natively serve /v1/responses). - needsCanonicalization = !h.deps.CanonicalBridge.TargetNativelyServesResponsesAPI(targetFmt) + needsCanonicalization = !h.deps.CanonicalBridge.ServesResponses(targetFmt, primary.ServesResponsesAPI) case ingressKind == typology.EndpointKindChat, isEmbeddingsIngress: needsCanonicalization = s.resolved.BodyFormat != targetFmt } diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_context.go b/packages/ai-gateway/internal/ingress/proxy/stage_context.go index 224b3c48..0632dceb 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stage_context.go +++ b/packages/ai-gateway/internal/ingress/proxy/stage_context.go @@ -134,7 +134,13 @@ func (h *Handler) newProxyState(in Ingress, w http.ResponseWriter, r *http.Reque fmt.Sprintf("unknown body format %q; supported: openai, anthropic, gemini, azure-openai, minimax, glm, deepseek", raw), "")) return nil, false } + // endpoint_type is chat-KIND for routing / cache / hook dispatch (via + // KindFromWireShape), but the Responses API carries its own label so the + // persisted traffic_event distinguishes /v1/responses from chat completions. endpointType := string(typology.KindFromWireShape(resolved.WireShape)) + if resolved.WireShape == typology.WireShapeOpenAIResponses { + endpointType = string(typology.EndpointKindResponses) + } // Stamp the effective ingress on the request context so the // VK extractor (vkauth) and format-aware model extractor diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_hooks.go b/packages/ai-gateway/internal/ingress/proxy/stage_hooks.go index 40e3f7e8..bb52fa60 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stage_hooks.go +++ b/packages/ai-gateway/internal/ingress/proxy/stage_hooks.go @@ -204,7 +204,15 @@ func (h *Handler) runRequestHooks(r *http.Request, w http.ResponseWriter, rec *a h.deps.Metrics.RecordHookRequest(ingressFormat, "request", string(hookResult.Decision)) } - if hookResult.Decision == hookcore.RejectHard { + // Refuse on a folded BLOCK action (RejectHard OR a standalone BlockSoft): + // ActionFromDecision folds BlockSoft → block, so the request stage dispatches + // it exactly like RejectHard (a hard 403), matching the SSE/response path and + // error-taxonomy-architecture.md — there is no soft-block client response. + // The `!CarriesRedaction()` guard keeps a BlockSoft that MASKS a co-firing redact + // out of this arm so it falls through to the redaction arm below and redact-forwards + // (the #13 invariant): RejectHard never carries redaction, so its behavior is + // unchanged. + if hookcore.ActionFromDecision(hookResult.Decision) == hookcore.ActionBlock && !hookResult.CarriesRedaction() { // Write X-Nexus-Hook and via before writeError commits the status // line, so the client sees the marker even on hook-rejected 4xx responses. // X-Nexus-Mode is reserved as an empty position so an outer hop's @@ -225,7 +233,7 @@ func (h *Handler) runRequestHooks(r *http.Request, w http.ResponseWriter, rec *a // requested but not actionable here". Any other error (malformed, // unknown schema after Extract succeeded) indicates an internal // inconsistency and surfaces as 500. - if hookResult.Decision == hookcore.Modify && len(hookResult.ModifiedContent) > 0 { + if hookResult.CarriesRedaction() { rewriteStart := time.Now() rewriteContent := rewriteContentWithToolArgs(hookResult.ModifiedContent, normalized, hookResult.TransformSpans) rewritten, n, rErr := trafficAdapter.RewriteRequestBody(r.Context(), body, r.URL.Path, rewriteContent) @@ -272,6 +280,11 @@ func (h *Handler) runRequestHooks(r *http.Request, w http.ResponseWriter, rec *a default: rec.HookRewriteCount = n rec.HookRewritten = true + // Stamp the DISPOSITION action: an applied rewrite is a redact, even when + // the aggregate Decision is BlockSoft (a soft-block masking a co-firing + // redact) — without this the audit row would read Action=block + a masked + // request body. + rec.RequestAction = hookcore.ActionRedact // The redacted wire copy is what the raw storage policy // persists under action=redact (rec.RequestBody holds // the pre-hook bytes for normalization and must never reach diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_hooks_test.go b/packages/ai-gateway/internal/ingress/proxy/stage_hooks_test.go index 39a061f8..6ad02586 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stage_hooks_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/stage_hooks_test.go @@ -127,6 +127,52 @@ func TestServeProxy_RequestHookModify_ForwardsRewrittenBodyUpstream(t *testing.T } } +// TestServeProxy_RequestHookBlockSoftMaskedRedact_ForwardsRedactedUpstream is the +// end-to-end #13 leak-#4 regression: a redact hook co-firing with a soft-block hook +// resolves to Decision=BLOCK_SOFT, which the pre-#13 `Decision==Modify` gate skipped — +// forwarding the caller's ORIGINAL request body (PII intact) upstream. With the gate on +// CarriesRedaction(), the upstream provider must receive the masked body and never the +// original email. +func TestServeProxy_RequestHookBlockSoftMaskedRedact_ForwardsRedactedUpstream(t *testing.T) { + var mu sync.Mutex + var upstreamGot []byte + upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + mu.Lock() + upstreamGot = b + mu.Unlock() + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{ + "id":"x","object":"chat.completion","model":"gpt-4o", + "choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}], + "usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2} + }`)) + })) + defer upstream.Close() + + deps := makeOpenAIDeps(t, upstream.URL, newBlockSoftPlusRedactRequestHookCache(t)) + h := NewHandler(deps).ServeProxy(Ingress{ + WireShape: typology.WireShapeOpenAIChat, + BodyFormat: provcore.FormatOpenAI, + }) + w := httptest.NewRecorder() + h(w, freshChatRequest(t, `{"model":"gpt-4o","messages":[{"role":"user","content":"ping alice@example.com"}]}`)) + + if w.Code != http.StatusOK { + t.Fatalf("status=%d body=%s (co-firing soft-block+redact must redact-and-forward)", w.Code, w.Body.String()) + } + mu.Lock() + got := string(upstreamGot) + mu.Unlock() + if !strings.Contains(got, "[REDACTED_EMAIL]") { + t.Errorf("upstream body=%s want the redacted placeholder forwarded", got) + } + if strings.Contains(got, "alice@example.com") { + t.Errorf("upstream body=%s must NOT carry the original email (the co-firing BlockSoft leak)", got) + } +} + // TestRunRequestHooks_RewriteUnsupported_FailsClosed pins the degraded // Modify path: when the traffic adapter cannot reverse-encode // (ErrRewriteUnsupported) the request is rejected (fail CLOSED) rather than diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_routing.go b/packages/ai-gateway/internal/ingress/proxy/stage_routing.go index 172fe126..b9e66f3b 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stage_routing.go +++ b/packages/ai-gateway/internal/ingress/proxy/stage_routing.go @@ -177,8 +177,9 @@ func (st routingStage) run() bool { if s.resolved.BodyFormat == provcore.FormatOpenAIResponses && len(routeResult.Targets) > 0 && h.deps.CanonicalBridge != nil { - targetFormat := provcore.Format(routeResult.Targets[0].AdapterType) - if !h.deps.CanonicalBridge.TargetNativelyServesResponsesAPI(targetFormat) { + primary := routeResult.Targets[0] + targetFormat := provcore.Format(primary.AdapterType) + if !h.deps.CanonicalBridge.ServesResponses(targetFormat, primary.ServesResponsesAPI) { if rej := validateResponsesIngressForCrossFormat(s.body); rej != nil { h.writeResponsesFeatureRejection(s.w, s.rec, rej) return false diff --git a/packages/ai-gateway/internal/ingress/proxy/stream_context.go b/packages/ai-gateway/internal/ingress/proxy/stream_context.go index 53fa0644..76b893c6 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stream_context.go +++ b/packages/ai-gateway/internal/ingress/proxy/stream_context.go @@ -78,6 +78,12 @@ type streamState struct { streamMaxBufferBytes int transcoder canonicalbridge.StreamTranscoder ingressFormat provcore.Format + // allowVerbatim permits the live relay to forward a chunk's RawBytes + // byte-for-byte when the decode session marked it Verbatim (genuine + // Responses upstream on a /v1/responses ingress). Off on every enforcing + // scope and on the chat-ingress auto-upgrade path, so enforcement and + // cross-ingress re-encoding always win over passthrough. + allowVerbatim bool // Relay outputs (stream_relay.go), read by the accounting stage: // the SSE reader (terminal-error classification) and the usage diff --git a/packages/ai-gateway/internal/ingress/proxy/stream_pipeline_test.go b/packages/ai-gateway/internal/ingress/proxy/stream_pipeline_test.go index 2a1f63f7..9bd50aa6 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stream_pipeline_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/stream_pipeline_test.go @@ -212,7 +212,7 @@ func TestHandleStreamWithSubscription_UnmappedHitOrigin_KeepsDefaultTranscoderSe rec := &audit.Record{EndpointType: "chat", GatewayCacheStatus: audit.GatewayCacheHit} r := openAIIngressRequest(t) r = r.WithContext(WithStreamHitOrigin(r.Context(), StreamHitOrigin{ - WireShape: typology.WireShapeGeminiGenerateContent, // no Format mapping + WireShape: typology.WireShapeBedrockConverse, // no Format mapping (event-stream framing) })) w := httptest.NewRecorder() diff --git a/packages/ai-gateway/internal/ingress/proxy/stream_relay.go b/packages/ai-gateway/internal/ingress/proxy/stream_relay.go index 7e1124e8..8ef7e5f9 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stream_relay.go +++ b/packages/ai-gateway/internal/ingress/proxy/stream_relay.go @@ -55,15 +55,24 @@ func (st streamRelayStage) run() bool { // (which would drain the subscription into the wire reader). It streams // in real time behind a bounded tail + cheap union prescan, escalating to // canonical-buffer redaction on a confirmed hit. The prescan closes over - // the response probe's MayMatchRawContent (the cheap union prefilter). - s.sseReader = h.runModelAStream(r.Context(), s, tee, usageHolder, h.buildResponsePrescan(r.Context(), s)) + // the response probe's MayMatchRawContent (the cheap union prefilter); the + // derived lookahead sizes the flush-before-deliver boundary guard. + prescan, maxPattern := h.buildResponsePrescan(r.Context(), s) + s.sseReader = h.runModelAStream(r.Context(), s, tee, usageHolder, prescan, maxPattern) default: // Drain the subscription (replay or live broker pump) into an // io.Reader of SSE-formatted lines so the wire-consuming pipelines // (live / passthrough) can consume it unchanged. - sseReader := newChunkSSEReaderFromSubscription(r.Context(), s.sub, s.transcoder, s.ingressFormat) + sseReader := newChunkSSEReaderFromSubscription(r.Context(), s.sub, s.transcoder, s.ingressFormat, s.allowVerbatim) sseReader.usageSink = usageHolder + // respAcc folds the response-hook trace across every checkpoint scan into + // ONE record per hook (summed latency, latest decision). The live pipeline + // runs the response stage at each checkpoint + EOF, so without this the trace + // would either be missing (never appended) or N×-duplicated; appended once + // after dispatch (below). Lazy + single-writer — see responseHookAccumulator. + var respAcc responseHookAccumulator + hookCtx := &streaming.StreamHookContext{ RequestID: requestID, IngressType: "AI_GATEWAY", @@ -89,6 +98,10 @@ func (st streamRelayStage) run() bool { // Derived from the Decision so a no-match approve still stamps // ActionApprove (persist) rather than dropping the captured stream. rec.ResponseAction = hookcore.ActionFromDecision(res.Decision) + // Fold this checkpoint's per-hook latency into the single response + // trace appended once after dispatch (real scan CPU summed across + // checkpoints, no N× duplication). + respAcc.add(res.HookResults) }, } @@ -114,6 +127,11 @@ func (st streamRelayStage) run() bool { MaxBufferBytes: s.streamMaxBufferBytes, } dispatchStreamMode(r.Context(), s.streamMode, streamDeps) + // Record the response-hook trace ONCE, after all checkpoints have run. This + // is the live audit-only path's only response-trace write — previously the + // streamed response hooks never entered rec.HooksPipeline, leaving + // response_hooks_ms NULL. finalize() is nil (no-op) when no response hook ran. + rec.HooksPipeline = appendHookTrace(rec.HooksPipeline, "response", respAcc.finalize()) s.sseReader = sseReader } logger.Debug("stream response capture", diff --git a/packages/ai-gateway/internal/ingress/proxy/stream_shape.go b/packages/ai-gateway/internal/ingress/proxy/stream_shape.go index 151556eb..d5830876 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stream_shape.go +++ b/packages/ai-gateway/internal/ingress/proxy/stream_shape.go @@ -116,14 +116,7 @@ func (st streamShapeStage) run() bool { // encoder so the cached canonical chunks are re-encoded into // the request's wire SSE grammar. if originOK && transcoder == nil && originBodyFormat != ingress.BodyFormat { - switch ingress.BodyFormat { - case provcore.FormatOpenAIResponses: - transcoder = canonicalbridge.NewResponsesStreamEncoder(target.ModelCode) - default: - if ingress.BodyFormat.IsOpenAIFamily() { - transcoder = canonicalbridge.NewChatCompletionsStreamEncoder(target.ModelCode) - } - } + transcoder = canonicalbridge.IngressStreamEncoder(ingress.BodyFormat, target.ModelCode) } } } @@ -137,10 +130,22 @@ func (st streamShapeStage) run() bool { transcoder = canonicalbridge.NewChatCompletionsStreamEncoder(target.ModelCode) } + // Verbatim passthrough is allowed only on the non-enforced /v1/responses + // live lane and never on the chat-ingress auto-upgrade path: enforcement + // (which forces the canonical buffer) and cross-ingress re-encoding must + // always win over forwarding the upstream's original Responses frames. + // When allowed, the relay forwards a Verbatim chunk's RawBytes byte-for-byte + // so built-in-tool / audio events survive; otherwise the chunk's decoded + // canonical fields drive the transcoder. + allowVerbatim := ingressFormat == provcore.FormatOpenAIResponses && + !s.responseEnforcingBlock && !s.responseEnforcingRedact && + !ResponsesUpgradeFromContext(r.Context()) + s.emitDone = emitDone s.streamMode = streamMode s.streamMaxBufferBytes = streamMaxBufferBytes s.transcoder = transcoder s.ingressFormat = ingressFormat + s.allowVerbatim = allowVerbatim return true } diff --git a/packages/ai-gateway/internal/platform/audit/enums.go b/packages/ai-gateway/internal/platform/audit/enums.go index 9c43b1c7..8d3de01e 100644 --- a/packages/ai-gateway/internal/platform/audit/enums.go +++ b/packages/ai-gateway/internal/platform/audit/enums.go @@ -19,8 +19,11 @@ type HookExecRecord struct { Decision string `json:"decision"` Reason string `json:"reason,omitempty"` ReasonCode string `json:"reasonCode,omitempty"` - LatencyMs int `json:"latencyMs"` - Error string `json:"error,omitempty"` + // LatencyMs is the truncated integer-ms floor of LatencyUs (backward compat, + // never clamped); LatencyUs is the precise microsecond wall-clock per hook. + LatencyMs int `json:"latencyMs"` + LatencyUs int `json:"latencyUs"` + Error string `json:"error,omitempty"` } // CacheStatus is the UNIFIED rollup recorded on traffic_event.cache_status. diff --git a/packages/ai-gateway/internal/platform/audit/message.go b/packages/ai-gateway/internal/platform/audit/message.go index 955e159e..7b62e5aa 100644 --- a/packages/ai-gateway/internal/platform/audit/message.go +++ b/packages/ai-gateway/internal/platform/audit/message.go @@ -160,6 +160,8 @@ func (w *Writer) recordToMessage(rec *Record) *mq.TrafficEventMessage { UpstreamTotalMs: rec.UpstreamTotalMs, RequestHooksMs: firstNonNil(rec.RequestHooksMs, sumHookLatenciesMs(rec.HooksPipeline, "request", "connection")), ResponseHooksMs: firstNonNil(rec.ResponseHooksMs, sumHookLatenciesMs(rec.HooksPipeline, "response")), + RequestHooksUs: firstNonNil(rec.RequestHooksUs, sumHookLatenciesUs(rec.HooksPipeline, "request", "connection")), + ResponseHooksUs: firstNonNil(rec.ResponseHooksUs, sumHookLatenciesUs(rec.HooksPipeline, "response")), LatencyBreakdown: rec.LatencyBreakdown, } // Gateway response cache savings. diff --git a/packages/ai-gateway/internal/platform/audit/phase_g_batch_test.go b/packages/ai-gateway/internal/platform/audit/phase_g_batch_test.go index 742e9148..2ca00112 100644 --- a/packages/ai-gateway/internal/platform/audit/phase_g_batch_test.go +++ b/packages/ai-gateway/internal/platform/audit/phase_g_batch_test.go @@ -151,6 +151,159 @@ func TestConsumeLoop_CapsBatchAtMaxCount(t *testing.T) { } } +// TestConsumeLoop_LingerPublishesPartialBatch deterministically exercises the linger +// path: a partial batch (< batchMaxCount) that receives no more records publishes when +// the consumerLinger timer fires (the `case <-timer.C` branch + the arm branch). Without +// closing stopCh, the ONLY way the partial batch can publish is the linger timer, so a +// non-empty batch appearing proves that branch ran — no reliance on incidental timing. +func TestConsumeLoop_LingerPublishesPartialBatch(t *testing.T) { + bp := &batchMemProducer{} + w := NewWriter(bp, "q", nil, slog.Default()) + const n = 3 // < batchMaxCount, so it lingers rather than publishing on a full batch + w.recCh = make(chan *Record, n) + for i := range n { + w.recCh <- &Record{RequestID: fmt.Sprintf("r%d", i), Timestamp: time.Now()} + } + w.wg.Add(1) + go w.consumeLoop(0) + + // The partial batch can ONLY publish via the linger timer here (no full batch, no + // stop). Wait comfortably longer than consumerLinger (100ms) for it. + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + bp.mu.Lock() + got := len(bp.batches) + bp.mu.Unlock() + if got >= 1 { + break + } + time.Sleep(5 * time.Millisecond) + } + close(w.stopCh) + w.wg.Wait() + + bp.mu.Lock() + defer bp.mu.Unlock() + if len(bp.batches) == 0 { + t.Fatal("partial batch never published — the linger timer branch did not fire") + } + if len(bp.batches[0]) != n { + t.Fatalf("linger batch = %d records, want %d", len(bp.batches[0]), n) + } +} + +// TestDrainOnStop_PublishesRemainderBoundedByMaxCount deterministically covers the +// shutdown drain (consumeLoop's stopCh case): the remainder is published, bounded at +// batchMaxCount per publish. Driving drainOnStop directly avoids the main-select race +// that made this branch flaky under CI when exercised through consumeLoop. +func TestDrainOnStop_PublishesRemainderBoundedByMaxCount(t *testing.T) { + t.Run("partial remainder in one publish", func(t *testing.T) { + bp := &batchMemProducer{} + w := NewWriter(bp, "q", nil, slog.Default()) + const n = 5 + w.recCh = make(chan *Record, n) + for i := range n { + w.recCh <- &Record{RequestID: fmt.Sprintf("r%d", i), Timestamp: time.Now()} + } + w.drainOnStop(0, nil) + bp.mu.Lock() + defer bp.mu.Unlock() + if len(bp.batches) != 1 || len(bp.batches[0]) != n { + t.Fatalf("drainOnStop should publish one batch of %d, got %d batches", n, len(bp.batches)) + } + }) + + t.Run("queue over batchMaxCount splits, capped per publish", func(t *testing.T) { + bp := &batchMemProducer{} + w := NewWriter(bp, "q", nil, slog.Default()) + n := batchMaxCount + 88 // forces the in-drain full-batch publish + a final remainder + w.recCh = make(chan *Record, n) + for i := range n { + w.recCh <- &Record{RequestID: fmt.Sprintf("r%d", i), Timestamp: time.Now()} + } + w.drainOnStop(0, nil) + bp.mu.Lock() + defer bp.mu.Unlock() + total := 0 + for _, b := range bp.batches { + if len(b) > batchMaxCount { + t.Fatalf("a drain publish carried %d > cap %d", len(b), batchMaxCount) + } + total += len(b) + } + if total != n { + t.Fatalf("drainOnStop lost records: published %d, want %d", total, n) + } + if len(bp.batches) < 2 || len(bp.batches[0]) != batchMaxCount { + t.Fatalf("expected a full %d-record publish then the remainder, got batches=%d first=%d", batchMaxCount, len(bp.batches), len(bp.batches[0])) + } + }) + + t.Run("lingering batch published even with empty queue", func(t *testing.T) { + bp := &batchMemProducer{} + w := NewWriter(bp, "q", nil, slog.Default()) + w.recCh = make(chan *Record, 1) + w.drainOnStop(0, []*Record{{RequestID: "carried", Timestamp: time.Now()}}) + bp.mu.Lock() + defer bp.mu.Unlock() + if len(bp.batches) != 1 || len(bp.batches[0]) != 1 { + t.Fatalf("a lingering batch must publish on stop even with an empty queue, got %d batches", len(bp.batches)) + } + }) +} + +// TestConsumeLoop_FullBatchWhileArmedStopsTimer covers the publish()'s armed-cleanup +// branch: a partial batch arms the linger timer, then enough records arrive to fill a +// full batch — publishing it must Stop the armed timer (the `if armed` branch). The +// initial single record arms the timer (sleep ensures it armed before the flood), then a +// batchMaxCount flood drives the greedy-drain to a full publish while armed. +func TestConsumeLoop_FullBatchWhileArmedStopsTimer(t *testing.T) { + bp := &batchMemProducer{} + w := NewWriter(bp, "q", nil, slog.Default()) + w.recCh = make(chan *Record, batchMaxCount+8) + w.wg.Add(1) + go w.consumeLoop(0) + + // One record: the consumer absorbs it and ARMS the linger timer (1 < batchMaxCount). + w.recCh <- &Record{RequestID: "arm", Timestamp: time.Now()} + time.Sleep(30 * time.Millisecond) // << consumerLinger(100ms): armed, not yet fired + + // Flood a full batch: the recCh arm's greedy drain reaches batchMaxCount and publishes + // WHILE armed → publish() takes the `if armed { timer.Stop() }` branch. + for i := range batchMaxCount { + w.recCh <- &Record{RequestID: fmt.Sprintf("f%d", i), Timestamp: time.Now()} + } + deadline := time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + bp.mu.Lock() + var full bool + for _, b := range bp.batches { + if len(b) == batchMaxCount { + full = true + } + } + bp.mu.Unlock() + if full { + break + } + time.Sleep(5 * time.Millisecond) + } + close(w.stopCh) + w.wg.Wait() + + bp.mu.Lock() + defer bp.mu.Unlock() + var sawFull bool + for _, b := range bp.batches { + if len(b) == batchMaxCount { + sawFull = true + } + } + if !sawFull { + t.Fatalf("expected a full %d-record publish while the linger timer was armed", batchMaxCount) + } +} + // TestPublishBatch_ChunksByBytes: with a low byte cap the chunk boundary is // driven by accumulated marshaled bytes, not count. func TestPublishBatch_ChunksByBytes(t *testing.T) { diff --git a/packages/ai-gateway/internal/platform/audit/record.go b/packages/ai-gateway/internal/platform/audit/record.go index 83840ca6..ae17f02b 100644 --- a/packages/ai-gateway/internal/platform/audit/record.go +++ b/packages/ai-gateway/internal/platform/audit/record.go @@ -245,10 +245,15 @@ type Record struct { // LatencyBreakdown — Long-tail per-source phase durations (ms). // For ai-gateway: auth_ms, quota_ms, routing_ms, // cache_lookup_ms, req_adapter_ms, resp_adapter_ms. - UpstreamTtfbMs *int - UpstreamTotalMs *int - RequestHooksMs *int - ResponseHooksMs *int + UpstreamTtfbMs *int + UpstreamTotalMs *int + RequestHooksMs *int + ResponseHooksMs *int + // Microsecond-precision hook aggregates (additive; siblings of the _ms fields). + // Normally nil and derived from HooksPipeline at recordToMessage time; + // firstNonNil lets a caller stamp them explicitly. + RequestHooksUs *int + ResponseHooksUs *int LatencyBreakdown map[string]int // EmbeddingCostUsd is stamped on every L1 miss that triggered an embedding diff --git a/packages/ai-gateway/internal/platform/audit/record_message_helpers.go b/packages/ai-gateway/internal/platform/audit/record_message_helpers.go index 906bd3cb..31025efb 100644 --- a/packages/ai-gateway/internal/platform/audit/record_message_helpers.go +++ b/packages/ai-gateway/internal/platform/audit/record_message_helpers.go @@ -46,15 +46,13 @@ func firstNonNil(ps ...*int) *int { return nil } -// sumHookLatenciesMs returns the aggregate per-hook latency (ms) for the -// hook rows whose Stage matches one of `stages`. Returns nil when no hook -// in the requested stages ran — distinguished from zero so the resulting -// `request_hooks_ms` / `response_hooks_ms` columns stay NULL for bypass / -// no-hook requests (P95 queries should not count those as 0ms). -// -// Used by recordToMessage to populate the hook-aggregate columns from -// the existing per-hook latency data already in rec.HooksPipeline. -func sumHookLatenciesMs(in []HookExecRecord, stages ...string) *int { +// sumHookLatencies returns the aggregate per-hook latency for the hook rows whose +// Stage matches one of `stages`, reading the value via `pick`. Returns nil when no +// hook in the requested stages ran — distinguished from zero so the resulting +// aggregate columns stay NULL for bypass / no-hook requests (P95 queries should +// not count those as 0). Used by recordToMessage to populate the hook-aggregate +// columns from the per-hook latency already in rec.HooksPipeline. +func sumHookLatencies(in []HookExecRecord, pick func(HookExecRecord) int, stages ...string) *int { if len(in) == 0 || len(stages) == 0 { return nil } @@ -71,8 +69,8 @@ func sumHookLatenciesMs(in []HookExecRecord, stages ...string) *int { continue } ran = true - if r.LatencyMs > 0 { - total += r.LatencyMs + if v := pick(r); v > 0 { + total += v } } if !ran { @@ -81,6 +79,17 @@ func sumHookLatenciesMs(in []HookExecRecord, stages ...string) *int { return &total } +// sumHookLatenciesMs / sumHookLatenciesUs aggregate the truncated-millisecond and +// precise-microsecond per-hook latencies respectively. The _ms columns keep their +// shipped value (sum of truncated per-hook ms); the _us columns carry precision. +func sumHookLatenciesMs(in []HookExecRecord, stages ...string) *int { + return sumHookLatencies(in, func(r HookExecRecord) int { return r.LatencyMs }, stages...) +} + +func sumHookLatenciesUs(in []HookExecRecord, stages ...string) *int { + return sumHookLatencies(in, func(r HookExecRecord) int { return r.LatencyUs }, stages...) +} + // firstNonEmptyStr returns a if non-empty, else b. Used for // target_method / target_path stamping to fall back to the request-side // value when the gateway didn't set a distinct target (transparent path diff --git a/packages/ai-gateway/internal/platform/audit/record_message_helpers_test.go b/packages/ai-gateway/internal/platform/audit/record_message_helpers_test.go new file mode 100644 index 00000000..9860d8c9 --- /dev/null +++ b/packages/ai-gateway/internal/platform/audit/record_message_helpers_test.go @@ -0,0 +1,37 @@ +package audit + +import "testing" + +// TestSumHookLatencies_MsAndUs verifies the dual aggregation: the _ms aggregate +// keeps its shipped meaning (sum of truncated per-hook ms), while the _us aggregate +// carries precise microseconds — so a sub-millisecond hook that truncates to 0 ms +// is still visible in microseconds. The NULL-vs-0 distinction (no hook ran vs ran +// in <1ms) is preserved on both. +func TestSumHookLatencies_MsAndUs(t *testing.T) { + rows := []HookExecRecord{ + {Stage: "request", LatencyMs: 0, LatencyUs: 120}, + {Stage: "request", LatencyMs: 2, LatencyUs: 2100}, + {Stage: "response", LatencyMs: 0, LatencyUs: 300}, + } + + if got := sumHookLatenciesUs(rows, "request"); got == nil || *got != 2220 { + t.Fatalf("request us = %v want 2220", got) + } + if got := sumHookLatenciesMs(rows, "request"); got == nil || *got != 2 { + t.Fatalf("request ms = %v want 2", got) + } + // The precision win: a sub-ms response hook is 0 in ms but precise in us. + if got := sumHookLatenciesUs(rows, "response"); got == nil || *got != 300 { + t.Fatalf("response us = %v want 300", got) + } + if got := sumHookLatenciesMs(rows, "response"); got == nil || *got != 0 { + t.Fatalf("response ms = %v want 0 (present, sub-ms)", got) + } + // No matching stage → nil (NULL), distinguishing "no hook ran" from "ran in 0". + if got := sumHookLatenciesUs(rows, "connection"); got != nil { + t.Fatalf("connection us = %v want nil", got) + } + if got := sumHookLatenciesUs(nil, "request"); got != nil { + t.Fatalf("empty input us = %v want nil", got) + } +} diff --git a/packages/ai-gateway/internal/platform/audit/writer_publish.go b/packages/ai-gateway/internal/platform/audit/writer_publish.go index 2303fc53..eda14d25 100644 --- a/packages/ai-gateway/internal/platform/audit/writer_publish.go +++ b/packages/ai-gateway/internal/platform/audit/writer_publish.go @@ -26,10 +26,12 @@ func (w *Writer) consumeLoop(connIdx int) { timer := time.NewTimer(time.Hour) timer.Stop() armed := false + // publish flushes the current batch on connIdx. It is only ever invoked with a + // NON-EMPTY batch: the recCh arm fires it at len>=batchMaxCount, and the timer.C arm + // fires it only while armed (armed is set only after a record was appended and is + // cleared by every publish, so a fired linger timer always has a pending record). The + // shutdown path's empty-safe flush lives in drainOnStop. publish := func() { - if len(batch) == 0 { - return - } w.publishBatchOn(connIdx, batch) batch = batch[:0] if armed { @@ -64,18 +66,35 @@ func (w *Writer) consumeLoop(connIdx int) { armed = false publish() case <-w.stopCh: - for { - select { - case rec := <-w.recCh: - batch = append(batch, rec) - if len(batch) >= batchMaxCount { - publish() - } - default: - publish() - return - } + w.drainOnStop(connIdx, batch) + return + } + } +} + +// drainOnStop empties the bounded queue into batch on shutdown and publishes the +// remainder on connIdx, bounding each publish at batchMaxCount (same cap as the +// steady-state path). Split out of consumeLoop's stopCh case so the drain is +// unit-testable deterministically: its inner select is recCh-vs-default (records-ready +// wins over default), unlike consumeLoop's main select where stopCh-vs-recCh is random, +// so exercising the drain THROUGH consumeLoop cannot be forced. Logic-equivalent to the +// previous inline loop; timer/armed cleanup is irrelevant on the exit path, and the +// per-publish non-empty guard lives here (so the steady-state publish() never needs it). +// batch is the consumer's current (possibly lingering) batch, consumed by value. +func (w *Writer) drainOnStop(connIdx int, batch []*Record) { + for { + select { + case rec := <-w.recCh: + batch = append(batch, rec) + if len(batch) >= batchMaxCount { + w.publishBatchOn(connIdx, batch) + batch = batch[:0] } + default: + if len(batch) > 0 { + w.publishBatchOn(connIdx, batch) + } + return } } } diff --git a/packages/ai-gateway/internal/platform/store/provider.go b/packages/ai-gateway/internal/platform/store/provider.go index 6800b497..a88c10f4 100644 --- a/packages/ai-gateway/internal/platform/store/provider.go +++ b/packages/ai-gateway/internal/platform/store/provider.go @@ -31,6 +31,12 @@ type Provider struct { APIVersion *string Region *string Enabled bool + // ServesResponsesAPI is the per-provider override for whether this + // upstream natively serves OpenAI /v1/responses. nil = use the adapter + // RequestShapes default; the bridge treats an explicit value as a + // downgrade-only signal (false forces canonical(chat); true cannot + // exceed the adapter's declared capability). + ServesResponsesAPI *bool } // Model represents a model record. @@ -75,13 +81,13 @@ type Model struct { // GetProvider fetches a provider by ID. func (db *DB) GetProvider(ctx context.Context, id string) (*Provider, error) { row := db.pool.QueryRow(ctx, ` - SELECT id, name, "displayName", adapter_type, "baseUrl", "pathPrefix", "apiVersion", region, enabled + SELECT id, name, "displayName", adapter_type, "baseUrl", "pathPrefix", "apiVersion", region, enabled, serves_responses_api FROM "Provider" WHERE id = $1 `, id) var p Provider err := row.Scan(&p.ID, &p.Name, &p.DisplayName, &p.AdapterType, &p.BaseURL, - &p.PathPrefix, &p.APIVersion, &p.Region, &p.Enabled) + &p.PathPrefix, &p.APIVersion, &p.Region, &p.Enabled, &p.ServesResponsesAPI) if err != nil { return nil, fmt.Errorf("store: get provider: %w", err) } diff --git a/packages/ai-gateway/internal/platform/store/provider_test.go b/packages/ai-gateway/internal/platform/store/provider_test.go index 7a8aca0c..24042156 100644 --- a/packages/ai-gateway/internal/platform/store/provider_test.go +++ b/packages/ai-gateway/internal/platform/store/provider_test.go @@ -12,7 +12,7 @@ import ( var providerTestColumns = []string{ "id", "name", "displayName", "adapter_type", "baseUrl", "pathPrefix", - "apiVersion", "region", "enabled", + "apiVersion", "region", "enabled", "serves_responses_api", } func makeProviderRow(id string) []any { @@ -21,7 +21,7 @@ func makeProviderRow(id string) []any { region := "us-east-1" return []any{ id, "openai", &display, "openai", "https://api.openai.com", "/v1", - &apiV, ®ion, true, + &apiV, ®ion, true, (*bool)(nil), } } diff --git a/packages/ai-gateway/internal/providers/core/types.go b/packages/ai-gateway/internal/providers/core/types.go index 0a3c46ff..a34fcffc 100644 --- a/packages/ai-gateway/internal/providers/core/types.go +++ b/packages/ai-gateway/internal/providers/core/types.go @@ -159,6 +159,14 @@ type CallTarget struct { CredentialName string // human-readable name of the credential ProviderModelID string // vendor's model ID (e.g. "claude-3-5-sonnet-20241022") + // ServesResponsesAPI is the per-provider override for whether this + // upstream natively serves OpenAI /v1/responses. nil = adapter + // RequestShapes default; an explicit value is downgrade-only (false + // forces canonical(chat); true cannot exceed the adapter capability). + // Resolved once per target in the executor failover loop from the + // hydrated routing snapshot — never a per-request DB read. + ServesResponsesAPI *bool + // Extras carries provider-specific configuration that doesn't fit in // the universal fields above. Keys are dot-namespaced: "azure.apiVersion", // "aws.accessKey", "gcp.serviceAccountJSON", etc. @@ -302,6 +310,14 @@ type Chunk struct { // joiner, which then stamp usage_extraction_status="truncated" rather // than "ok". Unused on genuine streaming chunks. Truncated bool + // Verbatim marks a chunk whose RawBytes are already in the client's + // egress wire shape and MUST be forwarded byte-for-byte, not re-encoded. + // Set by the /v1/responses content copier for a genuine-Responses upstream + // so built-in-tool / audio events the canonical waist cannot represent + // survive. The live relay honours it only on the non-enforced passthrough + // lane; an enforcing scope forces the canonical buffer (decoded fields). + // Default false → existing decoders unaffected. + Verbatim bool } // ToolCallDelta is a partial OpenAI-canonical tool call patch within a diff --git a/packages/ai-gateway/internal/providers/specs/openai/responses/classify.go b/packages/ai-gateway/internal/providers/specs/openai/responses/classify.go new file mode 100644 index 00000000..188d7d19 --- /dev/null +++ b/packages/ai-gateway/internal/providers/specs/openai/responses/classify.go @@ -0,0 +1,58 @@ +// Package responses — classify.go detects the OpenAI response wire shape from +// the actual response bytes, so the egress side decides how to decode/encode a +// /v1/responses reply from what the upstream really returned rather than from +// the target provider's declared Format. A provider tagged OpenAI may serve the +// chat-completions wire; trusting the bytes (not the Format) is what keeps a +// chat-shaped reply from being forwarded verbatim to a Responses client. +package responses + +import ( + "strings" + + "github.com/tidwall/gjson" +) + +// WireClass is the detected OpenAI response wire shape. +type WireClass int + +const ( + // ClassUnknown means the bytes match neither known shape — a keep-alive + // comment, an empty frame, or garbage. Callers must fail closed (route to + // the canonical/transcode or enforced lane), never to verbatim passthrough. + ClassUnknown WireClass = iota + // ClassResponses is the Responses API shape: a non-stream body whose + // top-level object is "response", or a stream whose events are response.*. + ClassResponses + // ClassChat is the chat-completions shape: object "chat.completion" (non + // stream) or "chat.completion.chunk" (stream). + ClassChat +) + +// ClassifyNonStreamBody classifies a non-streaming response body by its +// top-level "object" discriminator. +func ClassifyNonStreamBody(body []byte) WireClass { + switch gjson.GetBytes(body, "object").String() { + case "response": + return ClassResponses + case "chat.completion": + return ClassChat + default: + return ClassUnknown + } +} + +// ClassifyFirstSSEFrame classifies a stream from its first decoded SSE event. +// eventType is the SSE "event:" line value (may be empty when the grammar +// carries the type inside the data payload); data is the "data:" payload bytes. +func ClassifyFirstSSEFrame(eventType string, data []byte) WireClass { + if strings.HasPrefix(eventType, "response.") { + return ClassResponses + } + if strings.HasPrefix(gjson.GetBytes(data, "type").String(), "response.") { + return ClassResponses + } + if gjson.GetBytes(data, "object").String() == "chat.completion.chunk" { + return ClassChat + } + return ClassUnknown +} diff --git a/packages/ai-gateway/internal/providers/specs/openai/responses/classify_test.go b/packages/ai-gateway/internal/providers/specs/openai/responses/classify_test.go new file mode 100644 index 00000000..01417551 --- /dev/null +++ b/packages/ai-gateway/internal/providers/specs/openai/responses/classify_test.go @@ -0,0 +1,41 @@ +package responses + +import "testing" + +func TestClassifyNonStreamBody(t *testing.T) { + cases := []struct { + name string + body string + want WireClass + }{ + {"responses", `{"object":"response","id":"resp_1"}`, ClassResponses}, + {"chat", `{"object":"chat.completion","id":"x"}`, ClassChat}, + {"unknown_object", `{"object":"embedding"}`, ClassUnknown}, + {"no_object", `{"foo":1}`, ClassUnknown}, + {"empty", ``, ClassUnknown}, + } + for _, c := range cases { + if got := ClassifyNonStreamBody([]byte(c.body)); got != c.want { + t.Fatalf("%s: got %v want %v", c.name, got, c.want) + } + } +} + +func TestClassifyFirstSSEFrame(t *testing.T) { + cases := []struct { + name, ev, data string + want WireClass + }{ + {"event_line", "response.created", `{"type":"response.created"}`, ClassResponses}, + {"type_in_data", "", `{"type":"response.output_text.delta","delta":"hi"}`, ClassResponses}, + {"builtin_tool_event", "response.web_search_call.in_progress", `{}`, ClassResponses}, + {"chat_chunk", "", `{"object":"chat.completion.chunk","choices":[]}`, ClassChat}, + {"keepalive", "", `: keep-alive`, ClassUnknown}, + {"empty", "", ``, ClassUnknown}, + } + for _, c := range cases { + if got := ClassifyFirstSSEFrame(c.ev, []byte(c.data)); got != c.want { + t.Fatalf("%s: got %v want %v", c.name, got, c.want) + } + } +} diff --git a/packages/ai-gateway/internal/providers/specs/openai/stream/stream.go b/packages/ai-gateway/internal/providers/specs/openai/stream/stream.go index dc254d9a..a40d623a 100644 --- a/packages/ai-gateway/internal/providers/specs/openai/stream/stream.go +++ b/packages/ai-gateway/internal/providers/specs/openai/stream/stream.go @@ -29,7 +29,10 @@ func NewStreamDecoder(log *slog.Logger) *StreamDecoder { } // Open wraps body in the right session for the endpoint: -// - EndpointResponsesAPI → responsesStreamSession (Responses-API SSE grammar) +// - EndpointResponsesAPI → responsesEgressSession (content-detecting: forwards +// genuine Responses frames verbatim, decodes a chat.completion upstream into +// canonical so the proxy re-shapes it to Responses events — never leaks the +// wrong wire shape to a /v1/responses client) // - everything else → openaiStreamSession (chat-completions SSE) // // The dispatch is necessary because /v1/responses sends a completely different @@ -41,7 +44,7 @@ func (d *StreamDecoder) Open(body io.ReadCloser, endpoint typology.WireShape) (p return nil, fmt.Errorf("openai: nil stream body") } if endpoint == typology.WireShapeOpenAIResponses { - return &responsesStreamSession{ + return &responsesEgressSession{ scanner: specutil.NewSSEScanner(body), log: d.log, }, nil @@ -94,6 +97,14 @@ func (s *openaiStreamSession) Next(ctx context.Context) (provcore.Chunk, error) break } + return chatChunkFromFrame(ev), nil +} + +// chatChunkFromFrame decodes one non-empty OpenAI chat-completions SSE data +// frame into a canonical Chunk. Extracted so the /v1/responses content copier +// can reuse the exact chat-decode path when an upstream tagged for Responses +// actually returns chat.completion.chunk frames (the bulletproof egress). +func chatChunkFromFrame(ev specutil.SSEEvent) provcore.Chunk { chunk := provcore.Chunk{ RawBytes: formatSSE(ev.Event, ev.Data), NativeEvent: ev.Event, @@ -164,7 +175,7 @@ func (s *openaiStreamSession) Next(ctx context.Context) (provcore.Chunk, error) } } - return chunk, nil + return chunk } func (s *openaiStreamSession) Close() error { diff --git a/packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses.go b/packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses.go index 69907cf5..46be18dc 100644 --- a/packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses.go +++ b/packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses.go @@ -26,10 +26,15 @@ // via the separate error_envelope path) // response.web_search_call.* / .file_search_call.* / .image_generation_call.* / // response.mcp_call_arguments.* / .code_interpreter_call_code.* / -// response.computer_call.* → informational built-in tool events; -// no canonical emission (these only fire -// on same-shape passthrough, which uses -// the raw byte copier, not this session) +// response.computer_call.* → built-in tool events; the canonical +// chat waist cannot represent them, so +// this Responses→canonical decoder +// drops them. They are preserved only +// on the verbatim passthrough lane +// (responsesEgressSession copier), +// which forwards the original frames; +// this session runs on the lossy +// cross-format / auto-upgrade decode. // // Unknown event types are logged once (slog WARN) and skipped — the // stream MUST never abort on a new event type the upstream introduces. @@ -229,9 +234,9 @@ func (s *responsesStreamSession) Next(ctx context.Context) (provcore.Chunk, erro NativeEvent: evType, }, nil default: - // Built-in tool events + unknown future events: silently drop - // from canonical emit (built-ins only fire on same-shape - // passthrough which doesn't go through this session). Log + // Built-in tool events + unknown future events: dropped from the + // canonical emit because the chat waist cannot carry them (they are + // preserved on the verbatim copier lane, not this lossy decode). Log // unknowns once per stream so operators see drift early. if !isBuiltinToolEvent(evType) { unknownResponsesEventOnce.Do(func() { @@ -261,9 +266,9 @@ func usagePtrOrNil(u provcore.Usage) *provcore.Usage { // isBuiltinToolEvent reports whether an event name belongs to one of the // OpenAI-native built-in tools (web_search, file_search, image_gen, -// computer, mcp, code_interpreter). These fire only on same-shape -// passthrough; the stream session is invoked only on the auto-upgrade / -// cross-format paths, so seeing one here is a no-op (not an error). +// computer, mcp, code_interpreter). On this lossy Responses→canonical decode +// (cross-format / auto-upgrade) they have no canonical representation and are +// dropped without a warning; the verbatim copier lane preserves them. func isBuiltinToolEvent(evType string) bool { prefixes := []string{ "response.web_search_call.", diff --git a/packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses_egress.go b/packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses_egress.go new file mode 100644 index 00000000..282a0c4a --- /dev/null +++ b/packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses_egress.go @@ -0,0 +1,179 @@ +// Package stream — stream_responses_egress.go is the response-side content +// detector for the /v1/responses ingress. The request may be sent upstream as +// Responses-shape (a target whose capability says it serves /v1/responses) yet +// the upstream can still answer with chat.completion.chunk frames (an +// OpenAI-compatible endpoint that only implements /v1/chat/completions, or a +// provider that mis-declares its capability). Trusting the declared Format +// would forward those chat frames straight to a /v1/responses client with no +// terminal event — the egress bug this detector exists to kill. +// +// The session performs exactly ONE classification, lazily on the first decoded +// SSE frame at the raw-byte boundary (reusing the shared SSEScanner buffer; one +// per-stream hold, zero per-chunk allocation): +// +// first frame is Responses-shape (event: response.* / data {"type":"response.*"}) +// → copier mode: forward every upstream frame VERBATIM (Chunk.Verbatim + +// RawBytes) so built-in-tool / audio events survive, AND decode the +// canonical Delta / tool-call / reasoning / usage fields onto the same +// chunk so the enforcement (buffer / Model A) lane — which ignores +// Verbatim and reads canonical — can still redact and account. +// first frame is chat.completion.chunk, or unclassifiable +// → chat mode: decode each frame as canonical chat (fail CLOSED — never +// verbatim) so the proxy's Responses encoder re-shapes it into the +// response.* event grammar with a terminal response.completed. +package stream + +import ( + "bytes" + "context" + "io" + "log/slog" + + provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/specs/openai/responses" + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/specutil" + "github.com/tidwall/gjson" +) + +// egressMode is the per-stream decode mode resolved from the first frame. +type egressMode int + +const ( + egressModeUnresolved egressMode = iota + egressModeCopier // genuine Responses upstream → verbatim + canonical tee + egressModeChat // chat upstream / unclassifiable → canonical chat decode (fail closed) +) + +// responsesEgressSession is the content-sniffing session returned for a +// /v1/responses ingress. It owns one SSEScanner and resolves egressMode once. +type responsesEgressSession struct { + scanner *specutil.SSEScanner + log *slog.Logger + mode egressMode + done bool + // sawToolCall records whether any function_call_arguments.delta was seen so + // the terminal response.completed reports finish_reason "tool_calls" (parity + // with the chat decoder), which a downstream re-encoder preserves. + sawToolCall bool +} + +func (s *responsesEgressSession) Next(ctx context.Context) (provcore.Chunk, error) { + if s.done { + return provcore.Chunk{}, io.EOF + } + for { + if err := ctx.Err(); err != nil { + return provcore.Chunk{}, err + } + ev, err := s.scanner.Next() + if err != nil { + return provcore.Chunk{}, err + } + // Responses-API never sends "[DONE]"; chat-completions does. Match it + // defensively for either wire and close the stream. + if bytes.Equal(bytes.TrimSpace(ev.Data), []byte("[DONE]")) { + s.done = true + return provcore.Chunk{ + Done: true, + RawBytes: []byte("data: [DONE]\n\n"), + NativeEvent: ev.Event, + Verbatim: s.mode == egressModeCopier, + }, nil + } + if len(ev.Data) == 0 { + // Keep-alive / comment frame — skipped by the scanner contract; it + // never participates in classification (fail-safe on garbage). + continue + } + + if s.mode == egressModeUnresolved { + evType := ev.Event + if evType == "" { + evType = gjson.GetBytes(ev.Data, "type").String() + } + switch responses.ClassifyFirstSSEFrame(evType, ev.Data) { + case responses.ClassResponses: + s.mode = egressModeCopier + default: + // ClassChat AND ClassUnknown both fail closed to the canonical + // chat lane — never verbatim. An unclassifiable first frame must + // not leak the wrong wire shape; the Responses encoder downstream + // produces a valid response.created → response.completed envelope. + s.mode = egressModeChat + } + } + + if s.mode == egressModeCopier { + chunk := s.copierChunk(ev) + if chunk.Done { + s.done = true + } + return chunk, nil + } + return chatChunkFromFrame(ev), nil + } +} + +func (s *responsesEgressSession) Close() error { + s.done = true + return s.scanner.Close() +} + +// copierChunk turns one Responses-API SSE frame into a chunk that is BOTH a +// verbatim copy (RawBytes + Verbatim, so the non-enforced live relay forwards +// the original bytes with zero loss of built-in-tool / audio events) AND a +// canonical decode of the delta / tool-call / reasoning / usage fields (so the +// enforcement lane, which reads canonical and ignores Verbatim, can still +// redact and account). The usage tee fires on response.completed / +// response.incomplete only — the same terminal frames that carry the token +// totals. +func (s *responsesEgressSession) copierChunk(ev specutil.SSEEvent) provcore.Chunk { + evType := ev.Event + if evType == "" { + evType = gjson.GetBytes(ev.Data, "type").String() + } + chunk := provcore.Chunk{ + RawBytes: formatSSE(ev.Event, ev.Data), + NativeEvent: evType, + Verbatim: true, + } + switch evType { + case "response.output_text.delta", "response.refusal.delta": + chunk.Delta = gjson.GetBytes(ev.Data, "delta").String() + case "response.function_call_arguments.delta": + s.sawToolCall = true + chunk.ToolCallDeltas = []provcore.ToolCallDelta{{ + Index: int(gjson.GetBytes(ev.Data, "output_index").Int()), + ID: gjson.GetBytes(ev.Data, "item_id").String(), + Arguments: gjson.GetBytes(ev.Data, "delta").String(), + }} + case "response.reasoning_summary_text.delta", "response.reasoning_text.delta": + chunk.ReasoningDelta = gjson.GetBytes(ev.Data, "delta").String() + case "response.completed", "response.incomplete": + chunk.Done = true + usage := specutil.ExtractOpenAIUsage(gjson.GetBytes(ev.Data, "response.usage")) + chunk.Usage = usagePtrOrNil(usage) + chunk.FinishReason = s.responsesTerminalFinishReason(evType, ev.Data) + case "response.failed", "response.error": + chunk.Done = true + } + return chunk +} + +// responsesTerminalFinishReason maps a Responses terminal event to the +// canonical OpenAI finish_reason so a downstream re-encoder preserves it. +// completed → stop (or tool_calls when function calls were seen); +// incomplete → length / content_filter per incomplete_details.reason +// (mirrors mapResponsesStatusToFinishReason). +func (s *responsesEgressSession) responsesTerminalFinishReason(evType string, data []byte) string { + if evType == "response.incomplete" { + if gjson.GetBytes(data, "response.incomplete_details.reason").String() == "content_filter" { + return "content_filter" + } + return "length" + } + if s.sawToolCall { + return "tool_calls" + } + return "stop" +} diff --git a/packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses_egress_test.go b/packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses_egress_test.go new file mode 100644 index 00000000..73cc3022 --- /dev/null +++ b/packages/ai-gateway/internal/providers/specs/openai/stream/stream_responses_egress_test.go @@ -0,0 +1,254 @@ +// Package stream_test — stream_responses_egress_test.go covers the +// content-sniffing /v1/responses egress session: it classifies the first SSE +// frame and either forwards genuine Responses frames verbatim (with a canonical +// tee for accounting) or decodes a chat.completion upstream into canonical so +// the proxy can re-shape it into Responses events with a terminal event. +package stream_test + +import ( + "context" + "errors" + "io" + "log/slog" + "strings" + "testing" + + provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" + ostream "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/specs/openai/stream" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/typology" +) + +func openResponsesEgress(t *testing.T, sse string) provcore.StreamSession { + t.Helper() + d := ostream.NewStreamDecoder(slog.Default()) + sess, err := d.Open(io.NopCloser(strings.NewReader(sse)), typology.WireShapeOpenAIResponses) + if err != nil { + t.Fatalf("Open: %v", err) + } + return sess +} + +func drainSession(t *testing.T, sess provcore.StreamSession) []provcore.Chunk { + t.Helper() + var out []provcore.Chunk + for { + c, err := sess.Next(context.Background()) + if errors.Is(err, io.EOF) { + break + } + if err != nil { + t.Fatalf("Next: %v", err) + } + out = append(out, c) + if c.Done { + break + } + } + return out +} + +// TestResponsesEgress_ChatFirst_DecodesCanonical is the core bug: a +// /v1/responses ingress whose upstream actually emits chat.completion.chunk +// frames must be DECODED into canonical chunks (not dropped, not forwarded +// verbatim) so the proxy's Responses encoder can re-shape them — the session +// resolves chat mode and never marks the chunks Verbatim. +func TestResponsesEgress_ChatFirst_DecodesCanonical(t *testing.T) { + sse := "data: " + `{"object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hel"}}]}` + "\n\n" + + "data: " + `{"object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"lo"}}]}` + "\n\n" + + "data: " + `{"object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":2,"total_tokens":5}}` + "\n\n" + + "data: [DONE]\n\n" + chunks := drainSession(t, openResponsesEgress(t, sse)) + + var text strings.Builder + sawDone := false + for _, c := range chunks { + if c.Verbatim { + t.Fatalf("chat-upstream chunks must NOT be marked Verbatim (would leak chat frames to a Responses client): %+v", c) + } + text.WriteString(c.Delta) + if c.Done { + sawDone = true + } + } + if text.String() != "Hello" { + t.Fatalf("decoded text = %q, want %q", text.String(), "Hello") + } + if !sawDone { + t.Fatal("stream must terminate with a Done chunk so the encoder emits response.completed") + } +} + +// TestResponsesEgress_ResponsesFirst_VerbatimAndUsage is the raw-byte copier: +// a genuine Responses upstream (incl. a built-in web_search_call event) is +// forwarded byte-for-byte (Verbatim + RawBytes) so built-in-tool events reach +// the client, while usage is still teed from response.completed. +func TestResponsesEgress_ResponsesFirst_VerbatimAndUsage(t *testing.T) { + webSearch := `{"type":"response.web_search_call.results","output_index":0,"results":[{"url":"https://x"}]}` + sse := "event: response.created\ndata: " + `{"type":"response.created","response":{"id":"resp_1"}}` + "\n\n" + + "event: response.output_text.delta\ndata: " + `{"type":"response.output_text.delta","delta":"hi"}` + "\n\n" + + "event: response.web_search_call.results\ndata: " + webSearch + "\n\n" + + "event: response.completed\ndata: " + `{"type":"response.completed","response":{"usage":{"input_tokens":7,"output_tokens":4,"total_tokens":11}}}` + "\n\n" + chunks := drainSession(t, openResponsesEgress(t, sse)) + + var rawAll strings.Builder + var finalUsage *provcore.Usage + sawText := false + for _, c := range chunks { + if !c.Verbatim { + t.Fatalf("genuine Responses frames must be Verbatim: %+v", c) + } + rawAll.Write(c.RawBytes) + if c.Delta == "hi" { + sawText = true + } + if c.Usage != nil { + finalUsage = c.Usage + } + } + if !sawText { + t.Fatal("output_text.delta must surface a canonical Delta for the enforcement lane") + } + if !strings.Contains(rawAll.String(), "response.web_search_call.results") { + t.Fatalf("built-in web_search_call event must reach the client byte-for-byte; raw=%s", rawAll.String()) + } + if finalUsage == nil || finalUsage.TotalTokens == nil || *finalUsage.TotalTokens != 11 { + t.Fatalf("usage must be teed from response.completed; got %+v", finalUsage) + } +} + +// TestResponsesEgress_Copier_CanonicalTee covers the canonical fields the +// enforcement lane reads off a verbatim copier chunk: reasoning + refusal +// deltas, function-call argument deltas, and a tool_calls finish_reason on the +// terminal frame once a function call was seen. +func TestResponsesEgress_Copier_CanonicalTee(t *testing.T) { + sse := "event: response.created\ndata: " + `{"type":"response.created"}` + "\n\n" + + "event: response.reasoning_summary_text.delta\ndata: " + `{"type":"response.reasoning_summary_text.delta","delta":"because"}` + "\n\n" + + "event: response.refusal.delta\ndata: " + `{"type":"response.refusal.delta","delta":"no"}` + "\n\n" + + "event: response.function_call_arguments.delta\ndata: " + `{"type":"response.function_call_arguments.delta","output_index":2,"item_id":"call_9","delta":"{\"a\":1}"}` + "\n\n" + + "event: response.completed\ndata: " + `{"type":"response.completed","response":{"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}` + "\n\n" + chunks := drainSession(t, openResponsesEgress(t, sse)) + + var reasoning, refusal, finish string + var tool *provcore.ToolCallDelta + for i := range chunks { + c := chunks[i] + reasoning += c.ReasoningDelta + if c.FinishReason != "" { + finish = c.FinishReason + } + if len(c.ToolCallDeltas) > 0 { + tool = &c.ToolCallDeltas[0] + } + if c.Delta == "no" { + refusal = c.Delta + } + } + if reasoning != "because" { + t.Fatalf("reasoning delta = %q, want %q", reasoning, "because") + } + if refusal != "no" { + t.Fatal("refusal.delta must surface a canonical Delta") + } + if tool == nil || tool.ID != "call_9" || tool.Index != 2 || tool.Arguments != `{"a":1}` { + t.Fatalf("function_call_arguments.delta must tee to a ToolCallDelta; got %+v", tool) + } + if finish != "tool_calls" { + t.Fatalf("terminal finish_reason after a function call = %q, want tool_calls", finish) + } +} + +// TestResponsesEgress_Copier_TerminalReasons covers the incomplete/failed +// terminal mappings: content_filter and length on response.incomplete, and a +// Done chunk on response.failed. +func TestResponsesEgress_Copier_TerminalReasons(t *testing.T) { + cases := []struct { + name, event, terminal, wantFinish string + }{ + {"incomplete_content_filter", "response.incomplete", `{"type":"response.incomplete","response":{"incomplete_details":{"reason":"content_filter"}}}`, "content_filter"}, + {"incomplete_length", "response.incomplete", `{"type":"response.incomplete","response":{"incomplete_details":{"reason":"max_output_tokens"}}}`, "length"}, + {"failed", "response.failed", `{"type":"response.failed","response":{}}`, ""}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + sse := "event: response.created\ndata: " + `{"type":"response.created"}` + "\n\n" + + "event: " + tc.event + "\ndata: " + tc.terminal + "\n\n" + chunks := drainSession(t, openResponsesEgress(t, sse)) + last := chunks[len(chunks)-1] + if !last.Done { + t.Fatalf("terminal event must produce a Done chunk; got %+v", last) + } + if last.FinishReason != tc.wantFinish { + t.Fatalf("FinishReason = %q, want %q", last.FinishReason, tc.wantFinish) + } + }) + } +} + +// TestResponsesEgress_Copier_EdgeFrames covers an event-less frame (type only +// in data), an empty-data keep-alive that is skipped, and the EOF returned when +// Next is called after the stream is Done. +func TestResponsesEgress_Copier_EdgeFrames(t *testing.T) { + sse := "event: response.created\ndata: " + `{"type":"response.created"}` + "\n\n" + + "data: " + `{"type":"response.output_text.delta","delta":"x"}` + "\n\n" + + "event: ping\ndata:\n\n" + + "event: response.completed\ndata: " + `{"type":"response.completed","response":{"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}}` + "\n\n" + sess := openResponsesEgress(t, sse) + chunks := drainSession(t, sess) + if _, err := sess.Next(context.Background()); !errors.Is(err, io.EOF) { + t.Fatalf("Next after Done must return io.EOF, got %v", err) + } + var text string + for _, c := range chunks { + if !c.Verbatim { + t.Fatalf("copier frames must be Verbatim: %+v", c) + } + text += c.Delta + } + if text != "x" { + t.Fatalf("event-less output_text.delta must still decode a canonical Delta; got %q", text) + } +} + +type errReadCloser struct{ err error } + +func (e errReadCloser) Read([]byte) (int, error) { return 0, e.err } +func (e errReadCloser) Close() error { return nil } + +// TestResponsesEgress_ScannerError: a non-EOF read error surfaces from Next. +func TestResponsesEgress_ScannerError(t *testing.T) { + boom := errors.New("boom") + d := ostream.NewStreamDecoder(slog.Default()) + sess, err := d.Open(errReadCloser{err: boom}, typology.WireShapeOpenAIResponses) + if err != nil { + t.Fatalf("Open: %v", err) + } + _, err = sess.Next(context.Background()) + if err == nil || errors.Is(err, io.EOF) { + t.Fatalf("a non-EOF read error must surface from Next, got %v", err) + } +} + +// TestResponsesEgress_ContextCancel: a cancelled context surfaces the error +// from Next rather than a chunk. +func TestResponsesEgress_ContextCancel(t *testing.T) { + sess := openResponsesEgress(t, "event: response.created\ndata: "+`{"type":"response.created"}`+"\n\n") + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if _, err := sess.Next(ctx); !errors.Is(err, context.Canceled) { + t.Fatalf("cancelled context must surface context.Canceled, got %v", err) + } +} + +// TestResponsesEgress_GarbageFirst_FailsClosedToChat: an unclassifiable first +// frame must NOT be forwarded verbatim — the session falls back to the chat +// decode lane (canonical, non-Verbatim) so the proxy re-encodes a valid +// Responses envelope rather than leaking garbage. +func TestResponsesEgress_GarbageFirst_FailsClosedToChat(t *testing.T) { + sse := "data: " + `{"unexpected":"garbage"}` + "\n\n" + "data: [DONE]\n\n" + chunks := drainSession(t, openResponsesEgress(t, sse)) + for _, c := range chunks { + if c.Verbatim { + t.Fatalf("unclassifiable frame must fail closed to chat mode (never Verbatim): %+v", c) + } + } +} diff --git a/packages/ai-gateway/internal/providers/target/resolver.go b/packages/ai-gateway/internal/providers/target/resolver.go index 20e06ef9..b6d79fa6 100644 --- a/packages/ai-gateway/internal/providers/target/resolver.go +++ b/packages/ai-gateway/internal/providers/target/resolver.go @@ -63,6 +63,11 @@ type ProviderRow struct { BaseURL string Extras map[string]string Disabled bool + // ServesResponsesAPI mirrors Provider.serves_responses_api. nil = use + // the adapter RequestShapes default; copied onto CallTarget so the + // executor/bridge can resolve the /v1/responses capability without a + // per-request DB read. + ServesResponsesAPI *bool } // ModelStore is the minimum model-catalog surface the default @@ -152,14 +157,15 @@ func (r *PgResolver) Resolve(ctx context.Context, providerID, modelID string, hi } target := provcore.CallTarget{ - ProviderID: pr.ID, - ProviderName: pr.Name, - Format: provcore.Format(pr.AdapterType), - BaseURL: pr.BaseURL, - APIKey: apiKey, - CredentialID: credID, - CredentialName: credName, - ProviderModelID: mr.ProviderModelID, + ProviderID: pr.ID, + ProviderName: pr.Name, + Format: provcore.Format(pr.AdapterType), + BaseURL: pr.BaseURL, + APIKey: apiKey, + CredentialID: credID, + CredentialName: credName, + ProviderModelID: mr.ProviderModelID, + ServesResponsesAPI: pr.ServesResponsesAPI, } if len(pr.Extras) > 0 { target.Extras = make(map[string]string, len(pr.Extras)) diff --git a/packages/ai-gateway/internal/routing/core/types.go b/packages/ai-gateway/internal/routing/core/types.go index 7c9c0dd6..43f73527 100644 --- a/packages/ai-gateway/internal/routing/core/types.go +++ b/packages/ai-gateway/internal/routing/core/types.go @@ -183,6 +183,12 @@ type RoutingTarget struct { BaseURL string Region string Source string // "primary", "fallback", "recovery" + // ServesResponsesAPI mirrors Provider.serves_responses_api (nil = adapter + // RequestShapes default). Carried on the routing snapshot so the proxy + // stages (cross-format guard, body canonicalization, egress reshape) and + // the executor resolve the /v1/responses capability identically without a + // per-request DB read. + ServesResponsesAPI *bool } // RoutingPlan is the output of route resolution. diff --git a/packages/ai-gateway/internal/routing/resolver.go b/packages/ai-gateway/internal/routing/resolver.go index a0237a3c..028abfdb 100644 --- a/packages/ai-gateway/internal/routing/resolver.go +++ b/packages/ai-gateway/internal/routing/resolver.go @@ -329,15 +329,16 @@ func (r *Resolver) lookupTarget(ctx context.Context, providerID, modelID string) region = *p.Region } return &core.RoutingTarget{ - ProviderID: p.ID, - ProviderName: p.Name, - AdapterType: p.AdapterType, - ModelID: m.ID, - ModelCode: m.Code, - ModelName: m.Name, - ProviderModelID: m.ProviderModelID, - BaseURL: p.BaseURL, - Region: region, + ProviderID: p.ID, + ProviderName: p.Name, + AdapterType: p.AdapterType, + ModelID: m.ID, + ModelCode: m.Code, + ModelName: m.Name, + ProviderModelID: m.ProviderModelID, + BaseURL: p.BaseURL, + Region: region, + ServesResponsesAPI: p.ServesResponsesAPI, }, nil } diff --git a/packages/compliance-proxy/internal/audit/event_message.go b/packages/compliance-proxy/internal/audit/event_message.go index 12e01486..e358f580 100644 --- a/packages/compliance-proxy/internal/audit/event_message.go +++ b/packages/compliance-proxy/internal/audit/event_message.go @@ -125,6 +125,8 @@ func toMessage(e AuditEvent, thingID, thingName string) mq.TrafficEventMessage { msg.UpstreamTotalMs = e.UpstreamTotalMs msg.RequestHooksMs = e.RequestHooksMs msg.ResponseHooksMs = e.ResponseHooksMs + msg.RequestHooksUs = e.RequestHooksUs + msg.ResponseHooksUs = e.ResponseHooksUs if len(e.LatencyBreakdown) > 0 { msg.LatencyBreakdown = e.LatencyBreakdown } diff --git a/packages/control-plane/internal/ai/providers/handler/helpers_test.go b/packages/control-plane/internal/ai/providers/handler/helpers_test.go index 33e93962..be928a25 100644 --- a/packages/control-plane/internal/ai/providers/handler/helpers_test.go +++ b/packages/control-plane/internal/ai/providers/handler/helpers_test.go @@ -165,10 +165,11 @@ func anonEchoCtx(req *http.Request, rec *httptest.ResponseRecorder) echo.Context // Test fixtures: column lists + row builders mirroring store/ -// providerCols mirrors the 13-column GetProvider / Create / Update projection. +// providerCols mirrors the 14-column GetProvider / Create / Update projection +// (serves_responses_api sits between enabled and headers). var providerCols = []string{ "id", "name", "displayName", "description", "adapter_type", "baseUrl", - "pathPrefix", "apiVersion", "region", "enabled", "headers", + "pathPrefix", "apiVersion", "region", "enabled", "serves_responses_api", "headers", "createdAt", "updatedAt", } @@ -178,6 +179,7 @@ var providerListCols = append(append([]string{}, providerCols...), "model_count" func strPtr(s string) *string { return &s } func intPtr(n int) *int { return &n } func f64Ptr(f float64) *float64 { return &f } +func boolPtr(b bool) *bool { return &b } // nowFixture returns a stable UTC timestamp suitable for fixture rows. func nowFixture() time.Time { return time.Now().UTC().Truncate(time.Second) } @@ -190,7 +192,7 @@ func makeProviderRow(now time.Time) []any { return []any{ "prov-1", "test-provider", &displayName, &desc, "openai", "https://api.test.com", "/test", &apiVer, ®ion, true, - json.RawMessage(`{"x":"y"}`), now, now, + boolPtr(false), json.RawMessage(`{"x":"y"}`), now, now, } } diff --git a/packages/control-plane/internal/ai/providers/handler/propagation_fail_test.go b/packages/control-plane/internal/ai/providers/handler/propagation_fail_test.go index 4b408607..b14d533b 100644 --- a/packages/control-plane/internal/ai/providers/handler/propagation_fail_test.go +++ b/packages/control-plane/internal/ai/providers/handler/propagation_fail_test.go @@ -215,7 +215,7 @@ func TestCreateProvider_HubFailure502(t *testing.T) { now := nowFixture() mock.ExpectBegin() mock.ExpectQuery(`INSERT INTO "Provider"`). - WithArgs(anyArgs(11)...). + WithArgs(anyArgs(12)...). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectCommit() hub := &hubSpy{invalidateErr: errors.New("hub down")} @@ -246,7 +246,7 @@ func TestUpdateProvider_HubFailure502(t *testing.T) { mock.ExpectQuery(`FROM "Provider"\s+WHERE id`).WithArgs("prov-1"). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectQuery(`UPDATE "Provider"`). - WithArgs(anyArgs(13)...). + WithArgs(anyArgs(15)...). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) hub := &hubSpy{invalidateErr: errors.New("hub down")} aud := &auditSpy{} diff --git a/packages/control-plane/internal/ai/providers/handler/providers.go b/packages/control-plane/internal/ai/providers/handler/providers.go index d4f96a52..56890a97 100644 --- a/packages/control-plane/internal/ai/providers/handler/providers.go +++ b/packages/control-plane/internal/ai/providers/handler/providers.go @@ -55,24 +55,27 @@ func (h *Handler) ListProviders(c echo.Context) error { } type listItem struct { - ID string `json:"id"` - Name string `json:"name"` - DisplayName *string `json:"displayName"` - Description *string `json:"description"` - AdapterType string `json:"adapterType"` - BaseURL string `json:"baseUrl"` - Region *string `json:"region"` - Enabled bool `json:"enabled"` - ModelCount *int `json:"modelCount"` - CreatedAt any `json:"createdAt"` + ID string `json:"id"` + Name string `json:"name"` + DisplayName *string `json:"displayName"` + Description *string `json:"description"` + AdapterType string `json:"adapterType"` + BaseURL string `json:"baseUrl"` + Region *string `json:"region"` + Enabled bool `json:"enabled"` + ServesResponsesAPI *bool `json:"servesResponsesApi,omitempty"` + ModelCount *int `json:"modelCount"` + CreatedAt any `json:"createdAt"` } data := make([]listItem, 0, len(providers)) for _, p := range providers { data = append(data, listItem{ ID: p.ID, Name: p.Name, DisplayName: p.DisplayName, Description: p.Description, AdapterType: p.AdapterType, BaseURL: p.BaseURL, - Region: p.Region, - Enabled: p.Enabled, ModelCount: p.ModelCount, CreatedAt: p.CreatedAt, + Region: p.Region, + Enabled: p.Enabled, + ServesResponsesAPI: p.ServesResponsesAPI, + ModelCount: p.ModelCount, CreatedAt: p.CreatedAt, }) } return c.JSON(http.StatusOK, map[string]any{"data": data, "total": total}) @@ -96,8 +99,10 @@ func (h *Handler) GetProvider(c echo.Context) error { "id": p.ID, "name": p.Name, "displayName": p.DisplayName, "description": p.Description, "adapterType": p.AdapterType, "baseUrl": p.BaseURL, "pathPrefix": p.PathPrefix, "apiVersion": p.APIVersion, - "region": p.Region, - "enabled": p.Enabled, "createdAt": p.CreatedAt, "updatedAt": p.UpdatedAt, + "region": p.Region, + "enabled": p.Enabled, + "servesResponsesApi": p.ServesResponsesAPI, + "createdAt": p.CreatedAt, "updatedAt": p.UpdatedAt, "models": models, } if p.Headers != nil { @@ -173,17 +178,18 @@ func conflictForUniqueViolation(pgErr *pgconn.PgError, providerName string) (mes func (h *Handler) CreateProvider(c echo.Context) error { var body struct { - Name string `json:"name"` - DisplayName string `json:"displayName"` - Description string `json:"description"` - BaseURL string `json:"baseUrl"` - AdapterType string `json:"adapterType"` - APIVersion string `json:"apiVersion"` - Region string `json:"region"` - Enabled *bool `json:"enabled"` - Headers json.RawMessage `json:"headers"` - Models []createProviderModelInput `json:"models"` - Credential *createProviderCredentialInput `json:"credential"` + Name string `json:"name"` + DisplayName string `json:"displayName"` + Description string `json:"description"` + BaseURL string `json:"baseUrl"` + AdapterType string `json:"adapterType"` + APIVersion string `json:"apiVersion"` + Region string `json:"region"` + Enabled *bool `json:"enabled"` + ServesResponsesAPI *bool `json:"servesResponsesApi"` + Headers json.RawMessage `json:"headers"` + Models []createProviderModelInput `json:"models"` + Credential *createProviderCredentialInput `json:"credential"` } if err := c.Bind(&body); err != nil { return c.JSON(http.StatusBadRequest, errJSON("Invalid request body", "validation_error", "")) @@ -320,17 +326,18 @@ func (h *Handler) CreateProvider(c echo.Context) error { p, insertedModels, insertedCred, err := h.providers.CreateProviderWithChildren( c.Request().Context(), providerstore.CreateParams{ - ID: newProviderID, // Matches the AAD when an inline credential is sealed; empty → store generates - Name: body.Name, - DisplayName: body.DisplayName, - Description: desc, - BaseURL: body.BaseURL, - PathPrefix: "/" + body.Name, - AdapterType: body.AdapterType, - APIVersion: apiVersion, - Region: region, - Enabled: enabled, - Headers: body.Headers, + ID: newProviderID, // Matches the AAD when an inline credential is sealed; empty → store generates + Name: body.Name, + DisplayName: body.DisplayName, + Description: desc, + BaseURL: body.BaseURL, + PathPrefix: "/" + body.Name, + AdapterType: body.AdapterType, + APIVersion: apiVersion, + Region: region, + Enabled: enabled, + ServesResponsesAPI: body.ServesResponsesAPI, + Headers: body.Headers, }, modelParams, credParams, @@ -377,20 +384,21 @@ func (h *Handler) CreateProvider(c echo.Context) error { h.audit.LogObserved(c.Request().Context(), ae) resp := map[string]any{ - "id": p.ID, - "name": p.Name, - "displayName": p.DisplayName, - "description": p.Description, - "adapterType": p.AdapterType, - "baseUrl": p.BaseURL, - "pathPrefix": p.PathPrefix, - "apiVersion": p.APIVersion, - "region": p.Region, - "enabled": p.Enabled, - "headers": p.Headers, - "createdAt": p.CreatedAt, - "updatedAt": p.UpdatedAt, - "models": insertedModels, + "id": p.ID, + "name": p.Name, + "displayName": p.DisplayName, + "description": p.Description, + "adapterType": p.AdapterType, + "baseUrl": p.BaseURL, + "pathPrefix": p.PathPrefix, + "apiVersion": p.APIVersion, + "region": p.Region, + "enabled": p.Enabled, + "servesResponsesApi": p.ServesResponsesAPI, + "headers": p.Headers, + "createdAt": p.CreatedAt, + "updatedAt": p.UpdatedAt, + "models": insertedModels, } if insertedCred != nil { resp["credential"] = insertedCred @@ -447,6 +455,7 @@ func (h *Handler) UpdateProvider(c echo.Context) error { } var regionParam **string var apiVersionParam **string + var servesResponsesParam **bool var updateHeaders bool var headersVal json.RawMessage { @@ -465,6 +474,16 @@ func (h *Handler) UpdateProvider(c echo.Context) error { apiVersionParam = &s } } + // servesResponsesApi is three-state like apiVersion: a present + // JSON null clears the override back to the adapter default, a + // present boolean sets/downgrades, and an absent key leaves the + // stored value untouched. + if rv, ok := rawFields["servesResponsesApi"]; ok { + var b *bool + if err := json.Unmarshal(rv, &b); err == nil { + servesResponsesParam = &b + } + } if rv, ok := rawFields["headers"]; ok { updateHeaders = true headersVal = rv @@ -474,16 +493,17 @@ func (h *Handler) UpdateProvider(c echo.Context) error { } updated, err := h.providers.UpdateProvider(c.Request().Context(), id, providerstore.UpdateParams{ - Name: body.Name, - DisplayName: body.DisplayName, - Description: body.Description, - BaseURL: body.BaseURL, - AdapterType: body.AdapterType, - Region: regionParam, - APIVersion: apiVersionParam, - UpdateHeaders: updateHeaders, - Headers: headersVal, - Enabled: body.Enabled, + Name: body.Name, + DisplayName: body.DisplayName, + Description: body.Description, + BaseURL: body.BaseURL, + AdapterType: body.AdapterType, + Region: regionParam, + APIVersion: apiVersionParam, + ServesResponsesAPI: servesResponsesParam, + UpdateHeaders: updateHeaders, + Headers: headersVal, + Enabled: body.Enabled, }) if err != nil { var pgErr *pgconn.PgError diff --git a/packages/control-plane/internal/ai/providers/handler/providers_test.go b/packages/control-plane/internal/ai/providers/handler/providers_test.go index d686179b..4af7e46f 100644 --- a/packages/control-plane/internal/ai/providers/handler/providers_test.go +++ b/packages/control-plane/internal/ai/providers/handler/providers_test.go @@ -283,10 +283,7 @@ func TestCreateProvider_HappyBareProvider(t *testing.T) { now := nowFixture() mock.ExpectBegin() mock.ExpectQuery(`INSERT INTO "Provider"`). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectCommit() // incrementConfigVersion → 2 system_metadata calls. @@ -324,10 +321,7 @@ func TestCreateProvider_HappyWithModelsAndCredential(t *testing.T) { now := nowFixture() mock.ExpectBegin() mock.ExpectQuery(`INSERT INTO "Provider"`). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) // 1 model insert — 18 params ($1..$18 including 4 capability cols). mock.ExpectQuery(`INSERT INTO "Model"`). @@ -391,10 +385,7 @@ func TestCreateProvider_HappyModelMissingCodeAndAliases(t *testing.T) { now := nowFixture() mock.ExpectBegin() mock.ExpectQuery(`INSERT INTO "Provider"`). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectQuery(`INSERT INTO "Model"`). WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), @@ -431,10 +422,7 @@ func TestCreateProvider_NameCollision409(t *testing.T) { mock, db := newMockStore(t) mock.ExpectBegin() mock.ExpectQuery(`INSERT INTO "Provider"`). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnError(&pgconn.PgError{Code: "23505", ConstraintName: "Provider_name_key"}) mock.ExpectRollback() h := newHandler(db, nil, &auditSpy{}, nil, nil, nil, ProxyConfig{}) @@ -457,10 +445,7 @@ func TestCreateProvider_ModelCollision409(t *testing.T) { now := nowFixture() mock.ExpectBegin() mock.ExpectQuery(`INSERT INTO "Provider"`). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectQuery(`INSERT INTO "Model"`). WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), @@ -494,10 +479,7 @@ func TestCreateProvider_GenericStoreError500(t *testing.T) { mock, db := newMockStore(t) mock.ExpectBegin() mock.ExpectQuery(`INSERT INTO "Provider"`). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnError(errors.New("disk full")) mock.ExpectRollback() h := newHandler(db, nil, &auditSpy{}, nil, nil, nil, ProxyConfig{}) @@ -519,10 +501,7 @@ func TestCreateProvider_EnabledFalseAndOptionalFieldsBranch(t *testing.T) { now := nowFixture() mock.ExpectBegin() mock.ExpectQuery(`INSERT INTO "Provider"`). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectQuery(`INSERT INTO "Credential"`). WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), @@ -627,11 +606,7 @@ func TestUpdateProvider_HappyAllFields(t *testing.T) { mock.ExpectQuery(`FROM "Provider"\s+WHERE id`).WithArgs("prov-1"). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectQuery(`UPDATE "Provider"`). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectQuery(`SELECT value FROM system_metadata`). WillReturnRows(pgxmock.NewRows([]string{"value"})) @@ -672,11 +647,7 @@ func TestUpdateProvider_EmptyStringsTreatedAsAbsent(t *testing.T) { mock.ExpectQuery(`FROM "Provider"\s+WHERE id`).WithArgs("prov-1"). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectQuery(`UPDATE "Provider"`). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectQuery(`SELECT value FROM system_metadata`). WillReturnRows(pgxmock.NewRows([]string{"value"})) @@ -703,11 +674,7 @@ func TestUpdateProvider_NameCollision409(t *testing.T) { mock.ExpectQuery(`FROM "Provider"\s+WHERE id`).WithArgs("prov-1"). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectQuery(`UPDATE "Provider"`). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnError(&pgconn.PgError{Code: "23505"}) h := newHandler(db, nil, &auditSpy{}, nil, nil, nil, ProxyConfig{}) req := httptest.NewRequest(http.MethodPut, "/", strings.NewReader(`{"name":"taken"}`)) @@ -727,11 +694,7 @@ func TestUpdateProvider_GenericError500(t *testing.T) { mock.ExpectQuery(`FROM "Provider"\s+WHERE id`).WithArgs("prov-1"). WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) mock.ExpectQuery(`UPDATE "Provider"`). - WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), - pgxmock.AnyArg()). + WithArgs(pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg(), pgxmock.AnyArg()). WillReturnError(errors.New("disk full")) h := newHandler(db, nil, &auditSpy{}, nil, nil, nil, ProxyConfig{}) req := httptest.NewRequest(http.MethodPut, "/", strings.NewReader(`{"name":"x"}`)) @@ -1164,5 +1127,108 @@ func TestListProviderCredentials_Happy(t *testing.T) { } } +// TestProvider_ServesResponsesApi_HTTPRoundTrip proves the new +// serves_responses_api capability flows admin-API-end to admin-API-end: +// the create request body reaches the INSERT arg, the GET projection echoes the +// stored value, and the update body reaches the UPDATE apply-flag + value. The +// silently-dropped field this fixes would leave the toggle non-persistent. +func TestProvider_ServesResponsesApi_HTTPRoundTrip(t *testing.T) { + t.Run("create threads false to INSERT and echoes it", func(t *testing.T) { + mock, db := newMockStore(t) + now := nowFixture() + insArgs := anyArgs(12) + // CreateProviderWithChildren INSERT arg order: id, name, displayName, + // description, baseUrl, pathPrefix, adapter_type, apiVersion, region, + // enabled, serves_responses_api, headers — so the column is arg[10]. + insArgs[10] = boolPtr(false) + mock.ExpectBegin() + mock.ExpectQuery(`INSERT INTO "Provider"`). + WithArgs(insArgs...). + WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) + mock.ExpectCommit() + mock.ExpectQuery(`SELECT value FROM system_metadata`). + WillReturnRows(pgxmock.NewRows([]string{"value"})) + mock.ExpectExec(`INSERT INTO system_metadata`). + WillReturnResult(pgxmock.NewResult("INSERT", 1)) + + h := newHandler(db, &hubSpy{}, &auditSpy{}, nil, nil, nil, ProxyConfig{}) + body := `{"name":"alpha","baseUrl":"https://x","adapterType":"openai","servesResponsesApi":false}` + req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c, _ := echoCtx(req, rec, "u-1") + if err := h.CreateProvider(c); err != nil { + t.Fatalf("CreateProvider: %v", err) + } + // 201 here is itself the proof the value was threaded: the INSERT + // expectation pins serves_responses_api=false at arg[10], so a dropped + // or wrong value would fail the ordered match → store error → 500. + if rec.Code != http.StatusCreated { + t.Fatalf("status=%d; want 201 (INSERT arg[10]=false unmatched ⇒ value dropped); body=%s", rec.Code, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), `"servesResponsesApi":false`) { + t.Errorf("create response must echo servesResponsesApi:false, got %s", rec.Body.String()) + } + }) + + t.Run("get echoes stored value", func(t *testing.T) { + mock, db := newMockStore(t) + now := nowFixture() + mock.ExpectQuery(`FROM "Provider"\s+WHERE id`).WithArgs("prov-1"). + WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) + mock.ExpectQuery(`FROM "Model" WHERE "providerId"`).WithArgs("prov-1"). + WillReturnRows(pgxmock.NewRows(modelCols).AddRow(makeModelRow(now)...)) + h := newHandler(db, nil, &auditSpy{}, nil, nil, nil, ProxyConfig{}) + req := httptest.NewRequest(http.MethodGet, "/", nil) + rec := httptest.NewRecorder() + c, _ := echoCtx(req, rec, "u-1") + c.SetParamNames("id") + c.SetParamValues("prov-1") + if err := h.GetProvider(c); err != nil { + t.Fatalf("GetProvider: %v", err) + } + // makeProviderRow stores serves_responses_api=false. + if !strings.Contains(rec.Body.String(), `"servesResponsesApi":false`) { + t.Errorf("get must echo stored servesResponsesApi, got %s", rec.Body.String()) + } + }) + + t.Run("update threads present false to UPDATE apply+value", func(t *testing.T) { + mock, db := newMockStore(t) + now := nowFixture() + mock.ExpectQuery(`FROM "Provider"\s+WHERE id`).WithArgs("prov-1"). + WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) + updArgs := anyArgs(15) + updArgs[13] = true // applyServesResponses + updArgs[14] = boolPtr(false) // servesResponsesVal + mock.ExpectQuery(`UPDATE "Provider"`). + WithArgs(updArgs...). + WillReturnRows(pgxmock.NewRows(providerCols).AddRow(makeProviderRow(now)...)) + mock.ExpectQuery(`SELECT value FROM system_metadata`). + WillReturnRows(pgxmock.NewRows([]string{"value"})) + mock.ExpectExec(`INSERT INTO system_metadata`). + WillReturnResult(pgxmock.NewResult("INSERT", 1)) + h := newHandler(db, &hubSpy{}, &auditSpy{}, nil, nil, nil, ProxyConfig{}) + req := httptest.NewRequest(http.MethodPut, "/", strings.NewReader(`{"servesResponsesApi":false}`)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + c, _ := echoCtx(req, rec, "u-1") + c.SetParamNames("id") + c.SetParamValues("prov-1") + if err := h.UpdateProvider(c); err != nil { + t.Fatalf("UpdateProvider: %v", err) + } + // 200 proves threading: the UPDATE expectation pins applyServesResponses + // (arg[13]=true) and the value (arg[14]=false); a no-change or wrong value + // would fail the ordered match → store error → 500. + if rec.Code != http.StatusOK { + t.Fatalf("status=%d; want 200 (UPDATE apply/value args unmatched ⇒ value dropped); body=%s", rec.Code, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), `"servesResponsesApi":false`) { + t.Errorf("update response must echo servesResponsesApi:false, got %s", rec.Body.String()) + } + }) +} + // quiet linter for occasional unused fixture helpers in this file var _ = time.Now diff --git a/packages/control-plane/internal/ai/providers/providerstore/provider.go b/packages/control-plane/internal/ai/providers/providerstore/provider.go index c598f4fd..2a982511 100644 --- a/packages/control-plane/internal/ai/providers/providerstore/provider.go +++ b/packages/control-plane/internal/ai/providers/providerstore/provider.go @@ -27,20 +27,27 @@ import ( // Nullable so existing providers keep working until an operator fills // it in; the runtime hook treats a missing region as "unknown". type Provider struct { - ID string `json:"id"` - Name string `json:"name"` - DisplayName *string `json:"displayName"` - Description *string `json:"description"` - AdapterType string `json:"adapterType"` - BaseURL string `json:"baseUrl"` - PathPrefix string `json:"pathPrefix"` - APIVersion *string `json:"apiVersion"` - Region *string `json:"region"` - Enabled bool `json:"enabled"` - Headers json.RawMessage `json:"headers"` - CreatedAt time.Time `json:"createdAt"` - UpdatedAt time.Time `json:"updatedAt"` - ModelCount *int `json:"modelCount,omitempty"` + ID string `json:"id"` + Name string `json:"name"` + DisplayName *string `json:"displayName"` + Description *string `json:"description"` + AdapterType string `json:"adapterType"` + BaseURL string `json:"baseUrl"` + PathPrefix string `json:"pathPrefix"` + APIVersion *string `json:"apiVersion"` + Region *string `json:"region"` + Enabled bool `json:"enabled"` + // ServesResponsesAPI mirrors Provider.serves_responses_api. nil = use the + // adapter type's RequestShapes default (FormatOpenAI → true; every other + // adapter → false); an explicit value only DOWNGRADES (false forces + // /v1/responses ingress through the canonical chat codec). The AI Gateway + // reads this column to decide whether to relay Responses-shape bytes + // upstream or downgrade to /v1/chat/completions. + ServesResponsesAPI *bool `json:"servesResponsesApi,omitempty"` + Headers json.RawMessage `json:"headers"` + CreatedAt time.Time `json:"createdAt"` + UpdatedAt time.Time `json:"updatedAt"` + ModelCount *int `json:"modelCount,omitempty"` } // ListParams holds filter/pagination for listing providers. @@ -78,7 +85,7 @@ func (s *Store) ListProviders(ctx context.Context, p ListParams) ([]Provider, in // Data with model count dataQuery := fmt.Sprintf(` SELECT p.id, p.name, p."displayName", p.description, p.adapter_type, p."baseUrl", - p."pathPrefix", p."apiVersion", p.region, p.enabled, p.headers, + p."pathPrefix", p."apiVersion", p.region, p.enabled, p.serves_responses_api, p.headers, p."createdAt", p."updatedAt", (SELECT COUNT(*) FROM "Model" m WHERE m."providerId" = p.id) AS model_count FROM "Provider" p @@ -100,7 +107,7 @@ func (s *Store) ListProviders(ctx context.Context, p ListParams) ([]Provider, in var mc int if err := rows.Scan( &pr.ID, &pr.Name, &pr.DisplayName, &pr.Description, &pr.AdapterType, &pr.BaseURL, - &pr.PathPrefix, &pr.APIVersion, &pr.Region, &pr.Enabled, &pr.Headers, + &pr.PathPrefix, &pr.APIVersion, &pr.Region, &pr.Enabled, &pr.ServesResponsesAPI, &pr.Headers, &pr.CreatedAt, &pr.UpdatedAt, &mc, ); err != nil { return nil, 0, fmt.Errorf("scan provider: %w", err) @@ -115,14 +122,14 @@ func (s *Store) ListProviders(ctx context.Context, p ListParams) ([]Provider, in func (s *Store) GetProvider(ctx context.Context, id string) (*Provider, error) { row := s.pool.QueryRow(ctx, ` SELECT id, name, "displayName", description, adapter_type, "baseUrl", - "pathPrefix", "apiVersion", region, enabled, headers, "createdAt", "updatedAt" + "pathPrefix", "apiVersion", region, enabled, serves_responses_api, headers, "createdAt", "updatedAt" FROM "Provider" WHERE id = $1 `, id) var p Provider err := row.Scan(&p.ID, &p.Name, &p.DisplayName, &p.Description, &p.AdapterType, &p.BaseURL, - &p.PathPrefix, &p.APIVersion, &p.Region, &p.Enabled, &p.Headers, &p.CreatedAt, &p.UpdatedAt) + &p.PathPrefix, &p.APIVersion, &p.Region, &p.Enabled, &p.ServesResponsesAPI, &p.Headers, &p.CreatedAt, &p.UpdatedAt) if errors.Is(err, pgx.ErrNoRows) { return nil, nil } @@ -148,20 +155,23 @@ type CreateParams struct { APIVersion *string Region *string Enabled bool - Headers json.RawMessage + // ServesResponsesAPI is the per-provider Responses-API capability override. + // nil persists SQL NULL (adapter default); a non-nil value downgrades. + ServesResponsesAPI *bool + Headers json.RawMessage } // CreateProvider inserts a new provider and returns it. func (s *Store) CreateProvider(ctx context.Context, p CreateParams) (*Provider, error) { row := s.pool.QueryRow(ctx, ` - INSERT INTO "Provider" (id, name, "displayName", description, "baseUrl", "pathPrefix", adapter_type, "apiVersion", region, enabled, headers, "createdAt", "updatedAt") - VALUES (gen_random_uuid(), $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, NOW(), NOW()) - RETURNING id, name, "displayName", description, adapter_type, "baseUrl", "pathPrefix", "apiVersion", region, enabled, headers, "createdAt", "updatedAt" - `, p.Name, p.DisplayName, p.Description, p.BaseURL, p.PathPrefix, p.AdapterType, p.APIVersion, p.Region, p.Enabled, p.Headers) + INSERT INTO "Provider" (id, name, "displayName", description, "baseUrl", "pathPrefix", adapter_type, "apiVersion", region, enabled, serves_responses_api, headers, "createdAt", "updatedAt") + VALUES (gen_random_uuid(), $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, NOW(), NOW()) + RETURNING id, name, "displayName", description, adapter_type, "baseUrl", "pathPrefix", "apiVersion", region, enabled, serves_responses_api, headers, "createdAt", "updatedAt" + `, p.Name, p.DisplayName, p.Description, p.BaseURL, p.PathPrefix, p.AdapterType, p.APIVersion, p.Region, p.Enabled, p.ServesResponsesAPI, p.Headers) var pr Provider err := row.Scan(&pr.ID, &pr.Name, &pr.DisplayName, &pr.Description, &pr.AdapterType, &pr.BaseURL, - &pr.PathPrefix, &pr.APIVersion, &pr.Region, &pr.Enabled, &pr.Headers, &pr.CreatedAt, &pr.UpdatedAt) + &pr.PathPrefix, &pr.APIVersion, &pr.Region, &pr.Enabled, &pr.ServesResponsesAPI, &pr.Headers, &pr.CreatedAt, &pr.UpdatedAt) if err != nil { return nil, fmt.Errorf("create provider: %w", err) } @@ -188,12 +198,12 @@ func (s *Store) CreateProviderWithChildren( } var pr Provider err = tx.QueryRow(ctx, ` - INSERT INTO "Provider" (id, name, "displayName", description, "baseUrl", "pathPrefix", adapter_type, "apiVersion", region, enabled, headers, "createdAt", "updatedAt") - VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, NOW(), NOW()) - RETURNING id, name, "displayName", description, adapter_type, "baseUrl", "pathPrefix", "apiVersion", region, enabled, headers, "createdAt", "updatedAt" - `, providerID, provider.Name, provider.DisplayName, provider.Description, provider.BaseURL, provider.PathPrefix, provider.AdapterType, provider.APIVersion, provider.Region, provider.Enabled, provider.Headers, + INSERT INTO "Provider" (id, name, "displayName", description, "baseUrl", "pathPrefix", adapter_type, "apiVersion", region, enabled, serves_responses_api, headers, "createdAt", "updatedAt") + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, NOW(), NOW()) + RETURNING id, name, "displayName", description, adapter_type, "baseUrl", "pathPrefix", "apiVersion", region, enabled, serves_responses_api, headers, "createdAt", "updatedAt" + `, providerID, provider.Name, provider.DisplayName, provider.Description, provider.BaseURL, provider.PathPrefix, provider.AdapterType, provider.APIVersion, provider.Region, provider.Enabled, provider.ServesResponsesAPI, provider.Headers, ).Scan(&pr.ID, &pr.Name, &pr.DisplayName, &pr.Description, &pr.AdapterType, &pr.BaseURL, - &pr.PathPrefix, &pr.APIVersion, &pr.Region, &pr.Enabled, &pr.Headers, &pr.CreatedAt, &pr.UpdatedAt) + &pr.PathPrefix, &pr.APIVersion, &pr.Region, &pr.Enabled, &pr.ServesResponsesAPI, &pr.Headers, &pr.CreatedAt, &pr.UpdatedAt) if err != nil { return nil, nil, nil, fmt.Errorf("insert provider: %w", err) } @@ -294,16 +304,19 @@ func (s *Store) CreateProviderWithChildren( // UpdateParams holds fields for updating a provider. type UpdateParams struct { - Name *string - DisplayName *string - Description *string - BaseURL *string - AdapterType *string - Region **string - APIVersion **string // nil=no change; &nil=clear; &s=set - UpdateHeaders bool - Headers json.RawMessage // used when UpdateHeaders=true; nil clears the column - Enabled *bool + Name *string + DisplayName *string + Description *string + BaseURL *string + AdapterType *string + Region **string + APIVersion **string // nil=no change; &nil=clear; &s=set + // ServesResponsesAPI follows the same three-state semantic as APIVersion: + // nil=no change; &nil=clear back to adapter default; &b=set the override. + ServesResponsesAPI **bool + UpdateHeaders bool + Headers json.RawMessage // used when UpdateHeaders=true; nil clears the column + Enabled *bool } // UpdateProvider updates a provider's mutable fields. @@ -318,6 +331,11 @@ func (s *Store) UpdateProvider(ctx context.Context, id string, p UpdateParams) ( if applyAPIVersion { apiVersionVal = *p.APIVersion } + applyServesResponses := p.ServesResponsesAPI != nil + var servesResponsesVal *bool + if applyServesResponses { + servesResponsesVal = *p.ServesResponsesAPI + } row := s.pool.QueryRow(ctx, ` UPDATE "Provider" SET name = COALESCE($2, name), @@ -329,16 +347,17 @@ func (s *Store) UpdateProvider(ctx context.Context, id string, p UpdateParams) ( enabled = COALESCE($9, enabled), "pathPrefix" = CASE WHEN $2 IS NOT NULL THEN '/' || $2 ELSE "pathPrefix" END, "apiVersion" = CASE WHEN $10::boolean THEN $11 ELSE "apiVersion" END, + serves_responses_api = CASE WHEN $14::boolean THEN $15 ELSE serves_responses_api END, headers = CASE WHEN $12::boolean THEN $13 ELSE headers END, "updatedAt" = NOW() WHERE id = $1 - RETURNING id, name, "displayName", description, adapter_type, "baseUrl", "pathPrefix", "apiVersion", region, enabled, headers, "createdAt", "updatedAt" + RETURNING id, name, "displayName", description, adapter_type, "baseUrl", "pathPrefix", "apiVersion", region, enabled, serves_responses_api, headers, "createdAt", "updatedAt" `, id, p.Name, p.DisplayName, p.Description, p.BaseURL, p.AdapterType, applyRegion, regionVal, p.Enabled, - applyAPIVersion, apiVersionVal, p.UpdateHeaders, p.Headers) + applyAPIVersion, apiVersionVal, p.UpdateHeaders, p.Headers, applyServesResponses, servesResponsesVal) var pr Provider err := row.Scan(&pr.ID, &pr.Name, &pr.DisplayName, &pr.Description, &pr.AdapterType, &pr.BaseURL, - &pr.PathPrefix, &pr.APIVersion, &pr.Region, &pr.Enabled, &pr.Headers, &pr.CreatedAt, &pr.UpdatedAt) + &pr.PathPrefix, &pr.APIVersion, &pr.Region, &pr.Enabled, &pr.ServesResponsesAPI, &pr.Headers, &pr.CreatedAt, &pr.UpdatedAt) if errors.Is(err, pgx.ErrNoRows) { return nil, nil } diff --git a/packages/control-plane/internal/ai/providers/providerstore/provider_pgxmock_test.go b/packages/control-plane/internal/ai/providers/providerstore/provider_pgxmock_test.go index 7a827054..b644280b 100644 --- a/packages/control-plane/internal/ai/providers/providerstore/provider_pgxmock_test.go +++ b/packages/control-plane/internal/ai/providers/providerstore/provider_pgxmock_test.go @@ -14,6 +14,7 @@ import ( ) func strptr(s string) *string { return &s } +func boolptr(b bool) *bool { return &b } func anyArgs(n int) []any { a := make([]any, n) for i := range a { @@ -23,7 +24,7 @@ func anyArgs(n int) []any { } var ( - provCols = []string{"id", "name", "displayName", "description", "adapter_type", "baseUrl", "pathPrefix", "apiVersion", "region", "enabled", "headers", "createdAt", "updatedAt"} + provCols = []string{"id", "name", "displayName", "description", "adapter_type", "baseUrl", "pathPrefix", "apiVersion", "region", "enabled", "serves_responses_api", "headers", "createdAt", "updatedAt"} provListCols = append(append([]string{}, provCols...), "model_count") modelCols = []string{ "id", "code", "name", "description", "providerId", "providerModelId", "type", "features", @@ -43,7 +44,13 @@ var ( var tNow = time.Date(2026, 5, 27, 0, 0, 0, 0, time.UTC) func provRow(id, name string) []any { - return []any{id, name, strptr("d"), strptr("desc"), "openai", "https://api.x", "/" + name, strptr("2024"), strptr("us-east-1"), true, []byte(`{}`), tNow, tNow} + return provRowServes(id, name, boolptr(false)) +} + +// provRowServes builds a Provider row with an explicit serves_responses_api +// value so the column round-trips can be asserted (nil = adapter default). +func provRowServes(id, name string, serves *bool) []any { + return []any{id, name, strptr("d"), strptr("desc"), "openai", "https://api.x", "/" + name, strptr("2024"), strptr("us-east-1"), true, serves, []byte(`{}`), tNow, tNow} } func provListRow(id, name string, mc int) []any { return append(provRow(id, name), mc) } func modelRow(id, code string) []any { @@ -144,13 +151,13 @@ func TestGetProvider(t *testing.T) { func TestCreateProvider(t *testing.T) { s, m := newMock(t) - m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(10)...). + m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(11)...). WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRow("p1", "openai")...)) p, err := s.CreateProvider(context.Background(), CreateParams{Name: "openai"}) if err != nil || p == nil || p.Name != "openai" { t.Fatalf("CreateProvider: %+v %v", p, err) } - m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(10)...).WillReturnError(errors.New("dup")) + m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(11)...).WillReturnError(errors.New("dup")) if _, err := s.CreateProvider(context.Background(), CreateParams{}); err == nil { t.Fatal("insert error should surface") } @@ -163,7 +170,7 @@ func TestCreateProvider(t *testing.T) { func TestCreateProviderWithChildren_Full(t *testing.T) { s, m := newMock(t) m.ExpectBegin() - m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(11)...). + m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(12)...). WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRow("p1", "openai")...)) m.ExpectQuery(`INSERT INTO "Model"`).WithArgs(anyArgs(20)...). WillReturnRows(pgxmock.NewRows(modelCols).AddRow(modelRow("m1", "gpt-4o")...)) @@ -189,7 +196,7 @@ func TestCreateProviderWithChildren_Full(t *testing.T) { func TestCreateProviderWithChildren_NoChildren(t *testing.T) { s, m := newMock(t) m.ExpectBegin() - m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(11)...). + m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(12)...). WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRow("p1", "openai")...)) m.ExpectCommit() p, models, cred, err := s.CreateProviderWithChildren(context.Background(), CreateParams{Name: "openai"}, nil, nil) @@ -209,7 +216,7 @@ func TestCreateProviderWithChildren_Errors(t *testing.T) { t.Run("provider insert", func(t *testing.T) { s, m := newMock(t) m.ExpectBegin() - m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(11)...).WillReturnError(errors.New("boom")) + m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(12)...).WillReturnError(errors.New("boom")) m.ExpectRollback() if _, _, _, err := s.CreateProviderWithChildren(context.Background(), CreateParams{}, nil, nil); err == nil { t.Fatal("provider insert error should surface") @@ -218,7 +225,7 @@ func TestCreateProviderWithChildren_Errors(t *testing.T) { t.Run("model insert", func(t *testing.T) { s, m := newMock(t) m.ExpectBegin() - m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(11)...). + m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(12)...). WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRow("p1", "o")...)) m.ExpectQuery(`INSERT INTO "Model"`).WithArgs(anyArgs(20)...).WillReturnError(errors.New("bad model")) m.ExpectRollback() @@ -230,7 +237,7 @@ func TestCreateProviderWithChildren_Errors(t *testing.T) { t.Run("credential insert", func(t *testing.T) { s, m := newMock(t) m.ExpectBegin() - m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(11)...). + m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(12)...). WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRow("p1", "o")...)) m.ExpectQuery(`INSERT INTO "Credential"`).WithArgs(anyArgs(9)...).WillReturnError(errors.New("bad cred")) m.ExpectRollback() @@ -242,7 +249,7 @@ func TestCreateProviderWithChildren_Errors(t *testing.T) { t.Run("commit", func(t *testing.T) { s, m := newMock(t) m.ExpectBegin() - m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(11)...). + m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(anyArgs(12)...). WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRow("p1", "o")...)) m.ExpectCommit().WillReturnError(errors.New("commit failed")) if _, _, _, err := s.CreateProviderWithChildren(context.Background(), CreateParams{}, nil, nil); err == nil { @@ -255,22 +262,121 @@ func TestUpdateProvider(t *testing.T) { s, m := newMock(t) region := strptr("eu-west-1") apiV := strptr("2025") - m.ExpectQuery(`UPDATE "Provider"`).WithArgs(anyArgs(13)...). + m.ExpectQuery(`UPDATE "Provider"`).WithArgs(anyArgs(15)...). WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRow("p1", "openai")...)) p, err := s.UpdateProvider(context.Background(), "p1", UpdateParams{Name: strptr("New"), Region: ®ion, APIVersion: &apiV, UpdateHeaders: true}) if err != nil || p == nil { t.Fatalf("UpdateProvider: %+v %v", p, err) } - m.ExpectQuery(`UPDATE "Provider"`).WithArgs(anyArgs(13)...).WillReturnError(pgx.ErrNoRows) + m.ExpectQuery(`UPDATE "Provider"`).WithArgs(anyArgs(15)...).WillReturnError(pgx.ErrNoRows) if p, err := s.UpdateProvider(context.Background(), "x", UpdateParams{}); err != nil || p != nil { t.Fatalf("missing → (nil,nil), got %+v %v", p, err) } - m.ExpectQuery(`UPDATE "Provider"`).WithArgs(anyArgs(13)...).WillReturnError(errors.New("db")) + m.ExpectQuery(`UPDATE "Provider"`).WithArgs(anyArgs(15)...).WillReturnError(errors.New("db")) if _, err := s.UpdateProvider(context.Background(), "x", UpdateParams{}); err == nil { t.Fatal("db error should surface") } } +// TestProvider_ServesResponsesAPI_RoundTrip pins the serves_responses_api +// admin write/read path: create persists an explicit override, GET returns it, +// update sets / clears / leaves it via the three-state pointer semantic, and the +// scanned value round-trips out of RETURNING. The arg-position assertions prove +// the column value (and the apply flag) reach the SQL, not just the response. +func TestProvider_ServesResponsesAPI_RoundTrip(t *testing.T) { + t.Run("create persists explicit false (downgrade)", func(t *testing.T) { + s, m := newMock(t) + // Arg $10 (0-based 9) is serves_responses_api: assert the store sends + // the caller's pointer value, and the row scans back as false. + args := anyArgs(11) + args[9] = boolptr(false) + m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(args...). + WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRowServes("p1", "openai", boolptr(false))...)) + p, err := s.CreateProvider(context.Background(), CreateParams{Name: "openai", ServesResponsesAPI: boolptr(false)}) + if err != nil || p == nil || p.ServesResponsesAPI == nil || *p.ServesResponsesAPI != false { + t.Fatalf("create did not round-trip serves=false: p=%+v err=%v", p, err) + } + if err := m.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet (serves value not sent to INSERT): %v", err) + } + }) + + t.Run("create with nil persists adapter default", func(t *testing.T) { + s, m := newMock(t) + args := anyArgs(11) + args[9] = (*bool)(nil) + m.ExpectQuery(`INSERT INTO "Provider"`).WithArgs(args...). + WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRowServes("p1", "openai", nil)...)) + p, err := s.CreateProvider(context.Background(), CreateParams{Name: "openai"}) + if err != nil || p == nil || p.ServesResponsesAPI != nil { + t.Fatalf("create with nil should keep adapter default (nil): p=%+v err=%v", p, err) + } + if err := m.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet: %v", err) + } + }) + + t.Run("get returns stored true override", func(t *testing.T) { + s, m := newMock(t) + m.ExpectQuery(`FROM "Provider"\s+WHERE id = \$1`).WithArgs("p1"). + WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRowServes("p1", "openai", boolptr(true))...)) + p, err := s.GetProvider(context.Background(), "p1") + if err != nil || p == nil || p.ServesResponsesAPI == nil || *p.ServesResponsesAPI != true { + t.Fatalf("get did not return serves=true: p=%+v err=%v", p, err) + } + }) + + t.Run("update set true sends apply=true, value=true", func(t *testing.T) { + s, m := newMock(t) + setTrue := boolptr(true) + args := anyArgs(15) + args[13] = true // applyServesResponses + args[14] = setTrue // servesResponsesVal + m.ExpectQuery(`UPDATE "Provider"`).WithArgs(args...). + WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRowServes("p1", "openai", boolptr(true))...)) + p, err := s.UpdateProvider(context.Background(), "p1", UpdateParams{ServesResponsesAPI: &setTrue}) + if err != nil || p == nil || p.ServesResponsesAPI == nil || *p.ServesResponsesAPI != true { + t.Fatalf("update set true did not round-trip: p=%+v err=%v", p, err) + } + if err := m.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet (set true not applied): %v", err) + } + }) + + t.Run("update clear sends apply=true, value=nil", func(t *testing.T) { + s, m := newMock(t) + var clear *bool // present-but-null → clear back to adapter default + args := anyArgs(15) + args[13] = true + args[14] = (*bool)(nil) + m.ExpectQuery(`UPDATE "Provider"`).WithArgs(args...). + WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRowServes("p1", "openai", nil)...)) + p, err := s.UpdateProvider(context.Background(), "p1", UpdateParams{ServesResponsesAPI: &clear}) + if err != nil || p == nil || p.ServesResponsesAPI != nil { + t.Fatalf("update clear should reset to adapter default (nil): p=%+v err=%v", p, err) + } + if err := m.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet (clear not applied): %v", err) + } + }) + + t.Run("update without field sends apply=false (no change)", func(t *testing.T) { + s, m := newMock(t) + args := anyArgs(15) + args[13] = false // applyServesResponses → CASE leaves the column untouched + args[14] = (*bool)(nil) + m.ExpectQuery(`UPDATE "Provider"`).WithArgs(args...). + WillReturnRows(pgxmock.NewRows(provCols).AddRow(provRowServes("p1", "openai", boolptr(true))...)) + p, err := s.UpdateProvider(context.Background(), "p1", UpdateParams{Name: strptr("New")}) + if err != nil || p == nil { + t.Fatalf("update no-change failed: p=%+v err=%v", p, err) + } + if err := m.ExpectationsWereMet(); err != nil { + t.Fatalf("unmet (no-change apply flag wrong): %v", err) + } + }) +} + func TestDeleteProvider(t *testing.T) { t.Run("ok", func(t *testing.T) { s, m := newMock(t) diff --git a/packages/control-plane/internal/fleet/handler/agent/helpers_test.go b/packages/control-plane/internal/fleet/handler/agent/helpers_test.go index bac5513e..0ddc0af6 100644 --- a/packages/control-plane/internal/fleet/handler/agent/helpers_test.go +++ b/packages/control-plane/internal/fleet/handler/agent/helpers_test.go @@ -395,10 +395,11 @@ func makeDeviceAssignmentRow(now time.Time) []any { } } -// providerListCols mirrors ListProviders's data-query projection. +// providerListCols mirrors ListProviders's data-query projection +// (serves_responses_api sits between enabled and headers). var providerListCols = []string{ "id", "name", "displayName", "description", "adapter_type", "baseUrl", - "pathPrefix", "apiVersion", "region", "enabled", "headers", + "pathPrefix", "apiVersion", "region", "enabled", "serves_responses_api", "headers", "createdAt", "updatedAt", "model_count", } @@ -407,9 +408,10 @@ func makeProviderRow(now time.Time) []any { desc := "primary provider" apiVer := "2025-01-01" region := "us-east-1" + serves := false return []any{ "prov-1", "openai", &dn, &desc, "openai", "https://api.openai.com", - "/v1", &apiVer, ®ion, true, json.RawMessage(`{}`), + "/v1", &apiVer, ®ion, true, &serves, json.RawMessage(`{}`), now, now, 5, } } diff --git a/packages/control-plane/internal/traffic/handler/traffic/traffic_cov_test.go b/packages/control-plane/internal/traffic/handler/traffic/traffic_cov_test.go index 2a31940c..65d69780 100644 --- a/packages/control-plane/internal/traffic/handler/traffic/traffic_cov_test.go +++ b/packages/control-plane/internal/traffic/handler/traffic/traffic_cov_test.go @@ -28,11 +28,11 @@ import ( // 89 base scan destinations plus the 6 payload-JOIN columns // (inline_request_body, inline_response_body, request_spill_ref, // response_spill_ref, inline_request_encoding, inline_response_encoding). Base -// columns 0/1/2/66 carry the non-nullable values (ID, Source, Timestamp, +// columns 0/1/2/68 carry the non-nullable values (ID, Source, Timestamp, // CreatedAt); every other base column is SQL NULL. The caller fills the 6 -// payload columns (indices 89..94) to drive the body/spill branches. +// payload columns (indices 91..96) to drive the body/spill branches. func trafficEventCovRow(id string, reqBody, respBody, reqSpill, respSpill []byte, reqEnc, respEnc string) *pgxmock.Rows { - const n = 89 + const n = 91 const extra = 6 cols := make([]string, n+extra) vals := make([]any, n+extra) @@ -43,13 +43,13 @@ func trafficEventCovRow(id string, reqBody, respBody, reqSpill, respSpill []byte vals[0] = id vals[1] = "ai-gateway" vals[2] = tNowCov - vals[66] = tNowCov - vals[89] = nullableBytes(reqBody) - vals[90] = nullableBytes(respBody) - vals[91] = nullableBytes(reqSpill) - vals[92] = nullableBytes(respSpill) - vals[93] = reqEnc - vals[94] = respEnc + vals[68] = tNowCov // CreatedAt — index +2 after request/response_hooks_us + vals[91] = nullableBytes(reqBody) + vals[92] = nullableBytes(respBody) + vals[93] = nullableBytes(reqSpill) + vals[94] = nullableBytes(respSpill) + vals[95] = reqEnc + vals[96] = respEnc return pgxmock.NewRows(cols).AddRow(vals...) } diff --git a/packages/control-plane/internal/traffic/handler/traffic/traffic_handler_test.go b/packages/control-plane/internal/traffic/handler/traffic/traffic_handler_test.go index de262577..56d004de 100644 --- a/packages/control-plane/internal/traffic/handler/traffic/traffic_handler_test.go +++ b/packages/control-plane/internal/traffic/handler/traffic/traffic_handler_test.go @@ -141,7 +141,10 @@ func (s *testSpillStore) Get(_ context.Context, ref sharedaudit.SpillRef) (io.Re return nil, s.getErr } if s.readErr != nil { - return io.NopCloser(errReader{err: s.readErr}), nil + // Get succeeds — the configured failure surfaces later on Read via + // errReader (a spill object that resolves but whose stream is truncated), + // so returning a nil error here is correct. + return io.NopCloser(errReader{err: s.readErr}), nil //nolint:nilerr } return io.NopCloser(strings.NewReader(string(s.data))), nil } @@ -2045,16 +2048,17 @@ func TestComplianceReport_HookHealthDBError_Returns500(t *testing.T) { // traffic.go — GetTrafficEvent success path (record != nil, spillStore branch) -// trafficEventGetCols lists the 71 columns scanned by GetTrafficEvent -// (trafficEventSelectColumns + 4 payload columns). +// trafficEventGetCols lists the 97 columns scanned by GetTrafficEvent +// (trafficEventSelectColumns 91 cols + 6 payload columns). var trafficEventGetCols = []string{ - // trafficEventSelectColumns (67 cols) + // trafficEventSelectColumns (91 cols) "id", "source", "timestamp", "source_ip", "target_host", "method", "path", "target_method", "target_path", "status_code", "latency_ms", "upstream_ttfb_ms", "upstream_total_ms", "request_hooks_ms", "response_hooks_ms", + "request_hooks_us", "response_hooks_us", "latency_breakdown", "trace_id", "external_request_id", "entity_type", "entity_id", "entity_name", @@ -2117,6 +2121,8 @@ func newTrafficEventRow() *pgxmock.Rows { nil, nil, // group 6: request_hooks_ms, response_hooks_ms (*int) nil, nil, + // group 6b: request_hooks_us, response_hooks_us (*int) + nil, nil, // group 7: latency_breakdown (json.RawMessage) []byte(`null`), // group 8: trace_id, external_request_id (*string) @@ -2236,6 +2242,8 @@ func TestGetTrafficEvent_WithSpillStore_ResolvesFailed_StillReturns200(t *testin nil, nil, // group 6: request_hooks_ms, response_hooks_ms nil, nil, + // group 6b: request_hooks_us, response_hooks_us + nil, nil, // group 7: latency_breakdown []byte(`null`), // group 8: trace_id, external_request_id diff --git a/packages/control-plane/internal/traffic/store/trafficstore/traffic_queries.go b/packages/control-plane/internal/traffic/store/trafficstore/traffic_queries.go index f2bfe532..23f763ac 100644 --- a/packages/control-plane/internal/traffic/store/trafficstore/traffic_queries.go +++ b/packages/control-plane/internal/traffic/store/trafficstore/traffic_queries.go @@ -36,10 +36,13 @@ type TrafficEvent struct { // row. NULL on historical rows where upstream_ttfb / hooks could not be // reconstructed. The UI computes `ourOverheadMs = latencyMs − upstream_total_ms` // at render time. - UpstreamTtfbMs *int `json:"upstreamTtfbMs,omitempty"` - UpstreamTotalMs *int `json:"upstreamTotalMs,omitempty"` - RequestHooksMs *int `json:"requestHooksMs,omitempty"` - ResponseHooksMs *int `json:"responseHooksMs,omitempty"` + UpstreamTtfbMs *int `json:"upstreamTtfbMs,omitempty"` + UpstreamTotalMs *int `json:"upstreamTotalMs,omitempty"` + RequestHooksMs *int `json:"requestHooksMs,omitempty"` + ResponseHooksMs *int `json:"responseHooksMs,omitempty"` + // Microsecond-precision hook aggregates (additive; siblings of the _ms fields). + RequestHooksUs *int `json:"requestHooksUs,omitempty"` + ResponseHooksUs *int `json:"responseHooksUs,omitempty"` LatencyBreakdown json.RawMessage `json:"latencyBreakdown,omitempty"` // Request tracing TraceID *string `json:"traceId,omitempty"` @@ -243,6 +246,7 @@ const trafficEventSelectColumns = ` a.status_code, a.latency_ms, a.upstream_ttfb_ms, a.upstream_total_ms, a.request_hooks_ms, a.response_hooks_ms, + a.request_hooks_us, a.response_hooks_us, a.latency_breakdown, a.trace_id, a.external_request_id, a.entity_type, a.entity_id, a.entity_name, @@ -339,6 +343,7 @@ func (store *Store) GetTrafficEvent(ctx context.Context, id string) (*TrafficEve &a.StatusCode, &a.LatencyMs, &a.UpstreamTtfbMs, &a.UpstreamTotalMs, &a.RequestHooksMs, &a.ResponseHooksMs, + &a.RequestHooksUs, &a.ResponseHooksUs, &a.LatencyBreakdown, &a.TraceID, &a.ExternalRequestID, &a.EntityType, &a.EntityID, &a.EntityName, @@ -394,6 +399,7 @@ func scanOneTrafficEvent(row interface{ Scan(dest ...any) error }, a *TrafficEve &a.StatusCode, &a.LatencyMs, &a.UpstreamTtfbMs, &a.UpstreamTotalMs, &a.RequestHooksMs, &a.ResponseHooksMs, + &a.RequestHooksUs, &a.ResponseHooksUs, &a.LatencyBreakdown, &a.TraceID, &a.ExternalRequestID, &a.EntityType, &a.EntityID, &a.EntityName, diff --git a/packages/control-plane/internal/traffic/store/trafficstore/trafficstore_pgxmock_test.go b/packages/control-plane/internal/traffic/store/trafficstore/trafficstore_pgxmock_test.go index 0f14100c..e36c3074 100644 --- a/packages/control-plane/internal/traffic/store/trafficstore/trafficstore_pgxmock_test.go +++ b/packages/control-plane/internal/traffic/store/trafficstore/trafficstore_pgxmock_test.go @@ -42,7 +42,7 @@ func newMock(t *testing.T) (*Store, pgxmock.PgxPoolIface) { // accept. extra appends additional nil columns for GetTrafficEvent's payload // JOIN (request/response body + spill refs). func trafficEventRow(extra int) *pgxmock.Rows { - const n = 89 + const n = 91 cols := make([]string, n+extra) vals := make([]any, n+extra) for i := range cols { @@ -52,7 +52,7 @@ func trafficEventRow(extra int) *pgxmock.Rows { vals[0] = "evt1" // ID (string) vals[1] = "ai-gateway" // Source (string) vals[2] = tNow // Timestamp (time.Time) - vals[66] = tNow // CreatedAt (time.Time) + vals[68] = tNow // CreatedAt (time.Time) — index +2 after request/response_hooks_us return pgxmock.NewRows(cols).AddRow(vals...) } diff --git a/packages/nexus-hub/internal/observability/consumer/alert_view.go b/packages/nexus-hub/internal/observability/consumer/alert_view.go index 340f7b04..ae3ddf45 100644 --- a/packages/nexus-hub/internal/observability/consumer/alert_view.go +++ b/packages/nexus-hub/internal/observability/consumer/alert_view.go @@ -154,7 +154,8 @@ func (r *recReader) skipField(id mq.FieldID) { case mq.FldPromptTokens, mq.FldCompletionTokens, mq.FldReasoningTokens, mq.FldCacheCreationTokens, mq.FldCacheReadTokens, mq.FldNormalizedStripCount, mq.FldNormalizedStripBytes, mq.FldCacheMarkerInjected, mq.FldUpstreamTtfbMs, - mq.FldUpstreamTotalMs, mq.FldRequestHooksMs, mq.FldResponseHooksMs: + mq.FldUpstreamTotalMs, mq.FldRequestHooksMs, mq.FldResponseHooksMs, + mq.FldRequestHooksUs, mq.FldResponseHooksUs: r.varint() case mq.FldReasoningCostUsd, mq.FldGatewayCacheSavingsUsd, mq.FldEmbeddingCostUsd, mq.FldAIGuardCostUsd, mq.FldCacheWriteCostUsd, mq.FldCacheReadSavingsUsd, diff --git a/packages/nexus-hub/internal/observability/consumer/binwire_decode.go b/packages/nexus-hub/internal/observability/consumer/binwire_decode.go index 797a3fdf..974f2e4f 100644 --- a/packages/nexus-hub/internal/observability/consumer/binwire_decode.go +++ b/packages/nexus-hub/internal/observability/consumer/binwire_decode.go @@ -254,6 +254,10 @@ func decodeBinaryRecordInto(e *TrafficEventMessage, data []byte, skipBody bool) e.RequestHooksMs = pInt(int(r.varint())) case mq.FldResponseHooksMs: e.ResponseHooksMs = pInt(int(r.varint())) + case mq.FldRequestHooksUs: + e.RequestHooksUs = pInt(int(r.varint())) + case mq.FldResponseHooksUs: + e.ResponseHooksUs = pInt(int(r.varint())) case mq.FldLatencyBreakdown: e.LatencyBreakdown = r.json() case mq.FldAttestationVerified: diff --git a/packages/nexus-hub/internal/observability/consumer/consumer_behaviors_test.go b/packages/nexus-hub/internal/observability/consumer/consumer_behaviors_test.go index 556f663d..0a08240e 100644 --- a/packages/nexus-hub/internal/observability/consumer/consumer_behaviors_test.go +++ b/packages/nexus-hub/internal/observability/consumer/consumer_behaviors_test.go @@ -653,7 +653,7 @@ func TestTrafficWriter_InsertTrafficEvents_FullRowSendBatch(t *testing.T) { mock.ExpectBegin() eb := mock.ExpectBatch() - eb.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) + eb.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) mock.ExpectCommit() tx, _ := mock.Begin(ctx) @@ -700,7 +700,7 @@ func TestTrafficWriter_InsertTrafficEvents_BatchExecErrorPropagates(t *testing.T ctx := context.Background() mock.ExpectBegin() eb := mock.ExpectBatch() - eb.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnError(errors.New("constraint-violation")) + eb.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnError(errors.New("constraint-violation")) tx, _ := mock.Begin(ctx) defer tx.Rollback(ctx) //nolint:errcheck diff --git a/packages/nexus-hub/internal/observability/consumer/flush_pgxmock_test.go b/packages/nexus-hub/internal/observability/consumer/flush_pgxmock_test.go index 77625c20..ac04af72 100644 --- a/packages/nexus-hub/internal/observability/consumer/flush_pgxmock_test.go +++ b/packages/nexus-hub/internal/observability/consumer/flush_pgxmock_test.go @@ -88,7 +88,7 @@ func TestTrafficWriter_Flush_HappyPath_AcksAllAndCommits(t *testing.T) { mock.ExpectBegin() eb1 := mock.ExpectBatch() - eb1.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) + eb1.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) eb2 := mock.ExpectBatch() eb2.ExpectExec(`INSERT INTO traffic_event_payload`).WithArgs(anyArgs(13)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) // Normalized sidecar fast path: SAVEPOINT (Begin) → pipelined batch → @@ -141,7 +141,7 @@ func TestTrafficWriter_Flush_PersistsEndpointType(t *testing.T) { mock.ExpectBegin() eb1 := mock.ExpectBatch() eb1.ExpectExec(`INSERT INTO traffic_event`). - WithArgs(append(anyArgs(89), "embeddings", "")...). + WithArgs(append(append(anyArgs(89), "embeddings", ""), anyArgs(2)...)...). WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) eb2 := mock.ExpectBatch() eb2.ExpectExec(`INSERT INTO traffic_event_payload`).WithArgs(anyArgs(13)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) @@ -192,7 +192,7 @@ func TestTrafficWriter_Flush_PersistsIngressFormat(t *testing.T) { mock.ExpectBegin() eb1 := mock.ExpectBatch() eb1.ExpectExec(`INSERT INTO traffic_event`). - WithArgs(append(anyArgs(90), "anthropic")...). + WithArgs(append(append(anyArgs(90), "anthropic"), anyArgs(2)...)...). WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) eb2 := mock.ExpectBatch() eb2.ExpectExec(`INSERT INTO traffic_event_payload`).WithArgs(anyArgs(13)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) @@ -255,7 +255,7 @@ func TestTrafficWriter_Flush_InternalOpsBreakdownNulStripped(t *testing.T) { eb := mock.ExpectBatch() // $88 (internal_ops_breakdown) must arrive NUL-free and stripped. args := append(anyArgs(87), nulFreeJSONArg{want: `{"raw":"ab","esc":"pq"}`}) - args = append(args, anyArgs(3)...) // $89 (l2 entry key), $90 (endpoint_type), $91 (ingress_format) + args = append(args, anyArgs(5)...) // $89 l2 key, $90 endpoint_type, $91 ingress_format, $92/$93 hooks_us eb.ExpectExec(`INSERT INTO traffic_event`).WithArgs(args...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) mock.ExpectCommit() @@ -295,14 +295,14 @@ func TestTrafficWriter_FlushItem_NormalizedFailureStillCommits(t *testing.T) { // Batched attempt fails on the traffic_event insert → triggers per-item. mock.ExpectBegin() ebBatch := mock.ExpectBatch() - ebBatch.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnError(transient) + ebBatch.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnError(transient) mock.ExpectRollback() // Per-item: traffic_event ok; body absent (no payload batch); normalized // sidecar fails (fast batch then row-by-row, both non-poison) but the outer // tx still commits. mock.ExpectBegin() ebTE := mock.ExpectBatch() - ebTE.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) + ebTE.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) mock.ExpectBegin() // normalized savepoint (fast path) ebN := mock.ExpectBatch() ebN.ExpectExec(`INSERT INTO traffic_event_normalized`).WithArgs(anyArgs(10)...).WillReturnError(errors.New("normalize-batch-failed")) @@ -387,12 +387,12 @@ func TestTrafficWriter_Flush_InsertPoisonPill22021AcksToSkip(t *testing.T) { // Batched fast-path attempt aborts on the poison row. mock.ExpectBegin() eb := mock.ExpectBatch() - eb.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnError(poison) + eb.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnError(poison) mock.ExpectRollback() // Per-item fallback re-runs the same row, hits the same typed poison, acks. mock.ExpectBegin() ebItem := mock.ExpectBatch() - ebItem.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnError(poison) + ebItem.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnError(poison) mock.ExpectRollback() var ackCount, nakCount int32 @@ -436,18 +436,18 @@ func TestTrafficWriter_Flush_PoisonRowIsolatedHealthyCommits(t *testing.T) { // the rest). mock.ExpectBegin() ebBatch := mock.ExpectBatch() - ebBatch.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnError(poison) - ebBatch.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnError(poison) + ebBatch.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnError(poison) + ebBatch.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnError(poison) mock.ExpectRollback() // Per-item: poison row first — Begin, insert(22021), Rollback, ack-to-skip. mock.ExpectBegin() ebP := mock.ExpectBatch() - ebP.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnError(poison) + ebP.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnError(poison) mock.ExpectRollback() // Per-item: healthy row — Begin, insert ok, Commit, ack. mock.ExpectBegin() ebH := mock.ExpectBatch() - ebH.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) + ebH.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) mock.ExpectCommit() var poisonAck, poisonNak int32 @@ -497,12 +497,12 @@ func TestTrafficWriter_Flush_InsertNonPoisonFailureNaksAll(t *testing.T) { transient := errors.New("unique_violation") mock.ExpectBegin() eb := mock.ExpectBatch() - eb.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnError(transient) + eb.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnError(transient) mock.ExpectRollback() // Per-item retry: same transient failure → nak. mock.ExpectBegin() ebItem := mock.ExpectBatch() - ebItem.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnError(transient) + ebItem.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnError(transient) mock.ExpectRollback() var ackCount, nakCount int32 @@ -539,14 +539,14 @@ func TestTrafficWriter_Flush_InsertPayloadsFailureNaksAll(t *testing.T) { payloadErr := errors.New("disk-full") mock.ExpectBegin() eb1 := mock.ExpectBatch() - eb1.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) + eb1.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) eb2 := mock.ExpectBatch() eb2.ExpectExec(`INSERT INTO traffic_event_payload`).WithArgs(anyArgs(13)...).WillReturnError(payloadErr) mock.ExpectRollback() // Per-item retry: traffic_event ok, payload fails again → nak. mock.ExpectBegin() eb1b := mock.ExpectBatch() - eb1b.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) + eb1b.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) eb2b := mock.ExpectBatch() eb2b.ExpectExec(`INSERT INTO traffic_event_payload`).WithArgs(anyArgs(13)...).WillReturnError(payloadErr) mock.ExpectRollback() @@ -586,7 +586,7 @@ func TestTrafficWriter_Flush_NormalizedFailureWarnsButCommits(t *testing.T) { mock.ExpectBegin() eb1 := mock.ExpectBatch() - eb1.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) + eb1.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) // Body absent → insertPayloads short-circuits, no traffic_event_payload // batch is sent. Normalized fast path: SAVEPOINT → batch(err) → ROLLBACK TO // SAVEPOINT, then row-by-row retry: SAVEPOINT → Exec(err) → ROLLBACK. The @@ -638,14 +638,14 @@ func TestTrafficWriter_Flush_CommitFailureNaksAll(t *testing.T) { commitErr := errors.New("commit-rejected") mock.ExpectBegin() eb1 := mock.ExpectBatch() - eb1.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) + eb1.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) // Body absent + no normalize → only the traffic_event insert; commit fails. mock.ExpectCommit().WillReturnError(commitErr) mock.ExpectRollback() // Per-item retry: insert ok, commit fails again → nak. mock.ExpectBegin() eb1b := mock.ExpectBatch() - eb1b.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) + eb1b.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) mock.ExpectCommit().WillReturnError(commitErr) mock.ExpectRollback() @@ -693,7 +693,7 @@ func TestTrafficWriter_Flush_NilRegistryHappyPath(t *testing.T) { mock.ExpectBegin() eb := mock.ExpectBatch() - eb.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(91)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) + eb.ExpectExec(`INSERT INTO traffic_event`).WithArgs(anyArgs(93)...).WillReturnResult(pgconn.NewCommandTag("INSERT 0 1")) mock.ExpectCommit() var ackCount, nakCount int32 diff --git a/packages/nexus-hub/internal/observability/consumer/message.go b/packages/nexus-hub/internal/observability/consumer/message.go index 09b3b065..51196722 100644 --- a/packages/nexus-hub/internal/observability/consumer/message.go +++ b/packages/nexus-hub/internal/observability/consumer/message.go @@ -211,10 +211,13 @@ type TrafficEventMessage struct { // Latency phase breakdown. Producers omit fields they did not measure. // See docs/developers/architecture "Latency Phase Taxonomy" for the // per-source population matrix and JSONB key schema. - UpstreamTtfbMs *int `json:"upstreamTtfbMs,omitempty"` - UpstreamTotalMs *int `json:"upstreamTotalMs,omitempty"` - RequestHooksMs *int `json:"requestHooksMs,omitempty"` - ResponseHooksMs *int `json:"responseHooksMs,omitempty"` + UpstreamTtfbMs *int `json:"upstreamTtfbMs,omitempty"` + UpstreamTotalMs *int `json:"upstreamTotalMs,omitempty"` + RequestHooksMs *int `json:"requestHooksMs,omitempty"` + ResponseHooksMs *int `json:"responseHooksMs,omitempty"` + // Microsecond-precision siblings (additive; the _ms fields are unchanged). + RequestHooksUs *int `json:"requestHooksUs,omitempty"` + ResponseHooksUs *int `json:"responseHooksUs,omitempty"` LatencyBreakdown json.RawMessage `json:"latencyBreakdown,omitempty"` // Agent attestation passthrough. Compliance-proxy sets both when a verified diff --git a/packages/nexus-hub/internal/observability/consumer/traffic_copy.go b/packages/nexus-hub/internal/observability/consumer/traffic_copy.go index 6f9954ee..8d11aa88 100644 --- a/packages/nexus-hub/internal/observability/consumer/traffic_copy.go +++ b/packages/nexus-hub/internal/observability/consumer/traffic_copy.go @@ -93,6 +93,8 @@ var trafficEventColumns = []string{ "gateway_cache_l2_entry_key", "endpoint_type", "ingress_format", + // Microsecond hook aggregates (appended last; siblings of *_hooks_ms). + "request_hooks_us", "response_hooks_us", } // trafficEventRowValues returns the column values for one traffic_event row in @@ -174,6 +176,7 @@ func appendTrafficEventRow(dst []any, e TrafficEventMessage) []any { l2EntryKey, stripNul(e.EndpointType), stripNul(e.IngressFormat), + e.RequestHooksUs, e.ResponseHooksUs, ) } diff --git a/packages/nexus-hub/internal/observability/consumer/traffic_copy_test.go b/packages/nexus-hub/internal/observability/consumer/traffic_copy_test.go index 6f3c5851..98fa0661 100644 --- a/packages/nexus-hub/internal/observability/consumer/traffic_copy_test.go +++ b/packages/nexus-hub/internal/observability/consumer/traffic_copy_test.go @@ -71,8 +71,8 @@ func sqlInsertColumns(t *testing.T, sql string) []string { // not here (or vice versa) fails the build's tests, not silently in production. func TestTrafficEventColumnsParity(t *testing.T) { want := sqlInsertColumns(t, insertTrafficEventSQL) - if len(want) != 91 { - t.Fatalf("parsed %d columns from insertTrafficEventSQL, want 91", len(want)) + if len(want) != 93 { + t.Fatalf("parsed %d columns from insertTrafficEventSQL, want 93", len(want)) } if len(trafficEventColumns) != len(want) { t.Fatalf("trafficEventColumns has %d, SQL has %d", len(trafficEventColumns), len(want)) diff --git a/packages/nexus-hub/internal/observability/consumer/traffic_inserts.go b/packages/nexus-hub/internal/observability/consumer/traffic_inserts.go index 809a1e3a..291c43fb 100644 --- a/packages/nexus-hub/internal/observability/consumer/traffic_inserts.go +++ b/packages/nexus-hub/internal/observability/consumer/traffic_inserts.go @@ -125,6 +125,9 @@ func (w *TrafficEventWriter) insertTrafficEvents(ctx context.Context, tx pgx.Tx, // ingress_format ($91). Client-facing wire format; '' for // compliance-proxy / agent. NOT NULL DEFAULT '', bind directly. stripNul(e.IngressFormat), + // request_hooks_us / response_hooks_us ($92, $93). Nil *int → SQL NULL, + // same NULL-vs-0 semantics as the _ms aggregates. + e.RequestHooksUs, e.ResponseHooksUs, ) } @@ -196,7 +199,10 @@ INSERT INTO traffic_event ( -- Client-facing wire format (provcore.Format) the request/response bodies -- are stored in. Drives the control-plane view-time normalize. '' for -- compliance-proxy / agent transparent forwards. - ingress_format + ingress_format, + -- Microsecond-precision hook aggregates (siblings of request_hooks_ms / + -- response_hooks_ms). Appended LAST so existing $N positions never shift. + request_hooks_us, response_hooks_us ) VALUES ( $1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, @@ -237,7 +243,8 @@ INSERT INTO traffic_event ( $87, $88, $89, $90, - $91 + $91, + $92, $93 ) ON CONFLICT (id) DO NOTHING ` diff --git a/packages/shared/audit/event_types.go b/packages/shared/audit/event_types.go index 5a12d622..f672dcd9 100644 --- a/packages/shared/audit/event_types.go +++ b/packages/shared/audit/event_types.go @@ -122,10 +122,13 @@ type AuditEvent struct { // traffic.PhaseSink attached to the upstream call's context and from // the hook pipeline's per-hook latencyMs. Nil pointers stay NULL on // the wire and on the DB column. - UpstreamTtfbMs *int - UpstreamTotalMs *int - RequestHooksMs *int - ResponseHooksMs *int + UpstreamTtfbMs *int + UpstreamTotalMs *int + RequestHooksMs *int + ResponseHooksMs *int + // Microsecond-precision hook aggregates (additive; siblings of the _ms fields). + RequestHooksUs *int + ResponseHooksUs *int LatencyBreakdown map[string]int // DomainRuleID is the matched interception_domain.id when the host diff --git a/packages/shared/policy/decision/types.go b/packages/shared/policy/decision/types.go index f452e377..5931ef0f 100644 --- a/packages/shared/policy/decision/types.go +++ b/packages/shared/policy/decision/types.go @@ -135,6 +135,47 @@ type CompliancePipelineResult struct { // BlockingRule is the rule-pack attribution that caused the pipeline's // (reject) decision. BlockingRule *BlockingRule `json:"blockingRule,omitempty"` + // RedactionApplicable reports whether an ENFORCING per-hook decision + // (Modify or BlockSoft) contributed an APPLICABLE redaction artifact that + // the aggregate actually carries — ModifiedContent (Modify hooks only; the + // merge does not carry a BlockSoft hook's ModifiedContent) or a span whose + // ContentAddress is not an audit-only sentinel. It is the single signal + // CarriesRedaction() consults for the BlockSoft branch, separating a real + // inflight redaction masked behind a co-firing soft-block from an advisory + // approve-webhook / modify-webhook span set that cannot be applied. Derived + // in pipeline.mergeResults where per-hook decisions are visible; it cannot + // be reconstructed post-merge (the aggregate span union mixes advisory and + // enforcing spans). json:"-" — it is an in-memory predicate input, not a + // persisted/serialized field. + RedactionApplicable bool `json:"-"` +} + +// CarriesRedaction reports whether this result requires applying a redaction before +// delivery — a Modify, OR a BlockSoft that masks a co-firing redact (StrictestDecision +// promotes the aggregate Decision to BlockSoft while the redact's artifact is carried). +// Keying on Decision==Modify alone would silently DROP such a redaction and +// deliver/forward the flagged content raw. +// +// The BlockSoft branch gates on RedactionApplicable, NOT on raw span/content presence: +// a co-firing approve-webhook (or a modify-webhook reconciled under an approve ceiling) +// emits AUDIT-ONLY spans (ContentAddress = the audit-only sentinel, which ApplySpans +// never resolves), so keying on len(TransformSpans)>0 would report a redaction the +// aggregate cannot apply and over-block (fail-closed) a stream that should soft-deliver. +// RedactionApplicable is true only when an enforcing hook contributed an artifact the +// aggregate actually carries (see its field doc). A standalone soft-block with no +// applicable artifact carries no redaction and folds to block per the dispatch sites. +// +// Consumers MUST gate redaction on this predicate, not on Decision==Modify, on EVERY +// path (request + response, stream + non-stream) so a co-firing redact+soft-block masks +// and delivers rather than leaking or failing closed. +func (r *CompliancePipelineResult) CarriesRedaction() bool { + if r == nil { + return false + } + if r.Decision == Modify { + return true + } + return r.Decision == BlockSoft && r.RedactionApplicable } // HookResult is the output produced by a single hook execution. @@ -146,7 +187,15 @@ type HookResult struct { Decision Decision `json:"decision"` Reason string `json:"reason,omitempty"` ReasonCode string `json:"reasonCode,omitempty"` - LatencyMs int `json:"latencyMs"` + // LatencyMs is the truncated integer-millisecond wall-clock for this hook, + // kept for backward compatibility. It is the integer-ms floor of LatencyUs and + // is NEVER clamped: per-hook latency is summed downstream, so clamping a + // sub-millisecond hook to ≥1 would inflate the aggregate. + LatencyMs int `json:"latencyMs"` + // LatencyUs is the precise microsecond wall-clock for this hook. Hooks run at + // microsecond scale, so LatencyMs truncates to 0 for sub-millisecond hooks; + // LatencyUs carries the real value for the audit aggregates and the UI. + LatencyUs int `json:"latencyUs"` // Tags emitted by this hook; merged into the pipeline-wide set. Tags []string `json:"tags,omitempty"` Error string `json:"error,omitempty"` // non-empty if the hook errored diff --git a/packages/shared/policy/decision/types_test.go b/packages/shared/policy/decision/types_test.go index f34a4593..46ebf24b 100644 --- a/packages/shared/policy/decision/types_test.go +++ b/packages/shared/policy/decision/types_test.go @@ -32,3 +32,29 @@ func TestActionValid(t *testing.T) { t.Fatal("unknown action must be invalid") } } + +func TestCarriesRedaction(t *testing.T) { + cases := []struct { + name string + r *CompliancePipelineResult + want bool + }{ + {"nil", nil, false}, + {"modify", &CompliancePipelineResult{Decision: Modify}, true}, + // The BlockSoft branch now keys on RedactionApplicable (set by mergeResults), + // NOT raw span/content presence — so a co-firing redact masked by soft-block + // (flag true) carries, while a soft-block whose only spans are advisory/audit- + // only (flag false) does not, even if a ModifiedContent slice is present on the + // struct. This documents that the merge-computed flag is authoritative. + {"blocksoft-applicable", &CompliancePipelineResult{Decision: BlockSoft, RedactionApplicable: true}, true}, + {"blocksoft-flag-false-ignores-raw-content", &CompliancePipelineResult{Decision: BlockSoft, ModifiedContent: []ContentBlock{{}}}, false}, + {"blocksoft-without-redaction", &CompliancePipelineResult{Decision: BlockSoft}, false}, + {"approve", &CompliancePipelineResult{Decision: Approve}, false}, + {"rejecthard", &CompliancePipelineResult{Decision: RejectHard}, false}, + } + for _, c := range cases { + if got := c.r.CarriesRedaction(); got != c.want { + t.Errorf("%s: CarriesRedaction()=%v want %v", c.name, got, c.want) + } + } +} diff --git a/packages/shared/policy/hooks/core/capability_profiles_test.go b/packages/shared/policy/hooks/core/capability_profiles_test.go index 9936b358..fc451889 100644 --- a/packages/shared/policy/hooks/core/capability_profiles_test.go +++ b/packages/shared/policy/hooks/core/capability_profiles_test.go @@ -31,6 +31,17 @@ func TestChatOnly_SupportsEndpoint_Embeddings(t *testing.T) { } } +func TestChatOnly_SupportsEndpoint_Responses(t *testing.T) { + var h ChatOnly + // The Responses API carries its own endpoint_type label but is chat-family; + // ChatOnly hooks (PII redaction etc.) MUST still apply to it — otherwise + // labelling /v1/responses as "responses" would silently drop hooks that ran + // while it was labelled "chat". + if !h.SupportsEndpoint(EndpointTypeResponses) { + t.Error("ChatOnly must support EndpointTypeResponses (chat-family)") + } +} + func TestChatOnly_SupportsModality_Text(t *testing.T) { var h ChatOnly if !h.SupportsModality(ModalityText) { diff --git a/packages/shared/policy/hooks/core/noop.go b/packages/shared/policy/hooks/core/noop.go index ab2bf4cd..89bcedee 100644 --- a/packages/shared/policy/hooks/core/noop.go +++ b/packages/shared/policy/hooks/core/noop.go @@ -21,11 +21,13 @@ func NewNoop(cfg *HookConfig) (Hook, error) { // Execute always returns a successful Approve result. func (n *NoopHook) Execute(_ context.Context, _ *HookInput) (*HookResult, error) { start := time.Now() + elapsed := time.Since(start) return &HookResult{ HookID: n.cfg.ID, ImplementationID: n.cfg.ImplementationID, HookName: n.cfg.Name, Decision: Approve, - LatencyMs: int(time.Since(start).Milliseconds()), + LatencyMs: int(elapsed.Milliseconds()), + LatencyUs: int(elapsed.Microseconds()), }, nil } diff --git a/packages/shared/policy/hooks/core/onmatch.go b/packages/shared/policy/hooks/core/onmatch.go index 6573e636..22e01c95 100644 --- a/packages/shared/policy/hooks/core/onmatch.go +++ b/packages/shared/policy/hooks/core/onmatch.go @@ -244,10 +244,11 @@ func StrictestAction(a, b decision.Action) decision.Action { // // Rationale: Approve lets traffic through unchanged; Modify rewrites // content inflight (the request still completes, just with a modified -// body); BlockSoft stops the request from reaching the upstream but -// returns a soft-block response to the caller (more disruptive than a -// silent rewrite — the caller sees the block); RejectHard stops the -// request entirely. Abstain ranks at 0 so Strictest(Abstain, X) == X. +// body); BlockSoft ranks above Modify as a stricter ceiling, but at the +// dispatch sites ActionFromDecision folds it to the block action, so it +// stops the request exactly like RejectHard (a hard refuse — there is no +// distinct soft-block client response); RejectHard stops the request +// entirely. Abstain ranks at 0 so Strictest(Abstain, X) == X. // // The BlockSoft > Modify ordering matches the pipeline aggregator in // packages/shared/policy/pipeline/pipeline.go (its mergeResults prefers diff --git a/packages/shared/policy/hooks/core/types.go b/packages/shared/policy/hooks/core/types.go index 60ba8c5b..a9db48c4 100644 --- a/packages/shared/policy/hooks/core/types.go +++ b/packages/shared/policy/hooks/core/types.go @@ -238,6 +238,7 @@ const ( EndpointTypeVideoGeneration = typology.EndpointKindVideoGeneration EndpointTypeBatch = typology.EndpointKindBatch EndpointTypeJob = typology.EndpointKindJob + EndpointTypeResponses = typology.EndpointKindResponses ) // Modality identifies the content modality carried by a request or response. @@ -294,6 +295,31 @@ type RawContentPrescanner interface { MayMatchRaw(body []byte) bool } +// RuntimeEscalatable is the optional capability a hook implements when its +// runtime decision can be MORE restrictive than the action declared in its +// onMatch config — it can block or redact even when its declarative ceiling +// says approve (or redact). +// +// The streaming-routing predicates (Pipeline.MayBlock / Pipeline.MayRedact) +// answer "could this resolved pipeline ever block / redact" by reading each +// bound hook's declared onMatch action, BEFORE any response bytes are read, to +// decide whether the response must run in buffered execution. A hook that can +// exceed its declared action at execution would be under-routed onto the +// audit-only live path on the strength of a permissive declaration, and its +// runtime enforcement would be silently dropped. Such a hook reports +// MayExceedOnMatch()==true so the predicates treat it as both may-block and +// may-redact regardless of its configured action — the safe over-route +// direction (a pipeline that buffers but never fires is harmless; one that +// streams an enforceable block is not). +// +// webhook-forward is the canonical implementer: its Execute reconciles the +// remote endpoint's suggested decision against the admin ceiling via the +// STRICTEST of the two, so a remote reject_hard / modify overrides an approve +// or redact ceiling at runtime. +type RuntimeEscalatable interface { + MayExceedOnMatch() bool +} + // ChatOnly is the applicability helper for hooks that exclusively apply to // text-based chat traffic. Embed it into a hook struct to satisfy // SupportsEndpoint and SupportsModality: SupportsEndpoint returns true only @@ -308,10 +334,13 @@ type RawContentPrescanner interface { // } type ChatOnly struct{} -// SupportsEndpoint returns true for EndpointTypeChat and for the empty string -// (backward-compatible default when the caller has not yet classified the -// endpoint). -func (ChatOnly) SupportsEndpoint(e EndpointType) bool { return e == EndpointTypeChat || e == "" } +// SupportsEndpoint returns true for EndpointTypeChat, EndpointTypeResponses +// (the Responses API is chat-family — same conversational text content), and +// the empty string (backward-compatible default when the caller has not yet +// classified the endpoint). +func (ChatOnly) SupportsEndpoint(e EndpointType) bool { + return e == EndpointTypeChat || e == EndpointTypeResponses || e == "" +} // SupportsModality returns true for ModalityText and for the empty string. func (ChatOnly) SupportsModality(m Modality) bool { return m == ModalityText || m == "" } diff --git a/packages/shared/policy/hooks/matcher/match_bound.go b/packages/shared/policy/hooks/matcher/match_bound.go new file mode 100644 index 00000000..89779ab9 --- /dev/null +++ b/packages/shared/policy/hooks/matcher/match_bound.go @@ -0,0 +1,162 @@ +package matcher + +import ( + "regexp/syntax" + "unicode" + "unicode/utf8" +) + +// MaxMatchBytes returns a conservative UPPER BOUND, in bytes, on the longest +// contiguous text a single match of `expr` can span, and whether that bound is +// FINITE. It feeds the Model-A streaming engine's flush-before-deliver lookahead +// (modela.Config.MaxPatternBytes): a value straddling a held-unit boundary must be +// scanned over a window at least this wide before the leading unit is delivered, or +// it can leak across two raw deliveries. Correctness direction is asymmetric: +// +// - OVER-estimating is safe — a wider lookahead only costs extra (sound) scans. +// - UNDER-estimating reopens the boundary leak. +// +// So every node rounds UP (char classes count utf8.UTFMax bytes), and any UNBOUNDED +// repeat — `*`, `+`, or `{n,}` — makes the whole match unbounded (`bounded=false`). +// Treating `{n,}` as unbounded (rather than mirroring the vectorscan detection cap of +// detectionRepeatCap) is deliberate: the cap is `//go:build vectorscan` only, so the +// re2 fallback build does NOT cap, and a value an unbounded pattern matches can be +// window-sized regardless — that is the disclosed over-window best-effort surface, for +// which the finite lookahead cannot help anyway. An unbounded pattern therefore does +// NOT contribute to MaxPatternBytes (the caller leaves such patterns best-effort). +// +// A parse error yields (0, false): an unparseable pattern is treated as unbounded +// (conservative — it never shrinks the lookahead). +func MaxMatchBytes(expr string) (n int, bounded bool) { + re, err := syntax.Parse(expr, syntax.Perl) + if err != nil { + return 0, false + } + return maxMatchBytes(re) +} + +func maxMatchBytes(re *syntax.Regexp) (int, bool) { + switch re.Op { + case syntax.OpLiteral: + // Sum each literal rune's UTF-8 width — the true byte length (over-estimating ASCII + // at utf8.UTFMax would inflate the lookahead toward the tail window). A case-folded + // literal needs the orbit walk: Go stores `(?i)k` as a single OpLiteral carrying only + // the MINIMUM fold rune ('K'/'k'), yet it matches U+212A (Kelvin, 3 bytes), so the + // width must be taken over the whole simple-fold orbit or the bound under-counts. + fold := re.Flags&syntax.FoldCase != 0 + n := 0 + for _, r := range re.Rune { + n += foldedRuneBytes(r, fold) + } + return n, true + case syntax.OpCharClass: + // One char from the class: its widest member's UTF-8 byte width (an ASCII-only + // class is 1 byte, NOT 4). UTF-8 width is monotonic in code point, so the widest + // rune in each [lo,hi] range is hi. + w := 0 + for i := 0; i+1 < len(re.Rune); i += 2 { + if b := runeBytes(re.Rune[i+1]); b > w { + w = b + } + } + if w == 0 { + w = utf8.UTFMax // empty/degenerate class: conservative + } + return w, true + case syntax.OpAnyChar, syntax.OpAnyCharNotNL: + return utf8.UTFMax, true // any rune → up to 4 bytes + case syntax.OpEmptyMatch, + syntax.OpBeginLine, syntax.OpEndLine, + syntax.OpBeginText, syntax.OpEndText, + syntax.OpWordBoundary, syntax.OpNoWordBoundary: + return 0, true // zero-width + case syntax.OpCapture: + return maxMatchBytes(re.Sub[0]) + case syntax.OpQuest: + // 0 or 1 occurrence → the max is one occurrence. + return maxMatchBytes(re.Sub[0]) + case syntax.OpStar, syntax.OpPlus: + return 0, false // unbounded + case syntax.OpConcat: + total := 0 + for _, sub := range re.Sub { + sn, sb := maxMatchBytes(sub) + if !sb { + return 0, false + } + total += sn + } + return total, true + case syntax.OpAlternate: + best := 0 + for _, sub := range re.Sub { + sn, sb := maxMatchBytes(sub) + if !sb { + return 0, false + } + if sn > best { + best = sn + } + } + return best, true + case syntax.OpRepeat: + // `{n,m}` is bounded by m; `{n,}` (Max == -1) is unbounded. + if re.Max < 0 { + return 0, false + } + sn, sb := maxMatchBytes(re.Sub[0]) + if !sb { + return 0, false + } + return sn * re.Max, true + default: + // OpNoMatch and any future op: treat as unbounded (conservative). + return 0, false + } +} + +// runeBytes is the UTF-8 byte width of r, conservatively utf8.UTFMax for an invalid or +// out-of-range rune (e.g. a surrogate) so the bound never under-counts. +func runeBytes(r rune) int { + if n := utf8.RuneLen(r); n > 0 { + return n + } + return utf8.UTFMax +} + +// foldedRuneBytes is the widest UTF-8 byte width r can match. When the literal is NOT +// case-folded it is just runeBytes(r); when it IS folded it is the max over r's +// simple-fold orbit (e.g. {'K','k',U+212A} → 3), because a folded OpLiteral stores only +// the orbit's minimum rune but matches any orbit member — taking only runeBytes(r) would +// under-count (the reopened-boundary-leak failure mode). +func foldedRuneBytes(r rune, fold bool) int { + w := runeBytes(r) + if !fold { + return w + } + for f := unicode.SimpleFold(r); f != r; f = unicode.SimpleFold(f) { + if b := runeBytes(f); b > w { + w = b + } + } + return w +} + +// MaxBoundedMatchBytes returns the largest finite MaxMatchBytes across exprs (0 when +// none are bounded), and whether ANY expr was unbounded. The bounded max sizes the +// Model-A lookahead so every bounded-length enforceable value straddling a unit +// boundary is caught; unbounded patterns are reported so the caller can disclose them +// as best-effort (a value they match can exceed the window). +func MaxBoundedMatchBytes(exprs []string) (maxBounded int, anyUnbounded bool) { + for _, e := range exprs { + n, bounded := MaxMatchBytes(e) + if !bounded { + anyUnbounded = true + continue + } + if n > maxBounded { + maxBounded = n + } + } + return maxBounded, anyUnbounded +} diff --git a/packages/shared/policy/hooks/matcher/match_bound_test.go b/packages/shared/policy/hooks/matcher/match_bound_test.go new file mode 100644 index 00000000..ec78aabe --- /dev/null +++ b/packages/shared/policy/hooks/matcher/match_bound_test.go @@ -0,0 +1,138 @@ +package matcher + +import "testing" + +// TestMaxMatchBytes pins the conservative upper-bound walk that sizes the Model-A +// flush-before-deliver lookahead. Under-estimation reopens the boundary leak, so the +// bound must be ≥ the true longest match for every BOUNDED pattern, and every +// UNBOUNDED repeat (*, +, {n,}) must report bounded=false. +func TestMaxMatchBytes(t *testing.T) { + const utfMax = 4 + cases := []struct { + name string + expr string + wantBounded bool + // wantMin is a lower bound the result must meet/exceed when bounded (the walk + // over-estimates char classes at utf8.UTFMax, so we assert ≥ the byte length of + // the longest ASCII match, which it must never go below). + wantMin int + }{ + {"literal", "abc", true, 3}, + {"char class one", "[0-9]", true, 1}, + {"fixed repeat card", `[0-9]{16}`, true, 16}, + // Go regexp caps a single repeat at 1000; a long BOUNDED match is a CONCAT of + // repeats (the only way past the 4096 floor) — the medium-2 long-pattern case. + {"long bounded concat (medium-2)", `[A-Za-z0-9]{1000}[A-Za-z0-9]{200}`, true, 1200}, + {"bounded range", `[a-z]{1,64}`, true, 64}, + {"two-sided range", `[a-z]{2,24}`, true, 24}, + {"concat", `abc[0-9]{16}`, true, 3 + 16}, + {"alternate takes max", `(abc|[0-9]{16})`, true, 16}, + {"quest one occurrence", "a?", true, 0}, + {"anchored literal", "^abc$", true, 3}, + {"optional then fixed", `x?[0-9]{8}`, true, 8}, + // A surrogate rune has utf8.RuneLen == -1; runeBytes falls back to utf8.UTFMax so + // the bound never under-counts. + {"surrogate literal conservative", `\x{d800}`, true, 4}, + + {"star unbounded", "a*", false, 0}, + {"plus unbounded", "a+", false, 0}, + {"open repeat unbounded", `[a-z]{20,}`, false, 0}, + {"concat with star unbounded", `abc.*def`, false, 0}, + {"alternate with plus unbounded", `(abc|x+)`, false, 0}, + {"bounded repeat over unbounded sub", `(a*){2,3}`, false, 0}, + {"parse error unbounded", `[unterminated`, false, 0}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + n, bounded := MaxMatchBytes(c.expr) + if bounded != c.wantBounded { + t.Fatalf("MaxMatchBytes(%q) bounded = %v, want %v (n=%d)", c.expr, bounded, c.wantBounded, n) + } + if !bounded { + return + } + // The walk counts each class char at utf8.UTFMax, so the bound must be ≥ the + // ASCII-byte length (utfMax× for class-heavy patterns). Never below wantMin. + if n < c.wantMin { + t.Fatalf("MaxMatchBytes(%q) = %d, want ≥ %d (must not under-estimate)", c.expr, n, c.wantMin) + } + // And never absurdly large: a finite bound for these stays within utfMax× the + // char count (sanity that we didn't return a runaway). + if n > c.wantMin*utfMax+utfMax { + t.Fatalf("MaxMatchBytes(%q) = %d, unexpectedly large vs wantMin %d", c.expr, n, c.wantMin) + } + }) + } +} + +// TestMaxMatchBytes_TightByteWidth pins that ASCII constructs bound to their EXACT byte +// length (1 byte/char), not a blanket utf8.UTFMax (4×) — a 4× inflation would push a +// realistic ASCII rule set toward the tail window and degrade the batching guard, while a +// genuinely multi-byte class member is still counted at its real (upper-bound) UTF-8 width. +func TestMaxMatchBytes_TightByteWidth(t *testing.T) { + if n, _ := MaxMatchBytes(`[a-z]{100}`); n != 100 { + t.Fatalf("ascii class {100} = %d, want exactly 100 (1 byte/char, not 4x)", n) + } + if n, _ := MaxMatchBytes(`abcdef`); n != 6 { + t.Fatalf("ascii literal = %d, want 6", n) + } + // U+1000–U+1010 are 3-byte runes; one char from the class is counted at 3 (its real + // width, an exact upper bound), not 4. + if n, _ := MaxMatchBytes(`[\x{1000}-\x{1010}]{10}`); n != 30 { + t.Fatalf("3-byte-rune class {10} = %d, want 30 (3 bytes/char)", n) + } +} + +// TestMaxMatchBytes_CaseFoldedLiteral pins the fold-orbit fix: Go stores a (?i) literal as a +// single OpLiteral carrying only the orbit's MINIMUM rune, yet it matches any orbit member, so +// the bound must take the widest member or it under-counts (the reopened-boundary-leak mode the +// first tightening introduced). (?i)k matches U+212A (Kelvin, 3B); (?i)password matches paſſword +// (long-s U+017F ×2). A non-wide orbit stays 1 byte/char. +func TestMaxMatchBytes_CaseFoldedLiteral(t *testing.T) { + if n, b := MaxMatchBytes(`(?i)k`); !b || n != 3 { + t.Fatalf("(?i)k = %d bounded=%v, want 3 (Kelvin sign U+212A)", n, b) + } + if n, b := MaxMatchBytes(`(?i)password`); !b || n < 10 { + t.Fatalf("(?i)password = %d bounded=%v, want >= 10 (long-s ×2 ⇒ paſſword)", n, b) + } + if n, _ := MaxMatchBytes(`(?i)abc`); n != 3 { + t.Fatalf("(?i)abc = %d, want 3 (no wide fold members)", n) + } +} + +// TestMaxBoundedMatchBytes pins the aggregate: the largest finite bound across exprs, +// and that ANY unbounded expr is flagged (so the caller discloses it best-effort). +func TestMaxBoundedMatchBytes(t *testing.T) { + // Mixed: a short PII pattern, a long bounded CONCAT pattern, and an unbounded one. + maxB, anyUnb := MaxBoundedMatchBytes([]string{ + `[0-9]{16}`, // 64 bytes bounded + `[A-Za-z0-9]{1000}[A-Za-z0-9]{1000}`, // 2000 chars bounded — the long-pattern case + `[A-Za-z0-9+/=]{20,}`, // unbounded (JWT/base64 blob) → best-effort + }) + if maxB < 2000 { + t.Fatalf("MaxBoundedMatchBytes max = %d, want ≥ 2000 chars (the long bounded concat)", maxB) + } + if !anyUnb { + t.Fatal("MaxBoundedMatchBytes anyUnbounded = false, want true (the {20,} pattern is unbounded)") + } + + // All bounded, all short → small max, no unbounded. + maxB2, anyUnb2 := MaxBoundedMatchBytes([]string{`[0-9]{16}`, `abc`}) + if anyUnb2 { + t.Fatal("anyUnbounded = true for all-bounded exprs") + } + if maxB2 == 0 || maxB2 > 64+8 { + t.Fatalf("max = %d, want a small bounded value (~64)", maxB2) + } + + // All unbounded → max 0, anyUnbounded true. + maxB3, anyUnb3 := MaxBoundedMatchBytes([]string{`a*`, `b+`}) + if maxB3 != 0 || !anyUnb3 { + t.Fatalf("all-unbounded: max=%d anyUnbounded=%v, want 0/true", maxB3, anyUnb3) + } + + // Empty input → 0, false. + if m, u := MaxBoundedMatchBytes(nil); m != 0 || u { + t.Fatalf("empty: max=%d anyUnbounded=%v, want 0/false", m, u) + } +} diff --git a/packages/shared/policy/hooks/webhook/webhook.go b/packages/shared/policy/hooks/webhook/webhook.go index 2da9c1a2..7954d2b0 100644 --- a/packages/shared/policy/hooks/webhook/webhook.go +++ b/packages/shared/policy/hooks/webhook/webhook.go @@ -30,8 +30,8 @@ import ( // shape (redactions[] of {start, end, replacement, action, reason}), // webhook-forward decodes each redaction into a normalize.TransformSpan. // The offsets index a flat joined projection; spans use -// ContentAddress = "webhook.flat" which ApplySpans does not resolve, so -// they land in the audit row but are not applied inflight. Inflight +// ContentAddress = normalize.AddressAuditOnlySentinel which ApplySpans does +// not resolve, so they land in the audit row but are not applied inflight. Inflight // redaction of AI-Guard suggestions requires the internal aiguard-classify // path in ai-gateway. // @@ -140,6 +140,17 @@ func NewWebhookForwardWithClient(cfg *core.HookConfig, client *http.Client) (cor }, nil } +// MayExceedOnMatch satisfies core.RuntimeEscalatable: webhook-forward can return +// a decision stricter than its declarative onMatch ceiling. Execute reconciles +// the remote endpoint's suggested decision against the ceiling via +// core.StrictestDecision, so a remote reject_hard or modify wins over an approve +// or redact ceiling. The streaming-routing predicates must therefore treat every +// webhook-forward hook as may-block AND may-redact regardless of its configured +// action, so its runtime enforcement is never under-routed onto the audit-only +// live path. Always true: the reconcile can escalate under any ceiling, so the +// safe over-route applies unconditionally. +func (w *WebhookForward) MayExceedOnMatch() bool { return true } + func (w *WebhookForward) Execute(ctx context.Context, input *core.HookInput) (*core.HookResult, error) { // Build payload. The envelope is always included; content visibility // is controlled by payloadMode. @@ -301,9 +312,9 @@ type webhookRedactionWire struct { // redact when the wire field is missing; explicit "strip" / "inject" / // "replace" values are honored so the audit record is faithful. // -// ContentAddress is "webhook.flat" because the wire offsets index a flat -// joined projection that webhook-forward did not construct (the -// compliance-webhook shim on the AI-Guard side builds it). ApplySpans does +// ContentAddress is normalize.AddressAuditOnlySentinel because the wire +// offsets index a flat joined projection that webhook-forward did not construct +// (the compliance-webhook shim on the AI-Guard side builds it). ApplySpans does // not resolve this sentinel, so spans land in the audit row but do not // mutate inflight bytes. func redactionsToTransformSpans(redactions []webhookRedactionWire, endpoint string) []normalize.TransformSpan { @@ -327,7 +338,7 @@ func redactionsToTransformSpans(redactions []webhookRedactionWire, endpoint stri Source: normalize.SourceHook, SourceID: endpoint, Action: action, - ContentAddress: "webhook.flat", + ContentAddress: normalize.AddressAuditOnlySentinel, Start: r.Start, End: r.End, Replacement: r.Replacement, diff --git a/packages/shared/policy/hooks/webhook/webhook_escalatable_test.go b/packages/shared/policy/hooks/webhook/webhook_escalatable_test.go new file mode 100644 index 00000000..195d1c94 --- /dev/null +++ b/packages/shared/policy/hooks/webhook/webhook_escalatable_test.go @@ -0,0 +1,37 @@ +package webhook + +import ( + "net/http" + "testing" + + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// TestWebhookForward_MayExceedOnMatch asserts webhook-forward satisfies +// core.RuntimeEscalatable and reports true, so the streaming-routing predicates +// (Pipeline.MayBlock / MayRedact) over-route it to buffer regardless of its +// declared onMatch ceiling. Without this, a webhook configured with an approve +// ceiling whose remote reply returns reject_hard / modify would be left on the +// audit-only live path and have its enforcement silently dropped (HIGH-1). +func TestWebhookForward_MayExceedOnMatch(t *testing.T) { + cfg := &core.HookConfig{ + ID: "wh-1", + ImplementationID: "webhook-forward", + Name: "test-webhook", + // No onMatch block → factory sets the advisory approve ceiling, the exact + // configuration HIGH-1 leaks under. + Config: map[string]any{"endpoint": "http://example.com"}, + } + h, err := NewWebhookForwardWithClient(cfg, http.DefaultClient) + if err != nil { + t.Fatalf("NewWebhookForwardWithClient: %v", err) + } + + esc, ok := h.(core.RuntimeEscalatable) + if !ok { + t.Fatal("webhook-forward does not satisfy core.RuntimeEscalatable; predicates cannot detect its runtime escalation") + } + if !esc.MayExceedOnMatch() { + t.Error("MayExceedOnMatch()=false; a webhook with an approve ceiling would leak its runtime block/redact onto the live path") + } +} diff --git a/packages/shared/policy/pipeline/audit_emitter.go b/packages/shared/policy/pipeline/audit_emitter.go index b3430a58..ffb48b6a 100644 --- a/packages/shared/policy/pipeline/audit_emitter.go +++ b/packages/shared/policy/pipeline/audit_emitter.go @@ -329,8 +329,8 @@ func (e *AuditEmitter) buildEvent( // populated by shared/traffic tracing during the forward roundtrip. // LatencyBreakdown carries compliance-proxy long-tail keys from // forward_handler (conn_setup_ms, tls_handshake_ms). - requestHooksMs := sumHooksPipelineLatencyMs(reqPipeline) - responseHooksMs := sumHooksPipelineLatencyMs(respPipeline) + requestHooksMs, requestHooksUs := sumHooksPipelineLatency(reqPipeline) + responseHooksMs, responseHooksUs := sumHooksPipelineLatency(respPipeline) var upstreamTtfb, upstreamTotal *int if info.PhaseSink != nil { upstreamTtfb = info.PhaseSink.TtfbMs() @@ -385,6 +385,8 @@ func (e *AuditEmitter) buildEvent( UpstreamTotalMs: upstreamTotal, RequestHooksMs: requestHooksMs, ResponseHooksMs: responseHooksMs, + RequestHooksUs: requestHooksUs, + ResponseHooksUs: responseHooksUs, LatencyBreakdown: latencyBreakdown, DomainRuleID: info.DomainRuleID, PathAction: info.PathAction, @@ -409,31 +411,37 @@ func stageAction(r *CompliancePipelineResult) decision.Action { return r.Action } -// sumHooksPipelineLatencyMs walks the hooks_pipeline JSONB blob (an array -// of {latencyMs: int, ...} objects produced by stagePayload) and returns -// the sum of per-hook latency in ms. Returns nil for an empty / unparseable -// blob so the resulting traffic_event column stays NULL — distinguishing -// "no hooks ran" from "hooks ran in 0ms". -func sumHooksPipelineLatencyMs(pipelineJSON []byte) *int { +// sumHooksPipelineLatency sums the per-hook latency in the hooks_pipeline JSONB +// blob, returning BOTH the truncated-millisecond and precise-microsecond totals +// from a SINGLE unmarshal. Hooks run at microsecond scale, so the ms total +// truncates sub-millisecond hooks to 0 while the us total carries the real value. +// Each total is nil for an empty / unparseable / empty-array blob so the resulting +// traffic_event columns stay NULL ("no hooks ran" vs "ran in 0"). Zero/negative +// per-hook values are skipped. +func sumHooksPipelineLatency(pipelineJSON []byte) (msSum, usSum *int) { if len(pipelineJSON) == 0 { - return nil + return nil, nil } var rows []struct { LatencyMs int `json:"latencyMs"` + LatencyUs int `json:"latencyUs"` } if err := json.Unmarshal(pipelineJSON, &rows); err != nil { - return nil + return nil, nil } if len(rows) == 0 { - return nil + return nil, nil } - total := 0 + ms, us := 0, 0 for _, r := range rows { if r.LatencyMs > 0 { - total += r.LatencyMs + ms += r.LatencyMs + } + if r.LatencyUs > 0 { + us += r.LatencyUs } } - return &total + return &ms, &us } // classifyComplianceError derives (errorCode, errorReason) from the available diff --git a/packages/shared/policy/pipeline/audit_emitter_helpers_test.go b/packages/shared/policy/pipeline/audit_emitter_helpers_test.go index a4d99932..09f93119 100644 --- a/packages/shared/policy/pipeline/audit_emitter_helpers_test.go +++ b/packages/shared/policy/pipeline/audit_emitter_helpers_test.go @@ -22,30 +22,36 @@ func TestNullableString_NonEmptyReturnsFreshPointer(t *testing.T) { } } -func TestSumHooksPipelineLatencyMs(t *testing.T) { +func TestSumHooksPipelineLatency(t *testing.T) { + eq := func(t *testing.T, got, want *int) { + t.Helper() + switch { + case got == nil && want == nil: + case got == nil || want == nil: + t.Errorf("got %v, want %v", got, want) + case *got != *want: + t.Errorf("got %d, want %d", *got, *want) + } + } cases := []struct { - name string - in []byte - want *int + name string + in []byte + wantMs, wantUs *int }{ - {"nil yields nil", nil, nil}, - {"empty yields nil", []byte{}, nil}, - {"malformed JSON yields nil", []byte(`{not-json`), nil}, - {"empty array yields nil", []byte(`[]`), nil}, - {"sums positive latencies", []byte(`[{"latencyMs":10},{"latencyMs":15}]`), intPtr(25)}, - {"skips zero / negative", []byte(`[{"latencyMs":10},{"latencyMs":0},{"latencyMs":-5}]`), intPtr(10)}, + {"nil yields nil/nil", nil, nil, nil}, + {"empty yields nil/nil", []byte{}, nil, nil}, + {"malformed JSON yields nil/nil", []byte(`{not-json`), nil, nil}, + {"empty array yields nil/nil", []byte(`[]`), nil, nil}, + {"sums both from one parse", []byte(`[{"latencyMs":10,"latencyUs":10200},{"latencyMs":15,"latencyUs":15400}]`), intPtr(25), intPtr(25600)}, + {"skips zero / negative", []byte(`[{"latencyMs":10,"latencyUs":10000},{"latencyMs":0,"latencyUs":0},{"latencyMs":-5,"latencyUs":-5}]`), intPtr(10), intPtr(10000)}, + // Sub-millisecond hooks: ms total is 0 (present, ran), us total is precise. + {"sub-ms visible only in us", []byte(`[{"latencyMs":0,"latencyUs":200},{"latencyMs":0,"latencyUs":150}]`), intPtr(0), intPtr(350)}, } for _, c := range cases { t.Run(c.name, func(t *testing.T) { - got := sumHooksPipelineLatencyMs(c.in) - switch { - case got == nil && c.want == nil: - // ok - case got == nil || c.want == nil: - t.Errorf("got %v, want %v", got, c.want) - case *got != *c.want: - t.Errorf("got %d, want %d", *got, *c.want) - } + gotMs, gotUs := sumHooksPipelineLatency(c.in) + eq(t, gotMs, c.wantMs) + eq(t, gotUs, c.wantUs) }) } } diff --git a/packages/shared/policy/pipeline/enforcement.go b/packages/shared/policy/pipeline/enforcement.go index a1ff3cbf..13166a81 100644 --- a/packages/shared/policy/pipeline/enforcement.go +++ b/packages/shared/policy/pipeline/enforcement.go @@ -5,8 +5,65 @@ import ( "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/decision" "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/matcher" ) +// MaxPatternBound returns the conservative upper bound + unbounded flag for the +// pipeline's resolved hook set. When the pipeline was built by the PolicyResolver +// (the production path) BuildPipeline pre-stamps the per-generation CACHED value +// (see maxPatternBoundFor), so this is O(1) on the streaming hot path; a pipeline +// built directly via NewPipeline (tests) computes it lazily. The underlying +// computeMaxPatternBound is O(total patterns) — it parses every hook regex — so it +// must NOT run per request; that is why the resolver memoises it per resolved set. +func (p *Pipeline) MaxPatternBound() (maxBounded int, anyUnbounded bool) { + if p.boundComputed { + return p.maxBounded, p.anyUnbounded + } + return computeMaxPatternBound(p.hooks) +} + +// computeMaxPatternBound returns a conservative UPPER BOUND, in bytes, on the longest +// contiguous enforceable match across all bound content hooks — sizing the Model-A +// streaming flush-before-deliver lookahead (modela.Config.MaxPatternBytes). The bound +// must never under-estimate (that reopens the boundary leak), so each pattern is walked +// with matcher.MaxMatchBytes (which rounds up and reports unbounded `*`/`+`/`{n,}` as +// not-bounded). anyUnbounded reports whether ANY bound content hook carries an unbounded +// pattern (or cannot export its patterns) — a value such a pattern matches can exceed +// the tail window, the disclosed best-effort surface the finite lookahead cannot cover; +// callers MAY surface it (e.g. a config-time operator signal). maxBounded is 0 when no +// bounded pattern exists (the engine then falls back to its package default). This is the +// expensive computation (one regexp/syntax walk per pattern); callers cache the result per +// resolved hook set rather than invoking it per request. +func computeMaxPatternBound(hooks []boundHook) (maxBounded int, anyUnbounded bool) { + for i := range hooks { + pre, ok := hooks[i].hook.(core.RawContentPrescanner) + if !ok || !pre.ScansContent() { + continue // metadata hook: no redactable patterns + } + src, ok := hooks[i].hook.(core.PrescanPatternSource) + if !ok { + anyUnbounded = true // scans content but can't export patterns → conservatively best-effort + continue + } + pats := src.PrescanPatterns() + if pats == nil { + anyUnbounded = true + continue + } + for _, pp := range pats { + n, bounded := matcher.MaxMatchBytes(pp.Expr) + if !bounded { + anyUnbounded = true + continue + } + if n > maxBounded { + maxBounded = n + } + } + } + return maxBounded, anyUnbounded +} + // enforcement.go — scope-derived, content-INDEPENDENT predicates the streaming // relay uses BEFORE any response bytes are read to decide whether a request must // run in buffered execution, plus the runtime fail-posture derivation for a hook @@ -24,7 +81,10 @@ import ( // resolved pipeline ever redact", never "does this body match a pattern". // Over-approximation (a redact-action hook whose rules happen not to fire still // routes to buffer) is the safe direction — it can never under-route a request -// that would have redacted onto the unbuffered streaming path. +// that would have redacted onto the unbuffered streaming path. A hook whose +// runtime decision can exceed its declared ceiling (core.RuntimeEscalatable, +// e.g. webhook-forward) or whose onMatch is unparseable is counted as redact- +// capable here regardless of its declared action — see anyOnMatchAction. func (p *Pipeline) MayRedact() bool { return p.anyOnMatchAction(decision.ActionRedact) } @@ -33,23 +93,43 @@ func (p *Pipeline) MayRedact() bool { // "block" (hard reject). Same scope-derived, content-independent contract as // MayRedact. ParseOnMatch resolves an absent onMatch to the "block" default, so a // content hook with no explicit action is counted as block — matching its real -// execution behaviour and keeping the buffer-routing direction safe. +// execution behaviour and keeping the buffer-routing direction safe. A hook whose +// runtime decision can exceed its declared ceiling (core.RuntimeEscalatable, e.g. +// webhook-forward) or whose onMatch is unparseable is counted as block-capable +// regardless of its declared action — see anyOnMatchAction. func (p *Pipeline) MayBlock() bool { return p.anyOnMatchAction(decision.ActionBlock) } -// anyOnMatchAction reports whether any bound hook's parsed onMatch action equals -// want. A malformed onMatch is skipped: the hook fail-closes at execution, and a -// parse error must not be read as an enforcing action. +// anyOnMatchAction reports whether any bound hook could enforce the wanted +// action on a match. It returns true in three cases, every one biased toward the +// safe over-route direction (route to buffer rather than leak onto the live +// path): +// +// - the hook's runtime decision can exceed its declarative onMatch ceiling +// (it implements core.RuntimeEscalatable and reports true — canonically +// webhook-forward, whose remote reply can block/redact under any ceiling). +// Such a hook is treated as may-block AND may-redact regardless of want, +// because reading only its declared action would under-route it and silently +// drop the runtime enforcement; +// - its parsed onMatch action equals want; +// - its onMatch is unparseable. A malformed config is treated conservatively +// as enforcing — mirroring hookIsEnforcing — so a misconfigured hook routes +// to buffer instead of leaking onto the unbuffered live path. (The two +// predicates previously disagreed here: anyOnMatchAction skipped a parse +// error while hookIsEnforcing counted it.) func (p *Pipeline) anyOnMatchAction(want decision.Action) bool { for i := range p.hooks { + if esc, ok := p.hooks[i].hook.(core.RuntimeEscalatable); ok && esc.MayExceedOnMatch() { + return true + } cfg := p.hooks[i].config if cfg == nil { continue } om, err := core.ParseOnMatch(cfg.Config) if err != nil { - continue + return true } if om.Action == want { return true @@ -80,25 +160,32 @@ func (p *Pipeline) SetStrictFailClosed(strict bool) { // "Enforcing" means the hook's onMatch action is redact or block (it would // rewrite or reject on a match), derived via core.ParseOnMatch on the hook's // declarative config. A non-enforcing (approve-scope) hook never fails closed. -func failClosedOnError(strict bool, cfg *core.HookConfig) bool { +func failClosedOnError(strict bool, hook core.Hook, cfg *core.HookConfig) bool { switch strings.ToLower(strings.TrimSpace(cfg.FailBehavior)) { case "fail-closed": return true case "fail-open": return false default: - return strict && hookIsEnforcing(cfg) + return strict && hookIsEnforcing(hook, cfg) } } // hookIsEnforcing reports whether the hook would redact or block on a match — -// the scope whose guaranteed-execution contract a transient error breaks. The -// onMatch action is parsed from the hook's declarative config; a hook with no -// explicit onMatch defaults to block (the security "block on match" default in -// core.ParseOnMatch), so a match-only enforcer counts as enforcing. An -// unparseable onMatch is treated conservatively as enforcing so a strict caller -// fails closed rather than leaking on a misconfigured rule. -func hookIsEnforcing(cfg *core.HookConfig) bool { +// the scope whose guaranteed-execution contract a transient error breaks. A hook +// whose runtime decision can EXCEED its declared onMatch ceiling (a +// RuntimeEscalatable, e.g. a webhook whose remote verdict reconciles to a stricter +// action) counts as enforcing regardless of its declared action — mirroring the +// routing predicate (anyOnMatchAction) so a strict caller fails such a hook's +// transient error closed rather than leaking. Otherwise the onMatch action is parsed +// from the hook's declarative config; a hook with no explicit onMatch defaults to +// block (the security "block on match" default in core.ParseOnMatch), so a match-only +// enforcer counts as enforcing. An unparseable onMatch is treated conservatively as +// enforcing so a strict caller fails closed rather than leaking on a misconfigured rule. +func hookIsEnforcing(hook core.Hook, cfg *core.HookConfig) bool { + if esc, ok := hook.(core.RuntimeEscalatable); ok && esc.MayExceedOnMatch() { + return true + } oc, err := core.ParseOnMatch(cfg.Config) if err != nil { return true diff --git a/packages/shared/policy/pipeline/maxpatternbound_cache_test.go b/packages/shared/policy/pipeline/maxpatternbound_cache_test.go new file mode 100644 index 00000000..aac1b0ae --- /dev/null +++ b/packages/shared/policy/pipeline/maxpatternbound_cache_test.go @@ -0,0 +1,175 @@ +package pipeline + +import ( + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// TestMaxPatternBoundFor_CachePerGen pins the resolver memo: within one generation a +// repeat call is a cache hit, and a generation bump drops the previous generation's +// entries and recomputes — so a config reload can never serve a stale (too-small) bound. +func TestMaxPatternBoundFor_CachePerGen(t *testing.T) { + r := NewPolicyResolver(nil, core.NewHookRegistry(), nil) + hooks := []boundHook{bh(newPrescanHook(t, `[0-9]{16}`), "h1")} + want, wantU := computeMaxPatternBound(hooks) + if want == 0 { + t.Fatal("setup: expected a bounded pattern (>0)") + } + + // First call (gen 0): computes + caches under gen 0. + if b, u := r.maxPatternBoundFor(hooks, 0); b != want || u != wantU { + t.Fatalf("gen0: got (%d,%v) want (%d,%v)", b, u, want, wantU) + } + r.boundMu.Lock() + n, g := len(r.boundByKey), r.boundGen + r.boundMu.Unlock() + if n != 1 || g != 0 { + t.Fatalf("gen0 cache: len=%d gen=%d, want 1 / 0", n, g) + } + + // Same gen: cache hit, identical values. + if b, u := r.maxPatternBoundFor(hooks, 0); b != want || u != wantU { + t.Fatalf("gen0 hit: got (%d,%v)", b, u) + } + + // Generation bump: the gen-0 map is dropped, recomputed under gen 1. + if b, _ := r.maxPatternBoundFor(hooks, 1); b != want { + t.Fatalf("gen1: got %d want %d", b, want) + } + r.boundMu.Lock() + g = r.boundGen + r.boundMu.Unlock() + if g != 1 { + t.Fatalf("boundGen did not advance to 1 on a gen bump: %d", g) + } +} + +// TestMaxPatternBoundFor_AmbiguousKeyBypassesCache: a duplicate config ID makes the +// cache key ambiguous, so the bound is computed directly (correct value) and NOT cached +// — soundness over the optimisation, mirroring the union memo's fallback. +func TestMaxPatternBoundFor_AmbiguousKeyBypassesCache(t *testing.T) { + r := NewPolicyResolver(nil, core.NewHookRegistry(), nil) + dup := []boundHook{bh(newPrescanHook(t, `[0-9]{4}`), "x"), bh(newPrescanHook(t, `[a-f]{2}`), "x")} + want, _ := computeMaxPatternBound(dup) + if b, _ := r.maxPatternBoundFor(dup, 0); b != want { + t.Fatalf("ambiguous-key direct compute: got %d want %d", b, want) + } + r.boundMu.Lock() + n := len(r.boundByKey) + r.boundMu.Unlock() + if n != 0 { + t.Fatalf("ambiguous key must not populate the cache, got %d entries", n) + } +} + +// TestMaxPatternBound_StampedVsLazy: a BuildPipeline-stamped pipeline returns the cached +// fields (O(1) hot path); a NewPipeline-built one computes lazily over its hooks. +func TestMaxPatternBound_StampedVsLazy(t *testing.T) { + stamped := &Pipeline{boundComputed: true, maxBounded: 1234, anyUnbounded: true} + if b, u := stamped.MaxPatternBound(); b != 1234 || !u { + t.Fatalf("stamped: got (%d,%v) want (1234,true)", b, u) + } + lazy := newTestPipeline([]boundHook{bh(newPrescanHook(t, `[0-9]{8}`), "h")}) + if b, _ := lazy.MaxPatternBound(); b == 0 { + t.Fatal("lazy path should compute a bounded value (>0)") + } +} + +// TestBuildPipeline_StampsBoundResponseStageOnly: the response-stage build pre-stamps the +// cached bound (gen threaded from the same atomic snapshot) and populates the cache; a +// request-stage build does NOT stamp (only Model-A streaming on the response stage reads +// the bound), so it stays lazy. +func TestBuildPipeline_StampsBoundResponseStageOnly(t *testing.T) { + reg := core.NewHookRegistry() + reg.Register("prescan-impl", func(_ *core.HookConfig) (core.Hook, error) { + return newPrescanHook(t, `[0-9]{16}`), nil + }) + mkCfg := func(stage string) core.HookConfig { + return core.HookConfig{ID: "h-" + stage, ImplementationID: "prescan-impl", Name: "h", Enabled: true, Stage: stage, ApplicableIngress: []string{"ALL"}} + } + r := NewPolicyResolver([]core.HookConfig{mkCfg("response"), mkCfg("request")}, reg, testLogger()) + + resp, err := r.BuildPipeline("response", "AI_GATEWAY", "", nil, time.Second, 5*time.Second, false, true, testLogger()) + if err != nil || resp == nil { + t.Fatalf("response BuildPipeline: err=%v nil=%v", err, resp == nil) + } + if !resp.boundComputed { + t.Fatal("response-stage build must pre-stamp the bound") + } + if b, _ := resp.MaxPatternBound(); b != mustBound(t, resp.hooks) { + t.Fatalf("stamped bound %d != computed %d", b, mustBound(t, resp.hooks)) + } + r.boundMu.Lock() + cached := len(r.boundByKey) + r.boundMu.Unlock() + if cached == 0 { + t.Fatal("response build should populate the bound cache") + } + + req, err := r.BuildPipeline("request", "AI_GATEWAY", "", nil, time.Second, 5*time.Second, false, true, testLogger()) + if err != nil || req == nil { + t.Fatalf("request BuildPipeline: err=%v nil=%v", err, req == nil) + } + if req.boundComputed { + t.Fatal("request-stage build must NOT pre-stamp the bound (only response consumes it)") + } +} + +func mustBound(t *testing.T, hooks []boundHook) int { + t.Helper() + b, _ := computeMaxPatternBound(hooks) + return b +} + +// TestMaxPatternBoundFor_ConcurrentGenBumpsAlwaysCorrect hammers maxPatternBoundFor with +// monotonically advancing generations from many goroutines (mirroring concurrent +// BuildPipeline calls racing config reloads). Run under -race, it exercises the +// reset-on-read + off-lock-compute + store-guard (r.boundGen == gen) interleavings and +// asserts every call returns the value its own hook set computes — a stale-gen result can +// never be served. The bound is a pure function of the hooks, so correctness is gen-independent. +func TestMaxPatternBoundFor_ConcurrentGenBumpsAlwaysCorrect(t *testing.T) { + r := NewPolicyResolver(nil, core.NewHookRegistry(), nil) + hooks := []boundHook{bh(newPrescanHook(t, `[0-9]{16}`), "h1")} + want, wantU := computeMaxPatternBound(hooks) + + var gen atomic.Uint64 + var wg sync.WaitGroup + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + for i := range 2000 { + g := gen.Load() + if i%50 == 0 { + g = gen.Add(1) // some goroutines advance the generation + } + if b, u := r.maxPatternBoundFor(hooks, g); b != want || u != wantU { + t.Errorf("concurrent maxPatternBoundFor returned (%d,%v), want (%d,%v)", b, u, want, wantU) + return + } + } + }() + } + wg.Wait() +} + +// TestConfigSnapshot_GenLockstepWithSwapGen guards the bundle invariant: after a Swap the +// atomically-published configSnapshot.gen equals the monotonic swapGen counter. A future +// edit that bumps one without the other would mis-tag cache entries; this pins it. +func TestConfigSnapshot_GenLockstepWithSwapGen(t *testing.T) { + r := NewPolicyResolver(nil, core.NewHookRegistry(), nil) + if g, _ := r.loadSnapshot(); g != r.swapGen.Load() { + t.Fatalf("initial: configSnapshot.gen=%d != swapGen=%d", g, r.swapGen.Load()) + } + for i := range 3 { + r.Swap([]core.HookConfig{{ID: "h", ImplementationID: "x", Name: "h", Enabled: true, Stage: "response"}}) + g, _ := r.loadSnapshot() + if sg := r.swapGen.Load(); g != sg { + t.Fatalf("after Swap %d: configSnapshot.gen=%d != swapGen=%d", i, g, sg) + } + } +} diff --git a/packages/shared/policy/pipeline/merge.go b/packages/shared/policy/pipeline/merge.go new file mode 100644 index 00000000..026a58b8 --- /dev/null +++ b/packages/shared/policy/pipeline/merge.go @@ -0,0 +1,179 @@ +package pipeline + +// merge.go — result aggregation: folding the per-hook results into one +// CompliancePipelineResult (decision strictness merge, tag union, redaction +// artifact carry + applicability flag). Split out of pipeline.go along the +// aggregation seam. + +import ( + "sort" + + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" + normalize "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/normalize/core" +) + +// mergeResults aggregates individual hook results into a single pipeline result. +// +// Decision merging: +// - First REJECT_HARD wins overall +// - Any BLOCK_SOFT (with no REJECT_HARD) => BLOCK_SOFT +// - All APPROVE/ABSTAIN => APPROVE +// +// "First" is by hook priority (HookResult.Order), not by arrival order. +// executeParallel appends results in goroutine-completion order, so the +// raw slice can be in any order; sort up front by Order so the BLOCK_SOFT +// / Modify Reason+ReasonCode tie-breaks are deterministic across runs. +// +// Tags: union of all hook-emitted tags, sorted alphabetically and deduplicated. +func (p *Pipeline) mergeResults(results []core.HookResult) *core.CompliancePipelineResult { + sort.SliceStable(results, func(i, j int) bool { + return results[i].Order < results[j].Order + }) + + pr := &core.CompliancePipelineResult{ + Decision: core.Approve, + HookResults: results, + } + + // Merge tags from every executed hook (set union, sorted, deduped) up front, + // so tags from earlier hooks survive even if a later hook short-circuits + // the decision loop below via REJECT_HARD. + tagSet := make(map[string]struct{}) + for i := range results { + for _, tag := range results[i].Tags { + if tag == "" { + continue + } + tagSet[tag] = struct{}{} + } + } + merged := make([]string, 0, len(tagSet)) + for tag := range tagSet { + merged = append(merged, tag) + } + sort.Strings(merged) + pr.Tags = merged + + hasSoftReject := false + var softRejectReason, softRejectCode string + var softBlockingRule *core.BlockingRule + hasModify := false + // redactionApplicable drives CarriesRedaction()'s BlockSoft branch: an enforcing hook + // contributed an artifact the AGGREGATE carries. ASYMMETRIC — the merge carries + // ModifiedContent only from Modify hooks, so a Modify hook qualifies on ModifiedContent + // OR an applicable span, a BlockSoft hook on an applicable span ONLY. Advisory + // Approve/Abstain spans (an approve-webhook's audit-only sentinel) never qualify. + redactionApplicable := false + // First Modify hook's Reason / ReasonCode wins so a specific + // reason (e.g. ReasonAIGuardSuggestedVsPolicy stamped at the + // webhook-forward reconcile) propagates to CompliancePipelineResult + // instead of being clobbered by the generic "CONTENT_MODIFIED" default. + var modifyReason, modifyReasonCode string + var lastModifiedContent []core.ContentBlock + var allSpans []normalize.TransformSpan + + for i := range results { + r := &results[i] + // Aggregate spans from every hook regardless of terminal decision — + // even Approve hooks may emit informational transforms (e.g. + // cache-normaliser strips through a hook integration). Storage + // rewrite at the audit-write stage walks this aggregate. + if len(r.TransformSpans) > 0 { + allSpans = append(allSpans, r.TransformSpans...) + } + + switch r.Decision { + case core.RejectHard: + pr.Decision = core.RejectHard + pr.Reason = r.Reason + pr.ReasonCode = r.ReasonCode + pr.BlockingRule = r.BlockingRule + pr.TransformSpans = allSpans + pr.Action = core.ActionFromDecision(core.RejectHard) + return pr + case core.BlockSoft: + hasSoftReject = true + softRejectReason = r.Reason + softRejectCode = r.ReasonCode + if softBlockingRule == nil { + softBlockingRule = r.BlockingRule + } + // Applicable only via spans — the merge does NOT carry a BlockSoft + // hook's ModifiedContent into the aggregate. + if !redactionApplicable && anyApplicableSpan(r.TransformSpans) { + redactionApplicable = true + } + case core.Modify: + if !hasModify { + // First Modify hook's reason wins so a hook-stamped + // ReasonCode (e.g. ReasonAIGuardSuggestedVsPolicy) is not + // silently replaced by the generic "CONTENT_MODIFIED". + modifyReason = r.Reason + modifyReasonCode = r.ReasonCode + } + hasModify = true + if len(r.ModifiedContent) > 0 { + lastModifiedContent = r.ModifiedContent + } + // Applicable via the carried ModifiedContent OR an applicable span. + if !redactionApplicable && (len(r.ModifiedContent) > 0 || anyApplicableSpan(r.TransformSpans)) { + redactionApplicable = true + } + case core.Approve: + if p.clearSoftOnApprove { + hasSoftReject = false + softRejectReason = "" + softRejectCode = "" + softBlockingRule = nil + } + } + } + + pr.TransformSpans = allSpans + // Carry the redaction payload UNCONDITIONALLY, symmetric with TransformSpans above. + // When a redact hook (Modify + ModifiedContent + spans) co-fires with a soft-block + // hook, StrictestDecision promotes the aggregate Decision to BlockSoft; previously + // ModifiedContent was set only in the `else if hasModify` branch and was therefore + // DROPPED, leaving spans without their replacement content. Downstream consumers + // then could not apply the redaction and fell back to fail-closed (block) OR + // replay-original (leak), depending on the path. Carrying it here lets every + // consumer that keys on CarriesRedaction (not Decision==Modify) redact-and-deliver. + // lastModifiedContent is non-nil only when a Modify hook captured content; Approve + // and RejectHard (which short-circuits above) never populate it, so no non-redact + // decision gains content. + pr.ModifiedContent = lastModifiedContent + + if hasSoftReject { + pr.Decision = core.BlockSoft + pr.Reason = softRejectReason + pr.ReasonCode = softRejectCode + pr.BlockingRule = softBlockingRule + } else if hasModify { + pr.Decision = core.Modify + if modifyReason != "" { + pr.Reason = modifyReason + } else { + pr.Reason = "content modified by hook pipeline" + } + if modifyReasonCode != "" { + pr.ReasonCode = modifyReasonCode + } else { + pr.ReasonCode = "CONTENT_MODIFIED" + } + } + pr.RedactionApplicable = redactionApplicable + pr.Action = core.ActionFromDecision(pr.Decision) + return pr +} + +// anyApplicableSpan reports whether any span carries an APPLICABLE redaction — one whose +// ContentAddress is not an audit-only sentinel (a DENYLIST: an unknown/future address counts +// as applicable, the fail-safe over-block-never-leak direction). Short-circuits on first hit. +func anyApplicableSpan(spans []normalize.TransformSpan) bool { + for i := range spans { + if !normalize.IsAuditOnlySentinelAddress(spans[i].ContentAddress) { + return true + } + } + return false +} diff --git a/packages/shared/policy/pipeline/pipeline.go b/packages/shared/policy/pipeline/pipeline.go index 224426d6..98fce6a0 100644 --- a/packages/shared/policy/pipeline/pipeline.go +++ b/packages/shared/policy/pipeline/pipeline.go @@ -58,6 +58,15 @@ type Pipeline struct { // PolicyResolver). MayMatchRawContent scans it ONCE instead of looping one // cgo scan per hook. nil => use the per-hook loop. See unionprescan.go. unionPrescan matcher.Matcher + + // boundComputed/maxBounded/anyUnbounded carry the pre-stamped MaxPatternBound + // result. BuildPipeline sets these from the resolver's per-generation cache + // (maxPatternBoundFor) so MaxPatternBound() is O(1) on the streaming hot path + // instead of re-walking every hook regex per request. boundComputed is false on + // a NewPipeline-built pipeline (tests), where MaxPatternBound computes lazily. + boundComputed bool + maxBounded int + anyUnbounded bool } // SetAllowModify enables MODIFY decision passthrough (for ai-gateway). @@ -342,6 +351,7 @@ func (p *Pipeline) executeOneHook(ctx context.Context, bh *boundHook, input *cor ImplementationID: bh.config.ImplementationID, HookName: hookName, LatencyMs: int(elapsed.Milliseconds()), + LatencyUs: int(elapsed.Microseconds()), Error: err.Error(), } @@ -350,7 +360,7 @@ func (p *Pipeline) executeOneHook(ctx context.Context, bh *boundHook, input *cor // win; otherwise a strict (non-packet-path) caller fails an ENFORCING // hook closed so a transient error cannot leak PII/secrets on the very // rule meant to stop them, while packet-path callers stay fail-open. - if failClosedOnError(p.strictFailClosed, bh.config) { + if failClosedOnError(p.strictFailClosed, bh.hook, bh.config) { hr.Decision = core.RejectHard hr.Reason = fmt.Sprintf("hook error (fail-closed): %v", err) hr.ReasonCode = "HOOK_ERROR_FAIL_CLOSED" @@ -380,7 +390,10 @@ func (p *Pipeline) executeOneHook(ctx context.Context, bh *boundHook, input *cor if result.HookName == "" { result.HookName = hookName } + // Both derive from the single captured elapsed: LatencyUs is precise, LatencyMs + // its integer-ms floor (sub-millisecond hooks → 0). Not clamped (see HookResult). result.LatencyMs = int(elapsed.Milliseconds()) + result.LatencyUs = int(elapsed.Microseconds()) // MODIFY passes through unconditionally. The downstream caller applies // TransformSpans via TrafficAdapter.RewriteRequestBody; protocols that @@ -397,130 +410,3 @@ func (p *Pipeline) executeOneHook(ctx context.Context, bh *boundHook, input *cor HookDecisionTotal.WithLabelValues(hookName, string(result.Decision)).Inc() return *result } - -// mergeResults aggregates individual hook results into a single pipeline result. -// -// Decision merging: -// - First REJECT_HARD wins overall -// - Any BLOCK_SOFT (with no REJECT_HARD) => BLOCK_SOFT -// - All APPROVE/ABSTAIN => APPROVE -// -// "First" is by hook priority (HookResult.Order), not by arrival order. -// executeParallel appends results in goroutine-completion order, so the -// raw slice can be in any order; sort up front by Order so the BLOCK_SOFT -// / Modify Reason+ReasonCode tie-breaks are deterministic across runs. -// -// Tags: union of all hook-emitted tags, sorted alphabetically and deduplicated. -func (p *Pipeline) mergeResults(results []core.HookResult) *core.CompliancePipelineResult { - sort.SliceStable(results, func(i, j int) bool { - return results[i].Order < results[j].Order - }) - - pr := &core.CompliancePipelineResult{ - Decision: core.Approve, - HookResults: results, - } - - // Merge tags from every executed hook (set union, sorted, deduped) up front, - // so tags from earlier hooks survive even if a later hook short-circuits - // the decision loop below via REJECT_HARD. - tagSet := make(map[string]struct{}) - for i := range results { - for _, tag := range results[i].Tags { - if tag == "" { - continue - } - tagSet[tag] = struct{}{} - } - } - merged := make([]string, 0, len(tagSet)) - for tag := range tagSet { - merged = append(merged, tag) - } - sort.Strings(merged) - pr.Tags = merged - - hasSoftReject := false - var softRejectReason, softRejectCode string - var softBlockingRule *core.BlockingRule - hasModify := false - // First Modify hook's Reason / ReasonCode wins so a specific - // reason (e.g. ReasonAIGuardSuggestedVsPolicy stamped at the - // webhook-forward reconcile) propagates to CompliancePipelineResult - // instead of being clobbered by the generic "CONTENT_MODIFIED" default. - var modifyReason, modifyReasonCode string - var lastModifiedContent []core.ContentBlock - var allSpans []normalize.TransformSpan - - for i := range results { - r := &results[i] - // Aggregate spans from every hook regardless of terminal decision — - // even Approve hooks may emit informational transforms (e.g. - // cache-normaliser strips through a hook integration). Storage - // rewrite at the audit-write stage walks this aggregate. - if len(r.TransformSpans) > 0 { - allSpans = append(allSpans, r.TransformSpans...) - } - - switch r.Decision { - case core.RejectHard: - pr.Decision = core.RejectHard - pr.Reason = r.Reason - pr.ReasonCode = r.ReasonCode - pr.BlockingRule = r.BlockingRule - pr.TransformSpans = allSpans - pr.Action = core.ActionFromDecision(core.RejectHard) - return pr - case core.BlockSoft: - hasSoftReject = true - softRejectReason = r.Reason - softRejectCode = r.ReasonCode - if softBlockingRule == nil { - softBlockingRule = r.BlockingRule - } - case core.Modify: - if !hasModify { - // First Modify hook's reason wins so a hook-stamped - // ReasonCode (e.g. ReasonAIGuardSuggestedVsPolicy) is not - // silently replaced by the generic "CONTENT_MODIFIED". - modifyReason = r.Reason - modifyReasonCode = r.ReasonCode - } - hasModify = true - if len(r.ModifiedContent) > 0 { - lastModifiedContent = r.ModifiedContent - } - case core.Approve: - if p.clearSoftOnApprove { - hasSoftReject = false - softRejectReason = "" - softRejectCode = "" - softBlockingRule = nil - } - } - } - - pr.TransformSpans = allSpans - - if hasSoftReject { - pr.Decision = core.BlockSoft - pr.Reason = softRejectReason - pr.ReasonCode = softRejectCode - pr.BlockingRule = softBlockingRule - } else if hasModify { - pr.Decision = core.Modify - if modifyReason != "" { - pr.Reason = modifyReason - } else { - pr.Reason = "content modified by hook pipeline" - } - if modifyReasonCode != "" { - pr.ReasonCode = modifyReasonCode - } else { - pr.ReasonCode = "CONTENT_MODIFIED" - } - pr.ModifiedContent = lastModifiedContent - } - pr.Action = core.ActionFromDecision(pr.Decision) - return pr -} diff --git a/packages/shared/policy/pipeline/pipeline_maxpattern_test.go b/packages/shared/policy/pipeline/pipeline_maxpattern_test.go new file mode 100644 index 00000000..917decde --- /dev/null +++ b/packages/shared/policy/pipeline/pipeline_maxpattern_test.go @@ -0,0 +1,85 @@ +package pipeline + +import ( + "testing" + + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// prescanStub is a content hook that exports a fixed PrescanPattern set, so a test can +// drive Pipeline.MaxPatternBound through the PrescanPatternSource path. +type prescanStub struct { + stubHook + pats []core.PrescanPattern +} + +func (p *prescanStub) ScansContent() bool { return true } +func (p *prescanStub) MayMatchRaw([]byte) bool { return true } +func (p *prescanStub) PrescanPatterns() []core.PrescanPattern { return p.pats } + +// scanOnlyStub scans content but does NOT export its patterns (not a +// PrescanPatternSource) — the pipeline must treat it conservatively as best-effort. +type scanOnlyStub struct { + stubHook +} + +func (*scanOnlyStub) ScansContent() bool { return true } +func (*scanOnlyStub) MayMatchRaw([]byte) bool { return true } + +func contentHook(exprs ...string) boundHook { + pats := make([]core.PrescanPattern, len(exprs)) + for i, e := range exprs { + pats[i] = core.PrescanPattern{Expr: e} + } + return boundHook{hook: &prescanStub{pats: pats}, config: &core.HookConfig{ID: "h", Name: "h"}} +} + +// TestPipeline_MaxPatternBound pins the Model-A lookahead derivation: the aggregate +// returns the largest finite bound across content hooks and flags any unbounded pattern +// (best-effort surface). Under-derivation would reopen the streaming boundary leak. +func TestPipeline_MaxPatternBound(t *testing.T) { + t.Run("short bounded pattern", func(t *testing.T) { + b, u := newTestPipeline([]boundHook{contentHook(`[0-9]{16}`)}).MaxPatternBound() + if b == 0 || u { + t.Fatalf("bound=%d unbounded=%v, want >0 / false", b, u) + } + }) + t.Run("long bounded concat exceeds the 4096 floor", func(t *testing.T) { + // Tight ASCII byte-width = 1/char, so exceeding the 4096 floor needs >4096 chars: + // five {1000} concats = 5000 bytes. This is the case the derivation exists for. + b, _ := newTestPipeline([]boundHook{contentHook(`[A-Za-z0-9]{1000}[A-Za-z0-9]{1000}[A-Za-z0-9]{1000}[A-Za-z0-9]{1000}[A-Za-z0-9]{1000}`)}).MaxPatternBound() + if b <= 4096 { + t.Fatalf("bound=%d, want > 4096 (the long concat must size the lookahead above the default floor)", b) + } + }) + t.Run("max across hooks", func(t *testing.T) { + b, _ := newTestPipeline([]boundHook{contentHook(`[0-9]{16}`), contentHook(`[A-Za-z]{500}`)}).MaxPatternBound() + if b < 500 { + t.Fatalf("bound=%d, want the max across hooks (≥ 500)", b) + } + }) + t.Run("unbounded pattern flagged", func(t *testing.T) { + b, u := newTestPipeline([]boundHook{contentHook(`[a-z]+`)}).MaxPatternBound() + if !u { + t.Fatalf("unbounded `+` pattern not flagged (bound=%d)", b) + } + }) + t.Run("non-content hook skipped", func(t *testing.T) { + b, u := newTestPipeline([]boundHook{boundWithAction("meta", "redact")}).MaxPatternBound() + if b != 0 || u { + t.Fatalf("non-content hook: bound=%d unbounded=%v, want 0 / false (no patterns to bound)", b, u) + } + }) + t.Run("content hook without pattern source flagged best-effort", func(t *testing.T) { + b, u := newTestPipeline([]boundHook{{hook: &scanOnlyStub{}, config: &core.HookConfig{ID: "x", Name: "x"}}}).MaxPatternBound() + if !u { + t.Fatalf("scan-only content hook not flagged unbounded (bound=%d) — it can't be bounded so must be best-effort", b) + } + }) + t.Run("nil pattern set flagged best-effort", func(t *testing.T) { + b, u := newTestPipeline([]boundHook{{hook: &prescanStub{pats: nil}, config: &core.HookConfig{ID: "n", Name: "n"}}}).MaxPatternBound() + if !u { + t.Fatalf("nil-pattern content hook not flagged unbounded (bound=%d)", b) + } + }) +} diff --git a/packages/shared/policy/pipeline/pipeline_mayredact_test.go b/packages/shared/policy/pipeline/pipeline_mayredact_test.go index 1ec142dd..5c9366b7 100644 --- a/packages/shared/policy/pipeline/pipeline_mayredact_test.go +++ b/packages/shared/policy/pipeline/pipeline_mayredact_test.go @@ -1,6 +1,8 @@ package pipeline import ( + "context" + "errors" "testing" "time" @@ -113,3 +115,111 @@ func TestPipeline_MayRedact_NilConfigSkipped(t *testing.T) { t.Error("nil-config hook must report neither redact nor block") } } + +// escalatableHook is a test double for a hook whose runtime decision can exceed +// its declarative onMatch ceiling (the webhook-forward shape). exceed controls +// whether it advertises that capability. +type escalatableHook struct { + core.AnyEndpointAnyModality + exceed bool + execErr error +} + +func (h *escalatableHook) Execute(context.Context, *core.HookInput) (*core.HookResult, error) { + if h.execErr != nil { + return nil, h.execErr + } + return &core.HookResult{Decision: core.Approve}, nil +} + +func (h *escalatableHook) MayExceedOnMatch() bool { return h.exceed } + +// malformedOnMatchCfg returns a config whose onMatch.action is an unknown value, +// so core.ParseOnMatch returns an error. Mirrors a misconfigured hook row. +func malformedOnMatchCfg() map[string]any { + return map[string]any{"onMatch": map[string]any{"action": "bogus-not-an-action"}} +} + +// TestPipeline_MayRedact_MalformedOnMatchEnforces covers MEDIUM-2: a hook whose +// onMatch cannot be parsed must be treated as enforcing (both block- and +// redact-capable) so it routes to buffer rather than leaking onto the live path. +// This mirrors hookIsEnforcing, which already over-routes the same parse error. +func TestPipeline_MayRedact_MalformedOnMatchEnforces(t *testing.T) { + p := newTestPipeline([]boundHook{{ + hook: &stubHook{decision: core.Approve}, + config: &core.HookConfig{ID: "bad", Name: "bad", Config: malformedOnMatchCfg()}, + }}) + if !p.MayBlock() { + t.Error("MayBlock()=false for a malformed-onMatch hook; it would stream instead of routing to buffer") + } + if !p.MayRedact() { + t.Error("MayRedact()=false for a malformed-onMatch hook; it would stream instead of routing to buffer") + } +} + +// TestPipeline_MayRedact_RuntimeEscalatable covers HIGH-1: a hook that can exceed +// its declared onMatch ceiling at runtime (webhook-forward, which reconciles a +// remote reply against the ceiling via the strictest of the two) must be treated +// as both block- and redact-capable even when its declared action is approve, so +// its runtime enforcement is never under-routed onto the audit-only live path. +func TestPipeline_MayRedact_RuntimeEscalatable(t *testing.T) { + // Declared action is approve (non-enforcing), but the hook advertises that it + // can exceed that ceiling at runtime. + p := newTestPipeline([]boundHook{{ + hook: &escalatableHook{exceed: true}, + config: &core.HookConfig{ID: "webhook", Name: "webhook", Config: onMatchCfg("approve")}, + }}) + if !p.MayBlock() { + t.Error("MayBlock()=false for a runtime-escalatable approve-ceiling hook; runtime block would be dropped onto the live path") + } + if !p.MayRedact() { + t.Error("MayRedact()=false for a runtime-escalatable approve-ceiling hook; runtime redact would be dropped onto the live path") + } +} + +// TestPipeline_MayRedact_EscalatableOptOut guards the negative side of HIGH-1: a +// hook that implements core.RuntimeEscalatable but reports false (cannot exceed +// its ceiling) is governed by its declared action like any ordinary hook — the +// over-route fires only for genuinely escalatable hooks, not for every +// implementer of the interface. +func TestPipeline_MayRedact_EscalatableOptOut(t *testing.T) { + p := newTestPipeline([]boundHook{{ + hook: &escalatableHook{exceed: false}, + config: &core.HookConfig{ID: "h", Name: "h", Config: onMatchCfg("approve")}, + }}) + if p.MayBlock() || p.MayRedact() { + t.Error("a non-escalating (exceed=false) approve-ceiling hook must enforce neither block nor redact") + } +} + +// TestPipeline_FailClosedOnError_RuntimeEscalatable covers the error-posture variant +// (review 17-LOW-1): a runtime-escalatable hook with an approve ceiling that ERRORS +// under strict posture must fail closed (RejectHard), because its runtime decision +// could have exceeded the declared ceiling — mirroring the routing predicate so the +// strict appliance does not fail open on the very hook whose enforcement it cannot see. +func TestPipeline_FailClosedOnError_RuntimeEscalatable(t *testing.T) { + p := newTestPipeline([]boundHook{{ + hook: &escalatableHook{exceed: true, execErr: errors.New("transient webhook failure")}, + config: &core.HookConfig{ID: "webhook", Name: "webhook", Config: onMatchCfg("approve")}, + }}) + p.SetStrictFailClosed(true) + res := p.Execute(context.Background(), &core.HookInput{Stage: "response"}) + if res.Decision != core.RejectHard { + t.Fatalf("a strict escalatable hook's transient error must fail closed (RejectHard), got %s", res.Decision) + } +} + +// TestPipeline_FailClosedOnError_NonEscalatable_StaysOpen is the negative side: a +// non-escalatable approve-ceiling hook erroring under strict stays fail-open — its +// declared ceiling means a transient error cannot have leaked any enforcement. +func TestPipeline_FailClosedOnError_NonEscalatable_StaysOpen(t *testing.T) { + p := newTestPipeline([]boundHook{{ + hook: &escalatableHook{exceed: false, execErr: errors.New("transient failure")}, + config: &core.HookConfig{ID: "h", Name: "h", Config: onMatchCfg("approve")}, + }}) + p.SetStrictFailClosed(true) + res := p.Execute(context.Background(), &core.HookInput{Stage: "response"}) + if res.Decision == core.RejectHard { + t.Fatalf("a non-escalatable approve-ceiling hook's error must NOT fail closed, got %s", res.Decision) + } +} diff --git a/packages/shared/policy/pipeline/pipeline_test.go b/packages/shared/policy/pipeline/pipeline_test.go index 01107667..0b262d8b 100644 --- a/packages/shared/policy/pipeline/pipeline_test.go +++ b/packages/shared/policy/pipeline/pipeline_test.go @@ -14,6 +14,7 @@ import ( "github.com/prometheus/client_golang/prometheus/testutil" "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" + normalize "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/normalize/core" ) // stubHook is a test hook with configurable behavior. @@ -143,6 +144,45 @@ func TestPipeline_ParallelExecution(t *testing.T) { } } +// TestPipeline_HookLatencyMicroseconds verifies executeOneHook stamps precise +// microsecond latency (LatencyUs) alongside the truncated millisecond value, and +// that LatencyMs is the exact integer-ms floor of LatencyUs — i.e. LatencyMs is +// NOT clamped. A clamp-to-≥1 would break the floor invariant for sub-millisecond +// hooks and N×-inflate the summed _ms aggregate. +func TestPipeline_HookLatencyMicroseconds(t *testing.T) { + hks := []boundHook{ + // Real delay: LatencyUs must reflect it. + {hook: &stubHook{decision: core.Approve, delay: 3 * time.Millisecond}, config: &core.HookConfig{ID: "slow", Name: "slow", Priority: 1, FailBehavior: "fail-open"}}, + // Near-instant: on a fast host LatencyUs < 1000, so the floor invariant + // forces LatencyMs == 0 (proves no clamp). + {hook: &stubHook{decision: core.Approve, delay: 0}, config: &core.HookConfig{ID: "fast", Name: "fast", Priority: 2, FailBehavior: "fail-open"}}, + } + p := NewPipeline(hks, 5*time.Second, 30*time.Second, false, testLogger()) + result := p.Execute(context.Background(), &core.HookInput{}) + if len(result.HookResults) != 2 { + t.Fatalf("expected 2 hook results, got %d", len(result.HookResults)) + } + var slow *core.HookResult + for i := range result.HookResults { + hr := &result.HookResults[i] + // Floor invariant: both fields derive from the same elapsed, so ms is the + // exact integer-ms floor of us. A clamp would violate this for us<1000. + if hr.LatencyMs != hr.LatencyUs/1000 { + t.Fatalf("hook %s: LatencyMs=%d is not the floor of LatencyUs=%d (us/1000=%d) — clamp or skew?", + hr.HookID, hr.LatencyMs, hr.LatencyUs, hr.LatencyUs/1000) + } + if result.HookResults[i].HookID == "slow" { + slow = hr + } + } + if slow == nil { + t.Fatal("missing 'slow' hook result") + } + if slow.LatencyUs < 2000 { + t.Fatalf("slow hook LatencyUs=%d, expected >= 2000 (3ms delay)", slow.LatencyUs) + } +} + func TestPipeline_PerHookTimeout(t *testing.T) { hks := []boundHook{ {hook: &stubHook{decision: core.Approve, delay: 5 * time.Second}, config: &core.HookConfig{ @@ -777,3 +817,180 @@ func TestPipeline_Sequential_AccumulatesUpstreamTags(t *testing.T) { t.Fatalf("hook2 observed UpstreamTags = %v, want %v", recorder.observedUpstream, want) } } + +// redactStubHook returns a Modify result carrying replacement content (a redact hook). +type redactStubHook struct { + core.AnyEndpointAnyModality + content []core.ContentBlock +} + +func (h *redactStubHook) Execute(context.Context, *core.HookInput) (*core.HookResult, error) { + return &core.HookResult{Decision: core.Modify, ModifiedContent: h.content}, nil +} + +// TestPipeline_CoFiringRedactSoftBlock_CarriesRedaction pins the #13 core invariant: when +// a redact hook (Modify + ModifiedContent) co-fires with a soft-block hook, mergeResults +// promotes the aggregate Decision to BlockSoft (the strictest) but MUST carry the redact's +// ModifiedContent — previously it was dropped, leaving consumers unable to apply the +// redaction (fail-closed or leak). CarriesRedaction() must report true so every consumer +// applies the mask instead of keying on Decision==Modify and forwarding the original raw. +func TestPipeline_CoFiringRedactSoftBlock_CarriesRedaction(t *testing.T) { + hks := []boundHook{ + {hook: &redactStubHook{content: []core.ContentBlock{{Type: "text", Text: "[REDACTED]"}}}, config: &core.HookConfig{ID: "redact", Name: "redact", Priority: 1, FailBehavior: "fail-open"}}, + {hook: &stubHook{decision: core.BlockSoft, reason: "soft flag"}, config: &core.HookConfig{ID: "soft", Name: "soft", Priority: 2, FailBehavior: "fail-open"}}, + } + p := NewPipeline(hks, 5*time.Second, 30*time.Second, false, testLogger()) + result := p.Execute(context.Background(), &core.HookInput{}) + + if result.Decision != core.BlockSoft { + t.Fatalf("co-firing redact+soft-block must aggregate to BlockSoft (the ceiling), got %s", result.Decision) + } + if len(result.ModifiedContent) == 0 { + t.Fatal("the redact ModifiedContent must be CARRIED under BlockSoft (not dropped) — the #13 core invariant") + } + if !result.CarriesRedaction() { + t.Fatal("CarriesRedaction() must be true for a BlockSoft masking a co-firing redact") + } +} + +// TestPipeline_StandaloneSoftBlock_NoRedaction is the negative: a soft-block with no +// co-firing redact carries no redaction → CarriesRedaction() false → consumers +// deliver-with-warning (the original), not fail-closed/masked. +func TestPipeline_StandaloneSoftBlock_NoRedaction(t *testing.T) { + hks := []boundHook{ + {hook: &stubHook{decision: core.BlockSoft, reason: "soft flag"}, config: &core.HookConfig{ID: "soft", Name: "soft", Priority: 1, FailBehavior: "fail-open"}}, + } + p := NewPipeline(hks, 5*time.Second, 30*time.Second, false, testLogger()) + result := p.Execute(context.Background(), &core.HookInput{}) + if result.Decision != core.BlockSoft { + t.Fatalf("expected BlockSoft, got %s", result.Decision) + } + if result.CarriesRedaction() { + t.Fatal("a standalone soft-block carries no redaction") + } +} + +// spanStubHook returns a configurable decision with optional ModifiedContent and +// TransformSpans — exercises the applicable-artifact discriminator (#14). +type spanStubHook struct { + core.AnyEndpointAnyModality + decision core.Decision + content []core.ContentBlock + spans []normalize.TransformSpan +} + +func (h *spanStubHook) Execute(context.Context, *core.HookInput) (*core.HookResult, error) { + return &core.HookResult{Decision: h.decision, ModifiedContent: h.content, TransformSpans: h.spans}, nil +} + +func auditOnlySpan() normalize.TransformSpan { + return normalize.TransformSpan{Source: normalize.SourceHook, Action: normalize.ActionRedact, ContentAddress: normalize.AddressAuditOnlySentinel, Start: 0, End: 4} +} + +func applicableSpan() normalize.TransformSpan { + return normalize.TransformSpan{Source: normalize.SourceHook, Action: normalize.ActionRedact, ContentAddress: "messages.0.content.0.toolUse.input.0", Start: 0, End: 4} +} + +// TestPipeline_ApproveWebhookAuditSpans_CoFiringSoftBlock_NoRedaction is the #14 fix: +// an Approve hook emitting AUDIT-ONLY sentinel spans (the approve-webhook+redactions +// shape) co-firing with a soft-block promotes the aggregate to BlockSoft, but those +// advisory spans are NOT an applicable redaction — CarriesRedaction() must be false so +// the appliance soft-delivers instead of failing closed (over-block). +func TestPipeline_ApproveWebhookAuditSpans_CoFiringSoftBlock_NoRedaction(t *testing.T) { + hks := []boundHook{ + {hook: &spanStubHook{decision: core.Approve, spans: []normalize.TransformSpan{auditOnlySpan()}}, config: &core.HookConfig{ID: "wh", Name: "webhook", Priority: 1, FailBehavior: "fail-open"}}, + {hook: &stubHook{decision: core.BlockSoft, reason: "soft flag"}, config: &core.HookConfig{ID: "soft", Name: "soft", Priority: 2, FailBehavior: "fail-open"}}, + } + p := NewPipeline(hks, 5*time.Second, 30*time.Second, false, testLogger()) + result := p.Execute(context.Background(), &core.HookInput{}) + if result.Decision != core.BlockSoft { + t.Fatalf("expected BlockSoft aggregate, got %s", result.Decision) + } + if result.RedactionApplicable { + t.Fatal("an Approve hook's audit-only spans must NOT set RedactionApplicable") + } + if result.CarriesRedaction() { + t.Fatal("approve-webhook audit-only spans masked by soft-block must NOT CarryRedaction (the #14 over-block fix)") + } + // Audit union still records the advisory spans — the fix changes only the inflight predicate. + if len(result.TransformSpans) != 1 { + t.Fatalf("advisory span must still be unioned into the audit record, got %d", len(result.TransformSpans)) + } +} + +// TestPipeline_ModifyWebhookFlatOnly_CoFiringSoftBlock_NoRedaction closes the sibling +// hole: a Modify hook (e.g. a webhook reconciled under an approve ceiling) carrying ONLY +// audit-only sentinel spans and NO ModifiedContent is not an applicable inflight +// redaction, so under a co-firing soft-block CarriesRedaction() must be false. +func TestPipeline_ModifyWebhookFlatOnly_CoFiringSoftBlock_NoRedaction(t *testing.T) { + hks := []boundHook{ + {hook: &spanStubHook{decision: core.Modify, spans: []normalize.TransformSpan{auditOnlySpan()}}, config: &core.HookConfig{ID: "wh", Name: "webhook", Priority: 1, FailBehavior: "fail-open"}}, + {hook: &stubHook{decision: core.BlockSoft, reason: "soft flag"}, config: &core.HookConfig{ID: "soft", Name: "soft", Priority: 2, FailBehavior: "fail-open"}}, + } + p := NewPipeline(hks, 5*time.Second, 30*time.Second, false, testLogger()) + result := p.Execute(context.Background(), &core.HookInput{}) + if result.Decision != core.BlockSoft { + t.Fatalf("expected BlockSoft aggregate, got %s", result.Decision) + } + if result.CarriesRedaction() { + t.Fatal("modify-webhook flat-only spans masked by soft-block must NOT CarryRedaction (sibling of the #14 fix)") + } +} + +// TestPipeline_ToolArgOnlyModify_CoFiringSoftBlock_CarriesRedaction is the leak-guard: +// a real redaction carried by an APPLICABLE span with NO ModifiedContent (the tool-call- +// argument redaction shape) masked by a co-firing soft-block MUST still CarryRedaction — +// dropping the span clause would leak it. +func TestPipeline_ToolArgOnlyModify_CoFiringSoftBlock_CarriesRedaction(t *testing.T) { + hks := []boundHook{ + {hook: &spanStubHook{decision: core.Modify, spans: []normalize.TransformSpan{applicableSpan()}}, config: &core.HookConfig{ID: "redact", Name: "redact", Priority: 1, FailBehavior: "fail-open"}}, + {hook: &stubHook{decision: core.BlockSoft, reason: "soft flag"}, config: &core.HookConfig{ID: "soft", Name: "soft", Priority: 2, FailBehavior: "fail-open"}}, + } + p := NewPipeline(hks, 5*time.Second, 30*time.Second, false, testLogger()) + result := p.Execute(context.Background(), &core.HookInput{}) + if result.Decision != core.BlockSoft { + t.Fatalf("expected BlockSoft aggregate, got %s", result.Decision) + } + if !result.RedactionApplicable || !result.CarriesRedaction() { + t.Fatal("a tool-arg-only Modify (applicable span, no ModifiedContent) masked by soft-block MUST CarryRedaction — else it leaks") + } +} + +// TestPipeline_SelfRedactingBlockSoft_ApplicableSpan_CarriesRedaction: a single hook +// emitting BlockSoft WITH an applicable span carries that redaction (the aggregate unions +// every hook's spans), so CarriesRedaction() is true. +func TestPipeline_SelfRedactingBlockSoft_ApplicableSpan_CarriesRedaction(t *testing.T) { + hks := []boundHook{ + {hook: &spanStubHook{decision: core.BlockSoft, spans: []normalize.TransformSpan{applicableSpan()}}, config: &core.HookConfig{ID: "soft", Name: "soft", Priority: 1, FailBehavior: "fail-open"}}, + } + p := NewPipeline(hks, 5*time.Second, 30*time.Second, false, testLogger()) + result := p.Execute(context.Background(), &core.HookInput{}) + if result.Decision != core.BlockSoft { + t.Fatalf("expected BlockSoft, got %s", result.Decision) + } + if !result.CarriesRedaction() { + t.Fatal("a self-redacting BlockSoft with an applicable span must CarryRedaction (leak-guard)") + } +} + +// TestPipeline_SelfRedactingBlockSoft_ModifiedContentOnly_NoRedaction pins the merge +// asymmetry: mergeResults carries ModifiedContent ONLY from Modify hooks, so a BlockSoft +// hook's ModifiedContent is dropped from the aggregate. CarriesRedaction() must therefore +// be false (the aggregate cannot apply it) — fold-to-block is the fail-safe disposition, +// not a claimed-but-unappliable redaction. +func TestPipeline_SelfRedactingBlockSoft_ModifiedContentOnly_NoRedaction(t *testing.T) { + hks := []boundHook{ + {hook: &spanStubHook{decision: core.BlockSoft, content: []core.ContentBlock{{Type: "text", Text: "[REDACTED]"}}}, config: &core.HookConfig{ID: "soft", Name: "soft", Priority: 1, FailBehavior: "fail-open"}}, + } + p := NewPipeline(hks, 5*time.Second, 30*time.Second, false, testLogger()) + result := p.Execute(context.Background(), &core.HookInput{}) + if result.Decision != core.BlockSoft { + t.Fatalf("expected BlockSoft, got %s", result.Decision) + } + if len(result.ModifiedContent) != 0 { + t.Fatal("the merge must NOT carry a BlockSoft hook's ModifiedContent into the aggregate") + } + if result.CarriesRedaction() { + t.Fatal("a BlockSoft hook's ModifiedContent is not carried, so CarriesRedaction must be false (fold-to-block fail-safe)") + } +} diff --git a/packages/shared/policy/pipeline/policy.go b/packages/shared/policy/pipeline/policy.go index 7dc30083..79570057 100644 --- a/packages/shared/policy/pipeline/policy.go +++ b/packages/shared/policy/pipeline/policy.go @@ -26,8 +26,20 @@ var errFailClosedUnbuildable = fmt.Errorf("fail-closed compliance hook could not // replace it concurrently with in-flight Resolve*/Has* calls — a reader // keeps its loaded snapshot for the remainder of its call while the next // caller sees the new one. Config invalidation is lazy and non-blocking. +// configSnapshot bundles the hook-config slice with the generation it was published +// under, so a reader gets a CONSISTENT (gen, configs) pair from ONE atomic load. The +// per-generation cache memos (union prefilter, MaxPatternBound) tag their entries with +// the gen the configs were read under — bundling closes the window where a request could +// read a new swapGen but the old config (or vice versa) and cache a bound/union derived +// from one generation's hooks under another's key. swapGen stays the monotonic counter +// (lockstep with .gen) that Prewarm's staleness recheck and the union close-dance use. +type configSnapshot struct { + gen uint64 + configs []core.HookConfig +} + type PolicyResolver struct { - hookConfigs atomic.Pointer[[]core.HookConfig] + hookConfigs atomic.Pointer[configSnapshot] registry *core.HookRegistry logger *slog.Logger @@ -53,6 +65,12 @@ type PolicyResolver struct { // hook's anchor-stripped patterns) per resolved hook set, tagged with the // swapGen it was built under. See unionprescan.go. union unionState + + // patternBound memoises MaxPatternBound per resolved hook set, tagged with the + // generation it was built under (see maxPatternBoundFor in unionprescan.go). + boundMu sync.Mutex + boundGen uint64 + boundByKey map[string]patternBound } // NewPolicyResolver creates a resolver with an initial hook config snapshot @@ -66,7 +84,7 @@ func NewPolicyResolver(configs []core.HookConfig, registry *core.HookRegistry, l hookCache: make(map[string]core.Hook), } snapshot := append([]core.HookConfig(nil), configs...) - r.hookConfigs.Store(&snapshot) + r.hookConfigs.Store(&configSnapshot{gen: 0, configs: snapshot}) return r } @@ -88,11 +106,13 @@ func NewPolicyResolver(configs []core.HookConfig, registry *core.HookRegistry, l func (r *PolicyResolver) Swap(configs []core.HookConfig) { prevGen := r.swapGen.Add(1) - 1 snapshot := append([]core.HookConfig(nil), configs...) - oldPtr := r.hookConfigs.Swap(&snapshot) + // Publish gen+configs as ONE atomic value (lockstep with swapGen) so cache + // consumers reading hookConfigs get a consistent pair — see configSnapshot. + oldPtr := r.hookConfigs.Swap(&configSnapshot{gen: prevGen + 1, configs: snapshot}) oldByID := map[string]*core.HookConfig{} if oldPtr != nil { - old := *oldPtr + old := oldPtr.configs for i := range old { oldByID[old[i].ID] = &old[i] } @@ -160,9 +180,9 @@ func (r *PolicyResolver) Swap(configs []core.HookConfig) { // hands out the same slice. Returns true if a swap occurred. func (r *PolicyResolver) SwapIfChanged(configs []core.HookConfig) bool { cur := r.hookConfigs.Load() - if cur != nil && len(*cur) == len(configs) && len(configs) > 0 { + if cur != nil && len(cur.configs) == len(configs) && len(configs) > 0 { // Fast pointer check: if the backing array is the same, skip. - if &(*cur)[0] == &configs[0] { + if &cur.configs[0] == &configs[0] { return false } } @@ -170,72 +190,6 @@ func (r *PolicyResolver) SwapIfChanged(configs []core.HookConfig) bool { return true } -// Prewarm eagerly builds and caches the hook for every enabled request/response -// config so the factory's compile cost (notably the Vectorscan database, ~100s -// of ms) runs OFF the request path — at startup and in the background after each -// Swap. It is best-effort and idempotent: -// -// - Hooks are built WITHOUT holding hookMu (the compile must not block -// resolve()), then inserted under the lock with a double-check. -// - It is guarded by swapGen: if a Swap races while Prewarm is building, the -// captured generation no longer matches and Prewarm aborts without caching, -// so it can never install a hook for a superseded config. The newer Swap -// spawns its own Prewarm for the current snapshot. -// - Connection-stage hooks are skipped (they are cheap metadata hooks and -// resolve() applies an extra connection-compat gate before caching them). -// -// A hook left unbuilt (factory error, or aborted prewarm) is simply built -// lazily by the next resolve(), exactly as before — prewarm only removes the -// first-request latency spike, never changes correctness. -func (r *PolicyResolver) Prewarm() { - gen := r.swapGen.Load() - configs := r.snapshot() - for i := range configs { - cfg := &configs[i] - if !cfg.Enabled || strings.EqualFold(cfg.Stage, "connection") { - continue - } - factory := r.registry.Get(cfg.ImplementationID) - if factory == nil { - continue - } - r.hookMu.RLock() - _, hit := r.hookCache[cfg.ID] - r.hookMu.RUnlock() - if hit { - continue - } - // Build outside the lock — the compile is the expensive part and must - // not stall concurrent resolve() readers. - hook, err := factory(cfg) - if err != nil { - continue // resolve() will log+skip per its fail posture - } - r.hookMu.Lock() - switch { - case r.swapGen.Load() != gen: - // A swap raced; this snapshot may be stale. Drop our build. - r.hookMu.Unlock() - closeHook(hook) - case r.hookCache[cfg.ID] != nil: - // Someone (resolve or a concurrent prewarm) already built it. - r.hookMu.Unlock() - closeHook(hook) - default: - r.hookCache[cfg.ID] = hook - r.hookMu.Unlock() - } - } -} - -// closeHook releases a built-but-unused hook's resources (e.g. a Vectorscan -// matcher's cgo memory) when prewarm discards it. -func closeHook(h core.Hook) { - if c, ok := h.(io.Closer); ok { - _ = c.Close() - } -} - // snapshot returns the current hook config slice. Callers MUST capture // the return value in a local variable and operate on that local slice // for the remainder of their call — re-reading via snapshot() mid-loop @@ -245,7 +199,19 @@ func (r *PolicyResolver) snapshot() []core.HookConfig { if p == nil { return nil } - return *p + return p.configs +} + +// loadSnapshot returns the current (gen, configs) pair in one atomic load. Callers +// that build a per-generation cache entry (BuildPipeline → union/bound memos) MUST use +// THIS gen — not a separate r.swapGen.Load() — so the entry is tagged with the gen its +// configs were actually read under (closing the cross-generation cache window). +func (r *PolicyResolver) loadSnapshot() (uint64, []core.HookConfig) { + p := r.hookConfigs.Load() + if p == nil { + return 0, nil + } + return p.gen, p.configs } // ResolveHooks returns hooks to run for the given stage and ingress type, sorted @@ -278,15 +244,19 @@ func (r *PolicyResolver) ResolveHooks(stage, ingressType string, strictFailClose } // resolve filters configs by stage, ingress, and enabled, then instantiates core. +// It captures the current snapshot once (so a concurrent Swap cannot change the set +// mid-call) and delegates to resolveFrom. Callers that ALSO need the generation the +// configs were read under (BuildPipeline, for the per-gen caches) call loadSnapshot + +// resolveFrom directly so the hooks and the cache gen come from the SAME atomic load. func (r *PolicyResolver) resolve(stage, ingressType string, strictFailClosed bool) ([]boundHook, error) { - var out []boundHook + return r.resolveFrom(r.snapshot(), stage, ingressType, strictFailClosed) +} - // Capture the current snapshot once so that a concurrent Swap does - // not change the set of configs we iterate over mid-call. Pointers - // taken into this slice remain valid for the lifetime of the - // returned boundHook slice because Go GC keeps the backing array - // alive as long as any pointer references it. - configs := r.snapshot() +// resolveFrom is resolve over an already-captured config snapshot. Pointers taken into +// `configs` remain valid for the lifetime of the returned boundHook slice (Go GC keeps +// the backing array alive as long as any pointer references it). +func (r *PolicyResolver) resolveFrom(configs []core.HookConfig, stage, ingressType string, strictFailClosed bool) ([]boundHook, error) { + var out []boundHook for i := range configs { cfg := &configs[i] @@ -410,7 +380,10 @@ func (r *PolicyResolver) BuildPipeline( strictFailClosed bool, logger *slog.Logger, ) (*Pipeline, error) { - candidates, err := r.ResolveHooks(stage, ingressType, strictFailClosed) + // Load the (gen, configs) pair ONCE so the hooks we resolve and the generation + // we tag the per-gen caches with come from the same atomic snapshot. + gen, configs := r.loadSnapshot() + candidates, err := r.resolveFrom(configs, stage, ingressType, strictFailClosed) if err != nil { return nil, err } @@ -467,8 +440,19 @@ func (r *PolicyResolver) BuildPipeline( // (non-packet-path) callers — matching the build-time UNBUILDABLE posture. p.SetStrictFailClosed(strictFailClosed) // Fold every content hook's raw-body prefilter into one shared scan for this - // resolved set (cached per generation). nil => use the per-hook loop. - p.unionPrescan = r.unionPrescanFor(filtered) + // resolved set (cached per generation). nil => use the per-hook loop. The gen is + // the one the configs were loaded under (above), not a fresh read. + p.unionPrescan = r.unionPrescanFor(filtered, gen) + // Pre-stamp the per-generation cached MaxPatternBound so the streaming hot path + // reads O(1) instead of re-walking every hook regex per request. Only the response + // stage consumes the bound (Model-A streaming), so skip the cache lookup + signature + // work on request/connection-stage builds that never read it. EqualFold for parity + // with resolveFrom's stage match (a non-canonical "Response" would otherwise silently + // fall back to the per-request lazy compute — the exact cost being optimized away). + if strings.EqualFold(stage, "response") { + p.maxBounded, p.anyUnbounded = r.maxPatternBoundFor(filtered, gen) + p.boundComputed = true + } return p, nil } diff --git a/packages/shared/policy/pipeline/resolver_prewarm.go b/packages/shared/policy/pipeline/resolver_prewarm.go new file mode 100644 index 00000000..3392ab18 --- /dev/null +++ b/packages/shared/policy/pipeline/resolver_prewarm.go @@ -0,0 +1,73 @@ +package pipeline + +import ( + "io" + "strings" + + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// Prewarm eagerly builds and caches the hook for every enabled request/response +// config so the factory's compile cost (notably the Vectorscan database, ~100s +// of ms) runs OFF the request path — at startup and in the background after each +// Swap. It is best-effort and idempotent: +// +// - Hooks are built WITHOUT holding hookMu (the compile must not block +// resolve()), then inserted under the lock with a double-check. +// - It is guarded by swapGen: if a Swap races while Prewarm is building, the +// captured generation no longer matches and Prewarm aborts without caching, +// so it can never install a hook for a superseded config. The newer Swap +// spawns its own Prewarm for the current snapshot. +// - Connection-stage hooks are skipped (they are cheap metadata hooks and +// resolve() applies an extra connection-compat gate before caching them). +// +// A hook left unbuilt (factory error, or aborted prewarm) is simply built +// lazily by the next resolve(), exactly as before — prewarm only removes the +// first-request latency spike, never changes correctness. +func (r *PolicyResolver) Prewarm() { + gen, configs := r.loadSnapshot() + for i := range configs { + cfg := &configs[i] + if !cfg.Enabled || strings.EqualFold(cfg.Stage, "connection") { + continue + } + factory := r.registry.Get(cfg.ImplementationID) + if factory == nil { + continue + } + r.hookMu.RLock() + _, hit := r.hookCache[cfg.ID] + r.hookMu.RUnlock() + if hit { + continue + } + // Build outside the lock — the compile is the expensive part and must + // not stall concurrent resolve() readers. + hook, err := factory(cfg) + if err != nil { + continue // resolve() will log+skip per its fail posture + } + r.hookMu.Lock() + switch { + case r.swapGen.Load() != gen: + // A swap raced; this snapshot may be stale. Drop our build. + r.hookMu.Unlock() + closeHook(hook) + case r.hookCache[cfg.ID] != nil: + // Someone (resolve or a concurrent prewarm) already built it. + r.hookMu.Unlock() + closeHook(hook) + default: + r.hookCache[cfg.ID] = hook + r.hookMu.Unlock() + } + } +} + +// closeHook releases a built-but-unused hook's resources (e.g. a Vectorscan +// matcher's cgo memory) when prewarm discards it. +func closeHook(h core.Hook) { + if c, ok := h.(io.Closer); ok { + _ = c.Close() + } +} diff --git a/packages/shared/policy/pipeline/unionprescan.go b/packages/shared/policy/pipeline/unionprescan.go index 74d27f0f..7a8b47b9 100644 --- a/packages/shared/policy/pipeline/unionprescan.go +++ b/packages/shared/policy/pipeline/unionprescan.go @@ -218,14 +218,16 @@ func buildUnionPrescan(hooks []boundHook) (matcher.Matcher, bool) { // per-hook loop (the original, always-correct path), so a cache miss never blocks // the request path on the ~100ms+ compile and a reset never triggers a herd of // concurrent compiles. -func (r *PolicyResolver) unionPrescanFor(hooks []boundHook) matcher.Matcher { +func (r *PolicyResolver) unionPrescanFor(hooks []boundHook, gen uint64) matcher.Matcher { if !unionPrescanEnabled { return nil // A/B opt-out: per-hook loop } if !stableHookIDs(hooks) { return nil // ambiguous cache key → fall back rather than risk a wrong-set union } - gen := r.swapGen.Load() + // gen is the generation the hooks were resolved under (passed by the caller from + // the same atomic snapshot), NOT a fresh r.swapGen.Load() — so the entry is tagged + // with the gen its patterns belong to, closing the cross-generation cache window. key := hookSetSignature(hooks) r.union.mu.Lock() @@ -265,6 +267,54 @@ func (r *PolicyResolver) unionPrescanFor(hooks []boundHook) matcher.Matcher { return m } +// patternBound is one cached MaxPatternBound result for a resolved hook set. +type patternBound struct { + maxBounded int + anyUnbounded bool +} + +// maxPatternBoundFor returns the cached MaxPatternBound for the resolved hook set, +// computing it ONCE per generation. The bound is a static property of the hook set +// (it changes only on a config reload), so caching it keeps the O(total patterns) +// computeMaxPatternBound OFF the per-request streaming hot path — it was profiled +// re-walking every hook regex per request at ~51% CPU. gen is passed by the caller +// from the same atomic snapshot the hooks were resolved under (BuildPipeline), so the +// entry is tagged with the gen its patterns belong to — closing the cross-generation +// window. Unlike the union prefilter (cgo memory that must be Closed on a generation +// change), the result is two ints: no single-flight/close machinery is needed, the rare +// concurrent first-miss computes are harmless (identical within a generation), and a +// generation change just drops the map on the next access. The key is the resolved +// hook-set signature (sorted config IDs), the same key the union memo uses; an ambiguous +// key (empty/duplicate IDs) computes directly without caching. +func (r *PolicyResolver) maxPatternBoundFor(hooks []boundHook, gen uint64) (int, bool) { + if !stableHookIDs(hooks) { + return computeMaxPatternBound(hooks) // ambiguous cache key → compute (rare) + } + key := hookSetSignature(hooks) + + r.boundMu.Lock() + if r.boundByKey == nil || r.boundGen != gen { + r.boundByKey = make(map[string]patternBound) // drop the previous generation's entries + r.boundGen = gen + } + if b, ok := r.boundByKey[key]; ok { + r.boundMu.Unlock() + return b.maxBounded, b.anyUnbounded + } + r.boundMu.Unlock() + + // Compute off the lock (the expensive regex walk), then cache iff this generation + // is still current (a racing Swap that advanced boundGen discards our stale result). + mb, au := computeMaxPatternBound(hooks) + + r.boundMu.Lock() + if r.boundGen == gen && r.boundByKey != nil { + r.boundByKey[key] = patternBound{maxBounded: mb, anyUnbounded: au} + } + r.boundMu.Unlock() + return mb, au +} + // resetUnionsLockedIfGenChanged re-initialises the cache when the generation // advanced (or it is uninitialised), returning the previous generation's matchers // for the caller to Close AFTER releasing r.union.mu. Caller holds r.union.mu. diff --git a/packages/shared/policy/pipeline/unionprescan_test.go b/packages/shared/policy/pipeline/unionprescan_test.go index 44932cd3..e55b618c 100644 --- a/packages/shared/policy/pipeline/unionprescan_test.go +++ b/packages/shared/policy/pipeline/unionprescan_test.go @@ -171,18 +171,18 @@ func TestUnion_CacheAndGeneration(t *testing.T) { r := NewPolicyResolver(nil, core.NewHookRegistry(), nil) hooks := []boundHook{bh(newPrescanHook(t, `secret`, `\btoken\b`), "a")} - m1 := r.unionPrescanFor(hooks) + m1 := r.unionPrescanFor(hooks, r.swapGen.Load()) if m1 == nil { t.Fatal("expected a union matcher") } - m2 := r.unionPrescanFor(hooks) + m2 := r.unionPrescanFor(hooks, r.swapGen.Load()) if m1 != m2 { t.Error("expected the cached matcher on the second call within one generation") } prev := r.swapGen.Add(1) - 1 // simulate a Swap (increment, capture pre-gen) r.closeUnionsIfGen(prev) // Swap closes only the previous generation's unions - m3 := r.unionPrescanFor(hooks) + m3 := r.unionPrescanFor(hooks, r.swapGen.Load()) if m3 == nil { t.Fatal("expected a rebuilt union after generation bump") } @@ -198,20 +198,20 @@ func TestUnion_EmptyOrDupIDFallsBack(t *testing.T) { r := NewPolicyResolver(nil, core.NewHookRegistry(), nil) emptyID := []boundHook{{hook: newPrescanHook(t, `secret`), config: &core.HookConfig{ID: ""}}} - if m := r.unionPrescanFor(emptyID); m != nil { + if m := r.unionPrescanFor(emptyID, r.swapGen.Load()); m != nil { t.Error("empty hook ID must fall back to the per-hook loop (nil union)") } nilCfg := []boundHook{{hook: newPrescanHook(t, `secret`), config: nil}} - if m := r.unionPrescanFor(nilCfg); m != nil { + if m := r.unionPrescanFor(nilCfg, r.swapGen.Load()); m != nil { t.Error("nil hook config must fall back (nil union)") } dupID := []boundHook{bh(newPrescanHook(t, `a`), "x"), bh(newPrescanHook(t, `b`), "x")} - if m := r.unionPrescanFor(dupID); m != nil { + if m := r.unionPrescanFor(dupID, r.swapGen.Load()); m != nil { t.Error("duplicate hook ID must fall back (nil union)") } // Sanity: a clean, unique-ID set still unions. ok := []boundHook{bh(newPrescanHook(t, `a`), "x"), bh(newPrescanHook(t, `b`), "y")} - if m := r.unionPrescanFor(ok); m == nil { + if m := r.unionPrescanFor(ok, r.swapGen.Load()); m == nil { t.Error("unique non-empty IDs should produce a union") } } @@ -254,7 +254,7 @@ func TestUnion_ConcurrentSwapNoPanic(t *testing.T) { defer builders.Done() hooks := mk(id) for range 5000 { - if m := r.unionPrescanFor(hooks); m != nil { + if m := r.unionPrescanFor(hooks, r.swapGen.Load()); m != nil { _ = m.Scan([]string{"benign body without any token here"}, true) } } diff --git a/packages/shared/transport/mq/binwire.go b/packages/shared/transport/mq/binwire.go index 9dd6fcd0..21a098f9 100644 --- a/packages/shared/transport/mq/binwire.go +++ b/packages/shared/transport/mq/binwire.go @@ -150,6 +150,12 @@ const ( FldAttestationAgentID FieldID = 100 FldSourceProcess FieldID = 101 FldAction FieldID = 102 + // Additive microsecond hook aggregates (siblings of FldRequestHooksMs(96) / + // FldResponseHooksMs(97)). New ids — see the FORWARD-INCOMPATIBLE note on + // AllFieldIDs: a producer emitting these requires a Hub that decodes them + // (deploy order: schema → Hub → producers). + FldRequestHooksUs FieldID = 103 + FldResponseHooksUs FieldID = 104 ) // AllFieldIDs returns every registered field-id in wire order. It exists so the @@ -185,6 +191,7 @@ func AllFieldIDs() []FieldID { FldCredentialID, FldThingID, FldThingName, FldUpstreamTtfbMs, FldUpstreamTotalMs, FldRequestHooksMs, FldResponseHooksMs, FldLatencyBreakdown, FldAttestationVerified, FldAttestationAgentID, FldSourceProcess, FldAction, + FldRequestHooksUs, FldResponseHooksUs, } } @@ -334,6 +341,8 @@ func (m *TrafficEventMessage) AppendBinary(dst []byte) []byte { dst = optPtrIntAsI64(dst, FldUpstreamTotalMs, m.UpstreamTotalMs) dst = optPtrIntAsI64(dst, FldRequestHooksMs, m.RequestHooksMs) dst = optPtrIntAsI64(dst, FldResponseHooksMs, m.ResponseHooksMs) + dst = optPtrIntAsI64(dst, FldRequestHooksUs, m.RequestHooksUs) + dst = optPtrIntAsI64(dst, FldResponseHooksUs, m.ResponseHooksUs) // Floats (emit if non-zero). dst = optF64(dst, FldReasoningCostUsd, m.ReasoningCostUsd) diff --git a/packages/shared/transport/mq/messages.go b/packages/shared/transport/mq/messages.go index 0118cc6e..7ffb5e69 100644 --- a/packages/shared/transport/mq/messages.go +++ b/packages/shared/transport/mq/messages.go @@ -265,10 +265,15 @@ type TrafficEventMessage struct { // LatencyBreakdown — Per-source phase durations (ms); key set is closed per source. // // our_overhead_ms is not on the wire — derived as LatencyMs - UpstreamTotalMs (≥0) at read time. - UpstreamTtfbMs *int `json:"upstreamTtfbMs,omitempty"` - UpstreamTotalMs *int `json:"upstreamTotalMs,omitempty"` - RequestHooksMs *int `json:"requestHooksMs,omitempty"` - ResponseHooksMs *int `json:"responseHooksMs,omitempty"` + UpstreamTtfbMs *int `json:"upstreamTtfbMs,omitempty"` + UpstreamTotalMs *int `json:"upstreamTotalMs,omitempty"` + RequestHooksMs *int `json:"requestHooksMs,omitempty"` + ResponseHooksMs *int `json:"responseHooksMs,omitempty"` + // Microsecond-precision siblings of RequestHooksMs / ResponseHooksMs. Hooks run + // at microsecond scale, so the _ms aggregates truncate sub-millisecond hooks to + // 0; these carry the real value. Additive — the _ms fields are unchanged. + RequestHooksUs *int `json:"requestHooksUs,omitempty"` + ResponseHooksUs *int `json:"responseHooksUs,omitempty"` LatencyBreakdown map[string]int `json:"latencyBreakdown,omitempty"` // AttestationVerified and AttestationAgentID are set by the compliance-proxy diff --git a/packages/shared/transport/normalize/core/spans.go b/packages/shared/transport/normalize/core/spans.go index afeb9a42..18422ef1 100644 --- a/packages/shared/transport/normalize/core/spans.go +++ b/packages/shared/transport/normalize/core/spans.go @@ -77,3 +77,32 @@ type TransformSpan struct { // directly so non-redact sources (cache normaliser, cache_control // inject) flow through the same audit channel. type RedactionSpan = TransformSpan + +// AddressAuditOnlySentinel is the ContentAddress a span carries when it +// records WHAT was flagged but cannot be applied to any payload — the +// webhook-forward hook emits it because the webhook reports wire offsets +// against a flat body the gateway never reconstructs. ApplySpans drops +// such a span (its address resolves to no content block), so it lands in +// the audit row only, never mutating delivered or stored bytes. +// +// It is the single source of truth for that sentinel: the webhook producer +// and IsAuditOnlySentinelAddress both reference this const so they cannot +// drift. A future audit-only span family adds its sentinel here. +const AddressAuditOnlySentinel = "webhook.flat" + +// IsAuditOnlySentinelAddress reports whether a ContentAddress is a known +// audit-only sentinel — a span that records a flag for the audit log but +// is NOT applicable to delivered/stored bytes. +// +// Callers deciding "does this result carry an APPLICABLE redaction" +// (decision.CompliancePipelineResult via the pipeline aggregator) use the +// NEGATION of this as a deliberate DENYLIST, not an allowlist of resolvable +// address grammars: an unknown address is treated as applicable (fail-safe +// over-block on a dormant path), whereas an allowlist would mis-classify a +// future resolvable address family as inapplicable and risk a leak. This is +// NOT a payload-aware resolvability check — a malformed/empty address also +// fails ApplySpans yet is not a sentinel; the predicate answers only "is +// this a recognised audit-only marker". +func IsAuditOnlySentinelAddress(addr string) bool { + return addr == AddressAuditOnlySentinel +} diff --git a/packages/shared/transport/normalize/core/spans_test.go b/packages/shared/transport/normalize/core/spans_test.go new file mode 100644 index 00000000..4dd67d02 --- /dev/null +++ b/packages/shared/transport/normalize/core/spans_test.go @@ -0,0 +1,27 @@ +package core + +import "testing" + +// TestIsAuditOnlySentinelAddress pins the denylist discriminator: only the known +// audit-only sentinel is recognised; every other address (a real resolvable address, +// the empty string, an unknown form) is reported applicable — the fail-safe direction +// for the CarriesRedaction predicate (unknown → applicable → over-block, never leak). +func TestIsAuditOnlySentinelAddress(t *testing.T) { + cases := []struct { + addr string + want bool + }{ + {AddressAuditOnlySentinel, true}, + {"webhook.flat", true}, // the sentinel's literal value, pinned for drift + {"messages.0.content.0", false}, + {"messages.0.content.0.toolUse.input.1", false}, + {"http.bodyView", false}, + {"", false}, + {"webhook.flatx", false}, + } + for _, c := range cases { + if got := IsAuditOnlySentinelAddress(c.addr); got != c.want { + t.Errorf("IsAuditOnlySentinelAddress(%q)=%v want %v", c.addr, got, c.want) + } + } +} diff --git a/packages/shared/transport/streaming/buffer.go b/packages/shared/transport/streaming/buffer.go index bdb863b2..0f3d879b 100644 --- a/packages/shared/transport/streaming/buffer.go +++ b/packages/shared/transport/streaming/buffer.go @@ -67,6 +67,13 @@ type BufferPipeline struct { // verbatim replay. See FrameRedactor godoc. Nil preserves the // backward-compatible degrade-and-replay-original behavior. frameRedactor FrameRedactor + // strictFailClosed governs the no-redactor degrade posture (and mirrors the + // per-host redactor's own posture from the caller): the compliance-proxy appliance + // (true) fails a redaction it cannot apply CLOSED (error frame, no original replay); + // the agent NE host-packet path (false) degrades to a disclosed replay of the + // original. Nothing is delivered live in buffer mode, so failing closed here does + // not endanger the host packet path; the flag only selects the no-redactor disposition. + strictFailClosed bool } // WithPreHook installs a callback that runs between Phase 1 (read full @@ -86,6 +93,15 @@ func (b *BufferPipeline) WithFrameRedactor(fr FrameRedactor) *BufferPipeline { return b } +// WithStrictFailClosed selects the no-redactor degrade posture (GAP B): true (the +// compliance-proxy appliance) fails a redaction it cannot apply CLOSED; false (the agent +// NE host-packet path, the default) degrades to a disclosed replay of the original. +// Mirrors the per-host redactor's own caller-driven posture. +func (b *BufferPipeline) WithStrictFailClosed(strict bool) *BufferPipeline { + b.strictFailClosed = strict + return b +} + // NewBufferPipeline creates a buffer mode pipeline. func NewBufferPipeline(config BufferConfig, pipeline PipelineExecutor, logger *slog.Logger) *BufferPipeline { if logger == nil { @@ -205,9 +221,15 @@ func (b *BufferPipeline) Process( } } - // Phase 3: Replay, redact, or reject. - switch result.Decision { - case core.RejectHard: + // Phase 3: Replay, redact, or reject. Gate the redaction arm on CarriesRedaction() + // (Modify OR a BlockSoft masking a co-firing redact), NOT Decision==Modify — a + // co-firing soft-block promotes the aggregate Decision to BlockSoft while still + // carrying the redact's spans + ModifiedContent, so keying on Modify alone would + // fall through to the `default` replay arm and deliver the original RAW (the + // appliance leak this fixes). A standalone soft-block (no redaction) keeps the + // allow-with-warning replay on `default`. + switch { + case result.Decision == core.RejectHard: b.logger.Info("buffer pipeline: content rejected", "decision", result.Decision, "reason", result.Reason, @@ -221,15 +243,28 @@ func (b *BufferPipeline) Process( } return result, nil - case core.Modify: - // Modify = inflight redact. With a FrameRedactor wired, rewrite the - // buffered timeline so the masked frames (not the original) reach the - // client. Without one, preserve the backward-compatible degrade: the - // rewrite is silently ignored and the body replays verbatim, surfaced - // via WARN + nexus_streaming_modify_degraded_total{reason="buffer_mode"} - // so a misconfigured hook is visible to admins. + case result.CarriesRedaction(): + // Inflight redact. With a FrameRedactor wired, rewrite the buffered timeline so + // the masked frames (not the original) reach the client. Without one, the + // posture decides (GAP B): the compliance-proxy appliance fails CLOSED (a + // redaction it cannot apply must never deliver the original); the agent NE path + // degrades to a disclosed replay of the original (host-packet safety). if b.frameRedactor == nil { - b.logger.Warn("buffer mode: Modify decision degraded to Approve (no frame redactor; rewrite ignored)", + if b.strictFailClosed { + b.logger.Warn("buffer mode: redaction required but no frame redactor; failing closed (appliance)", + "requestId", baseInput.RequestID, + "reason", result.Reason, + ) + RecordModifyDegraded(reasonRedactInflightUnsupported) + if err := writeErrorAndDone(client); err != nil { + return result, fmt.Errorf("buffer pipeline: write error response: %w", err) + } + if flusher, ok := client.(http.Flusher); ok { + flusher.Flush() + } + return result, nil + } + b.logger.Warn("buffer mode: redaction required but no frame redactor; degrading to replay (agent fail-open)", "requestId", baseInput.RequestID, "reason", result.Reason, ) @@ -238,18 +273,16 @@ func (b *BufferPipeline) Process( } redacted, rErr := b.frameRedactor.RedactReplay(events, result) if rErr != nil { - // Genuinely non-reconstructable wire (or an internal redactor - // error) — FAIL CLOSED. The original (unredacted) frames are - // never replayed; emit the policy error frame so zero - // unredacted content reaches the client, and bump the counter - // under the distinct unsupported reason. ResponseBodyRedacted - // stays unset (caller's stream-relay guard stores NULL). + // Non-reconstructable wire — FAIL CLOSED: the original frames are never replayed; + // emit the policy error frame so zero unredacted content reaches the client. No + // coarse counter here: the redactor already recorded the more-informative ROOT-CAUSE + // label in redactUnsupported on this same fault (bumping the coarse one too + // double-counted). ResponseBodyRedacted stays unset (stream-relay guard stores NULL). b.logger.Warn("buffer mode: Modify redaction unsupported; failing closed (no original replay)", "requestId", baseInput.RequestID, "reason", result.Reason, "error", rErr, ) - RecordModifyDegraded(reasonRedactInflightUnsupported) if err := writeErrorAndDone(client); err != nil { return result, fmt.Errorf("buffer pipeline: write error response: %w", err) } @@ -258,8 +291,12 @@ func (b *BufferPipeline) Process( } return result, nil } - // Supported redaction — replay the masked timeline. No degrade - // counter: the rewrite was honored. + // Disposition: stamp action=redact when the mask was genuinely applied (a co-firing + // BlockSoft that redact-delivered reads redact, not the block ceiling); skip it on the + // agent fail-open degrade (RedactReplay relayed the original + ReasonRedactInflightUnsupported). + if result.ReasonCode != core.ReasonRedactInflightUnsupported { + result.Action = core.ActionRedact + } return result, b.replay(ctx, client, redacted) default: diff --git a/packages/shared/transport/streaming/buffer_test.go b/packages/shared/transport/streaming/buffer_test.go index 34002950..e90b7442 100644 --- a/packages/shared/transport/streaming/buffer_test.go +++ b/packages/shared/transport/streaming/buffer_test.go @@ -10,6 +10,7 @@ import ( "github.com/prometheus/client_golang/prometheus/testutil" "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" + normalize "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/normalize/core" ) func TestBufferPipeline_AllApproved(t *testing.T) { @@ -143,6 +144,98 @@ func readCounter(_ *testing.T, reason string) float64 { return testutil.ToFloat64(modifyDegradedTotal.WithLabelValues(reason)) } +// TestRecordModelAEscalation_SplitsByCause pins #11: the Model A escalation counter +// increments under the bounded cause label, so a memory-pressure eviction (a buffer-ceiling +// tuning signal) is observably distinct from a confirmed enforcing hit. +func TestRecordModelAEscalation_SplitsByCause(t *testing.T) { + modelAEscalationTotal.DeleteLabelValues(ModelAEscalationConfirmed) + modelAEscalationTotal.DeleteLabelValues(ModelAEscalationMemoryPressure) + + RecordModelAEscalation(ModelAEscalationConfirmed) + RecordModelAEscalation(ModelAEscalationConfirmed) + RecordModelAEscalation(ModelAEscalationMemoryPressure) + + if got := testutil.ToFloat64(modelAEscalationTotal.WithLabelValues(ModelAEscalationConfirmed)); got != 2 { + t.Errorf("confirmed escalations = %v, want 2", got) + } + if got := testutil.ToFloat64(modelAEscalationTotal.WithLabelValues(ModelAEscalationMemoryPressure)); got != 1 { + t.Errorf("memory_pressure escalations = %v, want 1", got) + } +} + +// TestBufferPipeline_ApproveWebhookAuditSpansSoftBlock_NoOverBlock is the #14 fix at the +// layer the over-block actually occurs: on the strict appliance with NO frame redactor, a +// BlockSoft aggregate whose only spans are audit-only (RedactionApplicable=false — the +// approve-webhook+redactions shape) must NOT take the redaction arm (which would fail +// closed) — it soft-delivers the original via the default replay arm. +func TestBufferPipeline_ApproveWebhookAuditSpansSoftBlock_NoOverBlock(t *testing.T) { + mp := &mockPipeline{ + decideFn: func(ctx context.Context, input *core.HookInput) *core.CompliancePipelineResult { + return &core.CompliancePipelineResult{ + Decision: core.BlockSoft, + Reason: "advisory soft flag", + TransformSpans: []normalize.TransformSpan{{Source: normalize.SourceHook, Action: normalize.ActionRedact, ContentAddress: normalize.AddressAuditOnlySentinel}}, + RedactionApplicable: false, + } + }, + } + // Strict appliance + NO frame redactor: the pre-fix predicate (len(spans)>0) would + // route here to the redaction arm and fail closed (over-block). + bp := NewBufferPipeline(BufferConfig{}, mp, slog.Default()).WithStrictFailClosed(true) + + input := makeOpenAISSE("Hello", " World") + baseTx := &core.HookInput{Stage: "response", IngressType: "COMPLIANCE_PROXY", TargetHost: "api.openai.com"} + + var output bytes.Buffer + result, err := bp.Process(context.Background(), strings.NewReader(input), &output, baseTx) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result.CarriesRedaction() { + t.Fatal("audit-only spans under soft-block must NOT CarryRedaction") + } + outputStr := output.String() + if strings.Contains(outputStr, "blocked by policy") { + t.Fatal("appliance over-blocked a soft-deliver stream (the #14 bug) — must replay the original") + } + if !strings.Contains(outputStr, "Hello World") && !strings.Contains(outputStr, "Hello") { + t.Fatalf("expected the original content soft-delivered, got %q", outputStr) + } +} + +// TestBufferPipeline_ApplicableRedactSoftBlock_FailsClosed is the contrast/leak-guard: a +// BlockSoft aggregate carrying an APPLICABLE redaction (RedactionApplicable=true) on the +// strict appliance WITHOUT a frame redactor still fails closed — the predicate correctly +// routes a real redaction to the protective arm (never replays the original raw). +func TestBufferPipeline_ApplicableRedactSoftBlock_FailsClosed(t *testing.T) { + mp := &mockPipeline{ + decideFn: func(ctx context.Context, input *core.HookInput) *core.CompliancePipelineResult { + return &core.CompliancePipelineResult{ + Decision: core.BlockSoft, + Reason: "redact masked by soft-block", + TransformSpans: []normalize.TransformSpan{{Source: normalize.SourceHook, Action: normalize.ActionRedact, ContentAddress: "messages.0.content.0", Start: 0, End: 5}}, + RedactionApplicable: true, + } + }, + } + bp := NewBufferPipeline(BufferConfig{}, mp, slog.Default()).WithStrictFailClosed(true) + + input := makeOpenAISSE("secret payload") + baseTx := &core.HookInput{Stage: "response", IngressType: "COMPLIANCE_PROXY", TargetHost: "api.openai.com"} + + var output bytes.Buffer + result, err := bp.Process(context.Background(), strings.NewReader(input), &output, baseTx) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !result.CarriesRedaction() { + t.Fatal("an applicable redaction masked by soft-block MUST CarryRedaction") + } + if strings.Contains(output.String(), "secret payload") { + t.Fatal("strict appliance with no redactor must NOT replay the original (leak) — fail closed") + } +} + func TestBufferPipeline_EmptyStream(t *testing.T) { mp := &mockPipeline{} logger := slog.Default() diff --git a/packages/shared/transport/streaming/fold_hook_results_test.go b/packages/shared/transport/streaming/fold_hook_results_test.go new file mode 100644 index 00000000..33201174 --- /dev/null +++ b/packages/shared/transport/streaming/fold_hook_results_test.go @@ -0,0 +1,61 @@ +package streaming + +import ( + "testing" + + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// TestFoldHookResults_CollapsesPerCheckpointDuplicates reproduces the observed +// chunked_async "RESPONSE PIPELINE (63)" duplication: the same hook scanned at 63 +// checkpoints must fold to ONE result with summed microsecond latency and the +// latest decision — not 63 identical rows. +func TestFoldHookResults_CollapsesPerCheckpointDuplicates(t *testing.T) { + var acc []core.HookResult + for range 63 { + acc = append(acc, core.HookResult{HookID: "h1", HookName: "response-quality-signals", Decision: core.Approve, LatencyUs: 100}) + } + got := foldHookResults(acc) + if len(got) != 1 { + t.Fatalf("63 checkpoint scans must fold to 1 row, got %d", len(got)) + } + if got[0].LatencyUs != 6300 { + t.Fatalf("summed LatencyUs=%d want 6300", got[0].LatencyUs) + } +} + +// TestFoldHookResults_LatestWinsAndSums verifies summed latency + latest decision, +// and that distinct hooks stay separate in first-seen order. +func TestFoldHookResults_LatestWinsAndSums(t *testing.T) { + in := []core.HookResult{ + {HookID: "a", Decision: core.Approve, LatencyMs: 0, LatencyUs: 200}, + {HookID: "b", Decision: core.Approve, LatencyUs: 50}, + {HookID: "a", Decision: core.RejectHard, Reason: "blocked", LatencyMs: 1, LatencyUs: 1100}, + } + got := foldHookResults(in) + if len(got) != 2 { + t.Fatalf("expected 2 hooks, got %d", len(got)) + } + if got[0].HookID != "a" || got[1].HookID != "b" { + t.Fatalf("first-seen order not preserved: %+v", got) + } + if got[0].LatencyUs != 1300 || got[0].LatencyMs != 1 { + t.Fatalf("hook a latency not summed: %+v", got[0]) + } + if got[0].Decision != core.RejectHard || got[0].Reason != "blocked" { + t.Fatalf("hook a latest decision not kept: %+v", got[0]) + } +} + +// TestFoldHookResults_PassThroughSmall confirms the <=1 fast path returns input +// unchanged (no allocation for the common single-hook / empty case). +func TestFoldHookResults_PassThroughSmall(t *testing.T) { + if got := foldHookResults(nil); got != nil { + t.Fatalf("nil → %v", got) + } + one := []core.HookResult{{HookID: "x", LatencyUs: 9}} + got := foldHookResults(one) + if len(got) != 1 || got[0].HookID != "x" { + t.Fatalf("single passthrough altered: %+v", got) + } +} diff --git a/packages/shared/transport/streaming/frame_redactor.go b/packages/shared/transport/streaming/frame_redactor.go index 91ab532a..23c1289f 100644 --- a/packages/shared/transport/streaming/frame_redactor.go +++ b/packages/shared/transport/streaming/frame_redactor.go @@ -47,6 +47,13 @@ var ErrRewriteUnsupported = errors.New("streaming: frame redactor cannot reconst // - Any other non-nil error is treated as fail-closed as well (no // original replay) — a redactor that hit an internal error must not // leak the unredacted stream. +// - On ANY non-nil error return, the implementation MUST record a +// `nexus_streaming_modify_degraded_total` root-cause label (via +// streaming.RecordModifyDegraded) BEFORE returning — the buffer pipeline +// deliberately does NOT bump a coarse counter on the redactor's error arm +// (it would double-count, see buffer.go), so this metric is the only +// operator signal that a fail-closed redaction-unsupported event occurred. +// The production sseFrameRedactor satisfies this in redactUnsupported. // // The implementation MUST be re-entrant: BufferPipeline.Process can run // concurrently for many requests, each with its own redactor instance, diff --git a/packages/shared/transport/streaming/frame_redactor_test.go b/packages/shared/transport/streaming/frame_redactor_test.go index bc8e1824..d9aa378e 100644 --- a/packages/shared/transport/streaming/frame_redactor_test.go +++ b/packages/shared/transport/streaming/frame_redactor_test.go @@ -123,8 +123,13 @@ func TestBufferPipeline_ModifyWithRedactor_ReplaysMasked(t *testing.T) { // TestBufferPipeline_ModifyRedactorUnsupported_FailsClosed asserts that // when the FrameRedactor cannot reconstruct the wire (ErrRewriteUnsupported) // the buffer fails CLOSED: zero original (unredacted) content reaches the -// client, the policy error frame is delivered, and the -// redact_inflight_unsupported counter is bumped once. +// client and the policy error frame is delivered. The buffer pipeline does NOT +// itself bump the coarse redact_inflight_unsupported counter on this path (#21): +// the FrameRedactor owns the degrade signal and records the more-informative +// ROOT-CAUSE label (the production sseFrameRedactor bumps tlsbump_splice_divergence / +// tlsbump_tool_arg_undeliverable in redactUnsupported); double-bumping the coarse +// counter here over-counted the strict-buffer fault. This test's mock redactor does +// not emit a root-cause label, so the coarse counter stays 0. func TestBufferPipeline_ModifyRedactorUnsupported_FailsClosed(t *testing.T) { fr := &mockFrameRedactor{ fn: func(_ []*SSEEvent, _ *core.CompliancePipelineResult) ([]*SSEEvent, error) { @@ -154,8 +159,10 @@ func TestBufferPipeline_ModifyRedactorUnsupported_FailsClosed(t *testing.T) { if !strings.Contains(out, "blocked by policy") { t.Errorf("expected policy error frame on fail-closed, got: %q", out) } - if got := readCounter(t, reasonRedactInflightUnsupported); got != 1 { - t.Errorf("expected redact_inflight_unsupported counter == 1, got %v", got) + // #21: the buffer no longer double-bumps the coarse counter here — the redactor owns + // the (root-cause) degrade signal. This mock emits none, so the coarse stays 0. + if got := readCounter(t, reasonRedactInflightUnsupported); got != 0 { + t.Errorf("expected coarse redact_inflight_unsupported counter == 0 (the redactor owns the degrade label), got %v", got) } } @@ -183,8 +190,8 @@ func TestBufferPipeline_ModifyNoRedactor_LegacyDegrade(t *testing.T) { if !strings.Contains(out, "verbatim") || !strings.Contains(out, "bytes") { t.Errorf("expected verbatim replay (Modify ignored without redactor), got: %q", out) } - if !strings.Contains(logBuf.String(), "Modify decision degraded to Approve") { - t.Errorf("expected WARN log about degradation, got: %s", logBuf.String()) + if !strings.Contains(logBuf.String(), "degrading to replay") { + t.Errorf("expected WARN log about the no-redactor degrade, got: %s", logBuf.String()) } if !strings.Contains(logBuf.String(), "buf-legacy-1") { t.Errorf("expected requestId in degradation log, got: %s", logBuf.String()) @@ -235,3 +242,98 @@ func TestBufferPipeline_RedactorConcurrency(t *testing.T) { t.Errorf("redaction B cross-contaminated: %q", b) } } + +// TestBufferPipeline_CoFiringBlockSoft_RedactDelivers proves the #13 appliance-leak fix: +// a co-firing soft-block promotes the aggregate Decision to BlockSoft but carries the +// redact's content, so the buffer routes it through the FrameRedactor (gating on +// CarriesRedaction, not Decision==Modify) and delivers the MASKED stream — not the +// original, which the old Decision==Modify switch fell through to `default` and replayed +// raw (the leak this fixes). +func TestBufferPipeline_CoFiringBlockSoft_RedactDelivers(t *testing.T) { + fr := &mockFrameRedactor{ + fn: func(events []*SSEEvent, _ *core.CompliancePipelineResult) ([]*SSEEvent, error) { + out := make([]*SSEEvent, len(events)) + for i, e := range events { + masked := *e + masked.Data = strings.ReplaceAll(e.Data, "SECRET", "[REDACTED]") + out[i] = &masked + } + return out, nil + }, + } + mp := &mockPipeline{decideFn: func(_ context.Context, _ *core.HookInput) *core.CompliancePipelineResult { + // Models a real co-firing redact (a Modify hook's ModifiedContent) masked by a + // soft-block: mergeResults would set RedactionApplicable; this test bypasses the + // merge, so it sets the flag explicitly to stay faithful to the aggregate shape. + return &core.CompliancePipelineResult{Decision: core.BlockSoft, ModifiedContent: []core.ContentBlock{{}}, RedactionApplicable: true, Reason: "soft-block masking redact"} + }} + bp := NewBufferPipeline(BufferConfig{}, mp, slog.Default()).WithFrameRedactor(fr) + + var output bytes.Buffer + result, err := bp.Process(context.Background(), strings.NewReader(makeOpenAISSE("my SECRET token")), + &output, &core.HookInput{Stage: "response", RequestID: "buf-cofire-1"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + out := output.String() + if strings.Contains(out, "SECRET") { + t.Fatalf("co-firing BlockSoft leaked the original content: %q", out) + } + if !strings.Contains(out, "[REDACTED]") { + t.Fatalf("co-firing BlockSoft must redact-deliver the masked stream, got %q", out) + } + // #13 P3: the disposition is stamped redact (the mask WAS applied), not the BlockSoft + // ceiling — so the audit row reads action=redact for the redact-delivered stream. + if result == nil || result.Action != core.ActionRedact { + t.Fatalf("result.Action = %v, want redact (an applied splice is a redact disposition, not block)", result) + } +} + +// TestBufferPipeline_CoFiringBlockSoft_FailOpenDegrade_KeepsBlockAction pins the #13 P3 +// guard: when the redactor CANNOT splice and fails open (agent: returns the original frames +// with a nil error after stamping ReasonRedactInflightUnsupported), the original was relayed, +// so the disposition must NOT be re-stamped redact — it stays the aggregate BlockSoft/block. +func TestBufferPipeline_CoFiringBlockSoft_FailOpenDegrade_KeepsBlockAction(t *testing.T) { + fr := &mockFrameRedactor{ + fn: func(events []*SSEEvent, result *core.CompliancePipelineResult) ([]*SSEEvent, error) { + // Agent fail-open degrade: relay the ORIGINAL frames unchanged, nil error, + // and stamp the disclosed unsupported reason (mirrors sseFrameRedactor). + result.ReasonCode = core.ReasonRedactInflightUnsupported + return events, nil + }, + } + mp := &mockPipeline{decideFn: func(_ context.Context, _ *core.HookInput) *core.CompliancePipelineResult { + return &core.CompliancePipelineResult{Decision: core.BlockSoft, Action: core.ActionBlock, ModifiedContent: []core.ContentBlock{{}}, Reason: "soft-block masking redact"} + }} + bp := NewBufferPipeline(BufferConfig{}, mp, slog.Default()).WithFrameRedactor(fr) + + var output bytes.Buffer + result, err := bp.Process(context.Background(), strings.NewReader(makeOpenAISSE("token")), + &output, &core.HookInput{Stage: "response", RequestID: "buf-cofire-failopen"}) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result == nil || result.Action != core.ActionBlock { + t.Fatalf("result.Action = %v, want block — the fail-open degrade relayed the original, so it must NOT stamp redact", result) + } +} + +// TestBufferPipeline_NoRedactor_Strict_FailsClosed proves GAP B for the appliance: a +// redaction with no frame redactor wired fails CLOSED (error frame, no original) under +// strictFailClosed, instead of the agent's degrade-replay. +func TestBufferPipeline_NoRedactor_Strict_FailsClosed(t *testing.T) { + bp := NewBufferPipeline(BufferConfig{}, modifyPipeline(), slog.Default()).WithStrictFailClosed(true) + + var output bytes.Buffer + if _, err := bp.Process(context.Background(), strings.NewReader(makeOpenAISSE("verbatim SECRET")), + &output, &core.HookInput{Stage: "response", RequestID: "buf-strict-1"}); err != nil { + t.Fatalf("unexpected error: %v", err) + } + out := output.String() + if strings.Contains(out, "SECRET") { + t.Fatalf("strict appliance must not replay the original when no redactor is wired, got %q", out) + } + if !strings.Contains(out, `"error"`) { + t.Fatalf("strict appliance must fail closed with an error frame, got %q", out) + } +} diff --git a/packages/shared/transport/streaming/live.go b/packages/shared/transport/streaming/live.go index 56d30635..a2da0929 100644 --- a/packages/shared/transport/streaming/live.go +++ b/packages/shared/transport/streaming/live.go @@ -3,7 +3,6 @@ package streaming import ( "context" "errors" - "fmt" "github.com/goccy/go-json" "io" "log/slog" @@ -56,13 +55,6 @@ type PipelineExecutor interface { Execute(ctx context.Context, input *core.HookInput) *core.CompliancePipelineResult } -// approvedChunk is a batch of events approved by a checkpoint evaluation. -type approvedChunk struct { - events []*SSEEvent - err error // non-nil signals the writer should emit an error and stop - reason string // human-readable rejection reason, if any -} - // LivePipeline processes an SSE stream with checkpoint-based compliance core. type LivePipeline struct { config LiveConfig @@ -143,9 +135,42 @@ func (l *LivePipeline) CapturedTruncated() bool { return l.captureBuf.Truncated() } -// Process reads SSE events from upstream, applies checkpoint-based compliance -// hooks, and writes approved events to the client writer. -// Returns the aggregated compliance result. +// foldHookResults collapses the per-checkpoint scans of the same hook into ONE +// result: LatencyMs/LatencyUs are summed (the real scan CPU spent across the +// stream) and every other field is taken from the latest scan (the last checkpoint +// is authoritative). The chunked_async response pipeline runs once per checkpoint +// and appends each scan's results, so without this fold the same hook is emitted N +// times — observed as a "RESPONSE PIPELINE (63)" list of identical rows — and the +// response hook aggregates are N×-inflated. Keyed by stable identity +// (HookID+ImplementationID), not the volatile per-scan Order; first-seen order is +// preserved. +func foldHookResults(in []core.HookResult) []core.HookResult { + if len(in) <= 1 { + return in + } + type key struct{ hookID, implID string } + idx := make(map[key]int, len(in)) + out := make([]core.HookResult, 0, len(in)) + for _, r := range in { + k := key{r.HookID, r.ImplementationID} + if i, ok := idx[k]; ok { + r.LatencyMs += out[i].LatencyMs + r.LatencyUs += out[i].LatencyUs + out[i] = r // latest scan authoritative for all non-latency fields + } else { + idx[k] = len(out) + out = append(out, r) + } + } + return out +} + +// Process reads SSE events from upstream and relays them to the client in REAL TIME +// (write-through), running observe-only compliance checkpoints on the accumulated +// content for AUDIT. It never holds back, blocks, or rewrites the wire: an enforcing +// (block/redact) scope is routed to the buffer / Model A paths upstream, so this +// path carries only non-enforcing traffic and a checkpoint's decision is recorded +// but never gates delivery. Returns the aggregated (always-Approve) result. func (l *LivePipeline) Process( ctx context.Context, upstream io.Reader, @@ -184,13 +209,10 @@ func (l *LivePipeline) Process( } eventChan := make(chan *SSEEvent, l.config.ChannelSize) - approvedChan := make(chan approvedChunk, l.config.ChannelSize) var ( - wg sync.WaitGroup - readerErr error - writerErr error - finalResult *core.CompliancePipelineResult + wg sync.WaitGroup + readerErr error ) // --- Reader goroutine --- @@ -226,182 +248,88 @@ func (l *LivePipeline) Process( } }() - // --- Compliance goroutine --- - wg.Add(1) - go func() { - defer wg.Done() - defer close(approvedChan) - - var ( - pendingEvents []*SSEEvent - // accumulatedAll is the full text accumulated so far (fed to every - // checkpoint). pendingText is the text since the last checkpoint. - // Both use strings.Builder so appending each SSE delta is amortized - // O(1); naive `s += delta` in this per-event loop is O(n²) over the - // length of a long stream. pendingLen tracks the - // builder's current length so the checkpoint threshold check stays a - // cheap int compare instead of len(pendingText.String()). - accumulatedAll strings.Builder - pendingText strings.Builder - pendingLen int - totalBytes int - allResults []core.HookResult - ) - - flushCheckpoint := func() core.Decision { - if len(pendingEvents) == 0 && pendingLen == 0 { - return core.Approve - } - - checkpointInput := buildCheckpointInput(baseInput, accumulatedAll.String()) - - // Let caller swap in a Registry-normalized payload so - // hooks see structured chat content (model name, tool_calls, - // reasoning segments, etc.) instead of the flat-text fallback. - // Receives the cumulative raw SSE wire bytes seen so far so - // each checkpoint re-normalizes the full accumulated payload - // — matches the "hooks operate on what user has seen so far" - // chunked_async semantic the user binding specifies. - if l.preHook != nil && rawAcc != nil { - l.preHook(rawAcc.Snapshot(), checkpointInput) - } - - result := l.pipeline.Execute(ctx, checkpointInput) - if result == nil { - // Pipeline returned nil — treat as approve. - return core.Approve - } - - allResults = append(allResults, result.HookResults...) + // --- Delivery + observe-only audit (current goroutine) --- + // accumulatedAll grows with the cumulative response text (bounded by + // MaxBufferSize); strings.Builder keeps the per-event append amortized O(1). + // pendingLen counts chars since the last checkpoint so the cadence check stays a + // cheap int compare. + var ( + accumulatedAll strings.Builder + pendingLen int + totalBytes int + allResults []core.HookResult + auditCapped bool + writerErr error + ) - switch result.Decision { - case core.RejectHard: - finalResult = &core.CompliancePipelineResult{ - Decision: core.RejectHard, - Reason: result.Reason, - ReasonCode: result.ReasonCode, - HookResults: allResults, - } - select { - case approvedChan <- approvedChunk{err: fmt.Errorf("blocked by policy"), reason: result.Reason}: - case <-ctx.Done(): - } - cancel() - // cancel alone doesn't unblock a reader sitting in - // upstream.Read — close the upstream so the reader - // goroutine exits and wg.Wait() can return (same wedge - // as the writer-error and overflow branches). - CloseUpstreamOnExit(upstream) - return core.RejectHard - - default: - // Approve or Abstain — flush pending events. - select { - case approvedChan <- approvedChunk{events: pendingEvents}: - case <-ctx.Done(): - return core.Approve - } - pendingEvents = nil - pendingText.Reset() - pendingLen = 0 - return core.Approve - } + // runCheckpoint fires the compliance pipeline OBSERVE-ONLY: it records the hook + // results for the audit aggregate but never blocks or rewrites the wire. The + // PreHook re-normalizes the cumulative raw bytes so hooks see structured content. + runCheckpoint := func() { + // Audit-only relay with no executor or no base input has nothing to record — + // skip the checkpoint rather than deref a nil. Delivery is independent of the + // checkpoint, so the stream still writes through. (Reached when a caller builds + // a hookless live pipeline for usage accumulation only, e.g. AI traffic whose + // response stage binds no hooks.) + if l.pipeline == nil || baseInput == nil { + return } - - for evt := range eventChan { - if ctx.Err() != nil { - return - } - - deltaText := extractDeltaText(evt) - totalBytes += len(evt.Data) - - if totalBytes > l.config.MaxBufferSize { - l.logger.Error("live pipeline: max buffer size exceeded", "bytes", totalBytes) - select { - case approvedChan <- approvedChunk{err: fmt.Errorf("stream buffer exceeded maximum size")}: - case <-ctx.Done(): - } - cancel() - // cancel doesn't unblock a slow upstream.Read — close - // the upstream so the reader goroutine exits and - // wg.Wait() can return (same wedge as the writer-error - // path). - CloseUpstreamOnExit(upstream) - return - } - - pendingEvents = append(pendingEvents, evt) - pendingText.WriteString(deltaText) - accumulatedAll.WriteString(deltaText) - pendingLen += len(deltaText) - - // Check if we've reached the checkpoint threshold. - if pendingLen >= l.config.CheckpointChars { - decision := flushCheckpoint() - if decision == core.RejectHard { - return - } - } + checkpointInput := buildCheckpointInput(baseInput, accumulatedAll.String()) + if l.preHook != nil && rawAcc != nil { + l.preHook(rawAcc.Snapshot(), checkpointInput) } - - // Final checkpoint: flush remaining accumulated text. - flushCheckpoint() - - // Build final aggregate result if not already set by a rejection. - if finalResult == nil { - finalResult = &core.CompliancePipelineResult{ - Decision: core.Approve, - HookResults: allResults, - } + if result := l.pipeline.Execute(ctx, checkpointInput); result != nil { + allResults = append(allResults, result.HookResults...) } - }() + } - // --- Writer goroutine (runs on current goroutine) --- - // flusher / canFlush were resolved above against the original client - // before MultiWriter wrapping (see comment near captureBuf init). - for chunk := range approvedChan { - if chunk.err != nil { - writerErr = writeErrorAndDone(client) - if canFlush { - flusher.Flush() - } - // Drain remaining items so the compliance goroutine does not block. - for range approvedChan { - } - break - } - writeOK := true - for _, evt := range chunk.events { - if err := WriteSSEEvent(client, evt); err != nil { - writerErr = err - cancel() - // cancel alone does NOT unblock a reader goroutine - // sitting inside upstream.Read. If - // upstream is a slow / hung connection, the reader stays - // blocked until upstream actually delivers bytes or the - // caller's outer defer Close() runs — but Process can't - // return until wg.Wait() does, and wg.Wait() can't return - // until the reader exits. CloseUpstreamOnExit calls - // upstream.(io.Closer).Close synchronously to unblock - // the reader's parser.Next; caller's outer defer Close - // is idempotent on http.Body. - CloseUpstreamOnExit(upstream) - writeOK = false - break - } - } - if !writeOK { + for evt := range eventChan { + // AUDIT-ONLY: deliver every event in real time — delivery is NEVER gated on a + // checkpoint. A write error closes the upstream so the reader goroutine exits + // and wg.Wait() can return (the slow-upstream wedge guard). + if err := WriteSSEEvent(client, evt); err != nil { + writerErr = err + cancel() + CloseUpstreamOnExit(upstream) break } if canFlush { flusher.Flush() } + + if auditCapped { + continue + } + deltaText := extractDeltaText(evt) + accumulatedAll.WriteString(deltaText) + pendingLen += len(deltaText) + totalBytes += len(evt.Data) + if totalBytes > l.config.MaxBufferSize { + // Audit-accumulation cap: stop scanning further content to bound memory, + // but KEEP delivering — an audit-only relay must never break a + // non-enforcing stream just because it grew past the scan budget. + l.logger.Warn("live pipeline: audit accumulation capped at max buffer size", "bytes", totalBytes) + auditCapped = true + continue + } + if pendingLen >= l.config.CheckpointChars { + runCheckpoint() + pendingLen = 0 + } + } + + // Mandatory final checkpoint (observe-only): scan the trailing content not yet + // covered by a periodic checkpoint so a stream shorter than the cadence — or the + // tail after the last checkpoint — is still audited once. Skipped when a write + // error aborted delivery or when the last checkpoint already covered everything. + if writerErr == nil && pendingLen > 0 { + runCheckpoint() } - // Wait for reader and compliance goroutines to finish. + // Wait for the reader goroutine to finish. wg.Wait() + finalResult := &core.CompliancePipelineResult{Decision: core.Approve, HookResults: foldHookResults(allResults)} if writerErr != nil { return finalResult, writerErr } diff --git a/packages/shared/transport/streaming/live_test.go b/packages/shared/transport/streaming/live_test.go index 4336280e..d03f72cb 100644 --- a/packages/shared/transport/streaming/live_test.go +++ b/packages/shared/transport/streaming/live_test.go @@ -43,6 +43,25 @@ func makeOpenAISSE(deltas ...string) string { return sb.String() } +func TestLivePipeline_NilExecutorAndNilBase_NoPanic(t *testing.T) { + // A hookless live pipeline (nil executor) built for usage accumulation only, with a + // nil base input, must not panic at the checkpoint — it skips checkpoint recording + // and still delivers every event (audit-only write-through). + lp := NewLivePipeline(LiveConfig{CheckpointChars: 5}, nil, slog.Default()) + input := makeOpenAISSE("Hello", " ", "World", "!") + var output bytes.Buffer + result, err := lp.Process(context.Background(), strings.NewReader(input), &output, nil) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if result == nil || result.Decision != core.Approve { + t.Fatalf("audit-only relay must return Approve, got %v", result) + } + if out := output.String(); !strings.Contains(out, "Hello") || !strings.Contains(out, "World") { + t.Fatalf("audit-only relay must deliver every event, got %q", out) + } +} + func TestLivePipeline_AllApproved(t *testing.T) { mp := &mockPipeline{} logger := slog.Default() @@ -88,21 +107,21 @@ func TestLivePipeline_AllApproved(t *testing.T) { } } +// TestLivePipeline_RejectAtCheckpoint pins the AUDIT-ONLY contract: a checkpoint +// hook returning RejectHard on the live path NEVER blocks — every event is delivered +// write-through, no in-band error frame is emitted, and Process returns Approve. The +// live path carries only non-enforcing traffic (enforcing scopes route to buffer / +// Model A upstream), so a hook decision is observed for audit but never gates the wire. func TestLivePipeline_RejectAtCheckpoint(t *testing.T) { callCount := 0 mp := &mockPipeline{ decideFn: func(ctx context.Context, input *core.HookInput) *core.CompliancePipelineResult { callCount++ - // Reject on second checkpoint. + // Even an enforcing RejectHard must NOT block the audit-only live path. if callCount >= 2 { - return &core.CompliancePipelineResult{ - Decision: core.RejectHard, - Reason: "policy violation detected", - } - } - return &core.CompliancePipelineResult{ - Decision: core.Approve, + return &core.CompliancePipelineResult{Decision: core.RejectHard, Reason: "policy violation detected"} } + return &core.CompliancePipelineResult{Decision: core.Approve} }, } logger := slog.Default() @@ -128,14 +147,53 @@ func TestLivePipeline_RejectAtCheckpoint(t *testing.T) { t.Fatal("expected non-nil result") return } - if result.Decision != core.RejectHard { - t.Errorf("expected REJECT_HARD, got %s", result.Decision) + if result.Decision != core.Approve { + t.Errorf("audit-only live path must always return Approve, got %s", result.Decision) + } + if callCount == 0 { + t.Error("expected the audit checkpoint pipeline to run") } - // Output should contain the error message. outputStr := output.String() - if !strings.Contains(outputStr, "blocked by policy") { - t.Errorf("expected output to contain error message, got:\n%s", outputStr) + if strings.Contains(outputStr, "blocked by policy") { + t.Errorf("audit-only live path must NOT emit a block frame, got:\n%s", outputStr) + } + // All content is delivered write-through despite the RejectHard checkpoint. + if !strings.Contains(outputStr, "Hello") || !strings.Contains(outputStr, "policy") { + t.Errorf("expected the full stream delivered write-through, got:\n%s", outputStr) + } +} + +// TestLivePipeline_AuditOnly_PreHookAndCapture exercises the observe-only checkpoint +// with a PreHook installed (the raw-bytes tee) plus body capture: the PreHook fires +// at each checkpoint, the full stream is delivered write-through, and the captured +// body mirrors the delivered bytes. Decision is Approve (audit-only never enforces). +func TestLivePipeline_AuditOnly_PreHookAndCapture(t *testing.T) { + preHookFired := 0 + lp := NewLivePipeline(LiveConfig{CheckpointChars: 5}, &mockPipeline{}, slog.Default()) + lp.WithPreHook(func(_ []byte, _ *core.HookInput) { preHookFired++ }) + lp.WithBodyCapture(1 << 20) + + input := makeOpenAISSE("Hello", " World", " this", " is", " enough") + var out bytes.Buffer + res, err := lp.Process(context.Background(), strings.NewReader(input), &out, &core.HookInput{ + Stage: "response", + IngressType: "COMPLIANCE_PROXY", + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if res == nil || res.Decision != core.Approve { + t.Errorf("audit-only must return Approve, got %+v", res) + } + if preHookFired == 0 { + t.Error("PreHook never fired — the raw-bytes tee + checkpoint PreHook branch is unexercised") + } + if !strings.Contains(out.String(), "Hello") || !strings.Contains(out.String(), "enough") { + t.Errorf("expected the full stream delivered write-through, got %q", out.String()) + } + if cb := lp.CapturedBytes(); cb == nil || !strings.Contains(string(cb), "Hello") { + t.Errorf("body capture did not record the delivered bytes, got %q", string(cb)) } } @@ -327,53 +385,15 @@ func TestLivePipeline_WriterError_ClosesUpstream(t *testing.T) { } } -// TestLivePipeline_RejectHard_ClosesUpstream pins the RejectHard -// close fix: the RejectHard branch in the -// compliance goroutine also needs to call CloseUpstreamOnExit (the -// initial fix covered writer-error + overflow but missed RejectHard). -// Without this, a slow upstream + RejectHard mid-stream wedges -// indefinitely. blockingReader makes the test fail deterministically -// if the fix is reverted. -func TestLivePipeline_RejectHard_ClosesUpstream(t *testing.T) { - rejectMP := &mockPipeline{ - decideFn: func(_ context.Context, _ *core.HookInput) *core.CompliancePipelineResult { - return &core.CompliancePipelineResult{Decision: core.RejectHard, Reason: "test reject"} - }, - } - upstream := newBlockingReader([]byte(makeOpenAISSE("hello world is enough text to trip checkpoint"))) - var writer bytes.Buffer - lp := NewLivePipeline(LiveConfig{CheckpointChars: 5}, rejectMP, slog.Default()) - - done := make(chan struct{}) - go func() { - defer close(done) - _, _ = lp.Process(context.Background(), upstream, &writer, &core.HookInput{ - Stage: "response", - IngressType: "COMPLIANCE_PROXY", - }) - }() - select { - case <-done: - case <-time.After(2 * time.Second): - t.Fatal("Process did not return within 2s after RejectHard — wedge regression (R-1: CloseUpstreamOnExit missing at RejectHard branch)") - } - if upstream.closeCount() == 0 { - t.Errorf("expected upstream.Close on RejectHard (R-1 fix); got 0") - } -} - -// TestLivePipeline_CancelDuringCheckpoint_FinalResultRaceFree drives -// the cancel-mid-checkpoint scenario. A slow pipeline.Execute simulates -// a hook taking time; the test cancels ctx WHILE Execute is in -// flight. The compliance goroutine may write `finalResult` after the -// cancel, then return; the main goroutine reads `finalResult` only -// after wg.Wait(), so happens-before via wg makes the publish safe. -// -// Run with `go test -race` (CI default). A potential race on -// `finalResult` was flagged here; this test pins that the -// path is race-free under the race detector. Any future refactor that -// reorders the publish before wg.Wait() will trip this test. -func TestLivePipeline_CancelDuringCheckpoint_FinalResultRaceFree(t *testing.T) { +// TestLivePipeline_CancelDuringCheckpoint_NoDeadlock drives the +// cancel-mid-checkpoint scenario: a slow observe-only pipeline.Execute holds until +// ctx is cancelled WHILE the checkpoint runs inline on the (single) delivery +// goroutine. The test pins that Process returns promptly (no deadlock on wg.Wait, +// no hang on the cancelled Execute) and returns a sane decision shape. There is no +// cross-goroutine publish to race: runCheckpoint / allResults / the final result are +// all touched only by the main goroutine after the audit-only refactor — run under +// `go test -race` to confirm. +func TestLivePipeline_CancelDuringCheckpoint_NoDeadlock(t *testing.T) { // Slow pipeline that lets the test inject cancel during Execute. executing := make(chan struct{}, 1) mp := &mockPipeline{ diff --git a/packages/shared/transport/streaming/metrics.go b/packages/shared/transport/streaming/metrics.go index 17b823b3..f3953510 100644 --- a/packages/shared/transport/streaming/metrics.go +++ b/packages/shared/transport/streaming/metrics.go @@ -34,3 +34,35 @@ var modifyDegradedTotal = promauto.NewCounterVec(prometheus.CounterOpts{ func RecordModifyDegraded(reason string) { modifyDegradedTotal.WithLabelValues(reason).Inc() } + +// Bounded `cause` label values for modelAEscalationTotal. +const ( + // ModelAEscalationConfirmed: a prescan hit was confirmed as an enforcing + // (block/redact) action, so the Model A engine handed off to buffer-to-end + // redaction — the normal enforcement path. + ModelAEscalationConfirmed = "confirmed" + // ModelAEscalationMemoryPressure: the held buffer hit MaxBufferBytes while a + // content unit was still incomplete, so the engine escalated rather than flush + // it raw — a tuning signal (the stream out-grew the buffer ceiling), distinct + // from a real policy hit. + ModelAEscalationMemoryPressure = "memory_pressure" +) + +// modelAEscalationTotal counts Model A streaming escalations split by CAUSE so an +// operator can tell a real policy hit ("confirmed") apart from a buffer-ceiling +// eviction ("memory_pressure") — the latter signals MaxBufferBytes may need raising +// for that traffic. This is a log/metric dimension only: the cause is NOT persisted +// on traffic_event (no schema add) and never reuses ResponseHookReasonCode (which the +// canonical-buffer redaction overwrites). Shared across all three data planes; the +// Prometheus job/instance labels distinguish which service emitted it. +var modelAEscalationTotal = promauto.NewCounterVec(prometheus.CounterOpts{ + Name: "nexus_streaming_modela_escalation_total", + Help: "Count of Model A streaming escalations, split by cause (confirmed enforcing hit vs memory-pressure buffer eviction).", +}, []string{"cause"}) + +// RecordModelAEscalation bumps the Model A escalation counter for the given cause +// (use the ModelAEscalation* constants). Exported so both Model A substrates (the +// ai-gateway canonical relay and the tlsbump wire path) record without re-registering. +func RecordModelAEscalation(cause string) { + modelAEscalationTotal.WithLabelValues(cause).Inc() +} diff --git a/packages/shared/transport/streaming/modela/coverage_warn.go b/packages/shared/transport/streaming/modela/coverage_warn.go new file mode 100644 index 00000000..66d8f466 --- /dev/null +++ b/packages/shared/transport/streaming/modela/coverage_warn.go @@ -0,0 +1,47 @@ +package modela + +import ( + "log/slog" + "sync" +) + +// DefaultTailWindowBytes is the package default trailing-window budget (the value +// withDefaults applies when Config.TailWindowBytes is unset). Exported so a substrate can +// compare a rule set's derived MaxPatternBytes against the window the engine will clamp the +// flush-before-deliver lookahead below (see withDefaults) and warn the operator when a +// contiguous enforceable pattern meets or exceeds it. +const DefaultTailWindowBytes = defaultTailWindowBytes + +// coverageWarnSeen dedupes the streaming-coverage Warn by derived bound so a busy stream +// does not log on every setup. The key space is the set of distinct derived pattern bounds +// across loaded rule sets — admin-authored and finite — so no eviction is needed. +var coverageWarnSeen sync.Map // maxBounded int -> struct{} + +// WarnStreamingCoverageGap emits a one-time (per distinct maxBounded) operator warning when +// a rule set's longest CONTIGUOUS enforceable pattern (maxBounded — already derived by the +// substrate via pipeline.MaxPatternBound and threaded in; this never recomputes it) meets or +// exceeds the streaming tail window. At/above that size the engine clamps the +// flush-before-deliver lookahead below the window, so Model-A real-time streaming is only +// best-effort for such a pattern: a value that long straddling unit boundaries may leak a +// bounded fragment before its completion is observed (the engine's disclosed surface). +// +// The remediation is deliberately NOT "raise the tail window" — the window is not an admin +// knob. Route the affected policy through BUFFERED streaming mode (full coverage) or narrow +// the >window-bounded rule. +// +// Silent in normal operation (realistic PII / token / JWT patterns sit well under the +// window); fires only on a genuinely unusual rule set. It is observability-only — off the +// per-byte path (called once at stream setup) and changes no enforcement. The dedupe uses +// LoadOrStore so two concurrent first-sight streams warn exactly once. +func WarnStreamingCoverageGap(logger *slog.Logger, maxBounded, tailWindow int) { + if logger == nil || maxBounded < tailWindow { + return + } + if _, loaded := coverageWarnSeen.LoadOrStore(maxBounded, struct{}{}); loaded { + return + } + logger.Warn("streaming compliance: a loaded rule set's longest contiguous enforceable pattern meets or exceeds the streaming tail window; Model-A real-time streaming is best-effort for it — route the policy through buffered streaming mode for full coverage, or narrow the rule", + "maxPatternBytes", maxBounded, + "tailWindowBytes", tailWindow, + ) +} diff --git a/packages/shared/transport/streaming/modela/coverage_warn_test.go b/packages/shared/transport/streaming/modela/coverage_warn_test.go new file mode 100644 index 00000000..64a7595e --- /dev/null +++ b/packages/shared/transport/streaming/modela/coverage_warn_test.go @@ -0,0 +1,58 @@ +package modela + +import ( + "context" + "log/slog" + "testing" +) + +// countingHandler records how many log records at each level were emitted. +type countingHandler struct { + warns int + last map[string]any +} + +func (h *countingHandler) Enabled(context.Context, slog.Level) bool { return true } +func (h *countingHandler) WithAttrs([]slog.Attr) slog.Handler { return h } +func (h *countingHandler) WithGroup(string) slog.Handler { return h } +func (h *countingHandler) Handle(_ context.Context, r slog.Record) error { + if r.Level == slog.LevelWarn { + h.warns++ + h.last = map[string]any{} + r.Attrs(func(a slog.Attr) bool { h.last[a.Key] = a.Value.Any(); return true }) + } + return nil +} + +func TestWarnStreamingCoverageGap(t *testing.T) { + h := &countingHandler{} + logger := slog.New(h) + + // Below the window: no warning (the normal case). + WarnStreamingCoverageGap(logger, 100, DefaultTailWindowBytes) + if h.warns != 0 { + t.Fatalf("below-window must not warn, got %d", h.warns) + } + + // At/above the window, first sight: exactly one warning carrying the sizes. + const big = 987654 // unique, unlikely to collide with other tests' keys + WarnStreamingCoverageGap(logger, big, DefaultTailWindowBytes) + if h.warns != 1 { + t.Fatalf("at/above-window first sight must warn once, got %d", h.warns) + } + if h.last["maxPatternBytes"] != int64(big) { + t.Errorf("warn maxPatternBytes = %v, want %d", h.last["maxPatternBytes"], big) + } + if h.last["tailWindowBytes"] != int64(DefaultTailWindowBytes) { + t.Errorf("warn tailWindowBytes = %v, want %d", h.last["tailWindowBytes"], DefaultTailWindowBytes) + } + + // Same bound again: deduped, no second warning. + WarnStreamingCoverageGap(logger, big, DefaultTailWindowBytes) + if h.warns != 1 { + t.Fatalf("repeated same-bound call must be deduped, got %d warns", h.warns) + } + + // Nil logger must not panic. + WarnStreamingCoverageGap(nil, big+1, DefaultTailWindowBytes) +} diff --git a/packages/shared/transport/streaming/modela/engine.go b/packages/shared/transport/streaming/modela/engine.go new file mode 100644 index 00000000..b1d37076 --- /dev/null +++ b/packages/shared/transport/streaming/modela/engine.go @@ -0,0 +1,413 @@ +// Package modela implements the substrate-agnostic streaming compliance engine +// shared by every ingress that relays an SSE response under a redact-scope +// chunked_async policy (the AI Gateway canonical relay and the tlsbump transparent +// proxy used by the agent and compliance-proxy). +// +// The engine is the prescan-gated REAL-TIME streaming algorithm: it forwards the +// response to the client in real time EXCEPT a trailing bounded tail held +// undelivered, runs a cheap union prescan once a batch of new content accumulates +// (and ALWAYS before releasing any held unit and at EOF, so batching never delivers +// unscanned content), pays for one full compliance confirm only on a prescan hit, +// and on a confirmed redact/block hit ESCALATES to buffer-to-end redaction so the +// held tail (and the remainder) are delivered redacted, never raw. +// +// The bounded tail is the load-bearing safety guarantee: a sensitive value no longer +// than MaxPatternBytes is still HELD AND re-scanned when its pattern-completing bytes +// arrive, so the prescan hit + confirm fire while the value is undelivered and the +// escalation redacts it — the complete value is never delivered. The firm guarantee is +// scoped to MaxPatternBytes (the longest enforceable pattern), NOT the whole window: +// because the prescan is batched (see Config.PrescanBatchBytes), a unit is re-scanned +// for completion only up to MaxPatternBytes past its end (the flush-before-deliver +// lookahead), and a unit is evicted once the held window fills regardless. Disclosed +// best-effort surfaces (strong compliance for these is the buffer or block modes, not +// Model A): (1) a value LARGER than the tail window leaks a bounded fragment before its +// completion is observed. (2) a value in (MaxPatternBytes, TailWindowBytes] that +// straddles a unit boundary where the leading unit's size approaches the window — its +// completion may not be held-AND-scanned before the leading unit is evicted. Soundness +// for sub-MaxPatternBytes values requires TailWindowBytes > MaxPatternBytes + +// PrescanBatchBytes + maxUnitSize; for substrates whose units approach the window (the +// tlsbump large-frame path) the operator MUST size TailWindowBytes (and MaxPatternBytes, +// to cover the longest contiguous enforceable pattern — e.g. JWT/PEM) accordingly. (3) a +// memory-cap (MaxBufferBytes) eviction that would otherwise flush an incomplete content +// unit raw instead ESCALATES (see Run), so memory pressure never silently overrides the +// in-window redaction guarantee. +// +// Everything substrate-specific — what a "unit" is, how its redactable text is +// extracted, how a unit is delivered to the client, how the confirmed remainder is +// redacted, and whether an unproducible redaction fails OPEN (tlsbump host-packet +// path) or CLOSED (the gateway appliance) — lives behind the Substrate interface. +// The engine carries no transport, no canonical/wire knowledge, and no fail +// posture: it only decides WHEN to hold, confirm, release, and escalate. +package modela + +import ( + "context" + "errors" + "io" + + hookcore "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +const ( + // defaultTailWindowBytes bounds the trailing redactable content held + // undelivered. A sensitive value shorter than this window is fully retained + // when its completing bytes arrive, so the prescan hit + escalation redact it + // before delivery; longer values carry the disclosed bounded-fragment risk. + defaultTailWindowBytes = 8 * 1024 + // defaultMaxBufferBytes caps total held bytes so a long non-content (e.g. + // reasoning) phase that holds many units cannot grow memory without bound. A + // substrate's own Escalate drain cap SHOULD agree with this so the main-loop + // ceiling and the buffer-drain ceiling do not diverge. + defaultMaxBufferBytes = 8 * 1024 * 1024 + // defaultPrescanBatchBytes is the byte threshold of unscanned redactable content + // that must accumulate before the engine triggers a (cgo) prescan. Deferring the + // prescan to a batch collapses the per-unit scan count ~10-25x on small-unit + // streams (each unit's redactable delta is tens of bytes) — the prescan is already + // O(window) bounded, so the cost was the per-unit FREQUENCY, not the per-scan size. + // 1KB is the cost/latency knee: smaller wastes scans; larger only marginally cuts + // the scan count while widening detection lag. Soundness never depends on the + // trigger firing — see Run's flush-before-deliver + EOF flush. + defaultPrescanBatchBytes = 1024 + // defaultMaxPatternBytes bounds the longest enforceable pattern the flush-before- + // deliver guard assumes. A unit is not released until the prescan has covered it PLUS + // this lookahead, so a sub-window value STARTING in the released unit and completing in + // a later still-held unit is caught even when batching deferred the scan (the unit was + // prescanned ALONE before its completing successor arrived). 4KB covers realistic PII / + // token patterns (cards, SSNs, emails, API keys, JWTs); a rule set with longer CONTIGUOUS + // patterns must raise Config.MaxPatternBytes. The guard is steady-state silent while + // TailWindowBytes > MaxPatternBytes + PrescanBatchBytes (the "comfortably larger than the + // longest pattern" assumption the TailWindowBytes doc already states). + defaultMaxPatternBytes = 4 * 1024 +) + +// DefaultMaxPatternBytes is the package default flush-before-deliver lookahead, exported +// so a substrate that DERIVES MaxPatternBytes from its rule set can floor the derived +// bound at this proven-safe baseline (never shrinking below the default even if the +// derivation under-counts). See Config.MaxPatternBytes. +const DefaultMaxPatternBytes = defaultMaxPatternBytes + +// Config tunes the tail window and the held-bytes ceiling. Zero fields take the +// package defaults. +type Config struct { + // TailWindowBytes is the trailing redactable-content budget held undelivered. + // It MUST be set comfortably larger than the longest enforceable pattern so a + // value straddling unit boundaries is always fully inside the re-scanned window + // when its completing bytes arrive. + TailWindowBytes int + // MaxBufferBytes caps total held bytes (across all channels, content or not). + MaxBufferBytes int + // PrescanBatchBytes defers the cheap union prescan until this many NEW redactable + // bytes have accumulated since the last scan, collapsing the per-unit cgo scan count + // on small-unit streams. It is a perf knob: a confirmed hit is caught at the next + // batch trigger, the flush-before-deliver guard (gated by MaxPatternBytes), or the EOF + // flush — never delivering unscanned content. Zero/-ve takes the package default; set 1 + // for per-unit scanning. + PrescanBatchBytes int + // MaxPatternBytes is the longest CONTIGUOUS enforceable pattern (value) the rules can + // match. The flush-before-deliver guard ensures the prescan has covered each released + // unit PLUS this lookahead, so a sub-window value that STARTS inside the released unit + // and completes in a later still-held unit is detected before the start unit is + // delivered raw — the boundary case batching would otherwise leak when a large unit is + // prescanned alone and released before its completing successor is scanned with it. MUST + // be ≥ the longest contiguous enforceable pattern; over-setting only costs extra flush + // scans on large-unit streams. Zero/-ve takes the package default; clamped below + // TailWindowBytes. Perf note: a window below ~MaxPatternBytes+PrescanBatchBytes + // degrades to a scan on (nearly) every release — safe, but loses the batching win; keep + // TailWindowBytes comfortably above MaxPatternBytes+PrescanBatchBytes+maxUnitSize. + MaxPatternBytes int +} + +func (c Config) withDefaults() Config { + if c.TailWindowBytes <= 0 { + c.TailWindowBytes = defaultTailWindowBytes + } + if c.MaxBufferBytes <= 0 { + c.MaxBufferBytes = defaultMaxBufferBytes + } + if c.PrescanBatchBytes <= 0 { + c.PrescanBatchBytes = defaultPrescanBatchBytes + } + if c.MaxPatternBytes <= 0 { + c.MaxPatternBytes = defaultMaxPatternBytes + } + // Clamp below the window: the lookahead can never need to reach past the held window + // (a value wider than the window is the disclosed best-effort surface, not a guarantee). + if c.MaxPatternBytes >= c.TailWindowBytes { + c.MaxPatternBytes = c.TailWindowBytes - 1 + } + return c +} + +// Substrate adapts a concrete streaming transport to the engine. U is the +// substrate's unit type — a canonical provider chunk for the gateway relay, a raw +// SSE frame for the tlsbump wire path. The engine never inspects U; it only routes +// units through these operations. +// +// Method contract: +// - Next yields the next unit. io.EOF ends the stream normally; any other error +// is terminal — the engine calls OnError and stops. A substrate that records +// usage / finish-reason / per-unit metadata does so inside Next as it produces +// each unit (the engine does not surface units back to the substrate except via +// Deliver and Escalate). +// - AppendRedactableText appends THIS unit's prescan/confirm bytes onto dst and +// returns the grown slice: the assistant content plus tool-call arguments / +// name / id, with a field separator between channels so a pattern cannot span +// two unrelated fields (over-matching the prefilter only wastes a confirm; it +// never misses a match). Reasoning is omitted to match the canonical +// redaction's coverage. Appending onto the engine's buffer (rather than +// returning a fresh []byte) keeps the hot miss path allocation-free. +// - UnitBytes returns the unit's total transport size, used for the MaxBufferBytes +// held-bytes ceiling. +// - ContentBytes returns the unit's redactable-content size WITHOUT field +// separators. It MUST measure the SAME content AppendRedactableText emits (a +// non-empty append iff ContentBytes>0; never over-report) — the engine uses it +// symmetrically to admit and evict the unit from the tail window so reasoning / +// non-content units do not evict redactable content from the window early, and +// an over-report would evict the tail prematurely. +// - IsDone reports whether this unit is the stream terminator. +// - Deliver forwards one held/released unit's BODY to the client in real time. It +// MUST emit only the unit body — never the stream terminator framing, even when +// IsDone(u) is true; the terminal frame is DeliverTerminal's job. A delivery +// error is terminal; the substrate records its own terminal classification and +// the engine stops, returning the error. +// - DeliverTerminal emits the stream's terminal framing (finish reason, usage, and +// any wire sentinel) exactly once after the last unit is delivered. +// - Prescan is the cheap union prefilter over accumulated redactable content. It +// MUST be sound: return false only when no bound hook can possibly match the +// bytes. A substrate whose prefilter is unavailable returns true ("always +// confirm") so the gate never silently skips enforcement. +// - Confirm runs ONE full compliance evaluation over the accumulated content and +// returns the aggregate result. nil is treated as approve (resume streaming); a +// substrate that CANNOT evaluate (build/exec error) MUST return a RejectHard +// result, never nil, so a missing enforcer never streams unredacted. +// - Escalate takes over on a confirmed redact/block hit (res carries the +// triggering decision) OR a memory-pressure eviction of an incomplete content +// unit (res is nil). It owns draining the remainder, RE-EVALUATING the buffered +// remainder on its own waist (it does NOT trust res's spans — res is advisory / +// audit context only), delivering ONLY the redacted remainder (the +// already-delivered prefix is never re-delivered; the held tail is delivered +// redacted, never raw), and the terminal framing. Its fail posture — fail-open +// relay-original vs fail-closed error-frame — is the substrate's, not the +// engine's. After Escalate returns the engine stops. +// - OnConfirmApproved stamps the audit outcome of a confirm that did NOT enforce (a +// prescan false positive, or a nil/approve result) so a SIEM sees "hooks +// evaluated" plus any tags, distinguishing it from "no response hook +// configured". res may be nil (nil-as-approve). +// - OnApproveEOF stamps an approve outcome when the stream reached EOF without any +// confirm ever running — the sound prescan cleared the whole stream, so the +// response is provably approved (vs "no hook configured"). +// - OnError emits the substrate's terminal error framing for a non-EOF Next error +// and returns the error the engine propagates. +type Substrate[U any] interface { + Next(ctx context.Context) (U, error) + AppendRedactableText(dst []byte, u U) []byte + UnitBytes(u U) int + ContentBytes(u U) int + IsDone(u U) bool + Deliver(ctx context.Context, u U) error + DeliverTerminal(ctx context.Context) error + Prescan(content []byte) bool + Confirm(ctx context.Context, content string) *hookcore.CompliancePipelineResult + Escalate(ctx context.Context, held []U, res *hookcore.CompliancePipelineResult) error + OnConfirmApproved(res *hookcore.CompliancePipelineResult) + OnApproveEOF() + OnError(ctx context.Context, err error) error +} + +// Run drives the prescan-gated real-time streaming algorithm over sub until EOF, +// an escalation, or a terminal error. It returns nil on a clean stream, the +// substrate's OnError result on a non-EOF Next error, a Deliver error if a real-time +// write fails, or the Escalate result on a confirmed redact/block hit or a +// memory-pressure escalation. +func Run[U any](ctx context.Context, sub Substrate[U], cfg Config) error { + cfg = cfg.withDefaults() + + var ( + held []U // trailing units NOT yet delivered to the client + heldScanLens []int // per-held-unit scanBuf byte length (parallel to held) + heldBytes int // total transport bytes currently held (MaxBufferBytes ceiling) + contentBytes int // redactable-content bytes currently held (tail window) + scanBuf []byte // accumulated redactable content the prescan/confirm read + deliveredScanLen int // scanBuf offset where the oldest still-held unit begins + scannedLen int // high-water scanBuf offset already prescanned (hit or miss) + anyConfirm bool // whether any confirm ran (audit: approve vs no-hook) + ) + + // scanThrough advances the windowed prescan to the current end of scanBuf and, on a + // hit, pays for ONE full confirm. It returns a non-nil result ONLY when the caller + // must escalate (a confirmed block/redact action); a miss or a false-positive approve + // is handled in place. It is the SINGLE prescan primitive — invoked by the batch + // trigger AND before every release / at EOF — so the invariant "scanned-through ≥ + // delivered" holds at every release point and batching never delivers unscanned + // content. The prescan reads the still-HELD window [deliveredScanLen:] (already + // bounded to ~TailWindowBytes by releases, so each scan is O(window)); scannedLen + // high-waters to len(scanBuf) so a repeated call with no new content is a no-op. + scanThrough := func() *hookcore.CompliancePipelineResult { + if len(scanBuf) <= scannedLen { + return nil + } + scannedLen = len(scanBuf) + if !sub.Prescan(scanBuf[deliveredScanLen:]) { + return nil + } + // Confirm over the still-HELD window only ([deliveredScanLen:]), not the whole + // accumulated buffer: Escalate re-evaluates the held remainder anyway, and a + // pattern entirely in the already-delivered (un-redactable) prefix could only fire + // a redundant escalate. Windowing keeps the confirm O(window). + res := sub.Confirm(ctx, string(scanBuf[deliveredScanLen:])) + // A confirm executed (even a nil/approve result) — the response was evaluated, + // which the EOF stamp must distinguish from "no hook". + anyConfirm = true + // Gate on the enforcing ACTION, not the raw decision enum: the pipeline aggregator + // ranks a co-firing soft-block ABOVE a redact, so a Modify (redact) carrying real + // spans can surface as BlockSoft. ActionFromDecision maps RejectHard/BlockSoft→block + // and Modify→redact, so keying on the action escalates every enforcing outcome + // instead of leaking a redact masked by a soft-block. + if res != nil { + act := hookcore.ActionFromDecision(res.Decision) + if act == hookcore.ActionBlock || act == hookcore.ActionRedact { + return res + } + } + // Prescan false positive / nil-as-approve: record the outcome, resume streaming. + sub.OnConfirmApproved(res) + return nil + } + + for { + u, err := sub.Next(ctx) + if err != nil { + if errors.Is(err, io.EOF) { + break + } + return sub.OnError(ctx, err) + } + + // Accumulate the unit's redactable content (with field separators) for the + // prescan, then HOLD the unit. Track the per-unit scanBuf length so a release + // can advance deliveredScanLen, and the tail-window content accounting uses + // the separator-free ContentBytes so the scan's separators do not perturb the + // eviction math. + before := len(scanBuf) + scanBuf = sub.AppendRedactableText(scanBuf, u) + held = append(held, u) + heldScanLens = append(heldScanLens, len(scanBuf)-before) + heldBytes += sub.UnitBytes(u) + contentBytes += sub.ContentBytes(u) + + // Prescan GATE (BATCHED): defer the cheap union prescan until PrescanBatchBytes of + // NEW redactable content have accumulated since the last scan, collapsing the + // per-unit cgo scan count on small-unit streams. Soundness does NOT rely on this + // trigger firing: the flush-before-deliver guard in the release loop and the EOF + // flush force a scanThrough over any held content before it is delivered, so a + // sub-threshold tail is never delivered unscanned. A confirmed enforcing hit hands + // off to buffer-to-end redaction (the held unit completing the match is still + // undelivered because the scan precedes its release). + if len(scanBuf)-scannedLen >= cfg.PrescanBatchBytes { + if res := scanThrough(); res != nil { + return sub.Escalate(ctx, held, res) + } + } + + // Release the tail. Two eviction triggers: + // - contentBytes > TailWindowBytes: the oldest held content is now older than + // the window, so any sub-window value it carries is complete — safe to + // deliver raw (the normal real-time path). + // - heldBytes > MaxBufferBytes: a memory cap. Evicting a NON-content unit + // (reasoning, ContentBytes==0) raw is safe (reasoning is never redacted). + // But evicting a CONTENT unit raw under memory pressure alone — when the + // window has NOT filled with content (it filled with reasoning) — could + // leak a still-incomplete sub-window value, so that case ESCALATES instead. + for len(held) > 0 { + overWindow := contentBytes > cfg.TailWindowBytes + overBuf := heldBytes > cfg.MaxBufferBytes + if !overWindow && !overBuf { + break + } + front := held[0] + if overBuf && !overWindow && sub.ContentBytes(front) > 0 { + // Memory-pressure eviction of a content-bearing unit whose value may be + // incomplete: escalate (buffer + redact the remainder) rather than leak. + return sub.Escalate(ctx, held, nil) + } + // FLUSH-BEFORE-DELIVER (with a MaxPatternBytes lookahead): release front only once + // the prescan has covered front's end PLUS MaxPatternBytes. frontScanEnd is front's + // absolute scanBuf end offset (deliveredScanLen is where front begins, heldScanLens[0] + // its length). The lookahead is load-bearing: a sub-window value that STARTS inside + // front and completes in a LATER still-held unit must be scanned before front is + // delivered raw — and batching may have scanned front ALONE (a large unit batch- + // triggered before its completing successor arrived), so requiring only + // scannedLen ≥ frontScanEnd would deliver front's pattern-prefix raw and then scan + // the suffix over a window that no longer holds the prefix (a boundary leak). With + // the lookahead, any value of length ≤ MaxPatternBytes straddling front's boundary is + // caught. Steady-state silent: the scan leads front by ~(TailWindow-PrescanBatch) ≫ + // MaxPatternBytes, so this only fires for a large single unit (or a tiny window). A + // flush hit escalates instead of delivering raw; scanThrough scans the full held + // window [deliveredScanLen:]. + if frontScanEnd := deliveredScanLen + heldScanLens[0]; scannedLen < len(scanBuf) && scannedLen < frontScanEnd+cfg.MaxPatternBytes { + if res := scanThrough(); res != nil { + return sub.Escalate(ctx, held, res) + } + } + if derr := sub.Deliver(ctx, front); derr != nil { + return derr + } + heldBytes -= sub.UnitBytes(front) + contentBytes -= sub.ContentBytes(front) + deliveredScanLen += heldScanLens[0] + var zero U + held[0] = zero // release the delivered unit's payload to the GC immediately + held = held[1:] + heldScanLens = heldScanLens[1:] + } + + // Compact scanBuf: drop the already-delivered prefix so memory is bounded at + // O(window) instead of O(N) for the whole stream (the prefix grows unbounded even + // on a clean, prescan-miss stream). Threshold-gated on `> TailWindowBytes` so the + // move is amortized O(1)/byte — never a per-unit memmove on the hot relay path. + // Delivered bytes are past the window and can no longer be redacted, so dropping + // them loses nothing the windowed prescan/confirm didn't already forgo. Rebase the + // absolute offsets by the dropped length; heldScanLens are RELATIVE (untouched). + // scannedLen is clamped at 0: if the scanned prefix lay entirely within the dropped + // region it becomes 0 (re-scan the held window on the next trigger — the safe + // direction). Because we only deliver content that has been scanned (flush-before- + // deliver), scannedLen ≥ deliveredScanLen here, so the rebase stays non-negative. + if deliveredScanLen > cfg.TailWindowBytes { + drop := deliveredScanLen + n := copy(scanBuf, scanBuf[drop:]) + scanBuf = scanBuf[:n] + deliveredScanLen = 0 + // Defensive floor: scannedLen ≥ drop in practice (we only deliver — advancing + // deliveredScanLen — content that flush-before-deliver already scanned), so this + // never actually clamps. Even if it did, a smaller scannedLen only forces extra + // (sound) scans, never a leak; max keeps the rebase robust without an untested branch. + scannedLen = max(0, scannedLen-drop) + } + + if sub.IsDone(u) { + break + } + } + + // EOF FLUSH: scan any unscanned tail (the final < PrescanBatchBytes that never hit + // the batch trigger) BEFORE delivering the held tail raw — else a sub-threshold final + // content unit would be delivered unscanned. A hit escalates (redact/block); only a + // miss / approve-cleared tail is delivered raw below. + if res := scanThrough(); res != nil { + return sub.Escalate(ctx, held, res) + } + + // EOF without escalation → every prescan was a miss or a false-positive approve, + // so the held tail is provably benign. When NO confirm ever ran, the sound + // prescan cleared the whole stream, so stamp the response as approved (vs "no + // response hook configured"). A confirm that ran already stamped its own outcome. + if !anyConfirm { + sub.OnApproveEOF() + } + for _, u := range held { + if derr := sub.Deliver(ctx, u); derr != nil { + return derr + } + } + return sub.DeliverTerminal(ctx) +} diff --git a/packages/shared/transport/streaming/modela/engine_test.go b/packages/shared/transport/streaming/modela/engine_test.go new file mode 100644 index 00000000..51fc78f6 --- /dev/null +++ b/packages/shared/transport/streaming/modela/engine_test.go @@ -0,0 +1,919 @@ +package modela + +import ( + "context" + "errors" + "io" + "runtime" + "strings" + "testing" + + hookcore "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// fakeUnit is the engine's opaque unit for tests: it carries the redactable text +// the prescan/confirm see, the content/transport byte sizes the tail-window math +// uses, and a done flag. +type fakeUnit struct { + id int + text string + content int + bytes int + done bool +} + +// fakeSubstrate is a fully-recording Substrate[fakeUnit]. It feeds a fixed unit +// queue, lets a test inject prescan/confirm verdicts and delivery/next errors, and +// records every engine-driven operation so a test can assert the observable +// algorithm behaviour (delivery order, escalation hand-off, audit stamps). +type fakeSubstrate struct { + units []fakeUnit + idx int + + // nextErr, when set, is returned by Next once idx reaches nextErrAt (instead of + // the queued unit / EOF) to exercise the terminal-error path. + nextErr error + nextErrAt int + + prescanFn func(content []byte) bool + confirmFn func(content string) *hookcore.CompliancePipelineResult + deliverErrAt int // unit id whose Deliver returns deliverErr; -1 disables + deliverErr error + + // recorded observations + delivered []int + terminalCalled int + approveEOF int + confirmApproved int + escalated bool + escalateHeld []int + escalateRes *hookcore.CompliancePipelineResult + onErrorCalled int + onErrorErr error + confirmCalls []string + prescanCalls []string +} + +func newFake(units []fakeUnit) *fakeSubstrate { + return &fakeSubstrate{units: units, nextErrAt: -1, deliverErrAt: -1} +} + +func (f *fakeSubstrate) Next(_ context.Context) (fakeUnit, error) { + if f.nextErr != nil && f.idx == f.nextErrAt { + return fakeUnit{}, f.nextErr + } + if f.idx >= len(f.units) { + return fakeUnit{}, io.EOF + } + u := f.units[f.idx] + f.idx++ + return u, nil +} + +func (f *fakeSubstrate) AppendRedactableText(dst []byte, u fakeUnit) []byte { + return append(dst, u.text...) +} +func (f *fakeSubstrate) UnitBytes(u fakeUnit) int { return u.bytes } +func (f *fakeSubstrate) ContentBytes(u fakeUnit) int { return u.content } +func (f *fakeSubstrate) IsDone(u fakeUnit) bool { return u.done } + +func (f *fakeSubstrate) Deliver(_ context.Context, u fakeUnit) error { + if f.deliverErr != nil && u.id == f.deliverErrAt { + return f.deliverErr + } + f.delivered = append(f.delivered, u.id) + return nil +} + +func (f *fakeSubstrate) DeliverTerminal(_ context.Context) error { f.terminalCalled++; return nil } + +func (f *fakeSubstrate) Prescan(content []byte) bool { + f.prescanCalls = append(f.prescanCalls, string(content)) + if f.prescanFn == nil { + return false + } + return f.prescanFn(content) +} + +func (f *fakeSubstrate) Confirm(_ context.Context, content string) *hookcore.CompliancePipelineResult { + f.confirmCalls = append(f.confirmCalls, content) + if f.confirmFn == nil { + return &hookcore.CompliancePipelineResult{Decision: hookcore.Approve} + } + return f.confirmFn(content) +} + +func (f *fakeSubstrate) Escalate(_ context.Context, held []fakeUnit, res *hookcore.CompliancePipelineResult) error { + f.escalated = true + for _, u := range held { + f.escalateHeld = append(f.escalateHeld, u.id) + } + f.escalateRes = res + return nil +} + +func (f *fakeSubstrate) OnConfirmApproved(_ *hookcore.CompliancePipelineResult) { f.confirmApproved++ } +func (f *fakeSubstrate) OnApproveEOF() { f.approveEOF++ } +func (f *fakeSubstrate) OnError(_ context.Context, err error) error { + f.onErrorCalled++ + f.onErrorErr = err + return err +} + +func reject() *hookcore.CompliancePipelineResult { + return &hookcore.CompliancePipelineResult{Decision: hookcore.RejectHard} +} +func modify() *hookcore.CompliancePipelineResult { + return &hookcore.CompliancePipelineResult{Decision: hookcore.Modify} +} +func approve() *hookcore.CompliancePipelineResult { + return &hookcore.CompliancePipelineResult{Decision: hookcore.Approve} +} + +func eq(a, b []int) bool { + if len(a) != len(b) { + return false + } + for i := range a { + if a[i] != b[i] { + return false + } + } + return true +} + +// TestNoPrescanHit_AllDeliveredApproveAtEOF: with a prescan that never matches and +// content under the tail window, every unit is held until EOF, then flushed in +// order; no confirm runs, so the engine stamps OnApproveEOF and emits the terminal. +func TestNoPrescanHit_AllDeliveredApproveAtEOF(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "a", content: 1, bytes: 1}, + {id: 2, text: "b", content: 1, bytes: 1}, + {id: 3, text: "c", content: 1, bytes: 1, done: true}, + }) + if err := Run(context.Background(), f, Config{}); err != nil { + t.Fatalf("Run returned error: %v", err) + } + if !eq(f.delivered, []int{1, 2, 3}) { + t.Fatalf("delivered order = %v, want [1 2 3]", f.delivered) + } + if f.escalated { + t.Fatal("escalated unexpectedly") + } + if len(f.confirmCalls) != 0 { + t.Fatalf("confirm ran %d times, want 0", len(f.confirmCalls)) + } + if f.approveEOF != 1 { + t.Fatalf("OnApproveEOF called %d times, want 1", f.approveEOF) + } + if f.terminalCalled != 1 { + t.Fatalf("DeliverTerminal called %d times, want 1", f.terminalCalled) + } +} + +// TestEOFWithoutDoneMarker: a stream that ends by Next returning io.EOF (no +// explicit done unit — some wires omit a terminator) still flushes the held tail, +// stamps OnApproveEOF, and emits the terminal frame. +func TestEOFWithoutDoneMarker(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "a", content: 1, bytes: 1}, + {id: 2, text: "b", content: 1, bytes: 1}, // no done flag → loop exits via io.EOF + }) + if err := Run(context.Background(), f, Config{}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !eq(f.delivered, []int{1, 2}) { + t.Fatalf("delivered = %v, want [1 2]", f.delivered) + } + if f.approveEOF != 1 || f.terminalCalled != 1 { + t.Fatalf("approveEOF=%d terminal=%d, want 1/1", f.approveEOF, f.terminalCalled) + } +} + +// TestPrescanFalsePositive_ResumesStreaming: prescan hits but confirm approves, so +// the engine stamps OnConfirmApproved and resumes real-time streaming — no +// escalation, and because a confirm DID run, OnApproveEOF is NOT stamped at EOF. +func TestPrescanFalsePositive_ResumesStreaming(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "secret", content: 6, bytes: 6}, + {id: 2, text: "more", content: 4, bytes: 4, done: true}, + }) + f.prescanFn = func([]byte) bool { return true } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return approve() } + if err := Run(context.Background(), f, Config{TailWindowBytes: 1024, MaxBufferBytes: 1 << 20}); err != nil { + t.Fatalf("Run error: %v", err) + } + if f.escalated { + t.Fatal("escalated on an approve confirm") + } + if f.confirmApproved == 0 { + t.Fatal("OnConfirmApproved never stamped on a false positive") + } + if f.approveEOF != 0 { + t.Fatalf("OnApproveEOF stamped %d times after a confirm ran, want 0", f.approveEOF) + } + if !eq(f.delivered, []int{1, 2}) { + t.Fatalf("delivered = %v, want [1 2]", f.delivered) + } +} + +// TestConfirmedReject_Escalates: a prescan hit confirmed as RejectHard hands the +// held units to Escalate and stops the engine — the held tail is NOT delivered raw. +func TestConfirmedReject_Escalates(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "card 4111", content: 9, bytes: 9}, + {id: 2, text: "1111", content: 4, bytes: 4, done: true}, + }) + f.prescanFn = func([]byte) bool { return true } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return reject() } + // PrescanBatchBytes:1 exercises per-unit scan timing — the confirm fires as soon as + // the first unit's content is scanned, so only unit 1 is held at escalation. + if err := Run(context.Background(), f, Config{TailWindowBytes: 1024, MaxBufferBytes: 1 << 20, PrescanBatchBytes: 1}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !f.escalated { + t.Fatal("did not escalate on a confirmed reject") + } + if !eq(f.escalateHeld, []int{1}) { + t.Fatalf("escalate held = %v, want [1] (the single held unit)", f.escalateHeld) + } + if len(f.delivered) != 0 { + t.Fatalf("delivered %v before escalation, want none (held under window)", f.delivered) + } + if f.escalateRes == nil || f.escalateRes.Decision != hookcore.RejectHard { + t.Fatalf("escalate result = %+v, want RejectHard", f.escalateRes) + } + if f.terminalCalled != 0 { + t.Fatal("DeliverTerminal ran on the engine path after escalation handed off") + } +} + +// TestConfirmedModify_Escalates: a Modify decision escalates exactly like a block — +// a redact that cannot be applied losslessly on the live wire goes to buffer-to-end. +func TestConfirmedModify_Escalates(t *testing.T) { + f := newFake([]fakeUnit{{id: 1, text: "ssn 078051120", content: 13, bytes: 13, done: true}}) + f.prescanFn = func([]byte) bool { return true } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return modify() } + if err := Run(context.Background(), f, Config{}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !f.escalated || f.escalateRes.Decision != hookcore.Modify { + t.Fatalf("modify did not escalate: escalated=%v res=%+v", f.escalated, f.escalateRes) + } +} + +// TestTailWindowReleasesPrefix: with a small tail window and a never-matching +// prescan, units beyond the window are delivered in real time (only the trailing +// window stays held until EOF) — proving the bounded-tail real-time behaviour. +func TestTailWindowReleasesPrefix(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "aaa", content: 3, bytes: 3}, + {id: 2, text: "bbb", content: 3, bytes: 3}, + {id: 3, text: "ccc", content: 3, bytes: 3, done: true}, + }) + // window = 3 bytes: after unit 2 the held content (6) exceeds 3, so unit 1 is + // released; after unit 3 the held content (now units 2+3 = 6) exceeds 3 so unit 2 + // releases. The final held tail (unit 3) flushes at EOF. + if err := Run(context.Background(), f, Config{TailWindowBytes: 3, MaxBufferBytes: 1 << 20}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !eq(f.delivered, []int{1, 2, 3}) { + t.Fatalf("delivered = %v, want [1 2 3] in order", f.delivered) + } +} + +// TestMaxBufferReleasesOnByteCeiling: a non-content-heavy unit (content 0 but large +// bytes) does not trip the content window, but the held-bytes ceiling forces a +// release so memory stays bounded during a long reasoning phase. +func TestMaxBufferReleasesOnByteCeiling(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "", content: 0, bytes: 100}, + {id: 2, text: "", content: 0, bytes: 100, done: true}, + }) + if err := Run(context.Background(), f, Config{TailWindowBytes: 1 << 20, MaxBufferBytes: 150}); err != nil { + t.Fatalf("Run error: %v", err) + } + // After unit 2, heldBytes (200) > 150 → release unit 1; unit 2 flushes at EOF. + if !eq(f.delivered, []int{1, 2}) { + t.Fatalf("delivered = %v, want [1 2]", f.delivered) + } +} + +// TestWindowedConfirmDedup: a persistent prescan hit confirms only on NEW content +// since the last confirm (confirmedLen gate), not once per unit — so a benign +// false-positive trigger is not re-confirmed until fresh bytes arrive. +func TestWindowedConfirmDedup(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "x", content: 1, bytes: 1}, + {id: 2, text: "", content: 0, bytes: 0}, // no new scan content + {id: 3, text: "y", content: 1, bytes: 1, done: true}, + }) + f.prescanFn = func([]byte) bool { return true } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return approve() } + // PrescanBatchBytes:1 exercises per-unit scan timing so the dedup is observable: + // unit 2 adds no scan bytes, so scanThrough is a no-op there. + if err := Run(context.Background(), f, Config{TailWindowBytes: 1024, MaxBufferBytes: 1 << 20, PrescanBatchBytes: 1}); err != nil { + t.Fatalf("Run error: %v", err) + } + // Unit 2 added no scan bytes, so len(scanBuf) did not exceed scannedLen → no + // second confirm. Confirms fire only on units 1 and 3. + if len(f.confirmCalls) != 2 { + t.Fatalf("confirm ran %d times, want 2 (only on new content)", len(f.confirmCalls)) + } +} + +// TestNilConfirmTreatedAsApprove: a nil confirm result is treated as approve — no +// escalation. But a confirm DID execute, so the audit records it as an evaluated +// approve (OnConfirmApproved), NOT as "no hook ran" (OnApproveEOF), preserving the +// SIEM distinction between "evaluated-approved" and "no response hook configured". +func TestNilConfirmTreatedAsApprove(t *testing.T) { + f := newFake([]fakeUnit{{id: 1, text: "z", content: 1, bytes: 1, done: true}}) + f.prescanFn = func([]byte) bool { return true } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return nil } + if err := Run(context.Background(), f, Config{}); err != nil { + t.Fatalf("Run error: %v", err) + } + if f.escalated { + t.Fatal("nil confirm result must not escalate") + } + if f.confirmApproved != 1 { + t.Fatalf("OnConfirmApproved = %d on a nil-as-approve result, want 1 (a confirm ran)", f.confirmApproved) + } + // A confirm ran → anyConfirm true → EOF must NOT re-stamp as no-hook approve. + if f.approveEOF != 0 { + t.Fatalf("OnApproveEOF = %d after a confirm ran, want 0", f.approveEOF) + } +} + +// TestBlockSoftMaskedRedact_Escalates pins the security gate fix: the pipeline +// aggregator ranks a co-firing soft-block ABOVE a redact, so a confirm that found +// PII and computed redact spans can surface as a BlockSoft decision. Keying the +// escalate gate on the enforcing ACTION (BlockSoft→block) — not the raw decision +// enum — escalates it instead of streaming the masked redact raw. +func TestBlockSoftMaskedRedact_Escalates(t *testing.T) { + f := newFake([]fakeUnit{{id: 1, text: "pii here", content: 8, bytes: 8, done: true}}) + f.prescanFn = func([]byte) bool { return true } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { + // A redact hook's spans masked behind a soft-block aggregated decision. + return &hookcore.CompliancePipelineResult{Decision: hookcore.BlockSoft} + } + if err := Run(context.Background(), f, Config{}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !f.escalated { + t.Fatal("BlockSoft (block action) was not escalated — PII would stream raw") + } + if len(f.delivered) != 0 { + t.Fatalf("delivered %v before escalation on a block-action confirm, want none", f.delivered) + } +} + +// TestMaxBufferContentEviction_Escalates pins the second leak fix: a content unit +// whose value is not yet complete must NOT be force-delivered raw under memory +// pressure. Here a small content unit is followed by reasoning (content==0, large +// bytes) that trips the MaxBufferBytes ceiling while the content window is NOT +// exceeded — the engine escalates rather than flushing the content unit raw. +func TestMaxBufferContentEviction_Escalates(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "SSN 078-05-", content: 11, bytes: 11}, // incomplete sensitive value + {id: 2, text: "", content: 0, bytes: 500}, // reasoning: trips the byte ceiling + }) + f.prescanFn = func([]byte) bool { return false } // isolate the memory-pressure path + if err := Run(context.Background(), f, Config{TailWindowBytes: 1 << 20, MaxBufferBytes: 150}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !f.escalated { + t.Fatal("memory-pressure eviction of an incomplete content unit did not escalate — sub-window leak") + } + if !eq(f.escalateHeld, []int{1, 2}) { + t.Fatalf("escalate held = %v, want [1 2]", f.escalateHeld) + } + if f.escalateRes != nil { + t.Fatal("memory-pressure escalation must pass a nil triggering result") + } + if len(f.delivered) != 0 { + t.Fatalf("delivered %v under memory pressure, want none (content held)", f.delivered) + } +} + +// TestReasoningEvictionUnderMemoryPressure: a non-content (reasoning) unit at the +// front IS delivered raw under the byte ceiling — reasoning is never redacted, so +// evicting it raw is safe and keeps memory bounded without forcing escalation. +func TestReasoningEvictionUnderMemoryPressure(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "", content: 0, bytes: 100}, // reasoning at front + {id: 2, text: "", content: 0, bytes: 100, done: true}, + }) + f.prescanFn = func([]byte) bool { return false } + if err := Run(context.Background(), f, Config{TailWindowBytes: 1 << 20, MaxBufferBytes: 150}); err != nil { + t.Fatalf("Run error: %v", err) + } + if f.escalated { + t.Fatal("reasoning eviction must not escalate") + } + if !eq(f.delivered, []int{1, 2}) { + t.Fatalf("delivered = %v, want [1 2] (reasoning evicted raw)", f.delivered) + } +} + +// TestPrescanWindowedToHeldContent locks the O(N) windowed-prescan contract: after a +// unit is released past the tail window, the next prescan sees ONLY the still-held +// content (from the released boundary), not the whole accumulated buffer — proving +// the lookbehind anchors on the oldest held unit, not a growing full-buffer scan. +func TestPrescanWindowedToHeldContent(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "aaa", content: 3, bytes: 3}, + {id: 2, text: "bbb", content: 3, bytes: 3}, + {id: 3, text: "ccc", content: 3, bytes: 3, done: true}, + }) + f.prescanFn = func([]byte) bool { return false } + if err := Run(context.Background(), f, Config{TailWindowBytes: 3, MaxBufferBytes: 1 << 20, PrescanBatchBytes: 1}); err != nil { + t.Fatalf("Run error: %v", err) + } + // Prescans: u1 sees "aaa"; u2 sees "aaabbb"; after u1 releases (deliveredScanLen=3), + // u3 sees only "bbbccc" — NOT the whole "aaabbbccc". + if len(f.prescanCalls) != 3 { + t.Fatalf("prescan ran %d times, want 3: %v", len(f.prescanCalls), f.prescanCalls) + } + if f.prescanCalls[2] != "bbbccc" { + t.Fatalf("third prescan saw %q, want windowed %q (delivered prefix excluded)", f.prescanCalls[2], "bbbccc") + } +} + +// TestSeparatorAccounting_TailUsesContentNotTextLen pins that the tail window is +// sized by separator-free ContentBytes, NOT by appended scan-text length. Three +// units of content 2 / scan-text 3 (a "\n" channel separator) under window 5: held +// content after two units is 4 ≤ 5 (nothing released), but if the window were +// (wrongly) sized by scan-text length it would be 6 > 5 and evict unit 1 early. The +// third unit triggers an escalation; asserting the escalation still holds [1 2 3] +// (and nothing was delivered) discriminates the two accountings — a text-length +// window would have released unit 1 and the escalation would see only [2 3]. +func TestSeparatorAccounting_TailUsesContentNotTextLen(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "ab\n", content: 2, bytes: 3}, + {id: 2, text: "cd\n", content: 2, bytes: 3}, + {id: 3, text: "ef\n", content: 2, bytes: 3, done: true}, + }) + // Prescan misses until the third unit's content arrives, then confirm blocks. + f.prescanFn = func(b []byte) bool { return strings.Contains(string(b), "ef") } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return reject() } + // PrescanBatchBytes:1 scans per unit so this pins the content-vs-text-length window + // accounting (not batch timing): unit 1 stays held under the content-sized window. + if err := Run(context.Background(), f, Config{TailWindowBytes: 5, MaxBufferBytes: 1 << 20, PrescanBatchBytes: 1}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !f.escalated { + t.Fatal("third-unit confirm did not escalate") + } + if !eq(f.escalateHeld, []int{1, 2, 3}) { + t.Fatalf("escalate held = %v, want [1 2 3] — a text-length window would have evicted unit 1", f.escalateHeld) + } + if len(f.delivered) != 0 { + t.Fatalf("delivered %v, want none (all held under content-sized window)", f.delivered) + } +} + +// TestNextError_OnError: a non-EOF Next error routes through OnError and the engine +// returns that error without delivering a terminal frame. +func TestNextError_OnError(t *testing.T) { + sentinel := errors.New("upstream exploded") + f := newFake([]fakeUnit{{id: 1, text: "a", content: 1, bytes: 1}}) + f.nextErr = sentinel + f.nextErrAt = 1 // first unit succeeds, second Next errors + err := Run(context.Background(), f, Config{}) + if !errors.Is(err, sentinel) { + t.Fatalf("Run error = %v, want %v", err, sentinel) + } + if f.onErrorCalled != 1 || !errors.Is(f.onErrorErr, sentinel) { + t.Fatalf("OnError calls=%d err=%v", f.onErrorCalled, f.onErrorErr) + } + if f.terminalCalled != 0 { + t.Fatal("DeliverTerminal must not run after a terminal Next error") + } +} + +// TestDeliverError_StopsDuringRelease: a Deliver error during tail release stops the +// engine and surfaces the error (the substrate has recorded its own terminal state). +func TestDeliverError_StopsDuringRelease(t *testing.T) { + derr := errors.New("client gone") + f := newFake([]fakeUnit{ + {id: 1, text: "aaa", content: 3, bytes: 3}, + {id: 2, text: "bbb", content: 3, bytes: 3}, + {id: 3, text: "ccc", content: 3, bytes: 3, done: true}, + }) + f.deliverErr = derr + f.deliverErrAt = 1 // unit 1 is the first released by the small window + err := Run(context.Background(), f, Config{TailWindowBytes: 3, MaxBufferBytes: 1 << 20}) + if !errors.Is(err, derr) { + t.Fatalf("Run error = %v, want %v", err, derr) + } + if f.terminalCalled != 0 { + t.Fatal("DeliverTerminal must not run after a delivery error") + } +} + +// TestDeliverError_StopsAtFinalFlush: a Deliver error while flushing the held tail at +// EOF also stops the engine before the terminal frame. +func TestDeliverError_StopsAtFinalFlush(t *testing.T) { + derr := errors.New("client gone at flush") + f := newFake([]fakeUnit{{id: 7, text: "a", content: 1, bytes: 1, done: true}}) + f.deliverErr = derr + f.deliverErrAt = 7 + err := Run(context.Background(), f, Config{}) + if !errors.Is(err, derr) { + t.Fatalf("Run error = %v, want %v", err, derr) + } + if f.terminalCalled != 0 { + t.Fatal("DeliverTerminal must not run after a final-flush delivery error") + } +} + +// allocSubstrate is a zero-overhead Substrate[int] whose methods neither record nor +// allocate, so testing.AllocsPerRun over Run measures ONLY the engine's own +// allocations (not test-fake bookkeeping). It streams `units` one-byte content units. +type allocSubstrate struct { + units int + idx int +} + +func (a *allocSubstrate) Next(context.Context) (int, error) { + if a.idx >= a.units { + return 0, io.EOF + } + a.idx++ + return a.idx, nil +} +func (a *allocSubstrate) AppendRedactableText(dst []byte, _ int) []byte { return append(dst, 'x') } +func (a *allocSubstrate) UnitBytes(int) int { return 1 } +func (a *allocSubstrate) ContentBytes(int) int { return 1 } +func (a *allocSubstrate) IsDone(int) bool { return false } +func (a *allocSubstrate) Deliver(context.Context, int) error { return nil } +func (a *allocSubstrate) DeliverTerminal(context.Context) error { return nil } +func (a *allocSubstrate) Prescan([]byte) bool { return false } +func (a *allocSubstrate) Confirm(context.Context, string) *hookcore.CompliancePipelineResult { + return nil +} +func (a *allocSubstrate) Escalate(context.Context, []int, *hookcore.CompliancePipelineResult) error { + return nil +} +func (a *allocSubstrate) OnConfirmApproved(*hookcore.CompliancePipelineResult) {} +func (a *allocSubstrate) OnApproveEOF() {} +func (a *allocSubstrate) OnError(_ context.Context, err error) error { return err } + +// TestRunHotPathAllocationsAreAmortized locks the perf guarantee behind the +// append-style AppendRedactableText seam: per-unit work allocates nothing; only the +// engine's three growable buffers (scanBuf / held / heldScanLens) allocate, and only +// on amortized doubling — O(log units), far below one-per-unit. A regression to a +// per-unit fresh-[]byte return (the original RedactableText shape) would push the +// count toward `units`. With prescan false, no confirm runs, so string(scanBuf) is +// never paid here. +func TestRunHotPathAllocationsAreAmortized(t *testing.T) { + const units = 256 + avg := testing.AllocsPerRun(50, func() { + s := &allocSubstrate{units: units} + _ = Run(context.Background(), s, Config{TailWindowBytes: 1 << 20, MaxBufferBytes: 1 << 20}) + }) + if avg > 64 { + t.Fatalf("Run allocated %.0f times for %d units; expected amortized O(log n) growth, not per-unit", avg, units) + } +} + +// TestConfigDefaults: zero config takes the package defaults; explicit values win. +func TestConfigDefaults(t *testing.T) { + got := Config{}.withDefaults() + if got.TailWindowBytes != defaultTailWindowBytes || got.MaxBufferBytes != defaultMaxBufferBytes { + t.Fatalf("defaults = %+v, want tail=%d max=%d", got, defaultTailWindowBytes, defaultMaxBufferBytes) + } + if got.PrescanBatchBytes != defaultPrescanBatchBytes { + t.Fatalf("PrescanBatchBytes default = %d, want %d", got.PrescanBatchBytes, defaultPrescanBatchBytes) + } + if got.MaxPatternBytes != defaultMaxPatternBytes { + t.Fatalf("MaxPatternBytes default = %d, want %d", got.MaxPatternBytes, defaultMaxPatternBytes) + } + custom := Config{TailWindowBytes: 7, MaxBufferBytes: 9}.withDefaults() + if custom.TailWindowBytes != 7 || custom.MaxBufferBytes != 9 { + t.Fatalf("explicit config overwritten: %+v", custom) + } + // MaxPatternBytes is clamped below TailWindowBytes (the lookahead never needs to reach + // past the held window — a wider value is the disclosed over-window surface). + if custom.MaxPatternBytes != custom.TailWindowBytes-1 { + t.Fatalf("MaxPatternBytes clamp = %d, want %d (TailWindowBytes-1)", custom.MaxPatternBytes, custom.TailWindowBytes-1) + } +} + +// TestCompactionWindowedConfirm proves the #12 fix keeps detection sound: with a small +// tail window, delivering several units advances deliveredScanLen past the window so +// scanBuf is compacted (the delivered prefix dropped + offsets rebased). A later prescan +// hit must STILL fire a confirm (the `len(scanBuf) > confirmedLen` gate not wrongly +// skipped after the rebase — the leak class), and the confirm must see only the held +// window (not the compacted-away delivered prefix). +func TestCompactionWindowedConfirm(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 0, text: "aaa", bytes: 3, content: 3}, + {id: 1, text: "bbb", bytes: 3, content: 3}, + {id: 2, text: "ccc", bytes: 3, content: 3}, + {id: 3, text: "ddd", bytes: 3, content: 3, done: true}, + }) + f.prescanFn = func(c []byte) bool { return strings.Contains(string(c), "ddd") } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { + return &hookcore.CompliancePipelineResult{Decision: hookcore.Modify} + } + if err := Run(context.Background(), f, Config{TailWindowBytes: 3, MaxBufferBytes: 1 << 20}); err != nil { + t.Fatal(err) + } + if len(f.confirmCalls) == 0 { + t.Fatal("confirm did not run on the prescan hit — the gate was wrongly skipped after compaction (leak class)") + } + if last := f.confirmCalls[len(f.confirmCalls)-1]; strings.Contains(last, "aaa") { + t.Fatalf("confirm saw the delivered+compacted prefix %q; want the held window only", last) + } + if !f.escalated { + t.Fatal("expected escalation on the confirmed redact — detection must survive compaction") + } +} + +// TestWindowedConfirmExcludesDeliveredPrefix is the precise #12 windowing assertion the +// existing TestCompactionWindowedConfirm cannot make: with TailWindowBytes=3 and a prescan +// hit on the THIRD unit, exactly one unit ("aaa") has been delivered (deliveredScanLen=3) +// but scanBuf is NOT yet compacted (deliveredScanLen is not > the window), so the delivered +// prefix is still physically present in scanBuf at confirm time. The confirm MUST receive +// only the held window "bbbccc", never the full "aaabbbccc" — pinning Confirm to +// scanBuf[deliveredScanLen:]. This kills a "confirm the whole buffer" mutant that the +// compacting test (which rebases deliveredScanLen to 0, making windowed and full identical) +// lets survive. +func TestWindowedConfirmExcludesDeliveredPrefix(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "aaa", content: 3, bytes: 3}, + {id: 2, text: "bbb", content: 3, bytes: 3}, + {id: 3, text: "ccc", content: 3, bytes: 3, done: true}, + }) + f.prescanFn = func(c []byte) bool { return strings.Contains(string(c), "ccc") } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return modify() } + if err := Run(context.Background(), f, Config{TailWindowBytes: 3, MaxBufferBytes: 1 << 20, PrescanBatchBytes: 1}); err != nil { + t.Fatal(err) + } + // Only "aaa" releases before the third-unit hit; scanBuf is still the un-compacted + // "aaabbbccc" (deliveredScanLen=3, not > window=3) when the confirm fires. + if !eq(f.delivered, []int{1}) { + t.Fatalf("delivered = %v, want [1] (only the prefix released before the hit)", f.delivered) + } + if len(f.confirmCalls) != 1 { + t.Fatalf("confirm ran %d times, want exactly 1 (on the third-unit hit)", len(f.confirmCalls)) + } + if got := f.confirmCalls[0]; got != "bbbccc" { + t.Fatalf("confirm saw %q, want the held window %q (delivered prefix \"aaa\" excluded; full buffer was \"aaabbbccc\")", got, "bbbccc") + } + if !f.escalated || !eq(f.escalateHeld, []int{2, 3}) { + t.Fatalf("escalate held = %v (escalated=%v), want [2 3]", f.escalateHeld, f.escalated) + } +} + +// TestCompactionRebasePreservesPrescanGate pins the stale-high-confirmedLen leak class the +// existing compaction test lets survive: an early run of approve-confirms drives confirmedLen +// up to len(scanBuf), then a release past the window compacts scanBuf and REBASES confirmedLen +// down by the dropped length to a still-POSITIVE value (here 3, not clamped to 0). A later +// prescan hit on genuinely-new content ("ddd") must still satisfy the `len(scanBuf) > +// confirmedLen` gate and fire the confirm. A mutant that forgot to subtract the dropped length +// (leaving confirmedLen stale-high at 9) would fail the gate (6 > 9 is false), silently skip +// the confirm, and stream the fresh sensitive content unredacted. +func TestCompactionRebasePreservesPrescanGate(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "aaa", content: 3, bytes: 3}, + {id: 2, text: "bbb", content: 3, bytes: 3}, + {id: 3, text: "ccc", content: 3, bytes: 3}, + {id: 4, text: "ddd", content: 3, bytes: 3, done: true}, + }) + f.prescanFn = func([]byte) bool { return true } // every new unit re-arms the gate + f.confirmFn = func(c string) *hookcore.CompliancePipelineResult { + if strings.Contains(c, "ddd") { + return modify() // the fresh post-compaction content must reach this + } + return approve() + } + if err := Run(context.Background(), f, Config{TailWindowBytes: 3, MaxBufferBytes: 1 << 20, PrescanBatchBytes: 1}); err != nil { + t.Fatal(err) + } + // aaa+bbb deliver before the final hit; compaction fires after bbb's release (drop=6, + // confirmedLen rebased 9→3). + if !eq(f.delivered, []int{1, 2}) { + t.Fatalf("delivered = %v, want [1 2]", f.delivered) + } + last := f.confirmCalls[len(f.confirmCalls)-1] + if !strings.Contains(last, "ddd") { + t.Fatalf("final confirm = %q, want it to include the post-compaction content \"ddd\" (gate skipped on a stale-high confirmedLen?)", last) + } + if !f.escalated || !eq(f.escalateHeld, []int{3, 4}) { + t.Fatalf("escalate held = %v (escalated=%v), want [3 4]", f.escalateHeld, f.escalated) + } +} + +// TestSustainedPrescanFP_ConfirmWorkIsLinear gates the #12 O(N²)→O(N) confirm fix with an +// asserting bound (the benchmark only prints). Under a perpetual prescan false-positive every +// unit fires a confirm; windowed confirm copies only the bounded tail (O(window) bytes per +// confirm → O(N) total), while the pre-#12 full-buffer confirm copied the whole accumulated +// prefix (O(N) bytes per confirm → O(N²) total). testing.AllocsPerRun is BLIND to this: the +// alloc COUNT is ~N either way — only the alloc BYTES diverge. So this measures TotalAlloc +// bytes at two stream lengths and asserts that 4× the units grows bytes ~linearly (< 8×); a +// quadratic-confirm regression would be ≈16× and fail. +func TestSustainedPrescanFP_ConfirmWorkIsLinear(t *testing.T) { + run := func(n int) uint64 { + units := make([]fakeUnit, n) + for i := range units { + units[i] = fakeUnit{id: i, text: "x", bytes: 1, content: 1, done: i == n-1} + } + f := newFake(units) + f.prescanFn = func([]byte) bool { return true } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return approve() } + var a, b runtime.MemStats + runtime.GC() + runtime.ReadMemStats(&a) + _ = Run(context.Background(), f, Config{TailWindowBytes: 8, MaxBufferBytes: 1 << 20}) + runtime.ReadMemStats(&b) + return b.TotalAlloc - a.TotalAlloc + } + const small, large = 1000, 4000 + bSmall := run(small) + bLarge := run(large) + // 4× units → ~4× bytes under O(N). Generous 8× ceiling absorbs fixed per-unit overhead + // (held / heldScanLens growth, fake bookkeeping); a quadratic confirm (~16×) blows it. + if bLarge > bSmall*8 { + t.Fatalf("confirm work grew %.1f× for %d× units (small=%d large=%d bytes); want ~linear (<8×), quadratic regression?", + float64(bLarge)/float64(bSmall), large/small, bSmall, bLarge) + } +} + +// --- Batched prescan (#13 port) --- + +// TestBatchedPrescan_CollapsesScanCount is the perf win: a long clean stream of small units +// triggers the cheap union prescan only once per PrescanBatchBytes of new content (plus the +// EOF flush), NOT once per unit — collapsing the cgo scan count. 300 ten-byte units = 3000 +// scan bytes ⇒ ~2 batch scans + 1 EOF flush, far below 300 per-unit scans. +func TestBatchedPrescan_CollapsesScanCount(t *testing.T) { + const units = 300 + us := make([]fakeUnit, units) + for i := range us { + us[i] = fakeUnit{id: i, text: strings.Repeat("x", 10), content: 10, bytes: 10, done: i == units-1} + } + f := newFake(us) + f.prescanFn = func([]byte) bool { return false } // clean stream + // Default PrescanBatchBytes (1024); huge window so nothing releases mid-stream. + if err := Run(context.Background(), f, Config{TailWindowBytes: 1 << 20, MaxBufferBytes: 1 << 20}); err != nil { + t.Fatalf("Run error: %v", err) + } + if got := len(f.prescanCalls); got > 5 { + t.Fatalf("prescan ran %d times for %d units; batching should collapse it to ~3 (2 batch + 1 EOF)", got, units) + } + if len(f.delivered) != units { + t.Fatalf("delivered %d units, want all %d", len(f.delivered), units) + } + if f.escalated { + t.Fatal("clean stream must not escalate") + } +} + +// TestBatchedPrescan_EOFTailScannedNoLeak pins that a sub-threshold final content unit (one +// that never reaches the batch trigger) is still scanned at the EOF flush BEFORE the held +// tail is delivered raw — a confirmed hit there escalates rather than leaking. +func TestBatchedPrescan_EOFTailScannedNoLeak(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "clean", content: 5, bytes: 5}, + {id: 2, text: "data", content: 4, bytes: 4}, + {id: 3, text: "secret", content: 6, bytes: 6, done: true}, // total 15B < batch threshold + }) + f.prescanFn = func(b []byte) bool { return strings.Contains(string(b), "secret") } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return reject() } + if err := Run(context.Background(), f, Config{TailWindowBytes: 1 << 20, MaxBufferBytes: 1 << 20}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !f.escalated { + t.Fatal("sub-threshold tail with a confirmed hit was NOT caught at the EOF flush — leak") + } + if !eq(f.escalateHeld, []int{1, 2, 3}) { + t.Fatalf("escalate held = %v, want [1 2 3] (all held, nothing delivered)", f.escalateHeld) + } + if len(f.delivered) != 0 { + t.Fatalf("delivered %v before the EOF-flush escalation, want none", f.delivered) + } + if got := len(f.prescanCalls); got != 1 { + t.Fatalf("prescan ran %d times, want 1 (only the EOF flush; the batch trigger never fired)", got) + } +} + +// TestBatchedPrescan_MidStreamTriggerEscalates pins that once a unit pushes accumulated new +// content past PrescanBatchBytes, the batch trigger fires MID-stream (not only at EOF) and a +// confirmed hit escalates immediately with the still-held units. +func TestBatchedPrescan_MidStreamTriggerEscalates(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "secret" + strings.Repeat("x", 1020), content: 1026, bytes: 1026}, // > 1024 ⇒ batch fires + {id: 2, text: "later", content: 5, bytes: 5, done: true}, // never reached + }) + f.prescanFn = func(b []byte) bool { return strings.Contains(string(b), "secret") } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return reject() } + if err := Run(context.Background(), f, Config{TailWindowBytes: 1 << 20, MaxBufferBytes: 1 << 20}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !f.escalated || !eq(f.escalateHeld, []int{1}) { + t.Fatalf("mid-stream batch trigger did not escalate at unit 1: escalated=%v held=%v", f.escalated, f.escalateHeld) + } + if len(f.prescanCalls) != 1 { + t.Fatalf("prescan ran %d times, want 1 (the mid-stream batch trigger, before EOF)", len(f.prescanCalls)) + } +} + +// TestBatchedPrescan_FlushBeforeDeliverNoLeak exercises the flush-before-deliver guard — the +// load-bearing soundness property when PrescanBatchBytes exceeds the tail window. With a small +// window and the default (large) batch threshold, the window fills and triggers a RELEASE +// before the batch trigger ever fires; the release must force a scanThrough over the held +// window first, so a pattern straddling the about-to-be-released prefix is caught (escalate), +// never delivered raw. +func TestBatchedPrescan_FlushBeforeDeliverNoLeak(t *testing.T) { + f := newFake([]fakeUnit{ + {id: 1, text: "aaaaaaaa", content: 8, bytes: 8}, + {id: 2, text: "secretXX", content: 8, bytes: 8}, + {id: 3, text: "cccccccc", content: 8, bytes: 8, done: true}, // content 24 > window 20 ⇒ release + }) + f.prescanFn = func(b []byte) bool { return strings.Contains(string(b), "secret") } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return reject() } + // Window 20 < default batch threshold (1024): the batch trigger never fires; the + // release-driven flush-before-deliver scan is the ONLY scan, and it must catch the hit. + if err := Run(context.Background(), f, Config{TailWindowBytes: 20, MaxBufferBytes: 1 << 20}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !f.escalated { + t.Fatal("flush-before-deliver did NOT scan the held window before release — the secret would deliver raw") + } + if len(f.delivered) != 0 { + t.Fatalf("delivered %v before the flush-before-deliver escalation, want none", f.delivered) + } + if len(f.prescanCalls) != 1 { + t.Fatalf("prescan ran %d times, want 1 (the flush-before-deliver scan; batch threshold never reached)", len(f.prescanCalls)) + } +} + +// TestBatchedPrescan_CoFiringBlockSoftAtEOF pins that a co-firing BlockSoft-carrying-redaction, +// confirmed only at the EOF flush under batching, still escalates on the enforcing ACTION +// (ActionFromDecision(BlockSoft)=block) rather than streaming the masked redact raw. +func TestBatchedPrescan_CoFiringBlockSoftAtEOF(t *testing.T) { + f := newFake([]fakeUnit{{id: 1, text: "pii here", content: 8, bytes: 8, done: true}}) + f.prescanFn = func([]byte) bool { return true } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { + return &hookcore.CompliancePipelineResult{Decision: hookcore.BlockSoft} + } + if err := Run(context.Background(), f, Config{TailWindowBytes: 1 << 20, MaxBufferBytes: 1 << 20}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !f.escalated || f.escalateRes == nil || f.escalateRes.Decision != hookcore.BlockSoft { + t.Fatalf("co-firing BlockSoft at EOF flush did not escalate: escalated=%v res=%+v", f.escalated, f.escalateRes) + } + if len(f.delivered) != 0 { + t.Fatalf("delivered %v before the co-firing escalation, want none", f.delivered) + } +} + +// TestBatchedPrescan_BoundarySpanningValue_LargeUnitNoLeak is the flush-before-deliver +// LOOKAHEAD regression (arch review): a LARGE content unit whose tail begins a sensitive +// value, followed by a small unit that completes it across the boundary. The large unit is +// batch-scanned ALONE (before its successor exists), so a flush guard keyed only on "front +// itself unscanned" (scannedLen < frontScanEnd) would release it raw — leaking the value's +// prefix — and then scan the suffix over a window that no longer holds the prefix, never +// matching. The MaxPatternBytes lookahead forces a re-scan over the completed value and +// escalates. (Defaults: W=8192, B=1024, MaxPatternBytes=4096.) +func TestBatchedPrescan_BoundarySpanningValue_LargeUnitNoLeak(t *testing.T) { + u1text := strings.Repeat("a", 7196) + "secr" // 7200B, ends mid-"secret" + u2text := "et" + strings.Repeat("b", 991) // 993B ⇒ 7200+993 > 8192 ⇒ u1 released + f := newFake([]fakeUnit{ + {id: 1, text: u1text, content: len(u1text), bytes: len(u1text)}, + {id: 2, text: u2text, content: len(u2text), bytes: len(u2text), done: true}, + }) + f.prescanFn = func(b []byte) bool { return strings.Contains(string(b), "secret") } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { return reject() } + if err := Run(context.Background(), f, Config{}); err != nil { + t.Fatalf("Run error: %v", err) + } + if !f.escalated { + t.Fatal("boundary-spanning value across a large unit was NOT caught — the large unit leaked its prefix raw (flush lookahead missing/too small)") + } + if len(f.delivered) != 0 { + t.Fatalf("delivered %v before the boundary escalation, want none (the large unit must not flush raw)", f.delivered) + } +} + +// BenchmarkRun_SustainedPrescanFP guards the O(N²)→O(N) confirm fix: a stream of N tiny +// content units under a perpetual prescan false-positive. Before #12 each unit ran a +// full-buffer confirm (Σ length = O(N²)); after, each confirm is windowed to O(window) +// and scanBuf is compacted, so total work + retained memory are O(N). +func BenchmarkRun_SustainedPrescanFP(b *testing.B) { + units := make([]fakeUnit, 2000) + for i := range units { + units[i] = fakeUnit{id: i, text: "x", bytes: 1, content: 1, done: i == len(units)-1} + } + b.ReportAllocs() + for range b.N { + f := newFake(units) + f.prescanFn = func([]byte) bool { return true } + f.confirmFn = func(string) *hookcore.CompliancePipelineResult { + return &hookcore.CompliancePipelineResult{Decision: hookcore.Approve} + } + _ = Run(context.Background(), f, Config{TailWindowBytes: 8, MaxBufferBytes: 1 << 20}) + } +} diff --git a/packages/shared/transport/streaming/pipelines_test.go b/packages/shared/transport/streaming/pipelines_test.go index c450c702..b6a5211d 100644 --- a/packages/shared/transport/streaming/pipelines_test.go +++ b/packages/shared/transport/streaming/pipelines_test.go @@ -388,31 +388,42 @@ func TestLivePipeline_NoCapture_CapturedNilAndFalse(t *testing.T) { // TestLivePipeline_MaxBufferExceeded — totalBytes > MaxBufferSize aborts // with an error event to the client and a non-nil writerErr. +// TestLivePipeline_MaxBufferExceeded pins the AUDIT-ONLY overflow behaviour: when the +// accumulated content exceeds MaxBufferSize the pipeline CAPS the audit scan to bound +// memory but KEEPS delivering — an audit-only relay must never break a non-enforcing +// stream just because it outgrew the scan budget. No block frame is emitted. func TestLivePipeline_MaxBufferExceeded(t *testing.T) { lp := NewLivePipeline(LiveConfig{ CheckpointChars: 100000, MaxBufferSize: 30, // tiny }, &mockPipeline{}, nil) - in := makeOpenAISSE( - "this string easily exceeds the 30-byte buffer cap once accumulated across frames") + // Several frames: the cap trips early, and the LATER frames must still be + // delivered write-through (exercising the post-cap continue). + in := makeOpenAISSE("alpha", "bravo", "charlie", "delta") var client bytes.Buffer res, err := lp.Process(context.Background(), strings.NewReader(in), &client, &core.HookInput{IngressType: "X"}) - _ = res if err != nil { - t.Fatalf("Process err = %v (expect nil; chunk error converts to client write)", err) + t.Fatalf("Process err = %v (audit-only overflow must not error)", err) } - if !strings.Contains(client.String(), "blocked by policy") { - t.Errorf("expected error envelope on client; got %q", client.String()) + if res == nil || res.Decision != core.Approve { + t.Errorf("audit-only overflow must return Approve, got %+v", res) + } + if strings.Contains(client.String(), "blocked by policy") { + t.Errorf("audit-only overflow must NOT emit a block frame; got %q", client.String()) + } + // Every frame is delivered write-through despite the audit cap tripping early. + for _, want := range []string{"alpha", "bravo", "charlie", "delta"} { + if !strings.Contains(client.String(), want) { + t.Errorf("expected %q delivered despite the audit cap; got %q", want, client.String()) + } } } -// TestLivePipeline_NilPipelineResult_NoPanic — when checkpoint returns nil, -// flushCheckpoint short-circuits with Approve. Observed behavior today: -// pendingEvents are NOT forwarded to approvedChan on the nil-result branch -// (the comment claims "treat as approve" but the code only returns the -// decision — see live.go:212-215). This test pins the no-panic + final -// decision == Approve guarantee, the only contract the call site relies on. +// TestLivePipeline_NilPipelineResult_NoPanic — when an observe-only checkpoint +// returns nil, the audit aggregate skips it (no panic, no HookResults appended) and +// delivery proceeds write-through. The final decision is Approve and the content is +// delivered regardless of the nil checkpoint (audit-only never gates the wire). func TestLivePipeline_NilPipelineResult_NoPanic(t *testing.T) { mp := &mockPipeline{ decideFn: func(_ context.Context, _ *core.HookInput) *core.CompliancePipelineResult { @@ -430,6 +441,9 @@ func TestLivePipeline_NilPipelineResult_NoPanic(t *testing.T) { if res == nil || res.Decision != core.Approve { t.Errorf("decision = %+v, want Approve", res) } + if !strings.Contains(out.String(), "abcdef") || !strings.Contains(out.String(), "ghijkl") { + t.Errorf("audit-only must deliver write-through on a nil checkpoint; got %q", out.String()) + } } // TestLivePipeline_ReaderError_PropagatesToProcess — non-EOF reader error @@ -473,25 +487,6 @@ func TestLivePipeline_ReplayWithFlushableClient_FlushesEachChunk(t *testing.T) { } } -// TestLivePipeline_HardReject_WriterErrorOnError — when the hard-reject -// path tries to write the error envelope but the client errors, Process -// returns that error. -func TestLivePipeline_HardReject_WriterErrorOnError(t *testing.T) { - mp := &mockPipeline{ - decideFn: func(_ context.Context, _ *core.HookInput) *core.CompliancePipelineResult { - return &core.CompliancePipelineResult{Decision: core.RejectHard} - }, - } - lp := NewLivePipeline(LiveConfig{CheckpointChars: 1}, mp, nil) - fc := &failingClientWriter{failAfter: 0, failureErr: errors.New("client gone")} - - in := makeOpenAISSE("trigger") - _, err := lp.Process(context.Background(), strings.NewReader(in), fc, &core.HookInput{IngressType: "X"}) - if err == nil { - t.Fatal("expected write error from reject path") - } -} - // TestSSEParser_InvalidRetryValue_Skipped — retry: must parse as int; // invalid values log a warning and leave the field at -1. func TestSSEParser_InvalidRetryValue_Skipped(t *testing.T) { diff --git a/packages/shared/transport/tlsbump/forward_handler_pins_test.go b/packages/shared/transport/tlsbump/forward_handler_pins_test.go index 74542df6..277bbe47 100644 --- a/packages/shared/transport/tlsbump/forward_handler_pins_test.go +++ b/packages/shared/transport/tlsbump/forward_handler_pins_test.go @@ -13,8 +13,181 @@ import ( "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/domain" "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" compliance "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/pipeline" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/traffic" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/traffic/adapters/api/openai" ) +// bodyCapturingRoundTripper records the BODY of each outbound upstream request so a test +// can assert exactly what bytes the upstream provider received (the masked request body on +// the redact path, never the original). +type bodyCapturingRoundTripper struct { + mu sync.Mutex + bodies [][]byte + makeResp func() *http.Response +} + +func (rt *bodyCapturingRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) { + var b []byte + if r.Body != nil { + b, _ = io.ReadAll(r.Body) + } + rt.mu.Lock() + rt.bodies = append(rt.bodies, b) + rt.mu.Unlock() + return rt.makeResp(), nil +} + +func (rt *bodyCapturingRoundTripper) lastBody() string { + rt.mu.Lock() + defer rt.mu.Unlock() + if len(rt.bodies) == 0 { + return "" + } + return string(rt.bodies[len(rt.bodies)-1]) +} + +func (rt *bodyCapturingRoundTripper) calls() int { + rt.mu.Lock() + defer rt.mu.Unlock() + return len(rt.bodies) +} + +// softBlockReqHook always returns a soft-block decision — the co-firing partner of a redact +// hook on the request stage. +type softBlockReqHook struct{} + +func (softBlockReqHook) Execute(context.Context, *core.HookInput) (*core.HookResult, error) { + return &core.HookResult{HookID: "h-sb", HookName: "softblock-req", Decision: core.BlockSoft, Reason: "content flagged (soft)", ReasonCode: "SOFT_FLAG"}, nil +} +func (softBlockReqHook) SupportsEndpoint(core.EndpointType) bool { return true } +func (softBlockReqHook) SupportsModality(core.Modality) bool { return true } + +// coFiringRequestResolver wires a request-stage redact hook (fixedModifyHook, returning the +// masked ModifiedContent) AND a request-stage soft-block hook. mergeResults ranks BlockSoft +// above Modify (Decision→BlockSoft) while carrying the redact's ModifiedContent — the +// co-firing shape that, pre-#13, forwarded the ORIGINAL request body upstream because the +// rewrite arm gated on Decision==Modify. +func coFiringRequestResolver(modified []core.ContentBlock) *compliance.PolicyResolver { + reg := core.NewHookRegistry() + reg.Register("req-redact", func(_ *core.HookConfig) (core.Hook, error) { + return fixedModifyHook{modified: modified}, nil + }) + reg.Register("req-softblock", func(_ *core.HookConfig) (core.Hook, error) { + return softBlockReqHook{}, nil + }) + redactCfg := map[string]any{"onMatch": map[string]any{"action": "redact"}} + return compliance.NewPolicyResolver([]core.HookConfig{ + {ID: "r-redact", ImplementationID: "req-redact", Name: "req-redact", Stage: "request", Enabled: true, FailBehavior: "fail-open", ApplicableIngress: []string{"ALL"}, Config: redactCfg}, + {ID: "r-softblock", ImplementationID: "req-softblock", Name: "req-softblock", Stage: "request", Enabled: true, FailBehavior: "fail-open", ApplicableIngress: []string{"ALL"}, Config: redactCfg}, + }, reg, discardSlog()) +} + +// adapterDomainEngine matches api.example.com with a PROCESS path action and the given +// adapter id, so runRequestPhase resolves a concrete adapter for the request rewrite. +func adapterDomainEngine(t *testing.T, adapterID string) *domain.Engine { + t.Helper() + eng := domain.NewEngine() + if err := eng.Swap([]domain.InterceptionDomain{{ + ID: "dom-adapter", Name: "example", HostPattern: "api.example.com", + HostMatchType: domain.HostMatchExact, DefaultPathAction: domain.PathActionProcess, + AdapterID: adapterID, Enabled: true, Priority: 10, + }}); err != nil { + t.Fatalf("engine swap: %v", err) + } + return eng +} + +// TestForwardHandler_RequestHookBlockSoftMaskedRedact_ForwardsRedactedUpstream pins #13 leak +// #4 on the tlsbump REQUEST path (forward_request_phase, previously ZERO-coverage): a redact +// hook co-firing with a soft-block aggregates to Decision=BLOCK_SOFT while carrying the +// redact's ModifiedContent. The pre-#13 `case Modify` gate skipped the rewrite and forwarded +// the ORIGINAL body upstream; the CarriesRedaction() gate now rewrites it, so the upstream +// provider receives the MASKED body and never the original email. The audit row records the +// aggregate BLOCK_SOFT decision (the redact is the disposition, not the ceiling). +func TestForwardHandler_RequestHookBlockSoftMaskedRedact_ForwardsRedactedUpstream(t *testing.T) { + writer := &recordingAuditWriter{} + rt := &bodyCapturingRoundTripper{makeResp: jsonUpstream} + areg := traffic.NewAdapterRegistry("test") + if err := areg.Register("openai-compat", func() traffic.Adapter { return &openai.Adapter{} }); err != nil { + t.Fatalf("adapter register: %v", err) + } + bo := &bumpOptions{ + policyResolver: coFiringRequestResolver([]core.ContentBlock{{Type: "text", Text: "ping [REDACTED_EMAIL]"}}), + auditEmitter: compliance.NewAuditEmitter(writer, discardSlog()), + domainEngine: adapterDomainEngine(t, "openai-compat"), + adapterRegistry: areg, + } + h := buildForwardHandler(context.Background(), "api.example.com:443", &UpstreamTransport{transport: rt}, discardSlog(), bo) + + body := `{"model":"gpt-4o","messages":[{"role":"user","content":"ping alice@example.com"}]}` + req := httptest.NewRequest(http.MethodPost, "https://api.example.com/v1/chat/completions", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if got := rt.calls(); got != 1 { + t.Fatalf("upstream forwards = %d, want 1 (co-firing soft-block+redact must redact-and-forward, not reject)", got) + } + up := rt.lastBody() + if !strings.Contains(up, "[REDACTED_EMAIL]") { + t.Fatalf("upstream body = %q, want the masked email forwarded", up) + } + if strings.Contains(up, "alice@example.com") { + t.Fatalf("upstream body = %q, must NOT carry the original email (the request-side co-firing leak)", up) + } + events := writer.snapshot() + if len(events) != 1 { + t.Fatalf("audit events = %d, want 1", len(events)) + } + if events[0].RequestHookDecision != string(core.BlockSoft) { + t.Fatalf("RequestHookDecision = %q, want %q (aggregate ceiling; the redact is the disposition)", events[0].RequestHookDecision, core.BlockSoft) + } +} + +// standaloneSoftBlockResolver wires ONLY a request-stage soft-block hook (no co-firing +// redact), so the aggregate is a standalone BlockSoft carrying no applicable redaction. +func standaloneSoftBlockResolver() *compliance.PolicyResolver { + reg := core.NewHookRegistry() + reg.Register("req-softblock", func(_ *core.HookConfig) (core.Hook, error) { return softBlockReqHook{}, nil }) + return compliance.NewPolicyResolver([]core.HookConfig{ + {ID: "r-softblock", ImplementationID: "req-softblock", Name: "req-softblock", Stage: "request", Enabled: true, FailBehavior: "fail-open", ApplicableIngress: []string{"ALL"}, Config: map[string]any{"onMatch": map[string]any{"action": "block"}}}, + }, reg, discardSlog()) +} + +// TestForwardHandler_StandaloneRequestBlockSoft_Refuses pins the #15 fold-to-block fix on +// the tlsbump request path: a standalone soft-block (no co-firing redact → no applicable +// redaction) folds to the block action and REFUSES the request (403), never forwarding it +// upstream. Before the fix the switch keyed on Decision==RejectHard only, so a BlockSoft +// fell through and forwarded the original unredacted body — a request-side leak on the +// dormant BlockSoft path. The agent fail-open path (richReject=false) returns a bare 403. +func TestForwardHandler_StandaloneRequestBlockSoft_Refuses(t *testing.T) { + writer := &recordingAuditWriter{} + rt := &bodyCapturingRoundTripper{makeResp: jsonUpstream} + bo := &bumpOptions{ + policyResolver: standaloneSoftBlockResolver(), + auditEmitter: compliance.NewAuditEmitter(writer, discardSlog()), + domainEngine: adapterDomainEngine(t, "openai-compat"), + } + h := buildForwardHandler(context.Background(), "api.example.com:443", &UpstreamTransport{transport: rt}, discardSlog(), bo) + + body := `{"model":"gpt-4o","messages":[{"role":"user","content":"flag me"}]}` + req := httptest.NewRequest(http.MethodPost, "https://api.example.com/v1/chat/completions", strings.NewReader(body)) + req.Header.Set("Content-Type", "application/json") + rec := httptest.NewRecorder() + h.ServeHTTP(rec, req) + + if rec.Code != http.StatusForbidden { + t.Fatalf("status = %d, want 403 (standalone request BlockSoft folds to block)", rec.Code) + } + if got := rt.calls(); got != 0 { + t.Fatalf("upstream forwards = %d, want 0 (a folded block must NOT reach upstream)", got) + } + events := writer.snapshot() + if len(events) != 1 || events[0].RequestHookDecision != string(core.BlockSoft) { + t.Fatalf("audit: want 1 event with RequestHookDecision=BLOCK_SOFT, got %+v", events) + } +} + // These tests pin the forward handler's observable per-phase contract: // attestation passthrough, domain path-policy (BLOCK / PASSTHROUGH), // request-hook decisions (REJECT_HARD / BLOCK_SOFT), upstream-failure diff --git a/packages/shared/transport/tlsbump/forward_request_phase.go b/packages/shared/transport/tlsbump/forward_request_phase.go index 2360f1b1..bcd9c480 100644 --- a/packages/shared/transport/tlsbump/forward_request_phase.go +++ b/packages/shared/transport/tlsbump/forward_request_phase.go @@ -176,29 +176,18 @@ func (x *bumpedExchange) runRequestPhase() bool { // convert it into a HookOutcomeInput for downstream writers. x.reqHookResult = result - switch result.Decision { - case compliance.RejectHard: - logger.Info("request blocked by compliance (REJECT_HARD)", - "target", x.flow.targetHost, - "transactionId", x.txID, - "reason", result.Reason, - ) - if bo.auditEmitter != nil { - bo.auditEmitter.Emit(reqInput, auditInfo, result, "BUMP_SUCCESS", http.StatusForbidden, int(time.Since(x.requestStart).Milliseconds()), captureBodyIfEnabled(x.pcCfg.StoreRequestBody, bodyBytes), nil, traffic.UsageMeta{}) - } - stampRejectMarkers(x.w.Header(), bo.identity, x.txID, x.domainRuleID, cpHookOutcomeFromResult(result)) - if bo.richReject { - WriteRejectResponse(x.w, x.r, bo.rejectConfig, x.txID, result.Reason, result.ReasonCode, http.StatusForbidden) - } else { - // Agent on-host interceptor: minimal 403 with no attribution - // body (it does not synthesize a rich error page in the host - // outbound path). - http.Error(x.w, "Forbidden", http.StatusForbidden) - } - return true - - case compliance.Modify: - // Hook requested inflight redact. Try to rewrite the + // Gate the rewrite arm on CarriesRedaction() (Modify OR a BlockSoft masking a + // co-firing redact), NOT Decision==Modify — a co-firing soft-block promotes the + // aggregate Decision to BlockSoft while carrying the redact's spans/content, so + // keying on Modify alone would forward the original request body UNREDACTED to the + // upstream provider (the request-side leak this fixes). + // Order matters: CarriesRedaction() is checked BEFORE the folded-block refuse + // arm so a BlockSoft that masks a co-firing redact takes the redaction arm + // (redact-forward, the #13 invariant) instead of being refused. + switch { + case result.CarriesRedaction(): + // Hook requested inflight redact (Modify, or a BlockSoft masking a + // co-firing redact). Try to rewrite the // upstream body via the resolved adapter. If the adapter // declares ErrRewriteUnsupported, fall back to // "upstream sees original, audit log stores spans" and @@ -235,6 +224,11 @@ func (x *bumpedExchange) runRequestPhase() bool { result.ReasonCode = core.ReasonRedactInflightUnsupported default: bodyBytes = rewritten + // Stamp the DISPOSITION action: an applied rewrite is a redact even + // when the aggregate Decision is BlockSoft (a soft-block masking a + // co-firing redact) — otherwise the audit row reads action=block for a + // redact-delivered request. + result.Action = core.ActionRedact // The rewritten copy is the only raw bytes allowed into // the audit store under the redact action — the emitter // selects it via StorageRawBody at build time. @@ -246,6 +240,33 @@ func (x *bumpedExchange) runRequestPhase() bool { } } // MODIFY does NOT short-circuit upstream; fall through. + + case core.ActionFromDecision(result.Decision) == core.ActionBlock: + // Refuse on a folded BLOCK action: RejectHard OR a standalone BlockSoft. + // ActionFromDecision folds BlockSoft → block, so the request stage dispatches + // it exactly like RejectHard (a hard 403), matching the SSE/response path and + // error-taxonomy-architecture.md — there is no soft-block client response. + // RejectHard never carries redaction, so its behavior is unchanged; a BlockSoft + // masking a co-firing redact already took the CarriesRedaction arm above. + logger.Info("request blocked by compliance", + "target", x.flow.targetHost, + "transactionId", x.txID, + "decision", result.Decision, + "reason", result.Reason, + ) + if bo.auditEmitter != nil { + bo.auditEmitter.Emit(reqInput, auditInfo, result, "BUMP_SUCCESS", http.StatusForbidden, int(time.Since(x.requestStart).Milliseconds()), captureBodyIfEnabled(x.pcCfg.StoreRequestBody, bodyBytes), nil, traffic.UsageMeta{}) + } + stampRejectMarkers(x.w.Header(), bo.identity, x.txID, x.domainRuleID, cpHookOutcomeFromResult(result)) + if bo.richReject { + WriteRejectResponse(x.w, x.r, bo.rejectConfig, x.txID, result.Reason, result.ReasonCode, http.StatusForbidden) + } else { + // Agent on-host interceptor: minimal 403 with no attribution + // body (it does not synthesize a rich error page in the host + // outbound path). + http.Error(x.w, "Forbidden", http.StatusForbidden) + } + return true } // APPROVE / ABSTAIN / MODIFY-handled — continue to upstream. diff --git a/packages/shared/transport/tlsbump/sse.go b/packages/shared/transport/tlsbump/sse.go index ca79b932..b43a2ab9 100644 --- a/packages/shared/transport/tlsbump/sse.go +++ b/packages/shared/transport/tlsbump/sse.go @@ -228,7 +228,7 @@ func handleSSEResponse( // captureMax is the upper bound on bytes we'll keep for audit when // audCtx.storeResponseBody is true. Mirrors the streamingConfig's // MaxBufferSize so a single tunable governs both the live-pipeline - // hold-back budget and the audit-capture budget. + // audit-accumulation cap and the audit-capture budget. captureMax := 0 if audCtx != nil && audCtx.storeResponseBody { captureMax = bo.streamingConfig.MaxBufferSize @@ -279,6 +279,13 @@ func handleSSEResponse( return } + // Scope-derived routing (mirror ai-gateway stream_shape.go §B2): an enforcing + // response scope OVERRIDES the admin streaming mode — the audit-only live path + // cannot enforce, so enforcing traffic must route to buffer or Model A. The probe + // is reused as the Model A prescan/confirm pipeline. + var responseProbe *compliance.Pipeline + mode, responseProbe = scopeRouteSSEMode(bo, mode, respInput, audCtx, logger) + switch mode { case "passthrough": // No compliance inspection — just relay. Passthrough doesn't parse @@ -429,6 +436,7 @@ func handleSSEResponse( bufConfig := streaming.BufferConfig{MaxBufferSize: bo.streamingConfig.MaxBufferSize} bufPipeline := streaming.NewBufferPipeline(bufConfig, pipelineExec, logger) + bufPipeline.WithStrictFailClosed(bo.strictFailClosed) // no-redactor posture (GAP B) if acc != nil { bufPipeline.WithUsageAccumulator(acc) } @@ -442,11 +450,10 @@ func handleSSEResponse( if cb := buildSSEPreHookCallback(ctx, bo, audCtx, respInput, respContentType); cb != nil { bufPipeline.WithPreHook(cb) } - // Wire the per-host frame redactor so a Modify (redact) decision rewrites - // the buffered timeline — splicing the masked text into the per-host text - // frames and passing non-text frames byte-verbatim — instead of degrading - // to a verbatim replay of the unredacted stream (a PII leak). nil adapter - // leaves the redactor unset → backward-compatible Modify-degrade. + // Wire the per-host frame redactor so a redact decision rewrites the buffered + // timeline (masked text spliced in, non-text frames byte-verbatim) instead of a + // verbatim replay (a PII leak). nil adapter → no redactor → posture-aware degrade + // (WithStrictFailClosed). Redactor carries the per-caller posture (open/closed). var sseAdapter traffic.Adapter if audCtx != nil { sseAdapter = audCtx.adapter @@ -456,7 +463,7 @@ func handleSSEResponse( ssePath = respInput.Path } spliceWired := false - if fr := newSSEFrameRedactor(ctx, sseAdapter, ssePath, logger); fr != nil { + if fr := newSSEFrameRedactor(ctx, sseAdapter, ssePath, bo.strictFailClosed, logger); fr != nil { bufPipeline.WithFrameRedactor(fr) spliceWired = true } @@ -474,6 +481,14 @@ func handleSSEResponse( stampSpliceRedactedBody(auditInfo, result, spliceWired, bufPipeline) emitAudit(logger, audCtx, respInput, auditInfo, bo, result, resp.StatusCode, requestStart, finalizeUsage(ctx, acc), bufPipeline.CapturedBytes()) + case "modela": + // redact-scope under chunked_async with a per-host adapter: prescan-gated + // real-time streaming with a bounded tail + escalate-to-buffer redaction via + // the shared Model-A engine (fail-OPEN wire substrate). responseProbe + the + // adapter are non-nil here (the routing reached "modela" only on MayRedact + + // an available adapter). + runSSEModelA(ctx, w, resp, audCtx, respInput, auditInfo, bo, logger, requestStart, responseProbe, acc, captureMax) + default: logger.Warn("unknown streaming mode, using passthrough", "mode", mode) var captureBuf *streaming.CappedBuffer diff --git a/packages/shared/transport/tlsbump/sse_buffer_redact_med_test.go b/packages/shared/transport/tlsbump/sse_buffer_redact_med_test.go index 4734bc17..eaa5516f 100644 --- a/packages/shared/transport/tlsbump/sse_buffer_redact_med_test.go +++ b/packages/shared/transport/tlsbump/sse_buffer_redact_med_test.go @@ -114,7 +114,7 @@ func TestSSE_BufferMode_SuccessfulSplice_StampsRedactedBody(t *testing.T) { func TestSSEFrameRedactor_AnthropicRedacts_FencePasses(t *testing.T) { ctx := context.Background() path := "/v1/messages" - fr := newSSEFrameRedactor(ctx, &anthropic.Adapter{}, path, nil) + fr := newSSEFrameRedactor(ctx, &anthropic.Adapter{}, path, false, nil) events := []*streaming.SSEEvent{ anthropicFrame("card "), anthropicFrame("4111111111111111"), @@ -162,8 +162,8 @@ func modifyDegradedValue(t *testing.T, reason string) float64 { // degrading to forward-original. func TestSSEFrameRedactor_FailOpenBumpsCounter(t *testing.T) { ctx := context.Background() - fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", nil) - before := modifyDegradedValue(t, failOpenReasonSplice) + fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", false, nil) + before := modifyDegradedValue(t, failReasonSplice) // ModifiedContent that cannot be reconstructed from the wire (the fence // mismatch) forces the splice-divergence fail-open. @@ -182,7 +182,7 @@ func TestSSEFrameRedactor_FailOpenBumpsCounter(t *testing.T) { if serialize(out[0]) != serialize(oaFrame("hello")) { t.Fatal("fail-open must forward the original frame unchanged") } - if got := modifyDegradedValue(t, failOpenReasonSplice) - before; got != 1 { + if got := modifyDegradedValue(t, failReasonSplice) - before; got != 1 { t.Fatalf("fail-open counter delta = %v, want 1", got) } } @@ -191,8 +191,8 @@ func TestSSEFrameRedactor_FailOpenBumpsCounter(t *testing.T) { // for the tool-arg-undeliverable fail-open reason. func TestSSEFrameRedactor_ToolArgFailOpenBumpsCounter(t *testing.T) { ctx := context.Background() - fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", nil) - before := modifyDegradedValue(t, failOpenReasonToolArg) + fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", false, nil) + before := modifyDegradedValue(t, failReasonToolArg) res := &core.CompliancePipelineResult{ Decision: core.Modify, ModifiedContent: []core.ContentBlock{{Type: "text", Text: "x"}}, @@ -201,7 +201,7 @@ func TestSSEFrameRedactor_ToolArgFailOpenBumpsCounter(t *testing.T) { if _, err := fr.RedactReplay([]*streaming.SSEEvent{oaFrame("hi"), doneFrame()}, res); err != nil { t.Fatalf("fail-open must return nil error, got %v", err) } - if got := modifyDegradedValue(t, failOpenReasonToolArg) - before; got != 1 { + if got := modifyDegradedValue(t, failReasonToolArg) - before; got != 1 { t.Fatalf("tool-arg fail-open counter delta = %v, want 1", got) } } diff --git a/packages/shared/transport/tlsbump/sse_frame_redactor.go b/packages/shared/transport/tlsbump/sse_frame_redactor.go index 45b9ac85..2ee9fc49 100644 --- a/packages/shared/transport/tlsbump/sse_frame_redactor.go +++ b/packages/shared/transport/tlsbump/sse_frame_redactor.go @@ -18,17 +18,25 @@ import ( // is delegated to the matched interception adapter's ExtractStreamChunk, so // non-OpenAI wires redact correctly. // -// FAIL-OPEN posture (best-effort, matching the non-stream tlsbump request path -// and the NE host-packet fail-open safety rule): whenever the masked wire -// cannot be soundly reconstructed, RedactReplay forwards the ORIGINAL events -// with a nil error and stamps REDACT_INFLIGHT_UNSUPPORTED on the result so the -// audit trail is honest. It never returns streaming.ErrRewriteUnsupported — -// that signals the buffer pipeline to fail CLOSED, which is the wrong posture -// for tlsbump's best-effort packet path. +// Fail posture is PER-CALLER, carried by strictFailClosed (sourced from +// bumpOptions.strictFailClosed at the build site). Whenever the masked wire +// cannot be soundly reconstructed, REDACT_INFLIGHT_UNSUPPORTED is always +// stamped on the result so the audit trail is honest, and then: +// +// - strictFailClosed=false (agent NE host-packet path): forward the ORIGINAL +// events with a nil error (FAIL-OPEN). The host's outbound packet path must +// never fail closed (CLAUDE.md NE safety rule); a hung/blocked redaction +// would take down the Mac's whole network. +// - strictFailClosed=true (compliance-proxy appliance): return +// streaming.ErrRewriteUnsupported (FAIL-CLOSED). The buffer pipeline then +// emits the policy error frame and replays no original byte, so zero +// unredacted content reaches the client — matching the ai-gateway sibling's +// fail-closed posture on an unproducible redaction. type sseFrameRedactor struct { - codec streaming.WireTextCodec - adapter traffic.Adapter - logger *slog.Logger + codec streaming.WireTextCodec + adapter traffic.Adapter + logger *slog.Logger + strictFailClosed bool } // adapterWireCodec adapts a traffic.Adapter's ExtractStreamChunk to the @@ -54,14 +62,17 @@ func (c adapterWireCodec) ChunkText(data string) (string, bool) { // newSSEFrameRedactor builds the redactor for one SSE response. Returns nil // when no adapter is available (the buffer pipeline then keeps its // backward-compatible Modify-degrade behavior rather than redacting). -func newSSEFrameRedactor(ctx context.Context, adapter traffic.Adapter, path string, logger *slog.Logger) *sseFrameRedactor { +// strictFailClosed selects the unreconstructable-wire posture (see the type +// doc): false = fail-open forward-original (agent), true = fail-closed (CP). +func newSSEFrameRedactor(ctx context.Context, adapter traffic.Adapter, path string, strictFailClosed bool, logger *slog.Logger) *sseFrameRedactor { if adapter == nil { return nil } return &sseFrameRedactor{ - codec: adapterWireCodec{ctx: ctx, adapter: adapter, path: path}, - adapter: adapter, - logger: logger, + codec: adapterWireCodec{ctx: ctx, adapter: adapter, path: path}, + adapter: adapter, + logger: logger, + strictFailClosed: strictFailClosed, } } @@ -77,11 +88,12 @@ func (r *sseFrameRedactor) RedactReplay(events []*streaming.SSEEvent, result *co // owns tool-arg redaction). GuardToolArgMasking is the shared gate that // the non-stream request path also uses — for a per-host adapter that is // not a traffic.ToolArgMasker it reports ErrRewriteUnsupported. We - // additionally fail open on ANY tool-arg span, so even a ToolArgMasker - // adapter never leaks an unmasked tool argument through a verbatim frame. + // additionally treat ANY tool-arg span as unsupported, so even a + // ToolArgMasker adapter never leaks an unmasked tool argument through a + // verbatim frame (posture handled by redactUnsupported per caller). guard := traffic.NormalizedContent{ToolCallArgs: toolArgSentinel(result.TransformSpans)} if traffic.GuardToolArgMasking(r.adapter, guard) != nil || len(guard.ToolCallArgs) > 0 { - return r.failOpen(events, result, "tool-call argument masking is undeliverable on the streaming wire", failOpenReasonToolArg) + return r.redactUnsupported(events, result, "tool-call argument masking is undeliverable on the streaming wire", failReasonToolArg) } // (2) Splice the text frames; non-text frames pass byte-verbatim. A false @@ -89,33 +101,53 @@ func (r *sseFrameRedactor) RedactReplay(events []*streaming.SSEEvent, result *co // usually a divergence between the normalized text the spans were // computed against and the per-frame ExtractStreamChunk wire transcript // (the splice's divergence fence). The counter exposes how often this - // happens so operators can see real redaction degrading to forward-original. + // happens so operators can see real redaction could not be applied + // (then forwarded-original for the agent, or blocked for the appliance). redacted, ok := streaming.SpliceTextFrames(events, result, r.codec) if !ok { - return r.failOpen(events, result, "frame splice could not reconstruct the masked wire", failOpenReasonSplice) + return r.redactUnsupported(events, result, "frame splice could not reconstruct the masked wire", failReasonSplice) } return redacted, nil } -// failOpen reasons are bounded metric labels for nexus_streaming_modify_degraded_total, -// distinguishing the two best-effort degradations of the tlsbump packet path so +// fail reasons are bounded metric labels for nexus_streaming_modify_degraded_total, +// distinguishing the two root causes of an unreconstructable masked wire so // operators can tell a structurally-undeliverable tool-arg mask apart from a -// wire/normalized divergence (the actionable signal — it means real text -// redaction is silently degrading to forward-original for some host). +// wire/normalized divergence (the actionable signal — real text redaction could +// not be applied for some host). These root-cause labels are recorded in BOTH +// postures; in the fail-closed posture the buffer pipeline additionally bumps +// its coarse "redact_inflight_unsupported" counter on the returned error. const ( - failOpenReasonToolArg = "tlsbump_tool_arg_undeliverable" - failOpenReasonSplice = "tlsbump_splice_divergence" + failReasonToolArg = "tlsbump_tool_arg_undeliverable" + failReasonSplice = "tlsbump_splice_divergence" ) -// failOpen forwards the original events unchanged, stamps the disclosed degraded -// reason on the result so the audit row reflects that the inflight redaction was -// not applied (no original is dropped, no error is surfaced), and bumps the -// modify-degraded counter so the fail-open rate is observable. -func (r *sseFrameRedactor) failOpen(events []*streaming.SSEEvent, result *core.CompliancePipelineResult, reason, metricReason string) ([]*streaming.SSEEvent, error) { +// redactUnsupported handles an unreconstructable masked wire. It always stamps +// the disclosed degraded reason on the result (audit honesty in both postures) +// and bumps the modify-degraded counter under the root-cause label, then +// branches on the per-caller fail posture (see the type doc): +// +// - strictFailClosed=false (agent): forward the ORIGINAL events with a nil +// error (fail-open). No original byte is dropped; the host packet path +// stays open. +// - strictFailClosed=true (compliance-proxy): return ErrRewriteUnsupported so +// the buffer pipeline fails CLOSED — it emits the policy error frame and +// replays no original byte, so zero unredacted content reaches the client. +func (r *sseFrameRedactor) redactUnsupported(events []*streaming.SSEEvent, result *core.CompliancePipelineResult, reason, metricReason string) ([]*streaming.SSEEvent, error) { result.ReasonCode = core.ReasonRedactInflightUnsupported streaming.RecordModifyDegraded(metricReason) + if r.strictFailClosed { + if r.logger != nil { + r.logger.Warn("SSE buffer redaction unsupported; failing closed (strict appliance, no original forwarded)", + "reason", reason, + "reasonCode", core.ReasonRedactInflightUnsupported, + "metricReason", metricReason, + ) + } + return nil, streaming.ErrRewriteUnsupported + } if r.logger != nil { - r.logger.Warn("SSE buffer redaction unsupported; forwarding original body (fail-open)", + r.logger.Warn("SSE buffer redaction unsupported; forwarding original body (fail-open, agent host-packet path)", "reason", reason, "reasonCode", core.ReasonRedactInflightUnsupported, "metricReason", metricReason, diff --git a/packages/shared/transport/tlsbump/sse_frame_redactor_test.go b/packages/shared/transport/tlsbump/sse_frame_redactor_test.go index a4dbd988..2f214b21 100644 --- a/packages/shared/transport/tlsbump/sse_frame_redactor_test.go +++ b/packages/shared/transport/tlsbump/sse_frame_redactor_test.go @@ -3,6 +3,7 @@ package tlsbump import ( "bytes" "context" + "errors" "fmt" "strings" "sync" @@ -60,7 +61,7 @@ func serialize(e *streaming.SSEEvent) string { // nil adapter → no redactor (buffer keeps its degrade behavior). func TestNewSSEFrameRedactor_NilAdapter(t *testing.T) { - if fr := newSSEFrameRedactor(context.Background(), nil, "/v1/chat/completions", nil); fr != nil { + if fr := newSSEFrameRedactor(context.Background(), nil, "/v1/chat/completions", false, nil); fr != nil { t.Fatalf("expected nil redactor for nil adapter") } } @@ -68,7 +69,7 @@ func TestNewSSEFrameRedactor_NilAdapter(t *testing.T) { // Through the REAL OpenAI adapter: a buffered Modify redacts the text frames. func TestSSEFrameRedactor_OpenAIRedacts(t *testing.T) { ctx := context.Background() - fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", nil) + fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", false, nil) codec := adapterWireCodec{ctx: ctx, adapter: &openai.Adapter{}, path: "/v1/chat/completions"} events := []*streaming.SSEEvent{ @@ -105,7 +106,7 @@ func TestSSEFrameRedactor_OpenAIRedacts(t *testing.T) { func TestSSEFrameRedactor_AnthropicRedacts(t *testing.T) { ctx := context.Background() path := "/v1/messages" - fr := newSSEFrameRedactor(ctx, &anthropic.Adapter{}, path, nil) + fr := newSSEFrameRedactor(ctx, &anthropic.Adapter{}, path, false, nil) codec := adapterWireCodec{ctx: ctx, adapter: &anthropic.Adapter{}, path: path} events := []*streaming.SSEEvent{ @@ -140,7 +141,7 @@ func TestSSEFrameRedactor_ToolArgMaskFailsOpen(t *testing.T) { ctx := context.Background() // OpenAI IS a ToolArgMasker; the streaming path still cannot deliver tool // args through verbatim frames, so it must fail open. - fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", nil) + fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", false, nil) events := []*streaming.SSEEvent{oaFrame("hello"), doneFrame()} before := serialize(events[0]) res := &core.CompliancePipelineResult{ @@ -164,7 +165,7 @@ func TestSSEFrameRedactor_ToolArgMaskFailsOpen(t *testing.T) { // GuardToolArgMasking reports unsupported → forward original + stamp. func TestSSEFrameRedactor_NonMaskerToolArgFailsOpen(t *testing.T) { ctx := context.Background() - fr := newSSEFrameRedactor(ctx, &anthropic.Adapter{}, "/v1/messages", nil) + fr := newSSEFrameRedactor(ctx, &anthropic.Adapter{}, "/v1/messages", false, nil) events := []*streaming.SSEEvent{anthropicFrame("hi"), doneFrame()} res := &core.CompliancePipelineResult{ Decision: core.Modify, @@ -182,7 +183,7 @@ func TestSSEFrameRedactor_NonMaskerToolArgFailsOpen(t *testing.T) { // Fail-OPEN: splice cannot reconstruct (fence mismatch) → forward + stamp. func TestSSEFrameRedactor_SpliceUnsupportedFailsOpen(t *testing.T) { ctx := context.Background() - fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", nil) + fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", false, nil) events := []*streaming.SSEEvent{oaFrame("hello world"), doneFrame()} res := &core.CompliancePipelineResult{ Decision: core.Modify, @@ -203,7 +204,7 @@ func TestSSEFrameRedactor_SpliceUnsupportedFailsOpen(t *testing.T) { // nil result → no-op, no panic. func TestSSEFrameRedactor_NilResult(t *testing.T) { - fr := newSSEFrameRedactor(context.Background(), &openai.Adapter{}, "/v1/chat/completions", nil) + fr := newSSEFrameRedactor(context.Background(), &openai.Adapter{}, "/v1/chat/completions", false, nil) events := []*streaming.SSEEvent{oaFrame("x")} out, err := fr.RedactReplay(events, nil) if err != nil || len(out) != 1 { @@ -258,7 +259,7 @@ func TestSSEFrameRedactor_ConcurrentRedactReplay(t *testing.T) { ModifiedContent: []core.ContentBlock{{Type: "text", Text: "x [P] y"}}, TransformSpans: []normalize.TransformSpan{tspan("messages.0.content.0", 2, 2+len(secret), "[P]")}, } - fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", nil) + fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", false, nil) out, err := fr.RedactReplay(events, res) if err != nil { t.Errorf("goroutine %d error: %v", n, err) @@ -278,6 +279,126 @@ func TestSSEFrameRedactor_ConcurrentRedactReplay(t *testing.T) { } } +// Posture-follows-caller: an unreconstructable masked wire fails CLOSED for the +// strict (compliance-proxy appliance) caller and fails OPEN for the non-strict +// (agent NE host-packet) caller, across BOTH root causes (tool-arg undeliverable +// and splice divergence). Fail-closed returns streaming.ErrRewriteUnsupported +// with no events (the buffer pipeline then emits the error frame, replaying zero +// original byte); fail-open returns the ORIGINAL events with a nil error. Both +// postures always stamp REDACT_INFLIGHT_UNSUPPORTED for audit honesty. +func TestSSEFrameRedactor_PostureFollowsCaller(t *testing.T) { + ctx := context.Background() + // A tool-arg span forces the tool-arg-undeliverable cause; a fence-mismatch + // ModifiedContent forces the splice-divergence cause. + toolArgSpans := []normalize.TransformSpan{tspan("messages.0.content.0.toolUse.input.0", 0, 1, "x")} + spliceSpans := []normalize.TransformSpan{tspan("messages.0.content.0", 0, 5, "BYE")} + + cases := []struct { + name string + strict bool + spans []normalize.TransformSpan + modText string + }{ + {"toolArg/strict_fail_closed", true, toolArgSpans, "hello"}, + {"toolArg/nonstrict_fail_open", false, toolArgSpans, "hello"}, + {"splice/strict_fail_closed", true, spliceSpans, "TOTALLY DIFFERENT"}, + {"splice/nonstrict_fail_open", false, spliceSpans, "TOTALLY DIFFERENT"}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", tc.strict, nil) + events := []*streaming.SSEEvent{oaFrame("hello world"), doneFrame()} + before := serialize(events[0]) + res := &core.CompliancePipelineResult{ + Decision: core.Modify, + ModifiedContent: []core.ContentBlock{{Type: "text", Text: tc.modText}}, + TransformSpans: tc.spans, + } + + out, err := fr.RedactReplay(events, res) + + // Audit honesty: both postures stamp the unsupported reason. + if res.ReasonCode != core.ReasonRedactInflightUnsupported { + t.Fatalf("expected REDACT_INFLIGHT_UNSUPPORTED stamp, got %q", res.ReasonCode) + } + + if tc.strict { + // FAIL-CLOSED: ErrRewriteUnsupported, no events. The buffer + // pipeline turns this into an error frame with zero original + // replay — no unredacted byte can reach the client. + if !errors.Is(err, streaming.ErrRewriteUnsupported) { + t.Fatalf("strict caller must return ErrRewriteUnsupported, got %v", err) + } + if out != nil { + t.Fatalf("strict fail-closed must return no events (no original forwarded), got %d", len(out)) + } + return + } + + // FAIL-OPEN: nil error, ORIGINAL events forwarded unchanged. + if err != nil { + t.Fatalf("non-strict caller must fail open with nil error, got %v", err) + } + if len(out) == 0 || serialize(out[0]) != before { + t.Fatalf("fail-open must forward the original frame unchanged") + } + }) + } +} + +// Strict fail-closed end-to-end through the buffer pipeline: a strict redactor +// whose splice cannot reconstruct makes BufferPipeline.Process emit the policy +// error frame and replay ZERO original byte — the client never sees the +// unredacted PII. This pins the leak-prevention contract that the strict +// posture exists for. +func TestSSEFrameRedactor_StrictFailClosedThroughBuffer(t *testing.T) { + ctx := context.Background() + pii := "ssn 123-45-6789" + exec := redactExecutor{result: &core.CompliancePipelineResult{ + Decision: core.Modify, + // Fence mismatch: ModifiedContent text does not match the wire + // transcript, forcing splice divergence → unreconstructable (same shape + // the splice fail-open unit test uses, here driven end-to-end). + ModifiedContent: []core.ContentBlock{{Type: "text", Text: "TOTALLY DIFFERENT"}}, + TransformSpans: []normalize.TransformSpan{tspan("messages.0.content.0", 0, 5, "BYE")}, + }} + fr := newSSEFrameRedactor(ctx, &openai.Adapter{}, "/v1/chat/completions", true, nil) + + pipeline := streaming.NewBufferPipeline( + streaming.BufferConfig{MaxBufferSize: 1 << 20}, exec, nil, + ).WithFrameRedactor(fr) + + body := strings.NewReader( + "data: " + oaFrame(pii).Data + "\n\n" + + "data: [DONE]\n\n", + ) + var client bytes.Buffer + res, err := pipeline.Process(ctx, body, &client, &core.HookInput{Path: "/v1/chat/completions"}) + if err != nil { + t.Fatalf("buffer pipeline error: %v", err) + } + if res.ReasonCode != core.ReasonRedactInflightUnsupported { + t.Fatalf("strict fail-closed must stamp REDACT_INFLIGHT_UNSUPPORTED, got %q", res.ReasonCode) + } + if strings.Contains(client.String(), "123-45") || strings.Contains(client.String(), "6789") { + t.Fatalf("strict fail-closed leaked unredacted PII to client: %q", client.String()) + } + if !strings.Contains(client.String(), "error") { + t.Fatalf("strict fail-closed must emit an error frame, got %q", client.String()) + } +} + +// redactExecutor is a streaming.PipelineExecutor test double that returns a +// fixed Modify result so the buffer pipeline routes into the frame redactor. +type redactExecutor struct { + result *core.CompliancePipelineResult +} + +func (e redactExecutor) Execute(_ context.Context, _ *core.HookInput) *core.CompliancePipelineResult { + return e.result +} + // Guard: the OpenAI adapter satisfies traffic.ToolArgMasker (anchors the // tool-arg fail-open reasoning). func TestOpenAIIsToolArgMasker(t *testing.T) { diff --git a/packages/shared/transport/tlsbump/sse_modela.go b/packages/shared/transport/tlsbump/sse_modela.go new file mode 100644 index 00000000..03f11b85 --- /dev/null +++ b/packages/shared/transport/tlsbump/sse_modela.go @@ -0,0 +1,488 @@ +package tlsbump + +// sse_modela.go — the wire substrate that drives the shared streaming Model-A +// engine (shared/transport/streaming/modela) over raw SSE frames for the tlsbump +// transparent proxy (agent + compliance-proxy). The engine owns the prescan-gated +// tail-hold algorithm; this adapter supplies the wire seams: parsing frames off the +// upstream, extracting per-frame visible text via the matched adapter, delivering +// frames in real time, confirming via the response pipeline, and escalating to a +// buffer-the-remainder redaction on the frame timeline. +// +// Fail posture is fail-OPEN (the agent NE host-packet path must never fail closed): +// the frame redactor relays the original frames + stamps REDACT_INFLIGHT_UNSUPPORTED +// when the masked wire cannot be soundly reconstructed (see sseFrameRedactor). A +// hard-BLOCK decision is an admin-intended enforcement (not a machinery error), so +// it still writes an in-band error frame — fail-open governs redaction faults, not +// deliberate blocks. + +import ( + "context" + "errors" + "io" + "log/slog" + "net/http" + "strings" + "time" + + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" + compliance "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/pipeline" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/streaming" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/streaming/modela" +) + +// overrideStreamingModeForScope applies the scope-derived routing override (mirror +// ai-gateway stream_shape.go §B2). An enforcing response scope overrides the admin +// streaming mode because the live path is audit-only and cannot enforce: MayBlock +// forces "buffer" (zero-leak hard block); MayRedact forces "buffer" unless the flow +// is chunked_async ("live") with a per-host adapter, where it arms Model A +// ("modela") — prescan-gated streaming + escalate-to-buffer redaction. A redact +// under raw "passthrough" cannot be applied on the wire, so it buffers. Non-enforcing +// scopes pass through unchanged. +func overrideStreamingModeForScope(mode string, mayBlock, mayRedact, hasAdapter bool) string { + switch { + case mayBlock: + return "buffer" + case mayRedact: + switch mode { + case "live": + if hasAdapter { + return "modela" + } + return "buffer" + case "passthrough": + return "buffer" + } + } + return mode +} + +// scopeRouteSSEMode builds the response probe once and applies the scope-derived +// streaming-mode override. Returns the (possibly overridden) mode and the probe +// pipeline (non-nil only when it can be reused as the Model A prescan/confirm +// pipeline). A nil respInput, an unbuildable fail-closed probe (non-strict caller), +// or no response hooks leaves the live/passthrough flow routed conservatively. +func scopeRouteSSEMode(bo *bumpOptions, mode string, respInput *core.HookInput, audCtx *requestAuditCtx, logger *slog.Logger) (string, *compliance.Pipeline) { + if respInput == nil { + return mode, nil + } + probe, pErr := bo.policyResolver.BuildPipeline( + "response", "COMPLIANCE_PROXY", + "", nil, + bo.perHookTimeout, bo.totalTimeout, bo.parallelHooks, + bo.strictFailClosed, + logger, + ) + switch { + case pErr != nil: + // A fail-closed response hook is unbuildable. Strict (appliance) callers are + // already refused by the SSE-entry guard (clean 451); a non-strict (agent) + // caller forces BUFFER so enforcing traffic never streams uninspected on the + // audit-only live path. + if mode == "live" || mode == "passthrough" { + return "buffer", nil + } + return mode, nil + case probe != nil: + hasAdapter := audCtx != nil && audCtx.adapter != nil + return overrideStreamingModeForScope(mode, probe.MayBlock(), probe.MayRedact(), hasAdapter), probe + } + return mode, nil +} + +// runSSEModelA drives the wire Model-A path: it constructs the wire substrate over +// the SSE body and runs the shared engine, then emits the audit row with the +// response-stage compliance result, the captured (best-effort) wire body, and the +// finalized usage. probe and audCtx.adapter are non-nil (the routing guarantees it). +func runSSEModelA( + ctx context.Context, + w http.ResponseWriter, + resp *http.Response, + audCtx *requestAuditCtx, + respInput *core.HookInput, + auditInfo *compliance.AuditInfo, + bo *bumpOptions, + logger *slog.Logger, + requestStart time.Time, + probe *compliance.Pipeline, + acc streaming.UsageAccumulator, + captureMax int, +) { + probe.SetAllowModify(true) + probe.SetClearSoftOnApprove(true) + + var captureBuf *streaming.CappedBuffer + var dest io.Writer = w + if captureMax > 0 { + captureBuf = streaming.NewCappedBuffer(captureMax) + dest = io.MultiWriter(w, captureBuf) + } + flusher, canFlush := w.(http.Flusher) + sseAdapter := audCtx.adapter + ssePath := "" + if respInput != nil { + ssePath = respInput.Path + } + maxBuf := bo.streamingConfig.MaxBufferSize + if maxBuf <= 0 { + maxBuf = modelaDefaultMaxBufferBytes + } + sub := &sseWireSubstrate{ + ctx: ctx, + parser: streaming.NewSSEParserWithLogger(resp.Body, logger), + client: dest, + flusher: flusher, + canFlush: canFlush, + codec: adapterWireCodec{ctx: ctx, adapter: sseAdapter, path: ssePath}, + redactor: newSSEFrameRedactor(ctx, sseAdapter, ssePath, bo.strictFailClosed, logger), + pipeline: probe, + base: respInput, + acc: acc, + maxBuf: maxBuf, + logger: logger, + } + maxPattern := deriveModelAMaxPattern(probe) + // Config-time operator signal (#16): tlsbump leaves TailWindowBytes at the engine + // default, so compare the derived bound against that default. Off the per-byte path. + modela.WarnStreamingCoverageGap(logger, maxPattern, modela.DefaultTailWindowBytes) + if err := modela.Run(ctx, sub, modela.Config{MaxBufferBytes: maxBuf, MaxPatternBytes: maxPattern}); err != nil { + target := "" + if respInput != nil { + target = respInput.TargetHost + } + logger.Error("SSE Model-A pipeline error", + "target", target, + "error", err, + "cancel_cause", cancelCause(ctx), + "duration_ms", int(time.Since(requestStart).Milliseconds()), + ) + } + // Storage parity (T3): persist the wire capture only on a non-enforcing outcome — + // it is then the benign, fully-delivered original. On a redaction/block escalation + // the wire capture holds the real-time prefix raw (a sub-window value is fully + // masked; a value larger than the tail window already delivered a bounded raw + // prefix to the client — the disclosed best-effort WIRE surface). The durable, + // queryable copy is held to a STRICTER standard than the ephemeral wire: an + // enforcing outcome never persists a raw prefix. + // + // This is defense-in-depth over the canonical gate: emitAudit → StorageRawBody + // already drops the body to NULL unless the response ACTION is approve. The guard + // here keys on the DECISION (via modelaResultEnforcing → ActionFromDecision) — the + // same judgment the engine's escalate gate uses — so a Decision/Action divergence + // (a redact Decision masked behind an approve-shaped Action) can never persist the + // raw capture. (Storing an off-path FULLY-redacted authoritative copy instead of + // NULL — a complete record strictly stronger than the wire — is tracked + // enhancement work; NULL is the safe floor.) + var capturedBytes []byte + if captureBuf != nil && !modelaResultEnforcing(sub.result) { + capturedBytes = captureBuf.Bytes() + } + emitAudit(logger, audCtx, respInput, auditInfo, bo, sub.result, resp.StatusCode, requestStart, finalizeUsage(ctx, acc), capturedBytes) +} + +// modelaResultEnforcing reports whether the Model A outcome blocked or redacted — +// the cases where the best-effort wire capture must NOT be persisted (it may carry a +// raw over-window prefix). A nil/approve result is non-enforcing. +func modelaResultEnforcing(res *core.CompliancePipelineResult) bool { + if res == nil { + return false + } + act := core.ActionFromDecision(res.Decision) + return act == core.ActionBlock || act == core.ActionRedact +} + +// modelaDefaultMaxBufferBytes mirrors the shared engine's MaxBufferBytes default +// (8 MB) so the escalation drain ceiling agrees with the main-loop ceiling when the +// admin streaming policy leaves MaxBufferSize unset. +const modelaDefaultMaxBufferBytes = 8 * 1024 * 1024 + +// wireUnit is the engine's unit for the wire path: one parsed SSE frame plus its +// visible text, extracted ONCE in Next so the per-frame prescan/window seams never +// re-parse the frame. +type wireUnit struct { + evt *streaming.SSEEvent + text string +} + +// sseWireSubstrate adapts a raw SSE stream to the shared Model-A engine. It holds +// the parser (source + escalation drain), the client writer (+ flusher), the +// per-host text codec + frame redactor, the response pipeline (prescan + confirm + +// escalation re-eval), and the per-frame audit accumulators. +type sseWireSubstrate struct { + ctx context.Context + parser *streaming.SSEParser + client io.Writer + flusher http.Flusher + canFlush bool + + codec streaming.WireTextCodec + redactor *sseFrameRedactor // nil when no adapter → redaction degrades (fail-open) + pipeline *compliance.Pipeline + maxBuf int // escalation drain ceiling; <=0 disables the cap + + base *core.HookInput + acc streaming.UsageAccumulator + + logger *slog.Logger + + // result is the authoritative response-stage compliance result the emit site + // records: the escalation re-eval outcome on a confirmed/pressure escalation, the + // last false-positive confirm otherwise, or an Approve on a clean stream. + result *core.CompliancePipelineResult +} + +var _ modela.Substrate[wireUnit] = (*sseWireSubstrate)(nil) + +// Next parses the next SSE frame, feeds the usage accumulator, and extracts the +// frame's visible text once. io.EOF ends the stream; any other error is terminal. +func (s *sseWireSubstrate) Next(_ context.Context) (wireUnit, error) { + evt, err := s.parser.Next() + if err != nil { + return wireUnit{}, err + } + if s.acc != nil { + s.acc.Feed(evt) + } + text, _ := s.codec.ChunkText(evt.Data) + return wireUnit{evt: evt, text: text}, nil +} + +// AppendRedactableText appends the frame's visible text (extracted in Next) for the +// prescan/confirm. Non-text frames (tool_use / ping / [DONE] / reasoning-only) +// contributed empty text and add nothing. +func (s *sseWireSubstrate) AppendRedactableText(dst []byte, u wireUnit) []byte { + return append(dst, u.text...) +} + +// UnitBytes is the frame's transport size (held-bytes ceiling budget). +func (s *sseWireSubstrate) UnitBytes(u wireUnit) int { return len(u.evt.Data) } + +// ContentBytes is the frame's redactable-content size (tail-window budget) — the +// extracted visible text, never the framing, so non-text frames do not evict +// redactable content from the window. +func (s *sseWireSubstrate) ContentBytes(u wireUnit) int { return len(u.text) } + +// IsDone reports the SSE [DONE] terminator frame. +func (s *sseWireSubstrate) IsDone(u wireUnit) bool { return u.evt.Done } + +// Deliver writes one held frame to the client in real time. A write error is +// terminal; the engine stops and the emit site records the partial outcome. +func (s *sseWireSubstrate) Deliver(_ context.Context, u wireUnit) error { + if err := s.writeFrame(u.evt); err != nil { + return err + } + return nil +} + +// DeliverTerminal is a no-op for SSE: the terminator ([DONE]) is itself a frame +// delivered through Deliver as the last held unit, so there is no separate terminal +// framing to emit. +func (s *sseWireSubstrate) DeliverTerminal(_ context.Context) error { return nil } + +// Prescan is the cheap union prefilter over accumulated visible text. A nil +// pipeline (no response hooks) fails safe to "always confirm". +func (s *sseWireSubstrate) Prescan(content []byte) bool { + if s.pipeline == nil { + return true + } + return s.pipeline.MayMatchRawContent(content) +} + +// Confirm runs ONE full response-stage evaluation over the accumulated visible +// text. A nil pipeline (no hooks) is treated as approve; the engine resumes. +func (s *sseWireSubstrate) Confirm(ctx context.Context, content string) *core.CompliancePipelineResult { + if s.pipeline == nil { + return &core.CompliancePipelineResult{Decision: core.Approve, Action: core.ActionApprove} + } + return s.pipeline.Execute(ctx, s.checkpointInput(content)) +} + +// Escalate hands the live path off to a buffer-to-end redaction over the frame +// timeline. It drains the remaining frames into the held set, RE-EVALUATES the +// remainder text on its own (so the redact spans map onto the still-held frames, not +// the already-delivered prefix — the divergence fence in SpliceTextFrames requires +// the transcript to equal the evaluated content), and delivers ONLY the redacted +// remainder. res is advisory (confirmed-hit vs memory-pressure); the re-eval is +// authoritative. Redaction-fault posture follows the caller (strictFailClosed): the +// agent NE path is fail-OPEN (an unproducible mask relays the original frames + +// stamps the degraded reason); the compliance-proxy appliance is fail-CLOSED (an +// unproducible mask emits an in-band block frame). A deliberate hard block always +// emits an in-band error frame regardless of posture. +// deriveModelAMaxPattern sizes the Model-A flush-before-deliver lookahead from the +// resolved response pipeline's longest contiguous enforceable match, floored at +// modela.DefaultMaxPatternBytes so it never drops below the proven-safe baseline even if +// the derivation under-counts. Two limits stay best-effort (the engine's disclosed +// surface, not this floor): (1) UNBOUNDED patterns (*, +, {n,}) whose match can exceed the +// tail window; (2) a derived bound at/above the tail window — the engine clamps the +// lookahead below the window. For (2) the caller surfaces a config-time operator warning +// (modela.WarnStreamingCoverageGap); the operator remediation is buffered streaming mode +// (full coverage) or narrowing the rule, NOT a tail-window knob. +func deriveModelAMaxPattern(probe *compliance.Pipeline) int { + bound, _ := probe.MaxPatternBound() + if bound < modela.DefaultMaxPatternBytes { + bound = modela.DefaultMaxPatternBytes + } + return bound +} + +func (s *sseWireSubstrate) Escalate(ctx context.Context, held []wireUnit, trigger *core.CompliancePipelineResult) error { + // Record the escalation CAUSE as a metric dimension (NOT a persisted field): a nil + // trigger means a memory-pressure eviction of an incomplete content unit; a non-nil + // trigger means a confirmed enforcing hit. Lets operators tell a real policy hit apart + // from a buffer-ceiling eviction (which signals MaxBufferBytes may need raising). + cause := streaming.ModelAEscalationConfirmed + if trigger == nil { + cause = streaming.ModelAEscalationMemoryPressure + } + streaming.RecordModelAEscalation(cause) + + // Drain the remainder of the stream into the held set (the prefix beyond the tail + // window is already delivered and excluded from held). The drain is capped at + // maxBuf so an arbitrarily long post-hit tail cannot grow the buffer without + // bound (the engine's MaxBufferBytes ceiling, mirrored on the drain per the + // engine doc + the ai-gateway sibling). On overflow the enforcing remainder + // cannot be buffered to redact, so it is BLOCKED with an in-band error frame — + // never relayed raw (bounded + zero-leak, the same outcome as a hard block). + heldBytes := 0 + for i := range held { + heldBytes += len(held[i].evt.Data) + } + for len(held) == 0 || !held[len(held)-1].evt.Done { + evt, err := s.parser.Next() + if err != nil { + if errors.Is(err, io.EOF) { + break + } + return err + } + if s.acc != nil { + s.acc.Feed(evt) + } + text, _ := s.codec.ChunkText(evt.Data) + held = append(held, wireUnit{evt: evt, text: text}) + heldBytes += len(evt.Data) + if s.maxBuf > 0 && heldBytes > s.maxBuf { + s.result = &core.CompliancePipelineResult{Decision: core.RejectHard, Action: core.ActionBlock, Reason: "stream exceeded maximum buffer size", ReasonCode: "stream_buffer_exceeded"} + return s.writeErrorAndDone() + } + } + + events := make([]*streaming.SSEEvent, len(held)) + var remainder strings.Builder + for i, u := range held { + events[i] = u.evt + remainder.WriteString(u.text) + } + + if s.pipeline == nil { + // No response pipeline — nothing to enforce; deliver the held frames as-is. + s.result = &core.CompliancePipelineResult{Decision: core.Approve, Action: core.ActionApprove} + return s.writeFrames(events) + } + + res := s.pipeline.Execute(ctx, s.checkpointInput(remainder.String())) + s.result = res + + switch res.Decision { + case core.RejectHard: + // Deliberate hard block — emit an in-band error frame (zero remainder + // content). Fail-open governs redaction machinery faults, not admin blocks. + return s.writeErrorAndDone() + default: + // Modify (redact) or any decision carrying redaction work: splice the masked + // text into the held text frames. The per-host redactor's posture follows the + // caller (strictFailClosed): the agent NE host-packet path is fail-OPEN — on an + // unproducible mask it relays the original frames + stamps the degraded reason + // and returns nil; the compliance-proxy appliance is fail-CLOSED — it returns + // ErrRewriteUnsupported, which we surface as an in-band block frame rather than + // relay the unredacted remainder. + out := events + if s.redactor != nil { + redacted, err := s.redactor.RedactReplay(events, res) + if err != nil { + return s.writeErrorAndDone() + } + out = redacted + // Disposition: stamp action=redact ONLY when the mask was genuinely spliced — + // NOT the agent fail-open relay-original degrade (RedactReplay returns the + // original frames with a nil error after stamping ReasonRedactInflightUnsupported). + // On a real splice the masked frames WERE delivered, so the audit row reads + // redact even under a co-firing BlockSoft ceiling; on the fail-open degrade the + // original was relayed, so the action stays the aggregate. + if res.ReasonCode != core.ReasonRedactInflightUnsupported { + res.Action = core.ActionRedact + } + } + return s.writeFrames(out) + } +} + +// OnConfirmApproved records a non-enforcing confirm outcome (false positive) so the +// emit site reports an evaluated-approved response. +func (s *sseWireSubstrate) OnConfirmApproved(res *core.CompliancePipelineResult) { + s.result = res +} + +// OnApproveEOF records an evaluated-approved outcome when the stream reached EOF +// without any confirm — the sound prescan cleared the whole stream. +func (s *sseWireSubstrate) OnApproveEOF() { + s.result = &core.CompliancePipelineResult{Decision: core.Approve, Action: core.ActionApprove} +} + +// OnError ends the stream on a non-EOF parser error. The agent host-packet path is +// fail-open: no synthetic error frame is forced onto the wire (the upstream fault is +// recorded for audit and the partial stream is left as delivered). +func (s *sseWireSubstrate) OnError(_ context.Context, err error) error { + if s.logger != nil { + s.logger.Error("SSE Model-A wire substrate upstream error", "error", err) + } + return err +} + +// checkpointInput builds a response HookInput carrying the accumulated/visible text +// as the single content block — the same flat-text shape the per-frame codec +// reconstructs, so a Modify's ModifiedContent aligns with SpliceTextFrames' per-frame +// transcript (the divergence fence holds). +func (s *sseWireSubstrate) checkpointInput(content string) *core.HookInput { + in := &core.HookInput{ + Stage: "response", + Normalized: core.PayloadFromTextSegments([]string{content}), + IngressType: s.base.IngressType, + TargetHost: s.base.TargetHost, + Method: s.base.Method, + Path: s.base.Path, + ContentType: s.base.ContentType, + SourceIP: s.base.SourceIP, + } + return in +} + +// writeFrame writes one SSE frame to the client (teed into the audit capture buffer +// when enabled) and flushes so the client sees it in real time. +func (s *sseWireSubstrate) writeFrame(evt *streaming.SSEEvent) error { + if err := streaming.WriteSSEEvent(s.client, evt); err != nil { + return err + } + if s.canFlush { + s.flusher.Flush() + } + return nil +} + +// writeFrames writes a sequence of frames (the redacted/approved remainder). +func (s *sseWireSubstrate) writeFrames(events []*streaming.SSEEvent) error { + for _, evt := range events { + if err := s.writeFrame(evt); err != nil { + return err + } + } + return nil +} + +// writeErrorAndDone emits the in-band block error frame + the [DONE] terminator, +// mirroring the buffer pipeline's hard-block delivery. +func (s *sseWireSubstrate) writeErrorAndDone() error { + if err := s.writeFrame(&streaming.SSEEvent{Event: "message", Data: `{"error": "blocked by policy"}`, Retry: -1}); err != nil { + return err + } + return s.writeFrame(&streaming.SSEEvent{Event: "message", Data: "[DONE]", Done: true, Retry: -1}) +} diff --git a/packages/shared/transport/tlsbump/sse_modela_test.go b/packages/shared/transport/tlsbump/sse_modela_test.go new file mode 100644 index 00000000..56c92a4b --- /dev/null +++ b/packages/shared/transport/tlsbump/sse_modela_test.go @@ -0,0 +1,527 @@ +package tlsbump + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" + compliance "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/pipeline" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/traffic/adapters/api/openai" + normalize "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/normalize/core" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/streaming" + streampolicy "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/streaming/policy" +) + +// TestOverrideStreamingModeForScope exhaustively pins the scope-routing decision +// table (mirror ai-gateway stream_shape.go §B2). +func TestOverrideStreamingModeForScope(t *testing.T) { + cases := []struct { + name string + mode string + mayBlock, mayRedact, adapter bool + want string + }{ + {"block overrides live", "live", true, false, true, "buffer"}, + {"block overrides passthrough", "passthrough", true, false, true, "buffer"}, + {"block beats redact", "live", true, true, true, "buffer"}, + {"redact live + adapter → modela", "live", false, true, true, "modela"}, + {"redact live no adapter → buffer", "live", false, true, false, "buffer"}, + {"redact passthrough → buffer", "passthrough", false, true, true, "buffer"}, + {"redact already buffer stays buffer", "buffer", false, true, true, "buffer"}, + {"non-enforcing live unchanged", "live", false, false, true, "live"}, + {"non-enforcing passthrough unchanged", "passthrough", false, false, true, "passthrough"}, + {"non-enforcing buffer unchanged", "buffer", false, false, true, "buffer"}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if got := overrideStreamingModeForScope(c.mode, c.mayBlock, c.mayRedact, c.adapter); got != c.want { + t.Fatalf("override(%q, block=%v, redact=%v, adapter=%v) = %q, want %q", + c.mode, c.mayBlock, c.mayRedact, c.adapter, got, c.want) + } + }) + } +} + +// TestModelaResultEnforcing pins the storage gate: only block/redact outcomes are +// enforcing (their best-effort wire capture must NOT be persisted — it may hold a raw +// over-window prefix); approve / abstain / nil are non-enforcing (the benign original +// is safe to persist). +func TestModelaResultEnforcing(t *testing.T) { + cases := []struct { + name string + res *core.CompliancePipelineResult + want bool + }{ + {"nil → not enforcing", nil, false}, + {"approve → not enforcing", &core.CompliancePipelineResult{Decision: core.Approve}, false}, + {"abstain → not enforcing", &core.CompliancePipelineResult{Decision: core.Abstain}, false}, + {"reject-hard → enforcing", &core.CompliancePipelineResult{Decision: core.RejectHard}, true}, + {"modify → enforcing", &core.CompliancePipelineResult{Decision: core.Modify}, true}, + {"block-soft → enforcing", &core.CompliancePipelineResult{Decision: core.BlockSoft}, true}, + } + for _, c := range cases { + t.Run(c.name, func(t *testing.T) { + if got := modelaResultEnforcing(c.res); got != c.want { + t.Fatalf("modelaResultEnforcing(%+v) = %v, want %v", c.res, got, c.want) + } + }) + } +} + +// lastResponseBodyPayload returns the persisted response-body column payload of the +// most recent audit event (empty ⇒ stored NULL). +func lastResponseBodyPayload(t *testing.T, w *recordingAuditWriter) []byte { + t.Helper() + evs := w.snapshot() + if len(evs) == 0 { + t.Fatal("no audit event emitted") + } + payload, _ := evs[len(evs)-1].ResponseBody.ColumnPayload() + return payload +} + +// TestSSE_ModelA_Storage_CleanPersists_EnforcingNull is the T3 storage-parity +// end-to-end guard (asserts the PERSISTED body, not just the wire): an enforcing +// (redact) modela stream stores NULL — never the raw best-effort prefix — while a +// clean (prescan-miss, approve-at-EOF) stream persists the delivered original. +func TestSSE_ModelA_Storage_CleanPersists_EnforcingNull(t *testing.T) { + // Enforcing redact → NULL. + redact := modelATestHook{ + decision: core.Modify, + modified: []core.ContentBlock{{Type: "text", Text: "card [CARD]"}}, + spans: []normalize.TransformSpan{tspan("messages.0.content.0", 5, 21, "[CARD]")}, + matchRaw: true, + } + bo, w := modelABumpOptions(modelARedactResolver(redact)) + respInput, auditInfo, audCtx := modelAAudCtx() + handleSSEResponse(context.Background(), httptest.NewRecorder(), sseFrameUpstream(oaContentBody("card ", "4111111111111111")), audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + if body := lastResponseBodyPayload(t, w); len(body) != 0 { + t.Fatalf("enforcing modela stream must persist a NULL response body, got %d bytes: %q", len(body), body) + } + + // Clean prescan-miss → approve at EOF → the delivered original is persisted. + clean := modelATestHook{decision: core.Modify, matchRaw: false} + bo2, w2 := modelABumpOptions(modelARedactResolver(clean)) + respInput2, auditInfo2, audCtx2 := modelAAudCtx() + handleSSEResponse(context.Background(), httptest.NewRecorder(), sseFrameUpstream(oaContentBody("hello ", "world")), audCtx2, respInput2, auditInfo2, bo2, discardSlog(), time.Now()) + if body := lastResponseBodyPayload(t, w2); len(body) == 0 { + t.Fatal("clean modela stream must persist the delivered original, got NULL") + } +} + +// modelATestHook is a configurable response hook for the wire Model-A path: its +// declarative onMatch is "redact" (so MayRedact routes the flow to modela), it +// returns a fixed runtime decision, and it acts as a RawContentPrescanner whose +// MayMatchRaw verdict the test controls (to drive prescan HIT vs MISS). +type modelATestHook struct { + decision core.Decision + modified []core.ContentBlock + spans []normalize.TransformSpan + matchRaw bool +} + +func (h modelATestHook) Execute(_ context.Context, _ *core.HookInput) (*core.HookResult, error) { + return &core.HookResult{ + Decision: h.decision, + Reason: "redact", + ReasonCode: "PII", + ModifiedContent: h.modified, + TransformSpans: h.spans, + }, nil +} +func (modelATestHook) SupportsEndpoint(core.EndpointType) bool { return true } +func (modelATestHook) SupportsModality(core.Modality) bool { return true } +func (modelATestHook) ScansContent() bool { return true } +func (h modelATestHook) MayMatchRaw([]byte) bool { return h.matchRaw } + +// modelARedactResolver wires a single response-stage hook whose onMatch is "redact" +// (→ MayRedact true → modela routing) and whose runtime behaviour is h. +func modelARedactResolver(h modelATestHook) *compliance.PolicyResolver { + reg := core.NewHookRegistry() + reg.Register("modela-test", func(_ *core.HookConfig) (core.Hook, error) { return h, nil }) + return compliance.NewPolicyResolver([]core.HookConfig{{ + ID: "h-modela", + ImplementationID: "modela-test", + Name: "modela-test", + Stage: "response", + Enabled: true, + FailBehavior: "fail-open", + ApplicableIngress: []string{"ALL"}, + Config: map[string]any{"onMatch": map[string]any{"action": "redact"}}, + }}, reg, discardSlog()) +} + +// modelACoFiringResolver wires TWO response-stage hooks, BOTH with onMatch "redact" (so the +// static scope is MayRedact && !MayBlock → modela routing, NOT buffer): a redact hook that +// returns Modify + ModifiedContent/spans, and a soft-block hook that returns BlockSoft. The +// pipeline aggregator ranks BlockSoft above Modify (Decision→BlockSoft) while carrying the +// redact's ModifiedContent/spans — the genuine co-firing shape leak #2 needs. (A soft-block +// hook declaring onMatch "block" would set MayBlock and route to buffer, never reaching the +// modela wire, so both hooks MUST be redact-scoped.) +func modelACoFiringResolver(redactHook, softBlockHook modelATestHook) *compliance.PolicyResolver { + reg := core.NewHookRegistry() + reg.Register("modela-redact", func(_ *core.HookConfig) (core.Hook, error) { return redactHook, nil }) + reg.Register("modela-softblock", func(_ *core.HookConfig) (core.Hook, error) { return softBlockHook, nil }) + redactCfg := map[string]any{"onMatch": map[string]any{"action": "redact"}} + return compliance.NewPolicyResolver([]core.HookConfig{ + {ID: "h-redact", ImplementationID: "modela-redact", Name: "modela-redact", Stage: "response", Enabled: true, FailBehavior: "fail-open", ApplicableIngress: []string{"ALL"}, Config: redactCfg}, + {ID: "h-softblock", ImplementationID: "modela-softblock", Name: "modela-softblock", Stage: "response", Enabled: true, FailBehavior: "fail-open", ApplicableIngress: []string{"ALL"}, Config: redactCfg}, + }, reg, discardSlog()) +} + +func modelABumpOptions(resolver *compliance.PolicyResolver) (*bumpOptions, *recordingAuditWriter) { + w := &recordingAuditWriter{} + store := streampolicy.NewStore(streampolicy.Policy{ + Mode: streampolicy.ModeChunkedAsync, ChunkBytes: 1024, + HookTimeoutMs: 1000, MaxBufferBytes: 1 << 20, FailBehavior: streampolicy.FailOpen, + }) + return &bumpOptions{ + policyResolver: resolver, + streamingPolicyStore: store, + auditEmitter: compliance.NewAuditEmitter(w, discardSlog()), + streamingConfig: streaming.LiveConfig{MaxBufferSize: 1 << 20}, + }, w +} + +func modelAAudCtx() (*core.HookInput, *compliance.AuditInfo, *requestAuditCtx) { + respInput := &core.HookInput{Stage: "response", TargetHost: "api.example.com", Path: "/v1/chat/completions", IngressType: "COMPLIANCE_PROXY"} + auditInfo := &compliance.AuditInfo{TransactionID: "tx-modela"} + audCtx := &requestAuditCtx{input: respInput, info: *auditInfo, adapter: &openai.Adapter{}, storeResponseBody: true} + return respInput, auditInfo, audCtx +} + +func oaContentBody(parts ...string) string { + var sb strings.Builder + for _, p := range parts { + sb.WriteString(`data: {"choices":[{"delta":{"content":"`) + sb.WriteString(p) + sb.WriteString(`"}}]}` + "\n\n") + } + sb.WriteString("data: [DONE]\n\n") + return sb.String() +} + +// TestSSE_ModelA_ConfirmedRedact_NoLeak is the load-bearing wire-path guard: a +// redact-scope chunked_async flow routes to Model A; a confirmed hit escalates and +// the complete card number never reaches the client — the masked value is spliced. +func TestSSE_ModelA_ConfirmedRedact_NoLeak(t *testing.T) { + hook := modelATestHook{ + decision: core.Modify, + modified: []core.ContentBlock{{Type: "text", Text: "card [CARD]"}}, + spans: []normalize.TransformSpan{tspan("messages.0.content.0", 5, 21, "[CARD]")}, + matchRaw: true, + } + bo, _ := modelABumpOptions(modelARedactResolver(hook)) + respInput, auditInfo, audCtx := modelAAudCtx() + body := oaContentBody("card ", "4111111111111111") + rec := httptest.NewRecorder() + + handleSSEResponse(context.Background(), rec, sseFrameUpstream(body), audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + + out := rec.Body.String() + if strings.Contains(out, "4111111111111111") { + t.Fatalf("Model A leaked the complete card number raw: %q", out) + } + if !strings.Contains(out, "[CARD]") { + t.Fatalf("Model A escalation did not splice the mask onto the wire: %q", out) + } +} + +// TestSSE_ModelA_BlockSoftMaskedRedact_NoLeak is leak #2's missing WIRE evidence: a redact +// hook co-firing with a soft-block hook (BOTH redact-scoped, so the flow routes to modela) +// aggregates to Decision=BlockSoft while carrying the redact's ModifiedContent/spans. Before +// #13 the merge dropped ModifiedContent under BlockSoft, so the wire splice fence +// (maskedMatchesModified) failed → redactUnsupported → the agent fail-open relayed the +// ORIGINAL held tail. Now mergeResults carries the content and the engine escalates on the +// enforcing ACTION (ActionFromDecision(BlockSoft)=block), so SpliceTextFrames masks the wire: +// the complete card number never reaches the client. (The ai-gateway RedactDelivers test +// covers the CANONICAL locus; this pins the raw-SSE splice end-to-end.) +func TestSSE_ModelA_BlockSoftMaskedRedact_NoLeak(t *testing.T) { + redact := modelATestHook{ + decision: core.Modify, + modified: []core.ContentBlock{{Type: "text", Text: "card [CARD]"}}, + spans: []normalize.TransformSpan{tspan("messages.0.content.0", 5, 21, "[CARD]")}, + matchRaw: true, + } + softBlock := modelATestHook{decision: core.BlockSoft, matchRaw: true} + bo, _ := modelABumpOptions(modelACoFiringResolver(redact, softBlock)) + respInput, auditInfo, audCtx := modelAAudCtx() + body := oaContentBody("card ", "4111111111111111") + rec := httptest.NewRecorder() + + handleSSEResponse(context.Background(), rec, sseFrameUpstream(body), audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + + out := rec.Body.String() + if strings.Contains(out, "4111111111111111") { + t.Fatalf("Model A co-firing BlockSoft leaked the complete card number raw: %q", out) + } + if !strings.Contains(out, "[CARD]") { + t.Fatalf("Model A co-firing BlockSoft did not splice the mask onto the wire: %q", out) + } +} + +// TestSSE_ModelA_StrictAppliance_FailClosed_OnUnredactableMask pins the caller-driven +// posture on the Model A wire escalation: a tool-arg mask is undeliverable on the +// streaming wire, so the redactor reports unsupported. Under the compliance-proxy +// appliance (strictFailClosed) the escalation BLOCKS with an in-band error frame — +// it must not relay the unredacted remainder raw; the agent (non-strict) fail-open +// path relays the original (NE host-packet safety) instead. +func TestSSE_ModelA_StrictAppliance_FailClosed_OnUnredactableMask(t *testing.T) { + mkHook := func() modelATestHook { + return modelATestHook{ + decision: core.Modify, + modified: []core.ContentBlock{{Type: "text", Text: "x"}}, + spans: []normalize.TransformSpan{tspan("messages.0.content.0.toolUse.input.0", 0, 1, "x")}, + matchRaw: true, + } + } + body := oaContentBody("secret ", "payload") + + // Appliance (strict) → fail-CLOSED block frame, no raw remainder. + boStrict, _ := modelABumpOptions(modelARedactResolver(mkHook())) + boStrict.strictFailClosed = true + ri, ai, ac := modelAAudCtx() + recStrict := httptest.NewRecorder() + handleSSEResponse(context.Background(), recStrict, sseFrameUpstream(body), ac, ri, ai, boStrict, discardSlog(), time.Now()) + if out := recStrict.Body.String(); !strings.Contains(out, "blocked by policy") || strings.Contains(out, "payload") { + t.Fatalf("strict appliance must block on an unredactable mask (no raw remainder), got %q", out) + } + + // Agent (non-strict) → fail-OPEN relays the original (NE safety). + boOpen, _ := modelABumpOptions(modelARedactResolver(mkHook())) + ri2, ai2, ac2 := modelAAudCtx() + recOpen := httptest.NewRecorder() + handleSSEResponse(context.Background(), recOpen, sseFrameUpstream(body), ac2, ri2, ai2, boOpen, discardSlog(), time.Now()) + out := recOpen.Body.String() + if strings.Contains(out, "blocked by policy") { + t.Fatalf("agent fail-open must NOT block on an unredactable mask, got %q", out) + } + // Agent fail-open relays the original remainder (the disclosed best-effort posture). + if !strings.Contains(out, "payload") { + t.Fatalf("agent fail-open must relay the original remainder, got %q", out) + } +} + +// TestSSE_ModelA_FalsePositive_Approve: the prescan HITs but the confirm Approves, so +// Model A resumes real-time streaming and delivers the full (benign) body unchanged. +func TestSSE_ModelA_FalsePositive_Approve(t *testing.T) { + hook := modelATestHook{decision: core.Approve, matchRaw: true} + bo, _ := modelABumpOptions(modelARedactResolver(hook)) + respInput, auditInfo, audCtx := modelAAudCtx() + body := oaContentBody("hello ", "world") + rec := httptest.NewRecorder() + + handleSSEResponse(context.Background(), rec, sseFrameUpstream(body), audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + + out := rec.Body.String() + if !strings.Contains(out, "hello") || !strings.Contains(out, "world") { + t.Fatalf("false-positive approve must deliver the full body, got %q", out) + } +} + +// TestSSE_ModelA_PrescanMiss_StreamsAndApproves: a sound prescan that never matches +// streams the whole body in real time with zero confirms and an approve-at-EOF stamp. +func TestSSE_ModelA_PrescanMiss_StreamsAndApproves(t *testing.T) { + hook := modelATestHook{decision: core.Modify, matchRaw: false} // prescan never hits + bo, _ := modelABumpOptions(modelARedactResolver(hook)) + respInput, auditInfo, audCtx := modelAAudCtx() + body := oaContentBody("totally ", "benign text") + rec := httptest.NewRecorder() + + handleSSEResponse(context.Background(), rec, sseFrameUpstream(body), audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + + out := rec.Body.String() + if !strings.Contains(out, "totally") || !strings.Contains(out, "benign text") { + t.Fatalf("prescan miss must stream the full body, got %q", out) + } + if !strings.Contains(out, "[DONE]") { + t.Fatalf("stream must terminate with the [DONE] frame, got %q", out) + } +} + +// TestSSE_ModelA_HardBlock_ErrorFrame: a confirmed RejectHard escalation emits an +// in-band block error frame with zero remainder content (deliberate block, not a +// fail-open redaction fault). +func TestSSE_ModelA_HardBlock_ErrorFrame(t *testing.T) { + hook := modelATestHook{decision: core.RejectHard, matchRaw: true} + bo, _ := modelABumpOptions(modelARedactResolver(hook)) + respInput, auditInfo, audCtx := modelAAudCtx() + body := oaContentBody("secret ", "payload") + rec := httptest.NewRecorder() + + handleSSEResponse(context.Background(), rec, sseFrameUpstream(body), audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + + out := rec.Body.String() + if !strings.Contains(out, "blocked by policy") { + t.Fatalf("hard block must emit an in-band error frame, got %q", out) + } + if strings.Contains(out, "payload") { + t.Fatalf("hard block must deliver zero remainder content, got %q", out) + } +} + +// TestSSE_ModelA_UsageAccumulated exercises the usage-accumulator seam: when the +// request was detected as AI traffic (provider set), the substrate feeds every +// parsed frame to the accumulator during Next, and the redact escalation still +// masks the value (no leak). +func TestSSE_ModelA_UsageAccumulated(t *testing.T) { + hook := modelATestHook{ + decision: core.Modify, + modified: []core.ContentBlock{{Type: "text", Text: "card [CARD]"}}, + spans: []normalize.TransformSpan{tspan("messages.0.content.0", 5, 21, "[CARD]")}, + matchRaw: true, + } + bo, _ := modelABumpOptions(modelARedactResolver(hook)) + respInput, auditInfo, audCtx := modelAAudCtx() + auditInfo.RequestMeta.Provider = "openai" // → handleSSEResponse builds a UsageAccumulator + body := oaContentBody("card ", "4111111111111111") + rec := httptest.NewRecorder() + + handleSSEResponse(context.Background(), rec, sseFrameUpstream(body), audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + + if strings.Contains(rec.Body.String(), "4111111111111111") { + t.Fatalf("Model A leaked the card with usage accounting on: %q", rec.Body.String()) + } +} + +// TestSSE_ModelA_EscalationDrainError exercises the escalation drain path hitting a +// non-EOF upstream error: the confirm fires on the first frame and escalates, then +// the drain read errors — the substrate returns the error without panicking or +// delivering the held value raw. +func TestSSE_ModelA_EscalationDrainError(t *testing.T) { + hook := modelATestHook{ + decision: core.Modify, + modified: []core.ContentBlock{{Type: "text", Text: "card [CARD]"}}, + spans: []normalize.TransformSpan{tspan("messages.0.content.0", 5, 21, "[CARD]")}, + matchRaw: true, + } + bo, _ := modelABumpOptions(modelARedactResolver(hook)) + respInput, auditInfo, audCtx := modelAAudCtx() + // One content frame, no [DONE], then an upstream error during the escalation drain. + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + Body: io.NopCloser(&errAfterReader{data: []byte(`data: {"choices":[{"delta":{"content":"card 4111"}}]}` + "\n\n")}), + } + rec := httptest.NewRecorder() + + handleSSEResponse(context.Background(), rec, resp, audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + if strings.Contains(rec.Body.String(), "4111") { + t.Fatalf("a drain error must not flush the held value raw: %q", rec.Body.String()) + } +} + +// failingWriter is an http.ResponseWriter whose Write always errors — to exercise +// the frame-write error path (client disconnected mid-stream). +type failingWriter struct{ header http.Header } + +func (f *failingWriter) Header() http.Header { return f.header } +func (f *failingWriter) WriteHeader(int) {} +func (f *failingWriter) Write([]byte) (int, error) { return 0, io.ErrClosedPipe } + +// TestSSE_ModelA_ClientWriteError exercises the frame-write error path (client +// disconnected mid-stream): on the prescan-miss real-time delivery a failing writer +// makes Deliver/writeFrame return, and the handler stops without panicking. +func TestSSE_ModelA_ClientWriteError(t *testing.T) { + hook := modelATestHook{decision: core.Modify, matchRaw: false} // miss → real-time held delivery at EOF + bo, _ := modelABumpOptions(modelARedactResolver(hook)) + respInput, auditInfo, audCtx := modelAAudCtx() + body := oaContentBody("hello ", "world") + fw := &failingWriter{header: http.Header{}} + + handleSSEResponse(context.Background(), fw, sseFrameUpstream(body), audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + // The assertion is that the failing writer did not panic the handler; failingWriter + // discards all writes, so there is nothing to inspect on the wire. +} + +// TestSSE_ModelA_EscalationWriteError exercises a frame-write error WHILE delivering +// the redacted remainder on escalation (client gone mid-escalation): writeFrames +// returns and the handler stops without leaking the raw value. +func TestSSE_ModelA_EscalationWriteError(t *testing.T) { + hook := modelATestHook{ + decision: core.Modify, + modified: []core.ContentBlock{{Type: "text", Text: "card [CARD]"}}, + spans: []normalize.TransformSpan{tspan("messages.0.content.0", 5, 21, "[CARD]")}, + matchRaw: true, + } + bo, _ := modelABumpOptions(modelARedactResolver(hook)) + respInput, auditInfo, audCtx := modelAAudCtx() + body := oaContentBody("card ", "4111111111111111") + fw := &failingWriter{header: http.Header{}} + + handleSSEResponse(context.Background(), fw, sseFrameUpstream(body), audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + // failingWriter discards all writes; the assertion is no panic and no raw value + // could reach the (failing) wire. +} + +// TestSSE_ModelA_EscalationDrainCap pins the bounded-drain guard: after a confirmed +// hit, a post-hit remainder larger than the buffer cap is BLOCKED with an in-band +// error frame (bounded memory, zero leak) rather than buffered unbounded or relayed +// raw. The cap is forced small via the streaming policy MaxBufferBytes. +func TestSSE_ModelA_EscalationDrainCap(t *testing.T) { + hook := modelATestHook{ + decision: core.Modify, + modified: []core.ContentBlock{{Type: "text", Text: "card [CARD]"}}, + spans: []normalize.TransformSpan{tspan("messages.0.content.0", 5, 21, "[CARD]")}, + matchRaw: true, + } + bo, _ := modelABumpOptions(modelARedactResolver(hook)) + bo.streamingConfig.MaxBufferSize = 64 // tiny cap → the drained remainder overflows + respInput, auditInfo, audCtx := modelAAudCtx() + big := strings.Repeat("4", 200) + body := oaContentBody("card ", big, big) + rec := httptest.NewRecorder() + + handleSSEResponse(context.Background(), rec, sseFrameUpstream(body), audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + + out := rec.Body.String() + if strings.Contains(out, big) { + t.Fatalf("drain-cap overflow leaked the raw remainder: %q", out) + } + if !strings.Contains(out, "blocked by policy") { + t.Fatalf("drain-cap overflow must block with an in-band error frame, got %q", out) + } +} + +// errAfterReader yields data once, then a non-EOF error — to exercise the substrate's +// upstream-error (OnError) path. +type errAfterReader struct { + data []byte + done bool +} + +func (r *errAfterReader) Read(p []byte) (int, error) { + if r.done { + return 0, io.ErrUnexpectedEOF + } + r.done = true + n := copy(p, r.data) + return n, nil +} + +// TestSSE_ModelA_UpstreamError_NoPanic: a non-EOF upstream read error on the +// prescan-miss path routes through OnError without panicking or forcing a synthetic +// frame (agent fail-open). +func TestSSE_ModelA_UpstreamError_NoPanic(t *testing.T) { + hook := modelATestHook{decision: core.Modify, matchRaw: false} // miss → main-loop Next hits the error + bo, _ := modelABumpOptions(modelARedactResolver(hook)) + respInput, auditInfo, audCtx := modelAAudCtx() + resp := &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{"Content-Type": {"text/event-stream"}}, + Body: io.NopCloser(&errAfterReader{data: []byte(`data: {"choices":[{"delta":{"content":"partial"}}]}` + "\n\n")}), + } + rec := httptest.NewRecorder() + + handleSSEResponse(context.Background(), rec, resp, audCtx, respInput, auditInfo, bo, discardSlog(), time.Now()) + // No panic + no synthetic error frame forced (fail-open). The partial frame may or + // may not have flushed depending on the window; the assertion is the call returns. + if strings.Contains(rec.Body.String(), "blocked by policy") { + t.Fatalf("agent fail-open must not force a block frame on an upstream error: %q", rec.Body.String()) + } +} diff --git a/packages/shared/transport/typology/defaults_test.go b/packages/shared/transport/typology/defaults_test.go index f52f78ab..b6aed933 100644 --- a/packages/shared/transport/typology/defaults_test.go +++ b/packages/shared/transport/typology/defaults_test.go @@ -19,6 +19,11 @@ func TestDefaults_EveryKindHasAtLeastOneRule(t *testing.T) { exempt := map[EndpointKind]bool{ EndpointKindVideoGeneration: true, EndpointKindJob: true, + // Label-only refinement: the path table classifies /v1/responses as + // chat-kind (defaults.go) for routing / cache / hook dispatch; the + // traffic_event endpoint_type stamp overrides the label to "responses". + // There is intentionally no dedicated path→responses rule. + EndpointKindResponses: true, } for _, k := range AllEndpointKinds { if exempt[k] { diff --git a/packages/shared/transport/typology/endpointkind.go b/packages/shared/transport/typology/endpointkind.go index c6a58f4a..0892f116 100644 --- a/packages/shared/transport/typology/endpointkind.go +++ b/packages/shared/transport/typology/endpointkind.go @@ -63,6 +63,13 @@ const ( // /v1/models/{model}). Never carries user content; hook pipeline // and cost layer ignore these. EndpointKindModels EndpointKind = "models" + + // EndpointKindResponses is the OpenAI Responses API (/v1/responses). It is + // a chat-FAMILY endpoint for routing / cache / hook purposes (the wire-shape + // dispatch KindFromWireShape keeps it chat-kind), but carries its own + // endpoint_type label so analytics and the route simulator can distinguish + // Responses traffic from /v1/chat/completions. + EndpointKindResponses EndpointKind = "responses" ) // AllEndpointKinds is the closed enumeration of every defined @@ -78,6 +85,7 @@ var AllEndpointKinds = []EndpointKind{ EndpointKindBatch, EndpointKindJob, EndpointKindModels, + EndpointKindResponses, } // IsValid reports whether k is one of the defined EndpointKind constants. diff --git a/packages/shared/transport/typology/endpointkind_test.go b/packages/shared/transport/typology/endpointkind_test.go index 291479a9..ee8a0405 100644 --- a/packages/shared/transport/typology/endpointkind_test.go +++ b/packages/shared/transport/typology/endpointkind_test.go @@ -22,6 +22,7 @@ func TestEndpointKindConstants(t *testing.T) { {EndpointKindBatch, "batch"}, {EndpointKindJob, "job"}, {EndpointKindModels, "models"}, + {EndpointKindResponses, "responses"}, } for _, c := range cases { if string(c.k) != c.want { @@ -48,6 +49,7 @@ func TestAllEndpointKindsExhaustive(t *testing.T) { EndpointKindBatch, EndpointKindJob, EndpointKindModels, + EndpointKindResponses, } if len(AllEndpointKinds) != len(want) { t.Fatalf("len(AllEndpointKinds) = %d, want %d", len(AllEndpointKinds), len(want)) From 2b5a28c921841adc26eef790346b4e8ac206c7eb Mon Sep 17 00:00:00 2001 From: nexus Date: Tue, 30 Jun 2026 21:42:55 +0800 Subject: [PATCH 3/8] feat(ui): sync control-plane UI, shared UI components, and tests (sync 9) Frontend code and tests (TLS-bump streaming-compliance + provider-adapter cycle): - control-plane-ui: provider form + API service/types, traffic audit-drawer HookTimeline, new ProviderForm test - i18n: en/es/zh pages.json parity Package-internal Go tests ship with PR2. --- .../src/api/services/ai-gateway/providers.ts | 9 ++ packages/control-plane-ui/src/api/types.ts | 19 ++++ .../src/i18n/locales/en/pages.json | 5 + .../src/i18n/locales/es/pages.json | 5 + .../src/i18n/locales/zh/pages.json | 5 + .../providers/list/ProviderForm.test.tsx | 101 ++++++++++++++++++ .../providers/list/ProviderForm.tsx | 34 +++++- .../traffic/audit-drawer/HookTimeline.tsx | 62 +++++++++-- 8 files changed, 233 insertions(+), 7 deletions(-) create mode 100644 packages/control-plane-ui/src/pages/ai-gateway/providers/list/ProviderForm.test.tsx diff --git a/packages/control-plane-ui/src/api/services/ai-gateway/providers.ts b/packages/control-plane-ui/src/api/services/ai-gateway/providers.ts index 3c710f8d..a34af7fe 100644 --- a/packages/control-plane-ui/src/api/services/ai-gateway/providers.ts +++ b/packages/control-plane-ui/src/api/services/ai-gateway/providers.ts @@ -49,6 +49,15 @@ export interface CreateProviderInput { */ adapterType: string; enabled?: boolean; + /** + * Whether this provider's endpoint serves the OpenAI Responses API + * (`/v1/responses`). `null` (or omitted) = use the adapter type's default; + * `false` forces a downgrade to `/v1/chat/completions` for + * OpenAI-compatible endpoints that only implement chat completions. + * Downgrade-only — a `true` value cannot exceed the wire adapter's + * actual capabilities. + */ + servesResponsesApi?: boolean | null; region?: string; apiVersion?: string; pathPrefix?: string; diff --git a/packages/control-plane-ui/src/api/types.ts b/packages/control-plane-ui/src/api/types.ts index 5e993dc5..4070bc3b 100644 --- a/packages/control-plane-ui/src/api/types.ts +++ b/packages/control-plane-ui/src/api/types.ts @@ -180,6 +180,16 @@ export interface Provider { region?: string | null; headers?: Record | null; enabled: boolean; + /** + * Whether this provider's endpoint serves the OpenAI Responses API + * (`/v1/responses`). `null`/absent means "use the adapter type's default" + * (the gateway decides from the wire adapter's capabilities). `false` + * forces the gateway to downgrade to `/v1/chat/completions` for + * OpenAI-compatible endpoints that only implement chat completions. + * Downgrade-only: a `true` value cannot make an adapter serve Responses + * if its wire format does not support it. + */ + servesResponsesApi?: boolean | null; createdAt: string; updatedAt?: string; } @@ -531,6 +541,9 @@ export interface HookExecutionRecord { reason?: string; reasonCode?: string; latencyMs?: number; + // Microsecond-precision per-hook latency (hooks run at microsecond scale, so + // latencyMs truncates sub-millisecond hooks to 0). Additive; latencyMs kept. + latencyUs?: number; error?: string; // PascalCase keys (compliance-proxy producer — shared/hooks.HookResult // serializes with Go field names since the struct lacks json tags). @@ -541,6 +554,7 @@ export interface HookExecutionRecord { Reason?: string; ReasonCode?: string; LatencyMs?: number; + LatencyUs?: number; Order?: number; Error?: string; } @@ -755,6 +769,11 @@ export interface TrafficEvent { upstreamTotalMs?: number | null; requestHooksMs?: number | null; responseHooksMs?: number | null; + // Microsecond-precision hook aggregates (additive; the _ms fields are kept). + // Hooks run at microsecond scale, so the _ms aggregates floor sub-millisecond + // hooks to 0; these carry the real value for precise display. + requestHooksUs?: number | null; + responseHooksUs?: number | null; /** * Long-tail phase durations. Values are milliseconds EXCEPT keys ending in * `_us`, which are microseconds (the sub-millisecond hook framing segments, diff --git a/packages/control-plane-ui/src/i18n/locales/en/pages.json b/packages/control-plane-ui/src/i18n/locales/en/pages.json index 6176dae3..9a2beb95 100644 --- a/packages/control-plane-ui/src/i18n/locales/en/pages.json +++ b/packages/control-plane-ui/src/i18n/locales/en/pages.json @@ -406,6 +406,11 @@ "apiVersion": "API Version", "apiVersionHelp": "Optional. Provider API version string (e.g. 2024-02-01 for Azure OpenAI). Leave blank unless the upstream endpoint requires a version segment.", "apiVersionPlaceholder": "e.g. 2024-02-01", + "servesResponsesApi": "Serves OpenAI Responses API", + "servesResponsesApiHelp": "Leave unset to use the provider type's default. Disable for OpenAI-compatible endpoints that only implement /v1/chat/completions.", + "servesResponsesApiOption_default": "Use adapter default", + "servesResponsesApiOption_true": "Enabled", + "servesResponsesApiOption_false": "Disabled (chat completions only)", "modelCode": "Model ID", "modelCodeLabel": "Model ID *", "placeholderModelCode": "e.g. gpt-4o (defaults to Upstream Model ID)", diff --git a/packages/control-plane-ui/src/i18n/locales/es/pages.json b/packages/control-plane-ui/src/i18n/locales/es/pages.json index d0e64e7b..bd885c7f 100644 --- a/packages/control-plane-ui/src/i18n/locales/es/pages.json +++ b/packages/control-plane-ui/src/i18n/locales/es/pages.json @@ -406,6 +406,11 @@ "apiVersion": "Versión de API", "apiVersionHelp": "Opcional. Cadena de versión de API del proveedor (p. ej. 2024-02-01 para Azure OpenAI). Dejar en blanco si el endpoint no requiere un segmento de versión.", "apiVersionPlaceholder": "p. ej. 2024-02-01", + "servesResponsesApi": "Sirve la API de Responses de OpenAI", + "servesResponsesApiHelp": "Déjalo sin definir para usar el valor predeterminado del tipo de proveedor. Desactívalo para endpoints compatibles con OpenAI que solo implementan /v1/chat/completions.", + "servesResponsesApiOption_default": "Usar el valor predeterminado del adaptador", + "servesResponsesApiOption_true": "Activado", + "servesResponsesApiOption_false": "Desactivado (solo chat completions)", "modelCode": "Model ID", "modelCodeLabel": "Model ID *", "placeholderModelCode": "p. ej. gpt-4o (por defecto usa el ID de modelo upstream)", diff --git a/packages/control-plane-ui/src/i18n/locales/zh/pages.json b/packages/control-plane-ui/src/i18n/locales/zh/pages.json index 9e165cf2..808243a2 100644 --- a/packages/control-plane-ui/src/i18n/locales/zh/pages.json +++ b/packages/control-plane-ui/src/i18n/locales/zh/pages.json @@ -406,6 +406,11 @@ "apiVersion": "API 版本", "apiVersionHelp": "可选。上游提供商的 API 版本字符串（如 Azure OpenAI 的 2024-02-01）。若上游端点不要求版本号，请留空。", "apiVersionPlaceholder": "例如 2024-02-01", + "servesResponsesApi": "提供 OpenAI Responses API", + "servesResponsesApiHelp": "留空则使用该供应商类型的默认值。对于仅实现 /v1/chat/completions 的 OpenAI 兼容端点，请将其禁用。", + "servesResponsesApiOption_default": "使用适配器默认值", + "servesResponsesApiOption_true": "启用", + "servesResponsesApiOption_false": "禁用（仅 chat completions）", "modelCode": "模型ID", "modelCodeLabel": "Model ID *", "placeholderModelCode": "例如 gpt-4o（留空则使用上游模型ID）", diff --git a/packages/control-plane-ui/src/pages/ai-gateway/providers/list/ProviderForm.test.tsx b/packages/control-plane-ui/src/pages/ai-gateway/providers/list/ProviderForm.test.tsx new file mode 100644 index 00000000..aafe198d --- /dev/null +++ b/packages/control-plane-ui/src/pages/ai-gateway/providers/list/ProviderForm.test.tsx @@ -0,0 +1,101 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; +import { render, screen, fireEvent, waitFor } from '@testing-library/react'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import { ToastProvider } from '@/context/ToastContext'; +import type { Provider } from '@/api/types'; +import { ProviderForm } from './ProviderForm'; + +// t returns the key so assertions don't depend on i18n initialization; the +// English fallback (2nd arg) is ignored. +vi.mock('react-i18next', () => ({ + useTranslation: () => ({ t: (k: string) => k }), + initReactI18next: { type: '3rdParty', init: () => {} }, + // ProviderConnectivityTestButton renders ; render its i18nKey as text. + Trans: ({ i18nKey }: { i18nKey: string }) => i18nKey, +})); + +// Capture the payload the form submits to the admin API. Both create() and +// update() are mocked so the form's useMutation resolves cleanly. +const create = vi.fn(async (data: unknown) => ({ id: 'p_new', ...(data as object) })); +const update = vi.fn(async (_id: string, data: unknown) => ({ id: _id, ...(data as object) })); +vi.mock('@/api/services', () => ({ + providerApi: { + create: (data: unknown) => create(data), + update: (id: string, data: unknown) => update(id, data), + // ProviderConnectivityTestButton references these; unused in these tests. + testExisting: vi.fn(), + testConnection: vi.fn(), + }, +})); + +function renderForm(provider?: Provider) { + const qc = new QueryClient({ defaultOptions: { queries: { retry: false } } }); + render( + + + + + , + ); +} + +const baseProvider: Provider = { + id: 'p_1', + name: 'openai-prod', + adapterType: 'openai', + baseUrl: 'https://api.openai.com', + pathPrefix: '/v1', + enabled: true, + createdAt: '2026-06-29T00:00:00Z', +}; + +describe('ProviderForm — servesResponsesApi round-trip', () => { + beforeEach(() => { + create.mockClear(); + update.mockClear(); + }); + + it('sends servesResponsesApi: null on create when left at the adapter default', async () => { + renderForm(); + // Fill the required fields (name + baseUrl) so the Save button enables. + fireEvent.change(document.querySelector('input[name="name"]') as HTMLInputElement, { + target: { value: 'my-provider' }, + }); + fireEvent.change(document.querySelector('input[name="baseUrl"]') as HTMLInputElement, { + target: { value: 'https://api.example.com' }, + }); + + fireEvent.click(screen.getByRole('button', { name: 'common:save' })); + + await waitFor(() => expect(create).toHaveBeenCalledTimes(1)); + expect(create.mock.calls[0][0]).toMatchObject({ servesResponsesApi: null }); + }); + + it('initialises the override to false and round-trips false on update', async () => { + renderForm({ ...baseProvider, servesResponsesApi: false }); + + fireEvent.click(screen.getByRole('button', { name: 'common:save' })); + + await waitFor(() => expect(update).toHaveBeenCalledTimes(1)); + expect(update.mock.calls[0][0]).toBe('p_1'); + expect(update.mock.calls[0][1]).toMatchObject({ servesResponsesApi: false }); + }); + + it('initialises the override to true and round-trips true on update', async () => { + renderForm({ ...baseProvider, servesResponsesApi: true }); + + fireEvent.click(screen.getByRole('button', { name: 'common:save' })); + + await waitFor(() => expect(update).toHaveBeenCalledTimes(1)); + expect(update.mock.calls[0][1]).toMatchObject({ servesResponsesApi: true }); + }); + + it('keeps a null/absent capability as the adapter default (null) on update', async () => { + renderForm({ ...baseProvider, servesResponsesApi: null }); + + fireEvent.click(screen.getByRole('button', { name: 'common:save' })); + + await waitFor(() => expect(update).toHaveBeenCalledTimes(1)); + expect(update.mock.calls[0][1]).toMatchObject({ servesResponsesApi: null }); + }); +}); diff --git a/packages/control-plane-ui/src/pages/ai-gateway/providers/list/ProviderForm.tsx b/packages/control-plane-ui/src/pages/ai-gateway/providers/list/ProviderForm.tsx index c33e8457..f8e1223f 100644 --- a/packages/control-plane-ui/src/pages/ai-gateway/providers/list/ProviderForm.tsx +++ b/packages/control-plane-ui/src/pages/ai-gateway/providers/list/ProviderForm.tsx @@ -26,6 +26,16 @@ export function ProviderForm({ provider, onClose, onSaved }: ProviderFormProps) const [region, setRegion] = useState(provider?.region ?? ''); const [apiVersion, setApiVersion] = useState(provider?.apiVersion ?? ''); const [enabled, setEnabled] = useState(provider?.enabled ?? true); + // Tri-state: 'default' = inherit the adapter type's capability (null), + // 'true'/'false' = explicit per-provider override. Stored as a string so the + // shared Select component (string-valued) can drive it directly. + const [servesResponsesApi, setServesResponsesApi] = useState<'default' | 'true' | 'false'>( + provider?.servesResponsesApi == null + ? 'default' + : provider.servesResponsesApi + ? 'true' + : 'false', + ); const { mutate, loading } = useMutation( (data: unknown) => @@ -41,7 +51,15 @@ export function ProviderForm({ provider, onClose, onSaved }: ProviderFormProps) ); const handleSubmit = () => { - mutate({ name, displayName, description, baseUrl, adapterType, region: region || undefined, apiVersion: apiVersion || undefined, enabled }); + mutate({ + name, displayName, description, baseUrl, adapterType, + region: region || undefined, + apiVersion: apiVersion || undefined, + enabled, + // 'default' clears the override back to null (inherit the adapter default); + // an explicit choice sends the boolean. + servesResponsesApi: servesResponsesApi === 'default' ? null : servesResponsesApi === 'true', + }); }; return ( @@ -124,6 +142,20 @@ export function ProviderForm({ provider, onClose, onSaved }: ProviderFormProps) placeholder={t('pages:providers.apiVersionPlaceholder')} /> + + diff --git a/packages/control-plane-ui/src/pages/traffic/audit-drawer/HookTimeline.tsx b/packages/control-plane-ui/src/pages/traffic/audit-drawer/HookTimeline.tsx index 0fdcb228..f4f4ebf2 100644 --- a/packages/control-plane-ui/src/pages/traffic/audit-drawer/HookTimeline.tsx +++ b/packages/control-plane-ui/src/pages/traffic/audit-drawer/HookTimeline.tsx @@ -17,9 +17,55 @@ function hookFields(r: HookExecutionRecord) { const reason = r.reason ?? r.Reason ?? ''; const reasonCode = r.reasonCode ?? r.ReasonCode ?? ''; const latencyMs = r.latencyMs ?? r.LatencyMs; + const latencyUs = r.latencyUs ?? r.LatencyUs; const order = r.order ?? r.Order ?? 0; const error = r.error ?? r.Error ?? ''; - return { id, name, impl, decision, reason, reasonCode, latencyMs, order, error }; + return { id, name, impl, decision, reason, reasonCode, latencyMs, latencyUs, order, error }; +} + +type HookFields = ReturnType; +type CollapsedHook = HookFields & { count: number; latencyMsSum?: number; latencyUsSum?: number }; + +// formatLatency renders the most precise available value: microseconds when +// present (hooks run at microsecond scale, so the millisecond value floors a +// sub-millisecond hook to 0), otherwise the legacy millisecond value. Sub-ms +// totals show as "210 µs"; larger ones as fractional milliseconds. +function formatLatency(latencyMs?: number, latencyUs?: number): string | null { + if (latencyUs != null && latencyUs > 0) { + if (latencyUs < 1000) return `${latencyUs} µs`; + return `${(latencyUs / 1000).toFixed(latencyUs < 10000 ? 2 : 1)} ms`; + } + if (latencyMs != null) return `${latencyMs} ms`; + return null; +} + +// collapseDuplicates folds repeated executions of the same hook (same id + impl) +// into ONE entry carrying the execution count and the summed latency. The data +// plane already emits one row per hook going forward, but historical rows +// captured before that fix can list the same hook many times — a streamed +// response was scanned at every checkpoint. Collapsing keeps the drawer readable +// while still showing how many times the hook ran (×N) and the total time; the +// latest scan's decision/reason win. +function collapseDuplicates(items: HookFields[]): CollapsedHook[] { + const byKey = new Map(); + const order: string[] = []; + for (const it of items) { + const key = `${it.id}|${it.impl}|${it.name}`; + const cur = byKey.get(key); + if (cur) { + cur.count += 1; + if (it.latencyMs != null) cur.latencyMsSum = (cur.latencyMsSum ?? 0) + it.latencyMs; + if (it.latencyUs != null) cur.latencyUsSum = (cur.latencyUsSum ?? 0) + it.latencyUs; + cur.decision = it.decision; + cur.reason = it.reason; + cur.reasonCode = it.reasonCode; + if (it.error) cur.error = it.error; + } else { + byKey.set(key, { ...it, count: 1, latencyMsSum: it.latencyMs, latencyUsSum: it.latencyUs }); + order.push(key); + } + } + return order.map((k) => byKey.get(k)!); } export function PipelineTimeline({ @@ -39,9 +85,9 @@ export function PipelineTimeline({ ); } - const ordered = [...rows] - .map((r) => ({ raw: r, ...hookFields(r) })) - .sort((a, b) => a.order - b.order); + const ordered = collapseDuplicates( + [...rows].map((r) => hookFields(r)).sort((a, b) => a.order - b.order), + ); return (
{label} ({ordered.length})
@@ -49,6 +95,7 @@ export function PipelineTimeline({ {ordered.map((r, idx) => { const primary = r.name || r.id || 'hook'; const showId = !!r.name && r.id; + const latency = formatLatency(r.latencyMsSum, r.latencyUsSum); return (
{primary} + {r.count > 1 && ( + ×{r.count} + )} {r.impl && ( {r.impl} @@ -70,11 +120,11 @@ export function PipelineTimeline({ id: {r.id}
)} - {(r.reason || r.reasonCode || r.latencyMs != null || r.error) && ( + {(r.reason || r.reasonCode || latency || r.error) && (
{r.reasonCode && [{r.reasonCode}]} {r.reason && {r.reason}} - {r.latencyMs != null && {r.latencyMs} ms} + {latency && {latency}} {r.error && error: {r.error}}
)} From ba3e5585f5bfab79660316575206f9d9d979512b Mon Sep 17 00:00:00 2001 From: nexus Date: Fri, 3 Jul 2026 22:39:28 +0800 Subject: [PATCH 4/8] chore(repo): sync docs, tooling, and CI config from internal main (sync 10) Bring repository-level assets up to date. - Docs: virtual-keys architecture + OpenAPI (personal VK duplicate-name handling), AI-gateway routing feature doc, .env.example - Tooling: new scripts/check-no-benchmark-env.sh guard - DB: tools/db-migrate schema-extras.sql + gateway.prisma - README: fix CI badge URLs to github.com/AlphaBitCore/nexus-gateway Co-Authored-By: Claude Opus 4.7 (1M context) --- .env.example | 14 +++++++++++++ README.md | 4 ++-- ...p-ai-providers-virtualkeys-architecture.md | 10 ++++++++++ .../openapi/control-plane/virtual-keys.yaml | 19 ++++++++++++++++++ .../features/cp-ui/ai-gateway-routing.md | 2 +- scripts/check-no-benchmark-env.sh | 20 +++++++++++++++++++ tools/db-migrate/schema-extras.sql | 6 ++++++ tools/db-migrate/schema/gateway.prisma | 4 ++++ 8 files changed, 76 insertions(+), 3 deletions(-) create mode 100755 scripts/check-no-benchmark-env.sh diff --git a/.env.example b/.env.example index 2cbf9a07..44090069 100644 --- a/.env.example +++ b/.env.example @@ -323,6 +323,20 @@ NATS_URL=nats://localhost:4222 # AI_GATEWAY_AUDIT_SPILL_RECOVERY_INTERVAL_MS=2000 # AI_GATEWAY_AUDIT_SPILL_RECOVERY_PACE_MS=50 +# --------------------------------------------------------------------------- +# BENCHMARK / ABLATION SWITCHES — NOT FOR PRODUCTION +# These exist only to measure throughput against forward-only gateways. +# Do NOT copy these into any production or customer environment file. +# --------------------------------------------------------------------------- +# NEXUS_PERF_PURE_FORWARD — BENCHMARK ONLY. Set to 1 to make the AI Gateway a +# pure forwarding proxy: it SKIPS the entire audit tail, so NO traffic_event +# rows are written. This DESTROYS the audit trail (GDPR / audit-retention +# loss) and must NEVER be set in any production or customer environment. +# Default (unset) = audit stored, normal operation. Detection when active: +# the startup WARN banner "PURE-FORWARD BENCHMARK MODE ACTIVE" and the +# Prometheus gauge nexus_ai_gateway_pure_forward_mode == 1. +# Intentionally has NO assignable line here so it is never accidentally shipped. + # Service discovery — co-located services on localhost; on prod each is a # domain or LB. Every URL below is bare-named (no service prefix) because # it identifies a shared environment-level entity, not a service-private diff --git a/README.md b/README.md index 10eb28c4..641aa52a 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # Nexus Gateway -[![CI](https://github.com/itechchoice/abc-nexus-gateway/actions/workflows/ci.yml/badge.svg?branch=main)](.github/workflows/ci.yml) -[![Go CI](https://github.com/itechchoice/abc-nexus-gateway/actions/workflows/go-ci.yml/badge.svg?branch=main)](.github/workflows/go-ci.yml) +[![CI](https://github.com/AlphaBitCore/nexus-gateway/actions/workflows/ci.yml/badge.svg?branch=main)](.github/workflows/ci.yml) +[![Go CI](https://github.com/AlphaBitCore/nexus-gateway/actions/workflows/go-ci.yml/badge.svg?branch=main)](.github/workflows/go-ci.yml) [![Coverage gate](https://img.shields.io/badge/coverage-%E2%89%A595%25%20per%20package-brightgreen)](./scripts/check-go-coverage.sh) [![Status: 1.1.0](https://img.shields.io/badge/status-1.1.0-brightgreen)](./CHANGELOG.md) [![License: Apache 2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](./LICENSE) diff --git a/docs/developers/architecture/services/control-plane/cp-ai-providers-virtualkeys-architecture.md b/docs/developers/architecture/services/control-plane/cp-ai-providers-virtualkeys-architecture.md index 10782600..2fefa3dd 100644 --- a/docs/developers/architecture/services/control-plane/cp-ai-providers-virtualkeys-architecture.md +++ b/docs/developers/architecture/services/control-plane/cp-ai-providers-virtualkeys-architecture.md @@ -89,6 +89,16 @@ clear the expiry to never-expire and escape the re-approval cadence; create and renew additionally require the value. **Personal** keys are exempt (uncapped, and may clear their expiry). +Names are unique per scope: application keys share one deployment-wide +namespace, while personal keys collide only within their owner — user A's +"test" never blocks user B's, and personal names are never disclosed across +users. Enforcement is two partial unique indexes in +`tools/db-migrate/schema-extras.sql` (`VirtualKey_application_name_uniq`, +`VirtualKey_personal_owner_name_uniq`); both create handlers (admin and +personal self-service) translate a violation into 409 `duplicate_name`, so the +check is race-proof rather than a read-then-insert. Update paths never touch +`name` (rename is not supported), so create is the only enforcement point. + Propagation splits by transition. Update, delete, and regenerate push a targeted invalidate-by-hash through `NotifyConfigChange` under the `virtual_keys` shadow key — an `invalidate` op carrying the affected key hash — so the gateway evicts diff --git a/docs/users/api/openapi/control-plane/virtual-keys.yaml b/docs/users/api/openapi/control-plane/virtual-keys.yaml index c693cf0a..a25af40f 100644 --- a/docs/users/api/openapi/control-plane/virtual-keys.yaml +++ b/docs/users/api/openapi/control-plane/virtual-keys.yaml @@ -113,6 +113,10 @@ paths: resolved type is "application", `projectId` and `expiresAt` become required, `expiresAt` must not exceed three months from now, and the key is created in "pending" status awaiting approval; other types are created "active". + + Names must be unique: application keys share one deployment-wide namespace, + and personal keys are unique per owner. A taken name is rejected with 409 + `duplicate_name`. x-nexus-iam-action: iam.ResourceVirtualKey.Action(iam.VerbCreate) x-nexus-tier: confirm requestBody: @@ -183,6 +187,21 @@ paths: schema: type: object additionalProperties: {} + "409": + description: |- + Conflict — the name is already taken (`error.type` is `duplicate_name`). + Application virtual-key names are unique across the deployment; personal + virtual-key names are unique per owner. + content: + application/json: + schema: + type: object + additionalProperties: {} + example: + error: + message: An application virtual key with this name already exists + type: duplicate_name + code: "" "500": description: Internal Server Error content: diff --git a/docs/users/features/cp-ui/ai-gateway-routing.md b/docs/users/features/cp-ui/ai-gateway-routing.md index 257e00b9..c936eb2d 100644 --- a/docs/users/features/cp-ui/ai-gateway-routing.md +++ b/docs/users/features/cp-ui/ai-gateway-routing.md @@ -64,7 +64,7 @@ This feature is limited to providers whose adapter type speaks the OpenAI wire f **List page.** Columns: name, project (with its organization), a status badge, expiry, an enabled toggle, and actions — approve / reject for pending keys, revoke for active keys, and delete. A "Create Virtual Key" button opens the create form; the list has a search box plus project and status filters. This page lists application-type keys only. -**Create and detail.** Creation collects name, an optional project (which binds the key to that project and its organization), a source-app label, an allowed-models list (per provider-and-model reference; an empty list means all models are allowed), a requests-per-minute rate limit, an expiry (or a never-expires flag), and the enabled flag. The secret is shown once, immediately after creation. The detail page has three tabs — info, quota, access-log; the info tab regenerates the secret (displayed afterward as a key-prefix plus masked remainder) and edits the key's scope, and the quota tab shows the rate limit. +**Create and detail.** Creation collects name, an optional project (which binds the key to that project and its organization), a source-app label, an allowed-models list (per provider-and-model reference; an empty list means all models are allowed), a requests-per-minute rate limit, an expiry (or a never-expires flag), and the enabled flag. Names must be unique among application keys (personal keys are unique per owner); leaving the name field runs an advisory duplicate check that flags a taken name inline before the rest of the form is filled in, and submission is rejected with a conflict error if the name is taken (the authoritative check). The secret is shown once, immediately after creation. The detail page has three tabs — info, quota, access-log; the info tab regenerates the secret (displayed afterward as a key-prefix plus masked remainder) and edits the key's scope, and the quota tab shows the rate limit. **Key concepts.** `vkType` is `application` or `personal` — this section manages application keys; personal keys are issued by developers from their own account settings. `vkStatus` moves through `pending`, `active`, `expired`, `rejected`, and `revoked`. The create form does not link a quota policy directly; quota association is shown on the detail page's quota tab. diff --git a/scripts/check-no-benchmark-env.sh b/scripts/check-no-benchmark-env.sh new file mode 100755 index 00000000..052e7376 --- /dev/null +++ b/scripts/check-no-benchmark-env.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +# Fail if the pure-forward benchmark switch is activated in any tracked file. +# NEXUS_PERF_PURE_FORWARD=1 disables the audit trail and must never ship — it is +# a benchmark-only, host-set env var. Commented references (e.g. in .env.example) +# are allowed; only an ACTIVE assignment is a violation. +set -euo pipefail + +cd "$(git rev-parse --show-toplevel)" + +# Match active assignments, ignore comment lines (leading # after optional space). +# No "\b" — macOS git grep -E (POSIX ERE) does not support it and silently matches +# nothing. The trailing "[:=]" anchors this to a real assignment, not prose. +hits="$(git grep -nE '^[^#]*NEXUS_PERF_PURE_FORWARD[[:space:]]*[:=][[:space:]]*(1|true|"1"|"true")' -- . ':!scripts/check-no-benchmark-env.sh' || true)" + +if [[ -n "$hits" ]]; then + echo "ERROR: NEXUS_PERF_PURE_FORWARD is activated in a committed file (audit-disabling benchmark switch must never ship):" >&2 + echo "$hits" >&2 + exit 1 +fi +echo "OK: no committed file activates NEXUS_PERF_PURE_FORWARD" diff --git a/tools/db-migrate/schema-extras.sql b/tools/db-migrate/schema-extras.sql index 3c065c50..337aa706 100644 --- a/tools/db-migrate/schema-extras.sql +++ b/tools/db-migrate/schema-extras.sql @@ -128,6 +128,12 @@ CREATE INDEX IF NOT EXISTS traffic_event_passthrough_active_idx ON public.traffi -- so without this partial the page seq-scans the whole table per sub-query. CREATE INDEX IF NOT EXISTS traffic_event_routed_provider_ts_idx ON public.traffic_event USING btree (routed_provider_id, "timestamp" DESC) WHERE ((source = 'ai-gateway'::text) AND (routed_provider_id IS NOT NULL)); CREATE UNIQUE INDEX IF NOT EXISTS "DeviceAssignment_deviceId_active_uidx" ON public."DeviceAssignment" USING btree ("deviceId") WHERE ("releasedAt" IS NULL); +-- Virtual-key name uniqueness is scoped by vkType: application keys share one +-- fleet-wide admin namespace, while personal keys are self-service and only +-- collide within their owner (user A's "test" must not block user B's). +-- Handlers map violations of these two indexes to 409 duplicate_name. +CREATE UNIQUE INDEX IF NOT EXISTS "VirtualKey_application_name_uniq" ON public."VirtualKey" USING btree (name) WHERE ("vkType" = 'application'::text); +CREATE UNIQUE INDEX IF NOT EXISTS "VirtualKey_personal_owner_name_uniq" ON public."VirtualKey" USING btree ("ownerId", name) WHERE ("vkType" = 'personal'::text); CREATE UNIQUE INDEX IF NOT EXISTS exemption_request_pending_dedup_uniq ON public.exemption_request USING btree (target_host, requested_by) WHERE (status = 'PENDING'::public."ExemptionRequestStatus"); CREATE UNIQUE INDEX IF NOT EXISTS thing_type_physical_id_uniq ON public.thing USING btree (type, physical_id) WHERE ((type = 'agent'::text) AND (physical_id IS NOT NULL)); CREATE UNIQUE INDEX IF NOT EXISTS uq_ops_rollup_1d ON public.metric_ops_rollup_1d USING btree (bucket_start, COALESCE(thing_id, ''::text), metric_name, dimension_key); diff --git a/tools/db-migrate/schema/gateway.prisma b/tools/db-migrate/schema/gateway.prisma index 53702206..97660296 100644 --- a/tools/db-migrate/schema/gateway.prisma +++ b/tools/db-migrate/schema/gateway.prisma @@ -35,6 +35,10 @@ model VirtualKey { rejectedAt DateTime? @db.Timestamptz(3) rejectReason String? + /// Name uniqueness is enforced by two partial unique indexes in + /// schema-extras.sql (Prisma cannot express WHERE-scoped uniques): + /// `VirtualKey_application_name_uniq` (name, vkType='application') and + /// `VirtualKey_personal_owner_name_uniq` (ownerId+name, vkType='personal'). @@index([name]) @@index([keyHash]) @@index([projectId]) From 24fe8b534be9775fc85326ddc0cb96aa2bb67ea7 Mon Sep 17 00:00:00 2001 From: nexus Date: Fri, 3 Jul 2026 22:40:20 +0800 Subject: [PATCH 5/8] feat(services): sync backend services and shared libs from internal main (sync 10) Backend service implementations and shared Go libraries: - ai-gateway: pure-forward wiring, ingress perf flags + stage accounting, audit body-pool, Anthropic codec, routing strategy-node JSON - control-plane: virtual-keys handlers (personal VK duplicate-name), SAML/OIDC display-name + federated-store test fixtures Co-Authored-By: Claude Opus 4.7 (1M context) --- packages/ai-gateway/cmd/ai-gateway/main.go | 1 + .../cmd/ai-gateway/wiring/pureforward.go | 30 ++++++ .../cmd/ai-gateway/wiring/pureforward_test.go | 37 ++++++++ .../internal/ingress/proxy/perf_flags.go | 13 +++ .../internal/ingress/proxy/perf_flags_test.go | 17 ++++ .../ingress/proxy/stage_accounting.go | 13 +++ .../ingress/proxy/stage_accounting_test.go | 45 +++++++++ .../platform/audit/record_bodypool.go | 7 ++ .../platform/audit/record_bodypool_test.go | 31 +++++++ .../providers/specs/anthropic/codec/codec.go | 11 ++- .../specs/anthropic/codec/codec_test.go | 2 + .../routing/core/strategynode_json.go | 32 +++++++ .../core/strategynode_unmarshal_test.go | 92 +++++++++++++++++++ .../internal/ai/virtualkeys/handler/user.go | 3 + .../ai/virtualkeys/handler/user_test.go | 28 ++++++ .../internal/ai/virtualkeys/handler/vk.go | 23 +++++ .../ai/virtualkeys/handler/vk_test.go | 48 ++++++++++ .../authserver/login/displayname_test.go | 12 +-- .../store/federated_store_mock_test.go | 4 +- 19 files changed, 438 insertions(+), 11 deletions(-) create mode 100644 packages/ai-gateway/cmd/ai-gateway/wiring/pureforward.go create mode 100644 packages/ai-gateway/cmd/ai-gateway/wiring/pureforward_test.go create mode 100644 packages/ai-gateway/internal/ingress/proxy/perf_flags_test.go create mode 100644 packages/ai-gateway/internal/ingress/proxy/stage_accounting_test.go create mode 100644 packages/ai-gateway/internal/routing/core/strategynode_json.go create mode 100644 packages/ai-gateway/internal/routing/core/strategynode_unmarshal_test.go diff --git a/packages/ai-gateway/cmd/ai-gateway/main.go b/packages/ai-gateway/cmd/ai-gateway/main.go index df45e6fa..7b2f4b2b 100644 --- a/packages/ai-gateway/cmd/ai-gateway/main.go +++ b/packages/ai-gateway/cmd/ai-gateway/main.go @@ -56,6 +56,7 @@ func run() int { } slog.SetDefault(logger) slog.Info("shared initialized", slog.Int("sharedHooks", len(builtins.Registry.All()))) + wiring.RegisterPureForwardMode(prometheus.DefaultRegisterer, logger) ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) defer stop() diff --git a/packages/ai-gateway/cmd/ai-gateway/wiring/pureforward.go b/packages/ai-gateway/cmd/ai-gateway/wiring/pureforward.go new file mode 100644 index 00000000..5d3b922f --- /dev/null +++ b/packages/ai-gateway/cmd/ai-gateway/wiring/pureforward.go @@ -0,0 +1,30 @@ +// pureforward.go — startup surface for the NEXUS_PERF_PURE_FORWARD benchmark +// switch: a loud WARN banner and a self-identifying Prometheus gauge so a process +// running with audit disabled is obvious in logs and on dashboards. +package wiring + +import ( + "log/slog" + + "github.com/prometheus/client_golang/prometheus" + + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/ingress/proxy" +) + +// RegisterPureForwardMode registers the pure_forward_mode gauge and, when the +// NEXUS_PERF_PURE_FORWARD benchmark switch is active, sets it to 1 and logs a +// loud WARN banner. The gauge stays 0 in normal operation. Registering it always +// (not only when on) means dashboards and alerts can rely on the series existing. +func RegisterPureForwardMode(reg prometheus.Registerer, logger *slog.Logger) { + g := prometheus.NewGauge(prometheus.GaugeOpts{ + Namespace: "nexus", + Subsystem: "ai_gateway", + Name: "pure_forward_mode", + Help: "1 when NEXUS_PERF_PURE_FORWARD benchmark mode is active (traffic_event audit persistence DISABLED). MUST be 0 in production.", + }) + reg.MustRegister(g) + if proxy.PerfPureForward() { + g.Set(1) + logger.Warn("PURE-FORWARD BENCHMARK MODE ACTIVE: traffic_event audit persistence is DISABLED (NEXUS_PERF_PURE_FORWARD=1). This MUST NOT run in production — the audit trail is not being written.") + } +} diff --git a/packages/ai-gateway/cmd/ai-gateway/wiring/pureforward_test.go b/packages/ai-gateway/cmd/ai-gateway/wiring/pureforward_test.go new file mode 100644 index 00000000..33f9d8ac --- /dev/null +++ b/packages/ai-gateway/cmd/ai-gateway/wiring/pureforward_test.go @@ -0,0 +1,37 @@ +package wiring + +import ( + "log/slog" + "testing" + + "github.com/prometheus/client_golang/prometheus" + + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/ingress/proxy" +) + +func gaugeValue(t *testing.T, reg *prometheus.Registry) float64 { + t.Helper() + mfs, err := reg.Gather() + if err != nil { + t.Fatalf("gather: %v", err) + } + for _, mf := range mfs { + if mf.GetName() == "nexus_ai_gateway_pure_forward_mode" { + return mf.GetMetric()[0].GetGauge().GetValue() + } + } + t.Fatal("gauge nexus_ai_gateway_pure_forward_mode not registered") + return 0 +} + +func TestRegisterPureForwardMode_GaugeReflectsFlag(t *testing.T) { + regOff := prometheus.NewRegistry() + RegisterPureForwardMode(regOff, slog.Default()) + want := 0.0 + if proxy.PerfPureForward() { + want = 1.0 + } + if v := gaugeValue(t, regOff); v != want { + t.Errorf("gauge=%v want %v (flag=%v)", v, want, proxy.PerfPureForward()) + } +} diff --git a/packages/ai-gateway/internal/ingress/proxy/perf_flags.go b/packages/ai-gateway/internal/ingress/proxy/perf_flags.go index 279fa945..65183ab0 100644 --- a/packages/ai-gateway/internal/ingress/proxy/perf_flags.go +++ b/packages/ai-gateway/internal/ingress/proxy/perf_flags.go @@ -56,3 +56,16 @@ func perfParallelHooks() bool { }) return parallelHooksFlag } + +// pureForward is the NEXUS_PERF_PURE_FORWARD benchmark switch. When on, the proxy +// skips its entire audit tail (no traffic_event) so throughput can be compared +// against forward-only gateways (Bifrost, agentgateway). Default off ("" or any +// value != "1") = audit stored, byte-identical to normal operation. Read once at +// package init; a plain var (not sync.Once) so tests can toggle it. Env-only by +// design — this must never be reachable via config/yaml/shadow, and must never be +// set in production (it disables the audit trail). +var pureForward = os.Getenv("NEXUS_PERF_PURE_FORWARD") == "1" + +// PerfPureForward reports whether pure-forward benchmark mode is active. Exported +// for the wiring layer's startup WARN banner and self-identifying gauge. +func PerfPureForward() bool { return pureForward } diff --git a/packages/ai-gateway/internal/ingress/proxy/perf_flags_test.go b/packages/ai-gateway/internal/ingress/proxy/perf_flags_test.go new file mode 100644 index 00000000..3eb3dcd3 --- /dev/null +++ b/packages/ai-gateway/internal/ingress/proxy/perf_flags_test.go @@ -0,0 +1,17 @@ +package proxy + +import "testing" + +func TestPerfPureForward_ReflectsPackageVar(t *testing.T) { + orig := pureForward + t.Cleanup(func() { pureForward = orig }) + + pureForward = true + if !PerfPureForward() { + t.Error("PerfPureForward() = false, want true when pureForward set") + } + pureForward = false + if PerfPureForward() { + t.Error("PerfPureForward() = true, want false") + } +} diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_accounting.go b/packages/ai-gateway/internal/ingress/proxy/stage_accounting.go index 5f6976ce..9359a9f0 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stage_accounting.go +++ b/packages/ai-gateway/internal/ingress/proxy/stage_accounting.go @@ -15,6 +15,19 @@ import ( // ServeProxy immediately after the state is built, so it covers every // stage's exit path. func (s *proxyState) finalizeAudit() { + // Pure-forward benchmark mode: skip the entire audit tail (no traffic_event). + // The record carries pooled request/response body buffers that the async + // writer normally returns downstream of Enqueue; since we skip Enqueue, return + // them here so the measured hot path recycles them instead of re-allocating + // the ~64 KB request body every request. Everything before this point + // (forwarding, cross-spec conversion, hook extraction/enforcement) has already + // run — this only drops the audit record. + if pureForward { + if s.h != nil && s.h.deps != nil && s.h.deps.AuditWriter != nil { + s.h.deps.AuditWriter.ReclaimRecordBody(s.rec) + } + return + } deferStart := time.Now() s.rec.UpstreamTtfbMs = s.phaseSink.TtfbMs() s.rec.UpstreamTotalMs = s.phaseSink.TotalMs() diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_accounting_test.go b/packages/ai-gateway/internal/ingress/proxy/stage_accounting_test.go new file mode 100644 index 00000000..64ef5b38 --- /dev/null +++ b/packages/ai-gateway/internal/ingress/proxy/stage_accounting_test.go @@ -0,0 +1,45 @@ +package proxy + +import ( + "log/slog" + "testing" + + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/platform/audit" +) + +// When pure-forward is on, finalizeAudit must reclaim the pooled request body and +// return WITHOUT enqueuing — the async writer never sees the record, so a capture +// producer receives zero messages, and the pooled body is returned (RequestBody +// cleared) so the benchmark hot path recycles it. +func TestFinalizeAudit_PureForward_SkipsEnqueueAndReclaims(t *testing.T) { + orig := pureForward + t.Cleanup(func() { pureForward = orig }) + pureForward = true + + prod := &captureProducer{} + w := audit.NewWriter(prod, "nexus.event.ai-traffic", nil, slog.Default()) + t.Cleanup(func() { w.Close() }) + + body, handle := audit.AcquireRequestBody([]byte(`{"model":"gpt-4o"}`)) + rec := &audit.Record{RequestBody: body} + rec.AttachPooledRequestBody(handle) + + s := &proxyState{h: &Handler{deps: &Deps{AuditWriter: w}}, rec: rec} + s.finalizeAudit() + + if rec.RequestBody != nil { + t.Errorf("pure-forward did not reclaim pooled request body: %v", rec.RequestBody) + } + // Pure-forward never calls Enqueue, so the writer's async workers never start + // and nothing can reach the producer — assert zero captured messages directly. + prod.mu.Lock() + n := len(prod.messages) + prod.mu.Unlock() + if n != 0 { + t.Errorf("pure-forward emitted %d audit messages, want 0", n) + } + // The early return must not have stamped the latency snapshot. + if rec.LatencyBreakdown != nil { + t.Errorf("pure-forward stamped LatencyBreakdown=%v, want nil (early return)", rec.LatencyBreakdown) + } +} diff --git a/packages/ai-gateway/internal/platform/audit/record_bodypool.go b/packages/ai-gateway/internal/platform/audit/record_bodypool.go index 3a9bce4f..8ecca037 100644 --- a/packages/ai-gateway/internal/platform/audit/record_bodypool.go +++ b/packages/ai-gateway/internal/platform/audit/record_bodypool.go @@ -84,6 +84,13 @@ func (w *Writer) reclaimRecordBody(rec *Record) { } } +// ReclaimRecordBody returns rec's pooled request/response body buffers to their +// pools. The normal terminal reclaim runs inside the writer downstream of +// Enqueue; the pure-forward benchmark path skips Enqueue and must return the +// buffers itself, or the measured hot path re-allocates the ~64 KB request body +// every request. Idempotent; safe on a nil or un-pooled record. +func (w *Writer) ReclaimRecordBody(rec *Record) { w.reclaimRecordBody(rec) } + // responseBodyPool reuses the streaming-capture tee's backing array. The SSE // capture tee buffers the response body into this array (one alloc per stream // was ~3.7 GB/window — the second-largest streaming-relay allocator); pooling it diff --git a/packages/ai-gateway/internal/platform/audit/record_bodypool_test.go b/packages/ai-gateway/internal/platform/audit/record_bodypool_test.go index c2765c5c..79313b84 100644 --- a/packages/ai-gateway/internal/platform/audit/record_bodypool_test.go +++ b/packages/ai-gateway/internal/platform/audit/record_bodypool_test.go @@ -265,3 +265,34 @@ func TestResponseBodyPool_ReuseNoStaleBytes(t *testing.T) { t.Fatalf("reused response buffer leaked stale bytes: %q", *h2) } } + +func TestReclaimRecordBody_ReturnsPooledBuffersAndClearsHandles(t *testing.T) { + w := NewWriter(nil, "topic", nil, slog.Default()) + + body, handle := AcquireRequestBody([]byte(`{"model":"gpt-4o"}`)) + rec := &Record{RequestBody: body} + rec.AttachPooledRequestBody(handle) + + respHandle := AcquireResponseBuffer() + *respHandle = append(*respHandle, []byte("streamed-bytes")...) + rec.ResponseBody = *respHandle + rec.AttachPooledResponseBody(respHandle) + + w.ReclaimRecordBody(rec) + + // After reclaim the record must not reference the pooled bytes — the guard + // relies on this so the benchmark hot path recycles buffers instead of + // re-allocating the ~64 KB request body every request. + if rec.RequestBody != nil { + t.Errorf("RequestBody not cleared after reclaim: %v", rec.RequestBody) + } + if rec.ResponseBody != nil { + t.Errorf("ResponseBody not cleared after reclaim: %v", rec.ResponseBody) + } + if rec.reqBodyHandle != nil || rec.respBodyHandle != nil { + t.Errorf("handles not nil'd: req=%v resp=%v", rec.reqBodyHandle, rec.respBodyHandle) + } + // Idempotent: a second call and a nil record must not panic. + w.ReclaimRecordBody(rec) + w.ReclaimRecordBody(nil) +} diff --git a/packages/ai-gateway/internal/providers/specs/anthropic/codec/codec.go b/packages/ai-gateway/internal/providers/specs/anthropic/codec/codec.go index 22b5bd8c..39553d8a 100644 --- a/packages/ai-gateway/internal/providers/specs/anthropic/codec/codec.go +++ b/packages/ai-gateway/internal/providers/specs/anthropic/codec/codec.go @@ -70,8 +70,8 @@ func (Codec) EncodeRequest(endpoint typology.WireShape, canonicalBody []byte, ta // Claude 4.x; the older 3.x family accepts every combination // unchanged. // - // (a) claude-opus-4-7 deprecated temperature / top_p / top_k - // entirely. Any of them present → 400 + // (a) claude-opus-4-7 and claude-opus-4-8 deprecated temperature / + // top_p / top_k entirely. Any of them present → 400 // "`temperature` is deprecated for this model." Strip all three. // // (b) Every other claude-4.x model (haiku-4-5, opus-4-1, opus-4-5, @@ -365,7 +365,12 @@ func appendSystemInstruction(existing any, instruction string) any { // yields 400 " is deprecated for this model." (initial // incident: traffic d914275a-0dae-4d13-a811-69e4d432c441). func anthropicModelRejectsSamplingParams(model string) bool { - return strings.HasPrefix(model, "claude-opus-4-7") + // claude-opus-4-7 AND claude-opus-4-8 return 400 + // "`temperature` is deprecated for this model." on any of + // temperature / top_p / top_k (opus-4-8 observed in prod smoke + // 2026-07-03; a temperature-sending client 400s on every call). + return strings.HasPrefix(model, "claude-opus-4-7") || + strings.HasPrefix(model, "claude-opus-4-8") } // anthropicModelRejectsTempTopPTogether reports whether the model diff --git a/packages/ai-gateway/internal/providers/specs/anthropic/codec/codec_test.go b/packages/ai-gateway/internal/providers/specs/anthropic/codec/codec_test.go index 6bfbab21..12c35e47 100644 --- a/packages/ai-gateway/internal/providers/specs/anthropic/codec/codec_test.go +++ b/packages/ai-gateway/internal/providers/specs/anthropic/codec/codec_test.go @@ -1182,6 +1182,8 @@ func TestAnthropicModelRejectsSamplingParams(t *testing.T) { }{ {"claude-opus-4-7", true}, {"claude-opus-4-7-20260101", true}, // dated variant in the same family + {"claude-opus-4-8", true}, // prod smoke 2026-07-03: 400 "temperature is deprecated" + {"claude-opus-4-8-20260101", true}, // dated variant in the same family {"claude-opus-4-1-20250805", false}, {"claude-opus-4-5-20250929", false}, // not yet observed; needs empirical confirmation {"claude-sonnet-4-6", false}, diff --git a/packages/ai-gateway/internal/routing/core/strategynode_json.go b/packages/ai-gateway/internal/routing/core/strategynode_json.go new file mode 100644 index 00000000..e29429f3 --- /dev/null +++ b/packages/ai-gateway/internal/routing/core/strategynode_json.go @@ -0,0 +1,32 @@ +package core + +import "github.com/goccy/go-json" + +// UnmarshalJSON decodes a StrategyNode, then reconciles the ab_split target +// list. The admin UI (source of truth for the config wire shape) authors +// ab_split targets under the generic "targets" key — the same key fallback +// uses for its []StrategyNode list — rather than "abTargets". Every ab_split +// rule persisted through the UI therefore stores {type:"ab_split","targets":[ +// {providerId,modelId,weight}]}. Without this shim those targets decode into +// Targets ([]StrategyNode, which drops the weight) and ABTargets stays empty, +// so the ab_split strategy resolves zero targets and the rule silently routes +// nothing. Hydrating ABTargets from "targets" (only when "abTargets" was not +// supplied) makes existing UI-authored rules route correctly with no data +// migration, while still honoring an explicit "abTargets" key if present. +func (n *StrategyNode) UnmarshalJSON(data []byte) error { + type alias StrategyNode + var a alias + if err := json.Unmarshal(data, &a); err != nil { + return err + } + *n = StrategyNode(a) + if n.Type == "ab_split" && len(n.ABTargets) == 0 { + var probe struct { + ABTargets []ABTarget `json:"targets"` + } + if err := json.Unmarshal(data, &probe); err == nil { + n.ABTargets = probe.ABTargets + } + } + return nil +} diff --git a/packages/ai-gateway/internal/routing/core/strategynode_unmarshal_test.go b/packages/ai-gateway/internal/routing/core/strategynode_unmarshal_test.go new file mode 100644 index 00000000..059c1650 --- /dev/null +++ b/packages/ai-gateway/internal/routing/core/strategynode_unmarshal_test.go @@ -0,0 +1,92 @@ +package core + +import ( + "github.com/goccy/go-json" + "testing" +) + +// TestStrategyNodeUnmarshalABSplitTargets locks the compat shim that lets an +// ab_split rule authored by the admin UI actually route. The UI persists +// ab_split targets under the generic "targets" key; the resolver reads +// ABTargets. Decoding the real UI config blob must populate ABTargets (with +// weights preserved) so the ab_split strategy resolves real targets instead of +// "no abTargets configured". This is the JSON path the in-memory struct tests +// never exercised, which is why the mismatch shipped. +func TestStrategyNodeUnmarshalABSplitTargets(t *testing.T) { + // Verbatim shape written by buildRoutingApiConfig's ab_split branch. + raw := []byte(`{ + "type": "ab_split", + "targets": [ + {"weight": 70, "modelId": "m-a", "providerId": "p-a"}, + {"weight": 30, "modelId": "m-b", "providerId": "p-b"} + ] + }`) + + var node StrategyNode + if err := json.Unmarshal(raw, &node); err != nil { + t.Fatalf("unmarshal ab_split config: %v", err) + } + + if len(node.ABTargets) != 2 { + t.Fatalf("ABTargets not hydrated from \"targets\": got %d, want 2", len(node.ABTargets)) + } + if node.ABTargets[0].ProviderID != "p-a" || node.ABTargets[0].ModelID != "m-a" || node.ABTargets[0].Weight != 70 { + t.Errorf("target[0] = %+v, want {p-a m-a 70}", node.ABTargets[0]) + } + if node.ABTargets[1].ProviderID != "p-b" || node.ABTargets[1].ModelID != "m-b" || node.ABTargets[1].Weight != 30 { + t.Errorf("target[1] = %+v, want {p-b m-b 30}", node.ABTargets[1]) + } +} + +// TestStrategyNodeUnmarshalABSplitPrefersExplicitAbTargets ensures an explicit +// "abTargets" key still wins and is not overwritten by the "targets" shim. +func TestStrategyNodeUnmarshalABSplitPrefersExplicitAbTargets(t *testing.T) { + raw := []byte(`{ + "type": "ab_split", + "abTargets": [{"weight": 100, "modelId": "m-x", "providerId": "p-x"}], + "targets": [{"weight": 1, "modelId": "m-ignored", "providerId": "p-ignored"}] + }`) + + var node StrategyNode + if err := json.Unmarshal(raw, &node); err != nil { + t.Fatalf("unmarshal: %v", err) + } + if len(node.ABTargets) != 1 || node.ABTargets[0].ModelID != "m-x" { + t.Fatalf("explicit abTargets should win, got %+v", node.ABTargets) + } +} + +// TestStrategyNodeUnmarshalMalformed surfaces the decode error rather than +// silently yielding a zero node that would route nothing. +func TestStrategyNodeUnmarshalMalformed(t *testing.T) { + var node StrategyNode + if err := json.Unmarshal([]byte(`{"type": "ab_split", "targets": "not-an-array"}`), &node); err == nil { + t.Fatal("expected error decoding malformed ab_split config, got nil") + } +} + +// TestStrategyNodeUnmarshalFallbackTargetsUnaffected guards against the shim +// leaking into fallback nodes, whose "targets" are []StrategyNode. +func TestStrategyNodeUnmarshalFallbackTargetsUnaffected(t *testing.T) { + raw := []byte(`{ + "type": "fallback", + "targets": [ + {"type": "single", "providerId": "p-a", "modelId": "m-a"}, + {"type": "single", "providerId": "p-b", "modelId": "m-b"} + ] + }`) + + var node StrategyNode + if err := json.Unmarshal(raw, &node); err != nil { + t.Fatalf("unmarshal fallback config: %v", err) + } + if len(node.Targets) != 2 { + t.Fatalf("fallback Targets = %d, want 2", len(node.Targets)) + } + if len(node.ABTargets) != 0 { + t.Errorf("fallback node must not hydrate ABTargets, got %d", len(node.ABTargets)) + } + if node.Targets[0].ProviderID != "p-a" || node.Targets[1].ModelID != "m-b" { + t.Errorf("fallback targets mis-decoded: %+v", node.Targets) + } +} diff --git a/packages/control-plane/internal/ai/virtualkeys/handler/user.go b/packages/control-plane/internal/ai/virtualkeys/handler/user.go index cd6c6744..2ba00961 100644 --- a/packages/control-plane/internal/ai/virtualkeys/handler/user.go +++ b/packages/control-plane/internal/ai/virtualkeys/handler/user.go @@ -122,6 +122,9 @@ func (h *Handler) CreateUserVirtualKey(c echo.Context) error { VKStatus: "active", }) if err != nil { + if msg, ok := vkNameConflict(err); ok { + return c.JSON(http.StatusConflict, errJSON(msg, "duplicate_name", "")) + } h.logger.Error("create user virtual key", "error", err) return c.JSON(http.StatusInternalServerError, errJSON("Internal server error", "server_error", "")) } diff --git a/packages/control-plane/internal/ai/virtualkeys/handler/user_test.go b/packages/control-plane/internal/ai/virtualkeys/handler/user_test.go index 93832fa9..1a2cabdc 100644 --- a/packages/control-plane/internal/ai/virtualkeys/handler/user_test.go +++ b/packages/control-plane/internal/ai/virtualkeys/handler/user_test.go @@ -9,6 +9,7 @@ import ( "time" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" "github.com/labstack/echo/v4" "github.com/pashagolub/pgxmock/v4" ) @@ -230,6 +231,33 @@ func TestCreateUserVirtualKey_DBError(t *testing.T) { } } +// TestCreateUserVirtualKey_DuplicateName covers the 409 envelope when the +// INSERT trips the per-owner personal-name unique index: the user is told +// they already own a key with this name, and no audit entry fires. +func TestCreateUserVirtualKey_DuplicateName(t *testing.T) { + h, mock, _, aud := newHandlerWithMockDB(t) + mock.ExpectQuery(`INSERT INTO "VirtualKey"`). + WithArgs(anyN(14)...). + WillReturnError(&pgconn.PgError{Code: "23505", ConstraintName: "VirtualKey_personal_owner_name_uniq"}) + + c, rec := makeJSONReq(t, http.MethodPost, "/x", `{"name":"mine"}`) + if err := h.CreateUserVirtualKey(c); err != nil { + t.Fatalf("CreateUserVirtualKey: %v", err) + } + if rec.Code != http.StatusConflict { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "duplicate_name") { + t.Errorf("body missing duplicate_name type: %s", rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "You already have a virtual key with this name") { + t.Errorf("body missing per-owner message: %s", rec.Body.String()) + } + if aud.count() != 0 { + t.Errorf("expected no audit on rejected create") + } +} + // UpdateUserVirtualKey (personal) // TestUpdateUserVirtualKey_Unauthenticated covers the no-auth → 401 path. diff --git a/packages/control-plane/internal/ai/virtualkeys/handler/vk.go b/packages/control-plane/internal/ai/virtualkeys/handler/vk.go index e85e866c..685b50c2 100644 --- a/packages/control-plane/internal/ai/virtualkeys/handler/vk.go +++ b/packages/control-plane/internal/ai/virtualkeys/handler/vk.go @@ -2,11 +2,13 @@ package virtualkey import ( "bytes" + "errors" "github.com/goccy/go-json" "io" "net/http" "time" + "github.com/jackc/pgx/v5/pgconn" "github.com/labstack/echo/v4" "github.com/AlphaBitCore/nexus-gateway/packages/control-plane/internal/ai/virtualkeys/vkstore" @@ -71,6 +73,24 @@ func (h *Handler) GetVirtualKey(c echo.Context) error { return c.JSON(http.StatusOK, vk) } +// vkNameConflict maps a unique violation on one of the VirtualKey name +// partial indexes (see tools/db-migrate/schema-extras.sql) to a caller-facing +// message. Application VK names are unique fleet-wide; personal VK names are +// unique per owner. +func vkNameConflict(err error) (string, bool) { + var pgErr *pgconn.PgError + if !errors.As(err, &pgErr) || pgErr.Code != "23505" { + return "", false + } + switch pgErr.ConstraintName { + case "VirtualKey_application_name_uniq": + return "An application virtual key with this name already exists", true + case "VirtualKey_personal_owner_name_uniq": + return "You already have a virtual key with this name", true + } + return "", false +} + func (h *Handler) CreateVirtualKey(c echo.Context) error { var body struct { Name string `json:"name"` @@ -151,6 +171,9 @@ func (h *Handler) CreateVirtualKey(c echo.Context) error { VKStatus: vkStatus, }) if err != nil { + if msg, ok := vkNameConflict(err); ok { + return c.JSON(http.StatusConflict, errJSON(msg, "duplicate_name", "")) + } h.logger.Error("create virtual key", "error", err) return internalServerError(c, "Internal server error") } diff --git a/packages/control-plane/internal/ai/virtualkeys/handler/vk_test.go b/packages/control-plane/internal/ai/virtualkeys/handler/vk_test.go index 26009226..1fca4f6f 100644 --- a/packages/control-plane/internal/ai/virtualkeys/handler/vk_test.go +++ b/packages/control-plane/internal/ai/virtualkeys/handler/vk_test.go @@ -10,6 +10,7 @@ import ( "time" "github.com/jackc/pgx/v5" + "github.com/jackc/pgx/v5/pgconn" "github.com/labstack/echo/v4" "github.com/pashagolub/pgxmock/v4" @@ -549,6 +550,53 @@ func TestCreateVirtualKey_DBError(t *testing.T) { } } +// TestCreateVirtualKey_DuplicateName covers the 409 envelope when the INSERT +// trips the fleet-wide application-name partial unique index: the caller gets +// a duplicate_name conflict, and no audit entry fires. +func TestCreateVirtualKey_DuplicateName(t *testing.T) { + h, mock, _, aud := newHandlerWithMockDB(t) + mock.ExpectQuery(`INSERT INTO "VirtualKey"`). + WithArgs(anyN(14)...). + WillReturnError(&pgconn.PgError{Code: "23505", ConstraintName: "VirtualKey_application_name_uniq"}) + + future := time.Now().UTC().Add(30 * 24 * time.Hour) + body := `{"name":"taken","vkType":"application","projectId":"p-1","expiresAt":"` + future.Format(time.RFC3339) + `"}` + c, rec := makeJSONReq(t, http.MethodPost, "/api/admin/virtual-keys", body) + if err := h.CreateVirtualKey(c); err != nil { + t.Fatalf("CreateVirtualKey: %v", err) + } + if rec.Code != http.StatusConflict { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "duplicate_name") { + t.Errorf("body missing duplicate_name type: %s", rec.Body.String()) + } + if !strings.Contains(rec.Body.String(), "application virtual key with this name already exists") { + t.Errorf("body missing application-name message: %s", rec.Body.String()) + } + if aud.count() != 0 { + t.Errorf("expected no audit on rejected create") + } +} + +// TestCreateVirtualKey_UnrelatedUniqueViolation pins the helper's boundary: a +// 23505 on any constraint other than the two name indexes (here the keyHash +// unique, an internal random-collision fault, not caller input) stays a 500. +func TestCreateVirtualKey_UnrelatedUniqueViolation(t *testing.T) { + h, mock, _, _ := newHandlerWithMockDB(t) + mock.ExpectQuery(`INSERT INTO "VirtualKey"`). + WithArgs(anyN(14)...). + WillReturnError(&pgconn.PgError{Code: "23505", ConstraintName: "VirtualKey_keyHash_key"}) + + c, rec := makeJSONReq(t, http.MethodPost, "/x", `{"name":"n","vkType":"personal"}`) + if err := h.CreateVirtualKey(c); err != nil { + t.Fatalf("CreateVirtualKey: %v", err) + } + if rec.Code != http.StatusInternalServerError { + t.Fatalf("status=%d body=%s", rec.Code, rec.Body.String()) + } +} + // TestCreateVirtualKey_NoAuth exercises the aa==nil branch — ownerID stays // nil and the INSERT still succeeds. func TestCreateVirtualKey_NoAuth(t *testing.T) { diff --git a/packages/control-plane/internal/identity/authserver/login/displayname_test.go b/packages/control-plane/internal/identity/authserver/login/displayname_test.go index 35333d0b..39f43740 100644 --- a/packages/control-plane/internal/identity/authserver/login/displayname_test.go +++ b/packages/control-plane/internal/identity/authserver/login/displayname_test.go @@ -24,8 +24,8 @@ func TestSAMLNameWithFallback(t *testing.T) { {"given + family composed", []saml.Attribute{attr(givenURI, "Carol"), attr(surnameURI, "King")}, "", "Carol King"}, {"given only", []saml.Attribute{attr("givenName", "Dave")}, "", "Dave"}, {"no name + no email → empty", []saml.Attribute{attr("dept", "eng")}, "", ""}, - {"email-valued name falls to humanized nickname", []saml.Attribute{attr(fullNameURI, "steve.chen@itechchoice.com"), attr("http://schemas.auth0.com/nickname", "steve.chen")}, "steve.chen@itechchoice.com", "steve chen"}, - {"email-valued name, no nickname → humanized email local-part", []saml.Attribute{attr(fullNameURI, "steve.chen@itechchoice.com")}, "steve.chen@itechchoice.com", "steve chen"}, + {"email-valued name falls to humanized nickname", []saml.Attribute{attr(fullNameURI, "steve.chen@alphabitcore.com"), attr("http://schemas.auth0.com/nickname", "steve.chen")}, "steve.chen@alphabitcore.com", "steve chen"}, + {"email-valued name, no nickname → humanized email local-part", []saml.Attribute{attr(fullNameURI, "steve.chen@alphabitcore.com")}, "steve.chen@alphabitcore.com", "steve chen"}, {"nickname handle when no full/given name", []saml.Attribute{attr("nickname", "ada.lovelace")}, "", "ada lovelace"}, {"no name attribute → humanized email local-part", []saml.Attribute{attr("dept", "eng")}, "grace.hopper@navy.mil", "grace hopper"}, } @@ -45,7 +45,7 @@ func TestHumanizeHandle(t *testing.T) { want string }{ {"steve.chen", "steve chen"}, - {"steve.chen@itechchoice.com", "steve chen"}, + {"steve.chen@alphabitcore.com", "steve chen"}, {"steve_chen", "steve chen"}, {"john.doe.42", "john doe"}, {"gracehopper", "gracehopper"}, @@ -80,9 +80,9 @@ func TestOIDCDisplayName(t *testing.T) { {"preferred_username humanized when no name/given", map[string]any{"preferred_username": "grace.hopper"}, "g@x", "sub-3", "grace hopper"}, {"humanized email local-part when no name claims", map[string]any{}, "harry.styles@x.com", "sub-4", "harry styles"}, {"falls back to subject when no name + no email", map[string]any{}, "", "sub-5", "sub-5"}, - {"email-valued name falls to humanized preferred_username", map[string]any{"name": "steve.chen@itechchoice.com", "preferred_username": "steve.chen"}, "steve.chen@itechchoice.com", "sub-6", "steve chen"}, - {"email-valued name falls to humanized nickname", map[string]any{"name": "steve.chen@itechchoice.com", "nickname": "steve.chen"}, "steve.chen@itechchoice.com", "sub-7", "steve chen"}, - {"email-valued name, no handles → humanized email local-part", map[string]any{"name": "steve.chen@itechchoice.com"}, "steve.chen@itechchoice.com", "sub-8", "steve chen"}, + {"email-valued name falls to humanized preferred_username", map[string]any{"name": "steve.chen@alphabitcore.com", "preferred_username": "steve.chen"}, "steve.chen@alphabitcore.com", "sub-6", "steve chen"}, + {"email-valued name falls to humanized nickname", map[string]any{"name": "steve.chen@alphabitcore.com", "nickname": "steve.chen"}, "steve.chen@alphabitcore.com", "sub-7", "steve chen"}, + {"email-valued name, no handles → humanized email local-part", map[string]any{"name": "steve.chen@alphabitcore.com"}, "steve.chen@alphabitcore.com", "sub-8", "steve chen"}, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { diff --git a/packages/control-plane/internal/identity/authserver/store/federated_store_mock_test.go b/packages/control-plane/internal/identity/authserver/store/federated_store_mock_test.go index ac88b925..edb15bc6 100644 --- a/packages/control-plane/internal/identity/authserver/store/federated_store_mock_test.go +++ b/packages/control-plane/internal/identity/authserver/store/federated_store_mock_test.go @@ -269,10 +269,10 @@ func TestFederatedStore_RefreshUserProfile_Success(t *testing.T) { ctx := context.Background() mock.ExpectExec(`UPDATE "NexusUser"\s+SET "displayName" = COALESCE\(NULLIF\(\$2, ''\), "displayName"\),\s+email\s+= COALESCE\(NULLIF\(\$3, ''\), email\),\s+"updatedAt"\s+= NOW\(\)\s+WHERE id = \$1`). - WithArgs("user_1", "steve chen", "steve.chen@itechchoice.com"). + WithArgs("user_1", "steve chen", "steve.chen@alphabitcore.com"). WillReturnResult(pgxmock.NewResult("UPDATE", 1)) - if err := s.RefreshUserProfile(ctx, "user_1", "steve chen", "steve.chen@itechchoice.com"); err != nil { + if err := s.RefreshUserProfile(ctx, "user_1", "steve chen", "steve.chen@alphabitcore.com"); err != nil { t.Fatalf("RefreshUserProfile: %v", err) } } From 3f7456b265a04561078e7e0a381144c174a77aa3 Mon Sep 17 00:00:00 2001 From: nexus Date: Fri, 3 Jul 2026 22:41:06 +0800 Subject: [PATCH 6/8] feat(ui): sync control-plane UI, shared UI components, and tests (sync 10) Frontend code and repository-level test suites: - Personal VK create + duplicate-name-check hook, FormInput - AI-gateway routing rule create/detail wizards + shared config - Compliance rule-packs form - i18n locale bundles (en/es/zh) - Vitest suites for the above Co-Authored-By: Claude Opus 4.7 (1M context) --- .../public/locales/en/pages.json | 21 ++++- .../public/locales/es/pages.json | 21 ++++- .../public/locales/zh/pages.json | 21 ++++- .../ai-gateway/personalVirtualKeys.ts | 4 +- .../src/hooks/useDuplicateNameCheck.test.tsx | 86 +++++++++++++++++++ .../src/hooks/useDuplicateNameCheck.ts | 59 +++++++++++++ .../src/i18n/locales/en/pages.json | 16 +++- .../src/i18n/locales/es/pages.json | 16 +++- .../src/i18n/locales/zh/pages.json | 16 +++- .../src/lib/forms/FormInput.test.tsx | 35 ++++++++ .../src/lib/forms/FormInput.tsx | 11 +++ .../account/personal-vks/PersonalVKCreate.tsx | 18 +++- .../wizard/ProviderWizard.module.css | 10 ++- .../routing-rule-config.splitWeights.test.ts | 37 ++++++++ .../routing/_shared/routing-rule-config.ts | 26 ++++++ .../create/RoutingRuleCreate.module.css | 18 ++++ .../routing/create/RoutingRuleCreatePage.tsx | 54 +++++++++++- .../routing/create/useRoutingRuleCreate.ts | 74 +++++++++------- .../detail/RoutingRuleDetail.module.css | 7 ++ .../routing/detail/RoutingRuleEditForm.tsx | 14 ++- .../routing/detail/useRoutingRuleDetail.ts | 13 ++- .../virtual-keys/VirtualKeyCreate.tsx | 21 ++++- .../rule-packs/form/RulePackCreatePage.tsx | 45 +++++++--- .../rule-packs/form/rulePackRules.test.ts | 49 +++++++++++ .../rule-packs/form/rulePackRules.ts | 30 +++++++ .../ai-gateway/personalVirtualKeys.test.ts | 4 +- .../create/RoutingRuleCreatePage.test.tsx | 25 +++--- 27 files changed, 663 insertions(+), 88 deletions(-) create mode 100644 packages/control-plane-ui/src/hooks/useDuplicateNameCheck.test.tsx create mode 100644 packages/control-plane-ui/src/hooks/useDuplicateNameCheck.ts create mode 100644 packages/control-plane-ui/src/lib/forms/FormInput.test.tsx create mode 100644 packages/control-plane-ui/src/pages/ai-gateway/routing/_shared/routing-rule-config.splitWeights.test.ts create mode 100644 packages/control-plane-ui/src/pages/compliance/rule-packs/form/rulePackRules.test.ts diff --git a/packages/control-plane-ui/public/locales/en/pages.json b/packages/control-plane-ui/public/locales/en/pages.json index 6176dae3..50f4e28d 100644 --- a/packages/control-plane-ui/public/locales/en/pages.json +++ b/packages/control-plane-ui/public/locales/en/pages.json @@ -406,6 +406,11 @@ "apiVersion": "API Version", "apiVersionHelp": "Optional. Provider API version string (e.g. 2024-02-01 for Azure OpenAI). Leave blank unless the upstream endpoint requires a version segment.", "apiVersionPlaceholder": "e.g. 2024-02-01", + "servesResponsesApi": "Serves OpenAI Responses API", + "servesResponsesApiHelp": "Leave unset to use the provider type's default. Disable for OpenAI-compatible endpoints that only implement /v1/chat/completions.", + "servesResponsesApiOption_default": "Use adapter default", + "servesResponsesApiOption_true": "Enabled", + "servesResponsesApiOption_false": "Disabled (chat completions only)", "modelCode": "Model ID", "modelCodeLabel": "Model ID *", "placeholderModelCode": "e.g. gpt-4o (defaults to Upstream Model ID)", @@ -525,6 +530,11 @@ "defaultModelFallback": "Default Model (fallback if router fails)", "providerModel": "Provider / Model", "splitPercent": "Split %", + "weightSumError": "Target percentages must add up to 100% (currently {{total}}%).", + "nameRequired": "Name is required.", + "fallbackIncomplete": "Finish or remove the partially-filled fallback target — set both provider and model.", + "smartMatchAutoRequired": "Smart routing must match the \"auto\" model literal — add it under Match conditions.", + "matchEmptyNote": "No match conditions set — this rule will apply to all requests.", "weight": "Weight", "remove": "Remove", "addTarget": "+ Add Target", @@ -878,6 +888,9 @@ "createNameRequired": "Name is required", "createVersionRequired": "Version is required", "createMaintainerRequired": "Maintainer is required", + "createNameInvalid": "Use \"/\" (lowercase), e.g. acme/pii-rules.", + "createVersionInvalid": "Use a v-prefixed semantic version, e.g. v1.0.0.", + "createNameHelp": "Format: / (lowercase), e.g. acme/pii-rules.", "createDescriptionHelp": "Optional. Shown to operators browsing the catalog.", "createRulesLabel": "Rules (JSON array)", "createRulesHelp": "Each rule requires ruleId, category, severity (hard|soft|info), pattern. Optional: flags, description, labels.", @@ -1054,6 +1067,7 @@ "virtualKeyCreatedMsg": "Virtual Key Created Msg", "nameHelpText": "Human-readable label for this virtual key. Used as the display name in dashboards and audit logs.", "namePlaceholder": "engineering-openai", + "nameTakenApplication": "An application virtual key with this name already exists. Choose a different name.", "projectTooltip": "Project this virtual key belongs to. Used for traffic attribution and quota calculation.", "expirationTooltip": "Application virtual keys require an expiration date capped at 3 months from now. The key cannot be used after it expires.", "disabledKeysTooltip": "Disabled Keys Tooltip", @@ -4322,6 +4336,7 @@ "createSubtitle": "Create a new personal virtual key", "name": "Name", "namePlaceholder": "my-api-key", + "nameTaken": "You already have a virtual key with this name. Choose a different name.", "sourceApp": "Source App", "status": "Status", "rpm": "Rate Limit (RPM)", @@ -6289,9 +6304,9 @@ "endpointExamplePlaceholder": "OpenAI (chat)" }, "rulePacks": { - "slugPlaceholder": "acme/rules", - "versionPlaceholder": "v1.0.0", - "sourcePlaceholder": "customer" + "slugPlaceholder": "e.g. acme/rules", + "versionPlaceholder": "e.g. v1.0.0", + "sourcePlaceholder": "e.g. customer" }, "idp": { "clientIdPlaceholder": "0oa1abc2def3ghi4j5k6" diff --git a/packages/control-plane-ui/public/locales/es/pages.json b/packages/control-plane-ui/public/locales/es/pages.json index d0e64e7b..ceee3a5a 100644 --- a/packages/control-plane-ui/public/locales/es/pages.json +++ b/packages/control-plane-ui/public/locales/es/pages.json @@ -406,6 +406,11 @@ "apiVersion": "Versión de API", "apiVersionHelp": "Opcional. Cadena de versión de API del proveedor (p. ej. 2024-02-01 para Azure OpenAI). Dejar en blanco si el endpoint no requiere un segmento de versión.", "apiVersionPlaceholder": "p. ej. 2024-02-01", + "servesResponsesApi": "Sirve la API de Responses de OpenAI", + "servesResponsesApiHelp": "Déjalo sin definir para usar el valor predeterminado del tipo de proveedor. Desactívalo para endpoints compatibles con OpenAI que solo implementan /v1/chat/completions.", + "servesResponsesApiOption_default": "Usar el valor predeterminado del adaptador", + "servesResponsesApiOption_true": "Activado", + "servesResponsesApiOption_false": "Desactivado (solo chat completions)", "modelCode": "Model ID", "modelCodeLabel": "Model ID *", "placeholderModelCode": "p. ej. gpt-4o (por defecto usa el ID de modelo upstream)", @@ -525,6 +530,11 @@ "defaultModelFallback": "Modelo predeterminado (respaldo si el enrutador falla)", "providerModel": "Proveedor / Modelo", "splitPercent": "% de división", + "weightSumError": "Los porcentajes de los objetivos deben sumar 100 % (actualmente {{total}} %).", + "nameRequired": "El nombre es obligatorio.", + "fallbackIncomplete": "Completa o elimina el objetivo de reserva a medio llenar: define proveedor y modelo.", + "smartMatchAutoRequired": "El enrutamiento inteligente debe coincidir con el literal de modelo «auto»: añádelo en Condiciones de coincidencia.", + "matchEmptyNote": "Sin condiciones de coincidencia: esta regla se aplicará a todas las solicitudes.", "weight": "Peso", "remove": "Eliminar", "addTarget": "+ Agregar objetivo", @@ -878,6 +888,9 @@ "createNameRequired": "El nombre es obligatorio", "createVersionRequired": "La versión es obligatoria", "createMaintainerRequired": "El mantenedor es obligatorio", + "createNameInvalid": "Usa «/»: solo minúsculas, números y guiones, p. ej. acme/pii-rules.", + "createVersionInvalid": "Usa una versión semántica con prefijo v, p. ej. v1.0.0.", + "createNameHelp": "Formato: / (minúsculas), p. ej. acme/pii-rules.", "createDescriptionHelp": "Opcional. Se muestra a los operadores que exploran el catálogo.", "createRulesLabel": "Reglas (array JSON)", "createRulesHelp": "Cada regla requiere ruleId, category, severity (hard|soft|info) y pattern. Opcional: flags, description, labels.", @@ -1054,6 +1067,7 @@ "virtualKeyCreatedMsg": "Mensaje de clave virtual creada", "nameHelpText": "Etiqueta legible para esta clave virtual. Se usa como nombre de visualización en paneles y registros de auditoría.", "namePlaceholder": "engineering-openai", + "nameTakenApplication": "Ya existe una clave virtual de aplicación con este nombre. Elige un nombre diferente.", "projectTooltip": "Proyecto al que pertenece esta clave virtual. Se utiliza para la atribución del tráfico y el cálculo de cuotas.", "expirationTooltip": "Las claves virtuales de aplicación requieren una fecha de vencimiento máxima de 3 meses. La clave no se puede usar después de su vencimiento.", "disabledKeysTooltip": "Tooltip de claves deshabilitadas", @@ -4322,6 +4336,7 @@ "createSubtitle": "Crear una nueva clave virtual personal", "name": "Nombre", "namePlaceholder": "my-api-key", + "nameTaken": "Ya tienes una clave virtual con este nombre. Elige un nombre diferente.", "sourceApp": "App de origen", "status": "Estado", "rpm": "Límite de tasa (RPM)", @@ -6289,9 +6304,9 @@ "endpointExamplePlaceholder": "OpenAI (chat)" }, "rulePacks": { - "slugPlaceholder": "acme/rules", - "versionPlaceholder": "v1.0.0", - "sourcePlaceholder": "customer" + "slugPlaceholder": "p. ej. acme/rules", + "versionPlaceholder": "p. ej. v1.0.0", + "sourcePlaceholder": "p. ej. customer" }, "idp": { "clientIdPlaceholder": "0oa1abc2def3ghi4j5k6" diff --git a/packages/control-plane-ui/public/locales/zh/pages.json b/packages/control-plane-ui/public/locales/zh/pages.json index 9e165cf2..86aa0627 100644 --- a/packages/control-plane-ui/public/locales/zh/pages.json +++ b/packages/control-plane-ui/public/locales/zh/pages.json @@ -406,6 +406,11 @@ "apiVersion": "API 版本", "apiVersionHelp": "可选。上游提供商的 API 版本字符串（如 Azure OpenAI 的 2024-02-01）。若上游端点不要求版本号，请留空。", "apiVersionPlaceholder": "例如 2024-02-01", + "servesResponsesApi": "提供 OpenAI Responses API", + "servesResponsesApiHelp": "留空则使用该供应商类型的默认值。对于仅实现 /v1/chat/completions 的 OpenAI 兼容端点，请将其禁用。", + "servesResponsesApiOption_default": "使用适配器默认值", + "servesResponsesApiOption_true": "启用", + "servesResponsesApiOption_false": "禁用（仅 chat completions）", "modelCode": "模型ID", "modelCodeLabel": "Model ID *", "placeholderModelCode": "例如 gpt-4o（留空则使用上游模型ID）", @@ -525,6 +530,11 @@ "defaultModelFallback": "默认模型（路由器失败时的回退）", "providerModel": "供应商 / 模型", "splitPercent": "分流百分比", + "weightSumError": "各目标的百分比之和必须为 100%(当前为 {{total}}%)。", + "nameRequired": "名称必填。", + "fallbackIncomplete": "请补全或删除半填的回退目标——供应商和模型都要选。", + "smartMatchAutoRequired": "智能路由必须匹配“auto”模型字面量——请在匹配条件里添加。", + "matchEmptyNote": "未设置匹配条件——该规则将应用于所有请求。", "weight": "权重", "remove": "移除", "addTarget": "+ 添加目标", @@ -878,6 +888,9 @@ "createNameRequired": "名称必填", "createVersionRequired": "版本必填", "createMaintainerRequired": "维护者必填", + "createNameInvalid": "请使用“/”格式——仅限小写字母、数字和连字符,例如 acme/pii-rules。", + "createVersionInvalid": "请使用带 v 前缀的语义化版本号,例如 v1.0.0。", + "createNameHelp": "格式:/(小写),例如 acme/pii-rules。", "createDescriptionHelp": "可选。在目录中展示给操作员。", "createRulesLabel": "规则 (JSON 数组)", "createRulesHelp": "每条规则需包含 ruleId、category、severity（hard|soft|info）、pattern；可选 flags、description、labels。", @@ -1054,6 +1067,7 @@ "virtualKeyCreatedMsg": "虚拟密钥已创建", "nameHelpText": "虚拟密钥的人读名称，在仪表板和审计日志中作为显示名使用。", "namePlaceholder": "engineering-openai", + "nameTakenApplication": "已存在同名的应用类型虚拟密钥，请更换名称。", "projectTooltip": "此虚拟密钥所属的项目，用于流量归属和配额计算。", "expirationTooltip": "应用类型虚拟密钥必须设置过期时间，且不得超过 3 个月。密钥过期后将无法用于认证。", "disabledKeysTooltip": "已禁用的密钥立即拒绝所有调用，但保留配置以供审计或重新启用。", @@ -4322,6 +4336,7 @@ "createSubtitle": "创建新的个人虚拟密钥", "name": "名称", "namePlaceholder": "my-api-key", + "nameTaken": "你已有同名的虚拟密钥，请更换名称。", "sourceApp": "来源应用", "status": "状态", "rpm": "速率限制 (RPM)", @@ -6289,9 +6304,9 @@ "endpointExamplePlaceholder": "OpenAI (chat)" }, "rulePacks": { - "slugPlaceholder": "acme/rules", - "versionPlaceholder": "v1.0.0", - "sourcePlaceholder": "customer" + "slugPlaceholder": "例如 acme/rules", + "versionPlaceholder": "例如 v1.0.0", + "sourcePlaceholder": "例如 customer" }, "idp": { "clientIdPlaceholder": "0oa1abc2def3ghi4j5k6" diff --git a/packages/control-plane-ui/src/api/services/ai-gateway/personalVirtualKeys.ts b/packages/control-plane-ui/src/api/services/ai-gateway/personalVirtualKeys.ts index 832727f0..3433feaa 100644 --- a/packages/control-plane-ui/src/api/services/ai-gateway/personalVirtualKeys.ts +++ b/packages/control-plane-ui/src/api/services/ai-gateway/personalVirtualKeys.ts @@ -11,8 +11,8 @@ export interface CreatePersonalVKInput { } export const personalVKApi = { - list: () => - api.get<{ data: VirtualKey[]; total: number }>('/api/my/virtual-keys'), + list: (params?: Record) => + api.get<{ data: VirtualKey[]; total: number }>('/api/my/virtual-keys', params), create: (data: CreatePersonalVKInput) => api.post('/api/my/virtual-keys', data), diff --git a/packages/control-plane-ui/src/hooks/useDuplicateNameCheck.test.tsx b/packages/control-plane-ui/src/hooks/useDuplicateNameCheck.test.tsx new file mode 100644 index 00000000..fdeff34b --- /dev/null +++ b/packages/control-plane-ui/src/hooks/useDuplicateNameCheck.test.tsx @@ -0,0 +1,86 @@ +import { describe, it, expect, vi } from 'vitest'; +import { renderHook, act, waitFor } from '@testing-library/react'; +import { useForm } from 'react-hook-form'; +import { useDuplicateNameCheck } from './useDuplicateNameCheck'; + +type Values = { name: string }; + +function setup(isTaken: (name: string) => Promise) { + return renderHook(() => { + const form = useForm({ defaultValues: { name: '' } }); + const check = useDuplicateNameCheck({ + form, + field: 'name', + message: 'name already taken', + isTaken, + }); + return { form, check }; + }); +} + +describe('useDuplicateNameCheck', () => { + it('sets the field error when the probe reports the name as taken', async () => { + const { result } = setup(async () => true); + act(() => result.current.form.setValue('name', 'dup')); + await act(() => result.current.check('dup')); + await waitFor(() => { + expect(result.current.form.getFieldState('name').error?.message).toBe('name already taken'); + }); + }); + + it('leaves the field clean when the name is available', async () => { + const { result } = setup(async () => false); + act(() => result.current.form.setValue('name', 'fresh')); + await act(() => result.current.check('fresh')); + expect(result.current.form.getFieldState('name').error).toBeUndefined(); + }); + + it('skips empty and whitespace-only values without probing', async () => { + const probe = vi.fn(async () => true); + const { result } = setup(probe); + await act(() => result.current.check('')); + await act(() => result.current.check(' ')); + expect(probe).not.toHaveBeenCalled(); + expect(result.current.form.getFieldState('name').error).toBeUndefined(); + }); + + it('stays silent when the probe fails (advisory only, server 409 is authoritative)', async () => { + const { result } = setup(async () => { throw new Error('network down'); }); + act(() => result.current.form.setValue('name', 'dup')); + await act(() => result.current.check('dup')); + expect(result.current.form.getFieldState('name').error).toBeUndefined(); + }); + + it('ignores a taken verdict if the field value changed while the probe was in flight', async () => { + let resolveProbe!: (v: boolean) => void; + const { result } = setup(() => new Promise(r => { resolveProbe = r; })); + act(() => result.current.form.setValue('name', 'old-name')); + let pending!: Promise; + act(() => { pending = result.current.check('old-name'); }); + // User re-focused and edited before the probe answered. + act(() => result.current.form.setValue('name', 'new-name')); + resolveProbe(true); + await act(() => pending); + expect(result.current.form.getFieldState('name').error).toBeUndefined(); + }); + + it('drops a stale probe response when a newer blur superseded it', async () => { + const resolvers: Array<(v: boolean) => void> = []; + const { result } = setup(() => new Promise(r => { resolvers.push(r); })); + act(() => result.current.form.setValue('name', 'first')); + let firstCheck!: Promise; + act(() => { firstCheck = result.current.check('first'); }); + + act(() => result.current.form.setValue('name', 'second')); + let secondCheck!: Promise; + act(() => { secondCheck = result.current.check('second'); }); + + // Second probe answers "available" first, then the stale first probe + // answers "taken" — the stale verdict must not surface an error. + resolvers[1](false); + await act(() => secondCheck); + resolvers[0](true); + await act(() => firstCheck); + expect(result.current.form.getFieldState('name').error).toBeUndefined(); + }); +}); diff --git a/packages/control-plane-ui/src/hooks/useDuplicateNameCheck.ts b/packages/control-plane-ui/src/hooks/useDuplicateNameCheck.ts new file mode 100644 index 00000000..f15c118e --- /dev/null +++ b/packages/control-plane-ui/src/hooks/useDuplicateNameCheck.ts @@ -0,0 +1,59 @@ +/** + * useDuplicateNameCheck — advisory blur-time duplicate-name probe for create + * forms whose names must be unique. + * + * Returns a `checkName(value)` callback to run when the user leaves the name + * field. It asks the caller-supplied `isTaken` probe and, on a hit, sets an + * inline field error so the user learns about the collision before filling in + * the rest of the form. The check is advisory only: probe failures are + * swallowed and never block input or submission — the server's 409 on create + * remains the authoritative (race-proof) duplicate rejection. + * + * The error is set via react-hook-form's `setError`, so it clears as soon as + * the user edits the field again (reValidateMode: 'onChange'). + */ +import { useCallback, useRef } from 'react'; +import type { FieldValues, Path, UseFormReturn } from 'react-hook-form'; + +export interface DuplicateNameCheckOptions { + form: UseFormReturn; + field: Path; + /** Inline error message shown when the name is already taken. */ + message: string; + /** Resolves true when `name` already exists in the relevant scope. */ + isTaken: (name: string) => Promise; +} + +export function useDuplicateNameCheck({ + form, + field, + message, + isTaken, +}: DuplicateNameCheckOptions) { + // Monotonic sequence: a slow probe for an old value must not clobber the + // state produced by a newer blur. + const seq = useRef(0); + + return useCallback( + async (raw: string) => { + const name = raw.trim(); + if (!name) return; + const mySeq = ++seq.current; + try { + const taken = await isTaken(name); + if (seq.current !== mySeq) return; // stale response — a newer blur ran + // The user may have re-focused and edited while the probe was in + // flight; only flag the value the field still holds. + const current = String(form.getValues(field) ?? '').trim(); + if (current !== name) return; + if (taken) { + form.setError(field, { type: 'duplicate', message }); + } + } catch { + // Advisory only — on probe failure stay silent; the create call's + // 409 still rejects duplicates. + } + }, + [form, field, message, isTaken], + ); +} diff --git a/packages/control-plane-ui/src/i18n/locales/en/pages.json b/packages/control-plane-ui/src/i18n/locales/en/pages.json index 9a2beb95..50f4e28d 100644 --- a/packages/control-plane-ui/src/i18n/locales/en/pages.json +++ b/packages/control-plane-ui/src/i18n/locales/en/pages.json @@ -530,6 +530,11 @@ "defaultModelFallback": "Default Model (fallback if router fails)", "providerModel": "Provider / Model", "splitPercent": "Split %", + "weightSumError": "Target percentages must add up to 100% (currently {{total}}%).", + "nameRequired": "Name is required.", + "fallbackIncomplete": "Finish or remove the partially-filled fallback target — set both provider and model.", + "smartMatchAutoRequired": "Smart routing must match the \"auto\" model literal — add it under Match conditions.", + "matchEmptyNote": "No match conditions set — this rule will apply to all requests.", "weight": "Weight", "remove": "Remove", "addTarget": "+ Add Target", @@ -883,6 +888,9 @@ "createNameRequired": "Name is required", "createVersionRequired": "Version is required", "createMaintainerRequired": "Maintainer is required", + "createNameInvalid": "Use \"/\" (lowercase), e.g. acme/pii-rules.", + "createVersionInvalid": "Use a v-prefixed semantic version, e.g. v1.0.0.", + "createNameHelp": "Format: / (lowercase), e.g. acme/pii-rules.", "createDescriptionHelp": "Optional. Shown to operators browsing the catalog.", "createRulesLabel": "Rules (JSON array)", "createRulesHelp": "Each rule requires ruleId, category, severity (hard|soft|info), pattern. Optional: flags, description, labels.", @@ -1059,6 +1067,7 @@ "virtualKeyCreatedMsg": "Virtual Key Created Msg", "nameHelpText": "Human-readable label for this virtual key. Used as the display name in dashboards and audit logs.", "namePlaceholder": "engineering-openai", + "nameTakenApplication": "An application virtual key with this name already exists. Choose a different name.", "projectTooltip": "Project this virtual key belongs to. Used for traffic attribution and quota calculation.", "expirationTooltip": "Application virtual keys require an expiration date capped at 3 months from now. The key cannot be used after it expires.", "disabledKeysTooltip": "Disabled Keys Tooltip", @@ -4327,6 +4336,7 @@ "createSubtitle": "Create a new personal virtual key", "name": "Name", "namePlaceholder": "my-api-key", + "nameTaken": "You already have a virtual key with this name. Choose a different name.", "sourceApp": "Source App", "status": "Status", "rpm": "Rate Limit (RPM)", @@ -6294,9 +6304,9 @@ "endpointExamplePlaceholder": "OpenAI (chat)" }, "rulePacks": { - "slugPlaceholder": "acme/rules", - "versionPlaceholder": "v1.0.0", - "sourcePlaceholder": "customer" + "slugPlaceholder": "e.g. acme/rules", + "versionPlaceholder": "e.g. v1.0.0", + "sourcePlaceholder": "e.g. customer" }, "idp": { "clientIdPlaceholder": "0oa1abc2def3ghi4j5k6" diff --git a/packages/control-plane-ui/src/i18n/locales/es/pages.json b/packages/control-plane-ui/src/i18n/locales/es/pages.json index bd885c7f..ceee3a5a 100644 --- a/packages/control-plane-ui/src/i18n/locales/es/pages.json +++ b/packages/control-plane-ui/src/i18n/locales/es/pages.json @@ -530,6 +530,11 @@ "defaultModelFallback": "Modelo predeterminado (respaldo si el enrutador falla)", "providerModel": "Proveedor / Modelo", "splitPercent": "% de división", + "weightSumError": "Los porcentajes de los objetivos deben sumar 100 % (actualmente {{total}} %).", + "nameRequired": "El nombre es obligatorio.", + "fallbackIncomplete": "Completa o elimina el objetivo de reserva a medio llenar: define proveedor y modelo.", + "smartMatchAutoRequired": "El enrutamiento inteligente debe coincidir con el literal de modelo «auto»: añádelo en Condiciones de coincidencia.", + "matchEmptyNote": "Sin condiciones de coincidencia: esta regla se aplicará a todas las solicitudes.", "weight": "Peso", "remove": "Eliminar", "addTarget": "+ Agregar objetivo", @@ -883,6 +888,9 @@ "createNameRequired": "El nombre es obligatorio", "createVersionRequired": "La versión es obligatoria", "createMaintainerRequired": "El mantenedor es obligatorio", + "createNameInvalid": "Usa «/»: solo minúsculas, números y guiones, p. ej. acme/pii-rules.", + "createVersionInvalid": "Usa una versión semántica con prefijo v, p. ej. v1.0.0.", + "createNameHelp": "Formato: / (minúsculas), p. ej. acme/pii-rules.", "createDescriptionHelp": "Opcional. Se muestra a los operadores que exploran el catálogo.", "createRulesLabel": "Reglas (array JSON)", "createRulesHelp": "Cada regla requiere ruleId, category, severity (hard|soft|info) y pattern. Opcional: flags, description, labels.", @@ -1059,6 +1067,7 @@ "virtualKeyCreatedMsg": "Mensaje de clave virtual creada", "nameHelpText": "Etiqueta legible para esta clave virtual. Se usa como nombre de visualización en paneles y registros de auditoría.", "namePlaceholder": "engineering-openai", + "nameTakenApplication": "Ya existe una clave virtual de aplicación con este nombre. Elige un nombre diferente.", "projectTooltip": "Proyecto al que pertenece esta clave virtual. Se utiliza para la atribución del tráfico y el cálculo de cuotas.", "expirationTooltip": "Las claves virtuales de aplicación requieren una fecha de vencimiento máxima de 3 meses. La clave no se puede usar después de su vencimiento.", "disabledKeysTooltip": "Tooltip de claves deshabilitadas", @@ -4327,6 +4336,7 @@ "createSubtitle": "Crear una nueva clave virtual personal", "name": "Nombre", "namePlaceholder": "my-api-key", + "nameTaken": "Ya tienes una clave virtual con este nombre. Elige un nombre diferente.", "sourceApp": "App de origen", "status": "Estado", "rpm": "Límite de tasa (RPM)", @@ -6294,9 +6304,9 @@ "endpointExamplePlaceholder": "OpenAI (chat)" }, "rulePacks": { - "slugPlaceholder": "acme/rules", - "versionPlaceholder": "v1.0.0", - "sourcePlaceholder": "customer" + "slugPlaceholder": "p. ej. acme/rules", + "versionPlaceholder": "p. ej. v1.0.0", + "sourcePlaceholder": "p. ej. customer" }, "idp": { "clientIdPlaceholder": "0oa1abc2def3ghi4j5k6" diff --git a/packages/control-plane-ui/src/i18n/locales/zh/pages.json b/packages/control-plane-ui/src/i18n/locales/zh/pages.json index 808243a2..86aa0627 100644 --- a/packages/control-plane-ui/src/i18n/locales/zh/pages.json +++ b/packages/control-plane-ui/src/i18n/locales/zh/pages.json @@ -530,6 +530,11 @@ "defaultModelFallback": "默认模型（路由器失败时的回退）", "providerModel": "供应商 / 模型", "splitPercent": "分流百分比", + "weightSumError": "各目标的百分比之和必须为 100%(当前为 {{total}}%)。", + "nameRequired": "名称必填。", + "fallbackIncomplete": "请补全或删除半填的回退目标——供应商和模型都要选。", + "smartMatchAutoRequired": "智能路由必须匹配“auto”模型字面量——请在匹配条件里添加。", + "matchEmptyNote": "未设置匹配条件——该规则将应用于所有请求。", "weight": "权重", "remove": "移除", "addTarget": "+ 添加目标", @@ -883,6 +888,9 @@ "createNameRequired": "名称必填", "createVersionRequired": "版本必填", "createMaintainerRequired": "维护者必填", + "createNameInvalid": "请使用“/”格式——仅限小写字母、数字和连字符,例如 acme/pii-rules。", + "createVersionInvalid": "请使用带 v 前缀的语义化版本号,例如 v1.0.0。", + "createNameHelp": "格式:/(小写),例如 acme/pii-rules。", "createDescriptionHelp": "可选。在目录中展示给操作员。", "createRulesLabel": "规则 (JSON 数组)", "createRulesHelp": "每条规则需包含 ruleId、category、severity（hard|soft|info）、pattern；可选 flags、description、labels。", @@ -1059,6 +1067,7 @@ "virtualKeyCreatedMsg": "虚拟密钥已创建", "nameHelpText": "虚拟密钥的人读名称，在仪表板和审计日志中作为显示名使用。", "namePlaceholder": "engineering-openai", + "nameTakenApplication": "已存在同名的应用类型虚拟密钥，请更换名称。", "projectTooltip": "此虚拟密钥所属的项目，用于流量归属和配额计算。", "expirationTooltip": "应用类型虚拟密钥必须设置过期时间，且不得超过 3 个月。密钥过期后将无法用于认证。", "disabledKeysTooltip": "已禁用的密钥立即拒绝所有调用，但保留配置以供审计或重新启用。", @@ -4327,6 +4336,7 @@ "createSubtitle": "创建新的个人虚拟密钥", "name": "名称", "namePlaceholder": "my-api-key", + "nameTaken": "你已有同名的虚拟密钥，请更换名称。", "sourceApp": "来源应用", "status": "状态", "rpm": "速率限制 (RPM)", @@ -6294,9 +6304,9 @@ "endpointExamplePlaceholder": "OpenAI (chat)" }, "rulePacks": { - "slugPlaceholder": "acme/rules", - "versionPlaceholder": "v1.0.0", - "sourcePlaceholder": "customer" + "slugPlaceholder": "例如 acme/rules", + "versionPlaceholder": "例如 v1.0.0", + "sourcePlaceholder": "例如 customer" }, "idp": { "clientIdPlaceholder": "0oa1abc2def3ghi4j5k6" diff --git a/packages/control-plane-ui/src/lib/forms/FormInput.test.tsx b/packages/control-plane-ui/src/lib/forms/FormInput.test.tsx new file mode 100644 index 00000000..cb2f6641 --- /dev/null +++ b/packages/control-plane-ui/src/lib/forms/FormInput.test.tsx @@ -0,0 +1,35 @@ +import { describe, it, expect, vi } from 'vitest'; +import { render, screen } from '@testing-library/react'; +import userEvent from '@testing-library/user-event'; +import { useForm } from 'react-hook-form'; + +vi.mock('react-i18next', () => ({ useTranslation: () => ({ t: (k: string) => k }) })); +// The '@/components/ui' barrel transitively imports '@/i18n' (Header/Sidebar +// language switcher), whose http-backend import has no place in a unit test. +vi.mock('@/i18n', () => ({ SUPPORTED_LANGUAGES: [], LANGUAGE_STORAGE_KEY: 'lang', default: {} })); + +import { FormInput } from './FormInput'; + +function Harness({ onValueBlur }: { onValueBlur?: (v: string) => void }) { + const form = useForm<{ name: string }>({ defaultValues: { name: '' } }); + return ; +} + +describe('FormInput onValueBlur', () => { + it('fires with the field value when the user leaves the field', async () => { + const onValueBlur = vi.fn(); + render(); + const input = screen.getByLabelText(/Name/); + await userEvent.type(input, 'my-key'); + await userEvent.tab(); + expect(onValueBlur).toHaveBeenCalledWith('my-key'); + }); + + it('renders and blurs normally when onValueBlur is not provided', async () => { + render(); + const input = screen.getByLabelText(/Name/); + await userEvent.type(input, 'x'); + await userEvent.tab(); + expect(input).toHaveValue('x'); + }); +}); diff --git a/packages/control-plane-ui/src/lib/forms/FormInput.tsx b/packages/control-plane-ui/src/lib/forms/FormInput.tsx index 7e425158..c3f81464 100644 --- a/packages/control-plane-ui/src/lib/forms/FormInput.tsx +++ b/packages/control-plane-ui/src/lib/forms/FormInput.tsx @@ -17,6 +17,12 @@ interface FormInputProps { placeholder?: string; disabled?: boolean; className?: string; + /** + * Called with the field's current value after RHF's own blur handling — + * for advisory async checks (e.g. duplicate-name probes) that should run + * when the user leaves the field. + */ + onValueBlur?: (value: string) => void; } export function FormInput({ @@ -30,6 +36,7 @@ export function FormInput({ placeholder, disabled, className, + onValueBlur, }: FormInputProps) { const { field, fieldState } = useController({ name, control: form.control }); @@ -45,6 +52,10 @@ export function FormInput({ { + field.onBlur(); + onValueBlur?.(e.target.value); + }} type={type} placeholder={placeholder} disabled={disabled} diff --git a/packages/control-plane-ui/src/pages/account/personal-vks/PersonalVKCreate.tsx b/packages/control-plane-ui/src/pages/account/personal-vks/PersonalVKCreate.tsx index 055a0ae4..9f5d04ff 100644 --- a/packages/control-plane-ui/src/pages/account/personal-vks/PersonalVKCreate.tsx +++ b/packages/control-plane-ui/src/pages/account/personal-vks/PersonalVKCreate.tsx @@ -1,9 +1,11 @@ +import { useCallback } from 'react'; import { useNavigate } from 'react-router-dom'; import { useTranslation } from 'react-i18next'; import { z } from 'zod'; import { personalVKApi } from '@/api/services/ai-gateway/personalVirtualKeys'; import type { CreatePersonalVKInput } from '@/api/services/ai-gateway/personalVirtualKeys'; import { useMutation } from '@/hooks/useMutation'; +import { useDuplicateNameCheck } from '@/hooks/useDuplicateNameCheck'; import { useZodForm, FormInput } from '@/lib/forms'; import { useUnsavedChangesWarning } from '@/hooks/useUnsavedChangesWarning'; import { @@ -38,6 +40,20 @@ export function PersonalVKCreate() { useUnsavedChangesWarning(form.formState.isDirty); + // Fail-fast duplicate probe on blur: personal VK names are unique per + // owner, and /api/my/virtual-keys is already scoped to the current user. + // Advisory only; the create 409 is authoritative. + const isNameTaken = useCallback(async (name: string) => { + const res = await personalVKApi.list({ q: name, limit: '100' }); + return (res.data ?? []).some(vk => vk.name === name); + }, []); + const checkName = useDuplicateNameCheck({ + form, + field: 'name', + message: t('pages:personalVks.nameTaken'), + isTaken: isNameTaken, + }); + const { mutate, loading } = useMutation( (data: CreatePersonalVKInput) => personalVKApi.create(data), { @@ -70,7 +86,7 @@ export function PersonalVKCreate() {
- + diff --git a/packages/control-plane-ui/src/pages/ai-gateway/providers/wizard/ProviderWizard.module.css b/packages/control-plane-ui/src/pages/ai-gateway/providers/wizard/ProviderWizard.module.css index 73075fc6..f1e7ef57 100644 --- a/packages/control-plane-ui/src/pages/ai-gateway/providers/wizard/ProviderWizard.module.css +++ b/packages/control-plane-ui/src/pages/ai-gateway/providers/wizard/ProviderWizard.module.css @@ -273,8 +273,9 @@ .templateCardSelected { composes: templateCard; - border-width: 2px; - box-shadow: var(--shadow-sm); + border-color: var(--color-primary); + background: var(--color-primary-light); + box-shadow: 0 0 0 3px var(--color-accent-bg); } .templateAvatar { @@ -340,8 +341,9 @@ .customCardSelected { composes: customCard; - border: 2px dashed var(--color-text); - background: var(--color-muted); + border: 2px dashed var(--color-primary); + background: var(--color-primary-light); + box-shadow: 0 0 0 3px var(--color-accent-bg); } .browseBtn { diff --git a/packages/control-plane-ui/src/pages/ai-gateway/routing/_shared/routing-rule-config.splitWeights.test.ts b/packages/control-plane-ui/src/pages/ai-gateway/routing/_shared/routing-rule-config.splitWeights.test.ts new file mode 100644 index 00000000..20cc01a4 --- /dev/null +++ b/packages/control-plane-ui/src/pages/ai-gateway/routing/_shared/routing-rule-config.splitWeights.test.ts @@ -0,0 +1,37 @@ +import { describe, it, expect } from 'vitest'; + +import { validateSplitWeights, type ProviderModelEntry } from './routing-rule-config'; + +const entry = (weight: string): ProviderModelEntry => ({ provider: 'p', model: 'm', weight }); + +// validateSplitWeights backs the sum-to-100 gate on ab_split "Split %" and +// loadbalance "Weight" targets. The resolvers treat weights as relative +// (weighted random over the sum), so the UI enforces 100 to keep the entered +// numbers truthful and the split predictable. +describe('validateSplitWeights', () => { + it('accepts weights that sum to exactly 100', () => { + expect(validateSplitWeights([entry('70'), entry('30')])).toEqual({ valid: true, total: 100 }); + expect(validateSplitWeights([entry('50'), entry('25'), entry('25')])).toEqual({ valid: true, total: 100 }); + expect(validateSplitWeights([entry('100')])).toEqual({ valid: true, total: 100 }); + }); + + it('rejects the 70 + 50 = 120 case and returns the running total', () => { + expect(validateSplitWeights([entry('70'), entry('50')])).toEqual({ valid: false, total: 120 }); + }); + + it('rejects sums below 100', () => { + expect(validateSplitWeights([entry('50')])).toEqual({ valid: false, total: 50 }); + expect(validateSplitWeights([entry('40'), entry('40')])).toEqual({ valid: false, total: 80 }); + }); + + it('rejects any target outside 0..100 even when the total is 100', () => { + // 150 + -50 = 100 but individual values are out of range → invalid. + expect(validateSplitWeights([entry('150'), entry('-50')]).valid).toBe(false); + expect(validateSplitWeights([entry('101'), entry('-1')]).valid).toBe(false); + }); + + it('treats non-numeric weights as invalid', () => { + expect(validateSplitWeights([entry(''), entry('100')]).valid).toBe(false); + expect(validateSplitWeights([entry('abc'), entry('100')]).valid).toBe(false); + }); +}); diff --git a/packages/control-plane-ui/src/pages/ai-gateway/routing/_shared/routing-rule-config.ts b/packages/control-plane-ui/src/pages/ai-gateway/routing/_shared/routing-rule-config.ts index 21e21b8c..54a398c6 100644 --- a/packages/control-plane-ui/src/pages/ai-gateway/routing/_shared/routing-rule-config.ts +++ b/packages/control-plane-ui/src/pages/ai-gateway/routing/_shared/routing-rule-config.ts @@ -35,6 +35,32 @@ export interface ProviderModelEntry { weight: string; } +/** + * Weighted-target strategies (ab_split "Split %" and loadbalance "Weight") + * present operators with per-target numbers that must sum to exactly 100, with + * every target in 0..100. Both resolvers treat weights as *relative* (weighted + * random over the sum), so 70+50 would silently become 58/42 — enforcing 100 + * makes the entered numbers truthful and the split predictable. Applied to any + * strategy that shows the weight column; single/fallback/conditional/etc. have + * no weights and are unaffected. + */ +export function validateSplitWeights(entries: ProviderModelEntry[]): { valid: boolean; total: number } { + let total = 0; + let allInRange = true; + for (const entry of entries) { + const raw = entry.weight.trim(); + const n = Number(raw); + // Number('') is 0, so guard empty/whitespace explicitly — a blank weight + // field is not a valid target, even though it would coerce to 0. + if (raw === '' || !Number.isFinite(n) || n < 0 || n > 100) { + allInRange = false; + continue; + } + total += n; + } + return { valid: allInRange && total === 100, total }; +} + export function mapLegacyStrategy(s: string): StrategyType { const map: Record = { priority: 'single', diff --git a/packages/control-plane-ui/src/pages/ai-gateway/routing/create/RoutingRuleCreate.module.css b/packages/control-plane-ui/src/pages/ai-gateway/routing/create/RoutingRuleCreate.module.css index 5e821c23..a4e88ced 100644 --- a/packages/control-plane-ui/src/pages/ai-gateway/routing/create/RoutingRuleCreate.module.css +++ b/packages/control-plane-ui/src/pages/ai-gateway/routing/create/RoutingRuleCreate.module.css @@ -26,6 +26,24 @@ margin-left: 2px; } +.weightSumError { + color: var(--color-danger); + font-size: var(--g-font-size-sm); + margin: 8px 0 0; +} + +.stepHint { + color: var(--color-danger); + font-size: var(--g-font-size-sm); + margin: 0 0 8px; +} + +.stepNote { + color: var(--color-text-muted); + font-size: var(--g-font-size-sm); + margin: 0 0 8px; +} + .helpBtn { display: inline-flex; align-items: center; diff --git a/packages/control-plane-ui/src/pages/ai-gateway/routing/create/RoutingRuleCreatePage.tsx b/packages/control-plane-ui/src/pages/ai-gateway/routing/create/RoutingRuleCreatePage.tsx index 764f76e2..f7f62b53 100644 --- a/packages/control-plane-ui/src/pages/ai-gateway/routing/create/RoutingRuleCreatePage.tsx +++ b/packages/control-plane-ui/src/pages/ai-gateway/routing/create/RoutingRuleCreatePage.tsx @@ -11,7 +11,7 @@ import { Switch, Tooltip } from '@/components/ui'; -import type { StrategyType } from '../_shared/routing-rule-config'; +import { type StrategyType, validateSplitWeights } from '../_shared/routing-rule-config'; import { ConditionalRoutingEditor } from '../editor/ConditionalRoutingEditor'; import { RoutingPrimaryWinnerCallout } from '../_shared/RoutingPrimaryWinnerCallout'; import { MatchConditionExtraFields } from '../editor/MatchConditionExtraFields'; @@ -41,6 +41,39 @@ export function RoutingRuleCreate() { const weightLabel = h.strategyType === 'ab_split' ? t('pages:routing.splitPercent') : t('pages:routing.weight'); const isLastStep = currentStep === WIZARD_TOTAL_STEPS - 1; + // Weighted targets (ab_split "Split %" + loadbalance "Weight") must sum to + // exactly 100 — see validateSplitWeights. + const hasWeightTargets = h.pipelineStage === '1' && h.showWeightColumn; + const weightCheck = validateSplitWeights(h.entries); + const weightSumInvalid = hasWeightTargets && !weightCheck.valid; + + // Smart routing must pin matchConditions to the "auto" literal (backend guard). + const smartNeedsAutoMatch = h.pipelineStage === '1' && h.strategyType === 'smart' && + !(h.matchRequestedModelLiterals.length > 0 && h.matchRequestedModelLiterals.every((l) => l === 'auto')); + // Non-smart rule with empty match conditions matches every request (soft note). + const matchIsEmpty = h.models.length === 0 && h.matchProviders.length === 0 && h.matchProjectIds.length === 0 && + h.matchRequestedModelLiterals.length === 0 && h.matchModelTypes.length === 0 && h.matchVirtualKeys.length === 0; + + // The blocking hint for a wizard step, or null when the step is complete. + const stepHint = (step: number): string | null => { + if (step === 0) { + return h.name.trim() === '' ? t('pages:routing.nameRequired', 'Name is required.') : null; + } + if (step === 1) { + return h.configValidity.ok ? null : h.configValidity.message; + } + if (step === 2) { + return h.fallbackIncomplete + ? t('pages:routing.fallbackIncomplete', 'Finish or remove the partially-filled fallback target — set both provider and model.') + : null; + } + return smartNeedsAutoMatch + ? t('pages:routing.smartMatchAutoRequired', 'Smart routing must match the "auto" model literal — add it under Match conditions.') + : null; + }; + const currentStepHint = stepHint(currentStep); + const allStepsValid = [0, 1, 2, 3].every((s) => stepHint(s) === null); + const goNext = () => setCurrentStep(s => Math.min(s + 1, WIZARD_TOTAL_STEPS - 1)); const goBack = () => { if (currentStep === 0) { h.navigate('/ai-gateway/routing'); return; } @@ -261,6 +294,11 @@ export function RoutingRuleCreate() {
))} + {weightSumInvalid && ( +

+ {t('pages:routing.weightSumError', { total: weightCheck.total })} +

+ )} )} @@ -340,17 +378,27 @@ export function RoutingRuleCreate() { + {/* ── Step validation hint / soft note ── */} + {currentStepHint && ( +

{currentStepHint}

+ )} + {isLastStep && matchIsEmpty && !smartNeedsAutoMatch && ( +

+ {t('pages:routing.matchEmptyNote', 'No match conditions set — this rule will apply to all requests.')} +

+ )} + {/* ── Wizard navigation footer ── */}
{isLastStep ? ( - ) : ( - )} diff --git a/packages/control-plane-ui/src/pages/ai-gateway/routing/create/useRoutingRuleCreate.ts b/packages/control-plane-ui/src/pages/ai-gateway/routing/create/useRoutingRuleCreate.ts index 1dbac64f..a9bb381f 100644 --- a/packages/control-plane-ui/src/pages/ai-gateway/routing/create/useRoutingRuleCreate.ts +++ b/packages/control-plane-ui/src/pages/ai-gateway/routing/create/useRoutingRuleCreate.ts @@ -19,6 +19,7 @@ import { resolveConditionalConfigFromEditor, DEFAULT_SMART_SYSTEM_PROMPT, buildMatchConditionsPayload, + validateSplitWeights, type ConditionalEditorHydration, type StrategyType, type ProviderModelEntry, @@ -115,36 +116,14 @@ export function useRoutingRuleCreate() { setStrategyType(next); }; - // ── Submit handler ──────────────────────────────────────────────── - const handleSubmit = () => { - const stageNum = pipelineStage === '0' ? 0 : 1; - if (stageNum === 0) { - const built = buildPolicyApiConfig(policyAllowM, policyDenyM, policyAllowP, policyDenyP); - if (!built.ok) { - addToast(built.message, 'error'); - return; - } - const matchConditions = buildMatchConditionsPayload({ - models, - requestedModelLiterals: matchRequestedModelLiterals, - modelTypes: matchModelTypes, - providers: matchProviders, - projects: matchProjectIds, - virtualKeys: matchVirtualKeys, - }); - mutate({ - name, - description, - strategyType: 'policy', - priority: Number(priority), - enabled, - pipelineStage: 0, - config: built.config, - matchConditions, - }); - return; + // Build the strategy config for the current form state. Shared by the + // wizard's Configuration-step gating (configValidity) and the final submit + // so what blocks Continue and what the API rejects can never drift. Returns + // the same {ok,message}|{ok,config} shape the build helpers use. + const buildConfig = (): { ok: true; config: unknown } | { ok: false; message: string } => { + if (pipelineStage === '0') { + return buildPolicyApiConfig(policyAllowM, policyDenyM, policyAllowP, policyDenyP); } - const built = strategyType === 'conditional' ? resolveConditionalConfigFromEditor(conditionalUi, providerGroups) @@ -159,6 +138,20 @@ export function useRoutingRuleCreate() { matchModelIds: models, preservedConditionalConfig: null, }); + if (!built.ok) return built; + // Weighted strategies additionally require the split to total 100. + if (strategyType === 'ab_split' || strategyType === 'loadbalance') { + const weightCheck = validateSplitWeights(entries); + if (!weightCheck.valid) { + return { ok: false, message: t('pages:routing.weightSumError', { total: weightCheck.total }) }; + } + } + return built; + }; + + // ── Submit handler ──────────────────────────────────────────────── + const handleSubmit = () => { + const built = buildConfig(); if (!built.ok) { addToast(built.message, 'error'); return; @@ -171,6 +164,19 @@ export function useRoutingRuleCreate() { projects: matchProjectIds, virtualKeys: matchVirtualKeys, }); + if (pipelineStage === '0') { + mutate({ + name, + description, + strategyType: 'policy', + priority: Number(priority), + enabled, + pipelineStage: 0, + config: built.config, + matchConditions, + }); + return; + } const fallbackChainApi = buildFallbackChainApi(fallbackEntries, providerGroups); mutate({ name, @@ -185,6 +191,14 @@ export function useRoutingRuleCreate() { }); }; + // Per-step required-field state for the wizard's Continue/submit gating. + const configValidity = buildConfig(); + // A fallback target is either fully filled or empty; a half-filled row + // (provider XOR model) would silently drop at buildFallbackChainApi. + const fallbackIncomplete = fallbackEntries.some( + (e) => (e.provider.trim() !== '') !== (e.model.trim() !== ''), + ); + // ── Derived values ──────────────────────────────────────────────── const showWeightColumn = strategyType === 'loadbalance' || strategyType === 'ab_split'; @@ -242,6 +256,8 @@ export function useRoutingRuleCreate() { // Derived showWeightColumn, configModelIds, + configValidity, + fallbackIncomplete, // Actions handleSubmit, diff --git a/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/RoutingRuleDetail.module.css b/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/RoutingRuleDetail.module.css index 86dd6780..8e2d8fab 100644 --- a/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/RoutingRuleDetail.module.css +++ b/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/RoutingRuleDetail.module.css @@ -458,6 +458,13 @@ margin-left: 2px; } +/* Weighted-target sum-to-100 validation message */ +.weightSumError { + color: var(--color-danger); + font-size: var(--g-font-size-sm); + margin: 8px 0 0; +} + /* Inline "Enabled" text next to the switch */ .enabledLabel { font-size: var(--g-font-size-sm); diff --git a/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/RoutingRuleEditForm.tsx b/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/RoutingRuleEditForm.tsx index f8ffb3b9..fa1ecd59 100644 --- a/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/RoutingRuleEditForm.tsx +++ b/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/RoutingRuleEditForm.tsx @@ -10,6 +10,7 @@ import { emptyConditionalFormState, hydrateConditionalEditorState, tryParseConditionalFormFromConfig, + validateSplitWeights, type StrategyType, } from '../_shared/routing-rule-config'; import { ConditionalRoutingEditor } from '../editor/ConditionalRoutingEditor'; @@ -76,6 +77,12 @@ export function RoutingRuleEditForm({ detail }: { detail: RoutingRuleDetailState const weightLabel = editStrategyType === 'ab_split' ? t('pages:routing.splitPercent') : t('pages:routing.weight'); + // Weighted targets (ab_split "Split %" + loadbalance "Weight") must sum to + // exactly 100 — see validateSplitWeights. + const hasWeightTargets = editPipelineStage === '1' && showWeightColumn; + const weightCheck = validateSplitWeights(entries); + const weightSumInvalid = hasWeightTargets && !weightCheck.valid; + const switchConditionalToJson = () => { if (conditionalUi.mode === 'json') return; const built = buildConditionalApiConfig(conditionalUi.form, providerGroups); @@ -297,6 +304,11 @@ export function RoutingRuleEditForm({ detail }: { detail: RoutingRuleDetailState
))} + {weightSumInvalid && ( +

+ {t('pages:routing.weightSumError', { total: weightCheck.total })} +

+ )} )} @@ -397,7 +409,7 @@ export function RoutingRuleEditForm({ detail }: { detail: RoutingRuleDetailState diff --git a/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/useRoutingRuleDetail.ts b/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/useRoutingRuleDetail.ts index b4077f71..e7e2aa6c 100644 --- a/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/useRoutingRuleDetail.ts +++ b/packages/control-plane-ui/src/pages/ai-gateway/routing/detail/useRoutingRuleDetail.ts @@ -39,6 +39,7 @@ import { DEFAULT_SMART_SYSTEM_PROMPT, parseMatchConditionsForm, buildMatchConditionsPayload, + validateSplitWeights, type ConditionalEditorHydration, type StrategyType, type ProviderModelEntry, @@ -331,7 +332,7 @@ export function useRoutingRuleDetail() { name: v.editName, description: v.editDescription, strategyType: 'policy', - priority: v.editPriority, + priority: Number(v.editPriority), enabled: v.editEnabled, pipelineStage: 0, config: built.config, @@ -341,6 +342,14 @@ export function useRoutingRuleDetail() { return; } + if (editStrategyType === 'ab_split' || editStrategyType === 'loadbalance') { + const weightCheck = validateSplitWeights(entries); + if (!weightCheck.valid) { + addToast(t('pages:routing.weightSumError', { total: weightCheck.total }), 'error'); + return; + } + } + const built = editStrategyType === 'conditional' ? resolveConditionalConfigFromEditor(conditionalUi, providerGroups) @@ -372,7 +381,7 @@ export function useRoutingRuleDetail() { name: v.editName, description: v.editDescription, strategyType: editStrategyType, - priority: v.editPriority, + priority: Number(v.editPriority), enabled: v.editEnabled, pipelineStage: 1, config: built.config, diff --git a/packages/control-plane-ui/src/pages/ai-gateway/virtual-keys/VirtualKeyCreate.tsx b/packages/control-plane-ui/src/pages/ai-gateway/virtual-keys/VirtualKeyCreate.tsx index cf284bb8..b9d08520 100644 --- a/packages/control-plane-ui/src/pages/ai-gateway/virtual-keys/VirtualKeyCreate.tsx +++ b/packages/control-plane-ui/src/pages/ai-gateway/virtual-keys/VirtualKeyCreate.tsx @@ -1,4 +1,4 @@ -import { useState, useMemo } from 'react'; +import { useState, useMemo, useCallback } from 'react'; import { useTranslation } from 'react-i18next'; import { useNavigate } from 'react-router-dom'; import { z } from 'zod'; @@ -6,6 +6,7 @@ import { useApi } from '@/hooks/useApi'; import { virtualKeyApi, projectApi, systemApi } from '@/api/services'; import type { CreateVirtualKeyInput } from '@/api/services'; import { useMutation } from '@/hooks/useMutation'; +import { useDuplicateNameCheck } from '@/hooks/useDuplicateNameCheck'; import { useZodForm, FormInput, FormSwitch } from '@/lib/forms'; import { useUnsavedChangesWarning } from '@/hooks/useUnsavedChangesWarning'; import { @@ -186,6 +187,22 @@ export function VirtualKeyCreate() { ['admin', 'projects', 'list', 'vk-create'], ); + // Fail-fast duplicate probe on blur: application VK names are unique + // deployment-wide, so ask the list endpoint (q is a substring filter — + // exact-compare client-side). Advisory only; the create 409 is + // authoritative, and a non-super-admin's list may be scoped to their own + // keys, so a silent miss here is possible and fine. + const isNameTaken = useCallback(async (name: string) => { + const res = await virtualKeyApi.list({ q: name, vkType: 'application', limit: '100' }); + return (res.data ?? []).some(vk => vk.name === name); + }, []); + const checkName = useDuplicateNameCheck({ + form, + field: 'name', + message: t('pages:virtualKeys.nameTakenApplication'), + isTaken: isNameTaken, + }); + const { mutate, loading } = useMutation( (data: CreateVirtualKeyInput) => virtualKeyApi.create(data) as Promise<{ key?: string; id?: string }>, { @@ -273,7 +290,7 @@ export function VirtualKeyCreate() { - + draftsToRules(ruleDrafts), [ruleDrafts]); const rulesValidation = rulesMode === 'json' ? parsedJson : parsedForm; + // Field-level format validation mirroring the backend contract. Only + // surfaced once the field is non-empty so the form does not shout at an + // untouched input; the empty case is covered by the required-field checks. + const nameError = + name.trim() !== '' && !isValidPackName(name.trim()) + ? t('pages:hooks.rulePacks.createNameInvalid', 'Use "/" (lowercase), e.g. acme/pii-rules.') + : undefined; + const versionError = + version.trim() !== '' && !isValidPackVersion(version.trim()) + ? t('pages:hooks.rulePacks.createVersionInvalid', 'Use a v-prefixed semantic version, e.g. v1.0.0.') + : undefined; + const { mutate: createPack, loading: creating } = useMutation( (body) => rulePacksApi.create(body), { @@ -74,10 +78,18 @@ export function RulePackCreatePage() { setFormError(t('pages:hooks.rulePacks.createNameRequired', 'Name is required')); return; } + if (nameError) { + setFormError(nameError); + return; + } if (version.trim() === '') { setFormError(t('pages:hooks.rulePacks.createVersionRequired', 'Version is required')); return; } + if (versionError) { + setFormError(versionError); + return; + } if (maintainer.trim() === '') { setFormError(t('pages:hooks.rulePacks.createMaintainerRequired', 'Maintainer is required')); return; @@ -141,6 +153,8 @@ export function RulePackCreatePage() { name.trim() !== '' && version.trim() !== '' && maintainer.trim() !== '' && + nameError === undefined && + versionError === undefined && rulesValidation.rules !== null && rulesValidation.error === null; @@ -160,10 +174,15 @@ export function RulePackCreatePage() {
- + / (lowercase), e.g. acme/pii-rules.')} + > setName(e.target.value)} placeholder={t('pages:rulePacks.slugPlaceholder')} /> - + setVersion(e.target.value)} placeholder={t('pages:rulePacks.versionPlaceholder')} />
diff --git a/packages/control-plane-ui/src/pages/compliance/rule-packs/form/rulePackRules.test.ts b/packages/control-plane-ui/src/pages/compliance/rule-packs/form/rulePackRules.test.ts new file mode 100644 index 00000000..7aa97026 --- /dev/null +++ b/packages/control-plane-ui/src/pages/compliance/rule-packs/form/rulePackRules.test.ts @@ -0,0 +1,49 @@ +import { describe, it, expect } from 'vitest'; + +import { isValidPackName, isValidPackVersion } from './rulePackRules'; + +// These mirror the backend contract in packages/shared/policy/rulepack/yaml.go +// (packNameRE / semverRE). If the backend regex changes, both must move +// together — the whole point is that the form rejects exactly what the API +// would reject, so an operator never sees the raw 400 detail. +describe('isValidPackName', () => { + it('accepts / with lowercase, digits and hyphens', () => { + expect(isValidPackName('acme/pii-rules')).toBe(true); + expect(isValidPackName('acme/rules')).toBe(true); + expect(isValidPackName('a1/b2')).toBe(true); + expect(isValidPackName('team-security/pci-dss-3')).toBe(true); + }); + + it('rejects a bare name with no namespace slash (the "test" case)', () => { + expect(isValidPackName('test')).toBe(false); + }); + + it('rejects uppercase, spaces, leading digits/hyphens, and empty segments', () => { + expect(isValidPackName('Acme/rules')).toBe(false); + expect(isValidPackName('acme/Rules')).toBe(false); + expect(isValidPackName('acme rules/x')).toBe(false); + expect(isValidPackName('1acme/rules')).toBe(false); + expect(isValidPackName('-acme/rules')).toBe(false); + expect(isValidPackName('acme/')).toBe(false); + expect(isValidPackName('/rules')).toBe(false); + expect(isValidPackName('acme/rules/extra')).toBe(false); + expect(isValidPackName('')).toBe(false); + }); +}); + +describe('isValidPackVersion', () => { + it('accepts v-prefixed semver with optional pre-release/build', () => { + expect(isValidPackVersion('v1.0.0')).toBe(true); + expect(isValidPackVersion('v10.20.30')).toBe(true); + expect(isValidPackVersion('v1.2.3-rc1')).toBe(true); + expect(isValidPackVersion('v1.2.3+build.5')).toBe(true); + }); + + it('rejects a missing v prefix, partial versions, and non-numeric parts', () => { + expect(isValidPackVersion('1.0.0')).toBe(false); + expect(isValidPackVersion('v1.0')).toBe(false); + expect(isValidPackVersion('v1')).toBe(false); + expect(isValidPackVersion('vx.y.z')).toBe(false); + expect(isValidPackVersion('')).toBe(false); + }); +}); diff --git a/packages/control-plane-ui/src/pages/compliance/rule-packs/form/rulePackRules.ts b/packages/control-plane-ui/src/pages/compliance/rule-packs/form/rulePackRules.ts index ff6d0340..fa1941a1 100644 --- a/packages/control-plane-ui/src/pages/compliance/rule-packs/form/rulePackRules.ts +++ b/packages/control-plane-ui/src/pages/compliance/rule-packs/form/rulePackRules.ts @@ -1,5 +1,23 @@ import type { RulePackRule } from '@/api/services'; +// Mirror the backend authoring contract in +// packages/shared/policy/rulepack/yaml.go (packNameRE / semverRE) so the JSON +// admin form rejects a malformed name/version inline — with a friendly hint — +// instead of round-tripping to the API and surfacing the raw 400 detail +// (`rulepack: name "test" must match /`). +const PACK_NAME_RE = /^[a-z][a-z0-9-]*\/[a-z][a-z0-9-]*$/; +const PACK_VERSION_RE = /^v\d+\.\d+\.\d+(?:[-+][A-Za-z0-9._-]+)?$/; + +/** True when name matches "/" (lowercase, digits, hyphens). */ +export function isValidPackName(name: string): boolean { + return PACK_NAME_RE.test(name); +} + +/** True when version is a v-prefixed semver, e.g. v1.0.0 or v1.2.3-rc1. */ +export function isValidPackVersion(version: string): boolean { + return PACK_VERSION_RE.test(version); +} + export interface RuleDraft { ruleId: string; category: string; @@ -22,6 +40,18 @@ export function emptyRuleDraft(): RuleDraft { }; } +/** Seed JSON shown in the create form's rules editor: one example PII rule. */ +export const DEFAULT_RULES_JSON = serializeRules([ + { + ruleId: 'example-email', + category: 'pii', + severity: 'hard', + pattern: '[\\w.+-]+@[\\w-]+\\.[\\w.-]+', + description: 'Blocks email addresses', + labels: ['pii:email'], + }, +]); + export function serializeRules(rules: RulePackRule[]): string { return JSON.stringify( rules.map((rule) => ({ diff --git a/packages/control-plane-ui/tests/api/services/ai-gateway/personalVirtualKeys.test.ts b/packages/control-plane-ui/tests/api/services/ai-gateway/personalVirtualKeys.test.ts index de0af91b..76bfa70b 100644 --- a/packages/control-plane-ui/tests/api/services/ai-gateway/personalVirtualKeys.test.ts +++ b/packages/control-plane-ui/tests/api/services/ai-gateway/personalVirtualKeys.test.ts @@ -7,11 +7,13 @@ beforeEach(() => Object.values(m).forEach((f) => f.mockClear())); describe('personalVKApi', () => { it('hits /api/my/virtual-keys', () => { personalVKApi.list(); + personalVKApi.list({ q: 'my-key' }); personalVKApi.create({} as never); personalVKApi.update('v1', {}); personalVKApi.delete('v1'); personalVKApi.regenerate('v1'); - expect(m.get).toHaveBeenCalledWith('/api/my/virtual-keys'); + expect(m.get).toHaveBeenCalledWith('/api/my/virtual-keys', undefined); + expect(m.get).toHaveBeenCalledWith('/api/my/virtual-keys', { q: 'my-key' }); expect(m.post).toHaveBeenCalledWith('/api/my/virtual-keys', {}); expect(m.put).toHaveBeenCalledWith('/api/my/virtual-keys/v1', {}); expect(m.delete).toHaveBeenCalledWith('/api/my/virtual-keys/v1'); diff --git a/packages/control-plane-ui/tests/pages/ai-gateway/routing/create/RoutingRuleCreatePage.test.tsx b/packages/control-plane-ui/tests/pages/ai-gateway/routing/create/RoutingRuleCreatePage.test.tsx index 4cfc3bf9..80e50fa5 100644 --- a/packages/control-plane-ui/tests/pages/ai-gateway/routing/create/RoutingRuleCreatePage.test.tsx +++ b/packages/control-plane-ui/tests/pages/ai-gateway/routing/create/RoutingRuleCreatePage.test.tsx @@ -92,20 +92,21 @@ describe('RoutingRuleCreate', () => { expect(navigate).toHaveBeenCalledWith('/ai-gateway/routing'); }); - it('Continue advances the wizard; Create on the last step (single, no model) surfaces an error', () => { + it('blocks Continue on step 0 until Name is filled', () => { wrap(); + const next = () => screen.getByRole('button', { name: i18n.t('pages:routing.wizardContinue', 'Continue') }); + expect(next()).toBeDisabled(); fireEvent.change(screen.getByTestId('routing-rule-name'), { target: { value: 'r' } }); - // walk to the last step - for (let i = 0; i < 6; i++) { - const next = screen.queryByRole('button', { name: i18n.t('pages:routing.wizardContinue', 'Continue') }); - if (!next) break; - fireEvent.click(next); - } - const create = screen.getByRole('button', { name: i18n.t('pages:routing.createRule') }); - expect(create).toBeEnabled(); // name present - fireEvent.click(create); - // single strategy with no provider/model resolved → validation toast, no mutation - expect(addToast).toHaveBeenCalledWith(expect.any(String), 'error'); + expect(next()).toBeEnabled(); + }); + + it('blocks Continue on the Configuration step when single strategy has no provider/model', () => { + wrap(); + fireEvent.change(screen.getByTestId('routing-rule-name'), { target: { value: 'r' } }); + // step 0 → Configuration + fireEvent.click(screen.getByRole('button', { name: i18n.t('pages:routing.wizardContinue', 'Continue') })); + // single strategy with no provider/model resolved → Continue gated, never submits + expect(screen.getByRole('button', { name: i18n.t('pages:routing.wizardContinue', 'Continue') })).toBeDisabled(); expect(mutate).not.toHaveBeenCalled(); }); }); From e97c218e06ae9e1c84199b35fb11683da7b070a2 Mon Sep 17 00:00:00 2001 From: nexus Date: Tue, 7 Jul 2026 15:59:32 +0800 Subject: [PATCH 7/8] chore(repo): sync docs, tooling, and CI config from internal main (sync 11) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bring repository-level docs and config up to date with internal main. - CHANGELOG.md: overload → retryable 429s (in-flight admission gate); hook-config reload-stampede fix; audit memory-budget entries. - .env.example: new AI_GATEWAY_AUDIT_MEM_MAX_BYTES (in-memory audit byte budget) and AI_GATEWAY_MAX_INFLIGHT (in-flight admission cap); revised AI_GATEWAY_AUDIT_MAX_QUEUED_RECORDS semantics. - README.md: doc refresh. - Architecture docs: audit-pipeline-architecture.md, ai-gateway/hook-architecture.md, provider-adapter-architecture.md. Co-Authored-By: Claude Opus 4.7 (1M context) --- .env.example | 37 ++++++++++++++++--- CHANGELOG.md | 35 ++++++++++++++++++ README.md | 4 +- .../audit-pipeline-architecture.md | 36 ++++++++++-------- .../services/ai-gateway/hook-architecture.md | 2 +- .../provider-adapter-architecture.md | 2 +- 6 files changed, 92 insertions(+), 24 deletions(-) diff --git a/.env.example b/.env.example index 44090069..0d628bee 100644 --- a/.env.example +++ b/.env.example @@ -280,13 +280,40 @@ NATS_URL=nats://localhost:4222 # default. # AI_GATEWAY_AUDIT_SPOOL_DIR=/var/lib/nexus/audit-spool -# In-heap audit record-buffer cap (overflow → durable spill above). Each queued -# record pins its pooled ~50 KB body until marshaled, so this bound is the primary -# control over the audit side-path's gw heap: 10000 (default) holds the body pool -# near ~1 GB under a slow-publish burst vs ~5 GB at the former 50000, same spill -# rate. Raise on a memory-rich box, lower on a constrained one. 0/unset → 10000. +# In-heap audit record-buffer cap (overflow → durable spill above). Bounds the +# queue's POINTER count only — the memory bound is AI_GATEWAY_AUDIT_MEM_MAX_BYTES +# below, which accounts each record's REAL body bytes. 0/unset keeps the default. # AI_GATEWAY_AUDIT_MAX_QUEUED_RECORDS=10000 +# In-memory audit byte budget — the memory half of the bounded audit queue. The +# audit writer reserves each record's REAL captured body bytes against this budget +# at admission; on a full budget the no-loss modes back-pressure the request path +# (RPS throttles to the audit drain rate, nothing dropped) and the lossy modes +# shed with a counted drop. Semantics mirror NEXUS_EVENTS_MAX_BYTES: unset/"auto" +# auto-sizes to ~15% of available RAM (2 GiB fallback off-Linux); an explicit +# human size ("8GB", "2048MB", raw bytes) pins it. A bigger budget buys BURST +# absorption (audit work deferred to quieter moments), not sustained throughput — +# under sustained overload the drain rate is the ceiling regardless. SIZING: the +# budget accounts raw body bytes only; REAL process memory runs ~2x the budget +# (Go GC headroom + marshal/compression copies), and the NATS memory stream holds +# another ~15% of RAM on the same box — keep a pin <=20% of RAM (a 10GB pin on a +# 32GB box OOM-killed the gateway at 15.9GB RSS in rig validation; a pin >25% +# logs a WARN). Observability: nexus_audit_mem_backpressure_total counts +# budget-full events; the resolved value is logged at startup. +# AI_GATEWAY_AUDIT_MEM_MAX_BYTES=auto + +# In-flight request admission cap. Bounds concurrent proxy requests so overload +# degrades into fast retryable 429s (Retry-After: 1, OpenAI-shaped error body) +# instead of unbounded in-heap queueing toward the GOMEMLIMIT collapse. Health, +# metrics, and admin endpoints are never gated. Shed requests are counted on +# nexus_ai_gateway_admission_shed_total (no audit record — auditing a shed storm +# would itself be load). DEFAULT unset/"auto" = 1024 x GOMAXPROCS — scaled by +# core count and sized generously so SSE streams (which hold a slot for the +# stream's full duration) do not false-shed a healthy streaming deployment; the +# bound caps heap growth under overload, not normal traffic. 0/negative +# disables; a non-numeric value falls back to auto with a warning. +# AI_GATEWAY_MAX_INFLIGHT=auto + # Audit overflow policy. Durable audit is a product promise + compliance # requirement. DEFAULT "spill" is spill-defer: the request path never # back-pressures — overflow goes to the durable on-disk spool and the diff --git a/CHANGELOG.md b/CHANGELOG.md index 1d9b014e..af420bd9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,41 @@ All notable changes to this project are documented here. The format follows ## [Unreleased] +### Changed — overload now degrades into retryable 429s (in-flight admission gate) + +- The AI Gateway bounds concurrent in-flight proxy requests (default + `1024 × GOMAXPROCS`; `AI_GATEWAY_MAX_INFLIGHT` overrides, `0` disables). At + arrival rates beyond the box's capacity, excess requests are rejected fast + with **429 + `Retry-After: 1`** in the caller's ingress error shape (OpenAI / + Anthropic / Gemini envelopes) instead of queueing in-heap until the Go memory + limit collapses throughput (measured pre-fix: 15.9s p99 at 1.5× capacity; + the pre-GOMEMLIMIT failure mode was an OOM kill). 429 was already part of the + data-plane contract (per-key rate limits and quota denials); SDK retry logic + engages unchanged. Health, metrics, and admin endpoints are never gated. Shed + requests are counted on `nexus_ai_gateway_admission_shed_total`. + +### Fixed — hook-config reload stampede at high load + +- Hook configuration freshness is now push-driven with a background TTL-backstop + ticker; the request path never loads configuration. Previously a TTL-stale + check on the request path could fan out one full rule-pack database load per + in-flight request while a slow load was running, collapsing the gateway at + high request rates (measured: p99 120s at 16k req/s with content hooks on; + fixed: p99 27ms at the same rate). Rule-pack install ordering also gained a + deterministic tiebreaker so no-change config reloads can no longer churn the + compiled matchers. + +### Performance — content-hook path allocation and CPU + +- Bodies-off deployments no longer allocate a fresh request-body buffer per + request (the pooled buffer is returned at request end; previously measured at + 52% of all gateway allocation under content-scan load). +- Redact-action rule packs skip re-localization entirely on benign traffic + (zero matches on a complete scan). +- Config snapshot loads expose `nexus_configcache_load_failures_total` and + `nexus_configcache_last_success_timestamp_seconds` for alerting on a frozen + config plane. + ### Fixed — Request/Response hook timing - **Streamed responses now record response-hook timing, exactly once per hook.** diff --git a/README.md b/README.md index 641aa52a..10eb28c4 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,7 @@ # Nexus Gateway -[![CI](https://github.com/AlphaBitCore/nexus-gateway/actions/workflows/ci.yml/badge.svg?branch=main)](.github/workflows/ci.yml) -[![Go CI](https://github.com/AlphaBitCore/nexus-gateway/actions/workflows/go-ci.yml/badge.svg?branch=main)](.github/workflows/go-ci.yml) +[![CI](https://github.com/itechchoice/abc-nexus-gateway/actions/workflows/ci.yml/badge.svg?branch=main)](.github/workflows/ci.yml) +[![Go CI](https://github.com/itechchoice/abc-nexus-gateway/actions/workflows/go-ci.yml/badge.svg?branch=main)](.github/workflows/go-ci.yml) [![Coverage gate](https://img.shields.io/badge/coverage-%E2%89%A595%25%20per%20package-brightgreen)](./scripts/check-go-coverage.sh) [![Status: 1.1.0](https://img.shields.io/badge/status-1.1.0-brightgreen)](./CHANGELOG.md) [![License: Apache 2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](./LICENSE) diff --git a/docs/developers/architecture/cross-cutting/observability/audit-pipeline-architecture.md b/docs/developers/architecture/cross-cutting/observability/audit-pipeline-architecture.md index 8dc70fbf..b71e3fd3 100644 --- a/docs/developers/architecture/cross-cutting/observability/audit-pipeline-architecture.md +++ b/docs/developers/architecture/cross-cutting/observability/audit-pipeline-architecture.md @@ -29,13 +29,15 @@ data-plane service MQ (NATS JetStream) Hub ───────────────── ───────────────────── ───── audit.Record (in-proc) │ - │ Enqueue (non-blocking) + │ Enqueue (byte-budget admission: no-loss modes + │ block on a full budget, lossy modes shed) ▼ -in-memory buffer (≤10000) - │ flushLoop tick (5s default) - │ batch of ≤100 → recordToMessage → JSON +bounded in-memory queue +(bytes ≤ ~15% RAM budget; count ≤ structural cap) + │ N consumer workers (one per producer connection) + │ batch of ≤512 (or 100ms linger) → marshal → frames ▼ -producer.Enqueue ──────► nexus.event.{ai-traffic|compliance|agent} +producer batch publish ──► nexus.event.{ai-traffic|compliance|agent} (NEXUS_EVENTS stream, InterestPolicy) │ hub-db-writer consumer group @@ -76,14 +78,13 @@ The same `Writer` is shared by the ai-gateway proxy handler and by ai-guard's cl ## 4. Writer — buffering, batching, retry, shutdown -The `Writer` (`audit.NewWriter`) owns the in-memory buffer and the background flush goroutine: +The `Writer` (`audit.NewWriter`) owns the bounded in-memory queue and the background consumer workers: -- `defaultBatchSize = 2000`, `defaultFlushInterval = 5s`, `maxQueueSize = 10000`. The audit insert is dominated by commit fsyncs, so the large drain batch amortizes them and pairs with the default-on COPY fast path; `NEXUS_HUB_CONSUMER_BATCH_SIZE` overrides it (a smaller batch trades throughput for a lower per-flush memory ceiling on tiny instances). Otherwise tunables are package constants, not yaml — there is one audit pipeline per process. -- `Enqueue` is **non-blocking**. When `len(buf) >= maxQueueSize` the record is dropped, a warning is logged with the `requestId`, and `nexus_audit_mq_dropped_total` is incremented. Dropping under sustained back-pressure is preferred over blocking the request path. +- **The queue is bounded twice: by BYTES (the real memory bound) and by pointer count (structural).** At `Enqueue` the writer reserves the record's **actual captured body size** (`len(RequestBody) + len(ResponseBody)`) against a lock-free soft byte budget (`shared/core/bytebudget`) sized by `AI_GATEWAY_AUDIT_MEM_MAX_BYTES` — same semantics as `NEXUS_EVENTS_MAX_BYTES`: unset/`auto` (default) = **~15% of `MemAvailable`** (2 GiB fallback where `/proc/meminfo` is unreadable), an explicit human size (`8GB`) pins it (a bigger budget buys **burst absorption**, not sustained throughput — the drain rate is the ceiling under sustained overload). Sizing: real process memory runs **~2× the budget** (Go GC headroom + marshal/compression copies) and the NATS memory stream holds another ~15% of RAM on a co-located box — keep a pin ≤20% of RAM (rig-validated: a 10 GB pin on a 32 GB box OOM-killed the gateway at 15.9 GB RSS; a pin >25% of available RAM logs a WARN, never clamps). The resolved value is logged at boot (`audit memory budget resolved`). The reservation is released **exactly once** when the record reaches a terminal — published, durably spilled, or counted-dropped. The no-loss modes (`block`, `spillblock`) **block at Enqueue** on a full budget (back-pressure to the request path); the lossy modes (`spill`, `drop`) never block — a full budget is an immediate counted drop. The channel caps (`200000` recCh / `100000` spillCh, `WithMaxQueuedRecords` override) bound pointer count only, a guard against a flood of tiny records. This replaced a count-derived sizing that assumed an average body size and let real ~100 KiB bodies OOM the process at 18 GiB. +- Enqueued records are drained by **N consumer workers** (one per producer pool connection), each batching up to 512 records (or a 100 ms linger for a partial batch) and publishing the batch framed on its own connection, so the per-connection ack barriers pipeline concurrently. - Embedding rows are coerced inside `audit.Writer.Enqueue`. When `EndpointType == EndpointTypeEmbeddings`, `coerceEmbeddingRow` (`packages/ai-gateway/internal/platform/audit/coerce.go`) zeroes chat-only fields (`completion_tokens`, `cache_read_tokens`, `cache_creation_tokens`, `reasoning_tokens`, `reasoning_cost_usd`) and emits a per-field warning naming the field, value, and request id. Every producer that publishes audit rows — proxy live response, proxy cache hits, and the ai-guard `WriterBackedTrafficSink` — emits through this single Enqueue path, so the coerce runs uniformly without each producer site re-implementing the rule. -- `flushLoop` ticks on `defaultFlushInterval`, snapshots the buffer under a mutex, marshals each record via `recordToMessage`, and calls `producer.Enqueue(ctx, queue, data)` with a per-record 5-second timeout. -- On a producer.Enqueue failure the record is **re-buffered** as long as space remains (`nexus_audit_mq_enqueue_errors_total` increments); if the buffer is at the cap the record is counted on `nexus_audit_mq_dropped_total` instead. There is no per-record back-off — the next flush tick is the retry. -- `Close()` drains the buffer through `drainBuffer` with a 15-second wall-time deadline and 200ms backoff between flush attempts. Records still in the buffer at the deadline are counted on `nexus_audit_mq_dropped_total` and logged so a sustained MQ outage at shutdown surfaces in monitoring instead of disappearing silently. +- On a publish failure the record is **re-queued for an in-memory retry up to 3 times** (`nexus_audit_mq_enqueue_errors_total` increments); past the cap — or when the queue is full — it is handed to the durable batched spill instead of circulating (the bounded-retry death-spiral fix). Never a silent loss. +- `Close()` stops intake, lets each consumer worker drain its remaining queued records and publish them (bounded at 512 per publish), drains the spill worker's channel, and durably spills any straggler that raced in after the final drain — a clean shutdown never leaves buffered audit unpublished or uncounted. - `WithThingIdentity(id, name)` stamps `ThingID` / `ThingName` onto every envelope so the consumer can attribute the row to the emitting Thing instance (gateway pod, proxy pod, agent host). Identity is set once at startup before any flush runs. - `WithSpillStore(store)` wires the out-of-band body backend; `WithPayloadCaptureStore(s)` wires the runtime cap snapshot so admin shadow updates take effect on the next flush without a restart; `WithNormalizer(fn)` wires a normalize closure that the write path does **not** invoke — the audit pipeline never persists a normalized projection. The closure exists only as the seam ai-gateway main uses to keep the normalize metrics series wired; `recordToMessage` does not call it. @@ -188,9 +189,10 @@ The `consumer.Manager` orchestrates the `hub-db-writer` group's consumers under Producer-side back-pressure: -- The in-memory buffer absorbs short bursts (10000 records). -- **Overflow is governed by `LossMode`** (`AI_GATEWAY_AUDIT_LOSS_MODE`); the config default is **`spillblock`** (zero-loss). On a full in-heap buffer the record is handed to the durable on-disk spool off the request path — identical to `spill` in the normal regime, so the request path is not blocked there — and only if the spool channel *itself* saturates does Enqueue back-pressure the request goroutine (bounded by `backpressureMaxWait`) instead of taking `spill`'s last-resort bounded drop. So durable audit stays a true zero-loss promise at `spill`'s throughput. `block` (the empty/unknown fallback) hard-back-pressures from the first full buffer; `spill`/`drop` are the explicit lossy opt-outs for non-compliance callers (a saturated `spill`, or any `drop`, counts on `nexus_audit_mq_dropped_total`). The resolved mode is logged at boot (`audit overflow policy resolved configured=… effective=…`). -- `Close()` retries draining for 15 seconds before counting the remainder as dropped, so a graceful rollout window cleanly flushes pending audit; a kill-9 path drops whatever is still in memory. On stop each bounded-queue consumer worker drains its remaining queued records and publishes them (bounded at `batchMaxCount` per publish, `drainOnStop`) before exiting, so a clean shutdown never leaves buffered audit unpublished. +- **The in-memory segment is byte-bounded** (§4): Enqueue admission reserves each record's real body bytes against the byte budget (`AI_GATEWAY_AUDIT_MEM_MAX_BYTES`, default `auto` ≈ 15% of RAM) and the no-loss modes block there when it is exhausted, so the audit heap can never grow past its RAM share no matter how large individual bodies are — the memory half of the bounded queue. `nexus_audit_mem_backpressure_total` counts those blocking events. The bounded channels absorb short bursts within the budget. +- **Overflow is governed by `LossMode`** (`AI_GATEWAY_AUDIT_LOSS_MODE`); the config default is **`spillblock`** (zero-loss). On a full in-heap buffer the record is handed to the durable on-disk spool off the request path — identical to `spill` in the normal regime, so the request path is not blocked there. If the spool channel *itself* saturates, the request goroutine **parks on the channel** (pure channel back-pressure — no per-goroutine disk write). And if the on-disk spool reaches its **total-size quota** (`AI_GATEWAY_AUDIT_SPOOL_MAX_TOTAL_MB`), `spillblock` still **never drops**: the single spill worker keeps the batch and retries (`ndjson.ErrSpoolQuotaExceeded` → back-pressure) while the recovery sweeper (§10.3) drains + deletes sealed files and frees space, so ingest self-throttles to the recovery-drain (PG-sustainable) rate instead of shedding records. The quota counts only reclaimable `audit-*.ndjson` content — `.poison` dead-letters are excluded, so undrainable content can never permanently wedge intake. The only `spillblock` drops are a genuine (non-quota) disk I/O error or shutdown. `block` (the empty/unknown fallback) hard-back-pressures from the first full buffer; `spill`/`drop` are the explicit lossy opt-outs for non-compliance callers (a saturated `spill`, or any `drop`, counts on `nexus_audit_mq_dropped_total`). The resolved mode is logged at boot (`audit overflow policy resolved configured=… effective=…`). +- **Single-writer spool.** Only the async spill worker writes to (and retries) the spool; request/consume goroutines back-pressure purely by blocking on the bounded in-heap channels. This keeps the per-write quota scan (a directory walk under the ndjson mutex) off every request goroutine, so a full-spool back-pressure cannot thrash the mutex the recovery sweeper needs to free space. +- `Close()` drains everything: each bounded-queue consumer worker drains its remaining queued records and publishes them (bounded at `batchMaxCount` per publish, `drainOnStop`), the spill worker drains its channel and flushes, and a final sweep durably spills any straggler left in either channel — a clean shutdown never leaves buffered audit unpublished or uncounted (each publish is individually time-bounded, so a wedged broker cannot hang the drain; a quota-full spool at shutdown is a one-shot counted drop). A kill-9 path drops whatever is still in memory. Consumer-side back-pressure: @@ -262,6 +264,8 @@ Producer side (each data-plane service exposes these on its own `/metrics`): | `nexus_audit_mq_reingested_total` | — | Spilled records replayed back into MQ by the spill-recovery sweeper (§10.3); `spilled − reingested − recovery_poisoned` ≈ records still on disk awaiting recovery | | `nexus_audit_mq_recovery_errors_total` | — | Spill-recovery sweeps left a file undrained (publish failure / read error); the sweep backs off and retries | | `nexus_audit_mq_recovery_poisoned_total` | — | Spilled records dead-lettered to a `.poison` sidecar because they exceed the broker `max_payload` and can never publish — a signal to lower the inline-body cap or enable out-of-band body spill | +| `nexus_audit_mq_spill_backpressure_total` | — | Times the spill worker entered a full-spool-quota back-pressure retry under `spillblock` (spool at `AI_GATEWAY_AUDIT_SPOOL_MAX_TOTAL_MB`). Rising **with `dropped_total` flat at 0** means the gateway is deliberately throttling the request path (not dropping, not hung) because the recovery drain (NATS/Hub/PG) < audit ingest — the distinct signal that tells an audit-spool wedge from an application hang | +| `nexus_audit_mem_backpressure_total` | — | Times `Enqueue` found the in-memory byte budget (~15% of RAM) exhausted. No-loss modes then block the request goroutine until drain frees bytes; lossy modes shed with a counted drop — this counter is what distinguishes a budget-full drop from a plain queue-full drop. Rising **with `dropped_total` flat at 0** = the designed bounded-memory overload behaviour: in-memory audit bytes hit their RAM share, ingest throttles to the drain rate instead of growing the heap | Consumer side (Hub `/metrics`): @@ -288,7 +292,9 @@ The cardinality on `nexus_mq_processed_total{queue}` is exactly three (the traff | Symptom | Where it shows up | Recovery | |--|--|--| | Producer-side burst overload | `nexus_audit_mq_dropped_total` climbs on the source service | None — drops are accepted to preserve the request path; investigate the back-pressure source (MQ outage? consumer wedge?) | -| Sustained MQ outage | `nexus_audit_mq_enqueue_errors_total` climbs; eventually `nexus_audit_mq_dropped_total` rises | Restart NATS; flush will retry on next tick. `Close()` deadline drops whatever is in-memory at shutdown | +| Audit spool at quota (`spillblock`) — request path deliberately throttled | `nexus_audit_mq_spill_backpressure_total` rises **while `nexus_audit_mq_dropped_total` stays 0**; request latency climbs / RPS falls to the audit-drain rate | Expected zero-loss behaviour under sustained overload where the recovery drain (NATS→Hub→PG) < audit ingest — the gateway throttles rather than dropping audit. Restore drain throughput (NATS/Hub/PG health) to relieve it. Operator levers if the throttle is unacceptable: raise `AI_GATEWAY_AUDIT_SPOOL_MAX_TOTAL_MB` (more disk buffer, still finite), or switch `AI_GATEWAY_AUDIT_LOSS_MODE` to `spill`/`drop` to trade audit completeness for request throughput. Distinguish from an app hang via this counter (a hang leaves it flat) | +| Audit memory budget exhausted — request path deliberately throttled | `nexus_audit_mem_backpressure_total` rises **while `nexus_audit_mq_dropped_total` stays 0**; gateway RSS plateaus near the budget instead of growing | Expected bounded-memory behaviour when audit ingest outruns the whole drain chain (publish + spill): in-memory audit bytes reached their ~15%-of-RAM budget and Enqueue blocks until drain frees bytes. Same relief levers as the spool-quota row — restore drain throughput; or trade completeness for throughput via `AI_GATEWAY_AUDIT_LOSS_MODE=spill`/`drop` (a full budget then sheds with a counted drop instead of throttling) | +| Sustained MQ outage | `nexus_audit_mq_enqueue_errors_total` climbs; records route to the durable spool (`spilled_total` rises); no-loss modes then back-pressure | Restart NATS; the recovery sweeper replays the spool. In no-loss modes nothing drops (ingest throttles instead); lossy modes count drops on `dropped_total` | | Consumer wedge | `nexus_mq_processed_total{queue}` flat while producer counters rise; eventually JetStream `MaxAge` / `MaxBytes` discards messages | Investigate `nexus_consumer_healthy{consumer="traffic-event-writer"}`; restart Hub; backfilled rows are lost past the discard window | | DB write failure (transient) | `nexus_mq_batch_flush_total{result="error"}` + `nexus_mq_traffic_errors_total{error_type=db_*}`; messages naked with backoff → redelivered (DB-backed DLQ at `redeliveryThresholdAttempts`, or on-disk DLQ if the DB is still down) | Self-heals once the DB recovers; investigate `nexus_mq_dlq_inserted_total{subject}` for repeat offenders, and `nexus_mq_disk_dlq_inserted_total{subject}` for rows that fell through to disk during a full DB outage (replay per §10.1) | | DB write failure (poison-pill, null bytes) | `nexus_mq_batch_flush_total{result="error"}` once per affected batch; warn-level "permanent encoding error, acking to skip poison batch" log | None needed — the batch is dropped and the next batch proceeds. The producer-side `stripNul` plumbing prevents this almost everywhere; a leak is a producer-side bug | diff --git a/docs/developers/architecture/services/ai-gateway/hook-architecture.md b/docs/developers/architecture/services/ai-gateway/hook-architecture.md index 3ed6ee68..e2706332 100644 --- a/docs/developers/architecture/services/ai-gateway/hook-architecture.md +++ b/docs/developers/architecture/services/ai-gateway/hook-architecture.md @@ -70,7 +70,7 @@ When a redacting hook (`Modify`) co-fires with a soft-block hook the aggregate i ## 6. Config flow -`HookConfigCache` is the bridge from stored config to the resolver: a loader reads the `HookConfig` rows and `Swap`s them into the `PolicyResolver`. On the server-side data planes it reloads when the Hub pushes a config change (via the thing-client `OnConfigChanged` callback) with a TTL backstop; the Agent has no direct database access, so it is push-only. Before the swap, `rulepack.Enrich` binds each installed rule pack into the relevant hook's config under `_rulePackInstalls`, so the `rulepack-engine` hook evaluates packs without holding a database handle inside `Execute`. +`HookConfigCache` is the bridge from stored config to the resolver: a loader reads the `HookConfig` rows and `Swap`s them into the `PolicyResolver`. On the server-side data planes it reloads when the Hub pushes a config change (via the thing-client `OnConfigChanged` callback), with a background TTL-backstop ticker covering a degraded push channel; the request path never loads — `Resolver()` is a pure snapshot getter, so a slow database cannot stall or stampede request goroutines. The Agent has no direct database access, so it is push-only. Before the swap, `rulepack.Enrich` binds each installed rule pack into the relevant hook's config under `_rulePackInstalls`, so the `rulepack-engine` hook evaluates packs without holding a database handle inside `Execute`. The AI Gateway invokes the pipeline at both the request and response stages: it builds a sequential pipeline for the stage, ingress, endpoint, and modality, enables `allowModify` and `clearSoftOnApprove`, and executes it against the `HookInput`. diff --git a/docs/developers/architecture/services/ai-gateway/provider-adapter-architecture.md b/docs/developers/architecture/services/ai-gateway/provider-adapter-architecture.md index 52a403f8..a0208a3f 100644 --- a/docs/developers/architecture/services/ai-gateway/provider-adapter-architecture.md +++ b/docs/developers/architecture/services/ai-gateway/provider-adapter-architecture.md @@ -71,7 +71,7 @@ The `shape` parameter tells the codec which of its native wire shapes the call t `specAdapter.Execute` runs `PrepareBody`, which chooses between two paths: -- **Passthrough.** When the caller's body is already in the adapter's `Format` (or both sides are OpenAI-family), `PassthroughRewrite` applies any in-place model rewrites and the body is forwarded. `stripNexusNamespace` deletes the `nexus` key from the body before it reaches the upstream, so extension metadata never leaks to the provider. +- **Passthrough.** When the caller's body is already in the adapter's `Format` (or both sides are OpenAI-family), the body is forwarded after the generic dispatcher rewrites the top-level `model` field to the resolved `CallTarget.ProviderModelID` (so a client-facing alias or a routing-changed model reaches the upstream as the vendor's real name, not the 404-inducing alias) and `PassthroughRewrite` applies any per-model wire quirks in place. The model rewrite fires for wires that carry the model at the JSON body root: OpenAI-family formats (`Format.IsOpenAIFamily`) and adapters that declare `AdapterSpec.PassthroughModelInBody` — today only Anthropic `/v1/messages`, whose same-format native passthrough skips the `EncodeRequest` that would otherwise stamp `ProviderModelID`. Wires carrying the model elsewhere leave it untouched here: Gemini (URL path, set by `Transport.BuildURL`), Bedrock (deleted from the body, encoded into the URL by its codec). `stripNexusNamespace` deletes the `nexus` key from the body before it reaches the upstream, so extension metadata never leaks to the provider. - **Codec.** Otherwise `SchemaCodec.EncodeRequest` translates the canonical body into the target wire. `canonicalbridge.Bridge` (`packages/ai-gateway/internal/execution/canonicalbridge`) holds the per-`Format` codecs and exposes `IngressChatToCanonical`; its `chatWireShapeForFormat` / `embeddingsWireShapeForFormat` helpers resolve the native `WireShape` for a `Format`. From 9f526c17053120b6fbccb3489363ad1830669c41 Mon Sep 17 00:00:00 2001 From: nexus Date: Tue, 7 Jul 2026 16:00:08 +0800 Subject: [PATCH 8/8] feat(services): sync backend services and shared libs from internal main (sync 11) Backend service implementations and shared Go libraries. Dominant theme: AI Gateway overload-resilience. - packages/ai-gateway (75): in-flight admission gate (retryable 429 + Retry-After under overload), audit memory-budget accounting (AI_GATEWAY_AUDIT_MEM_MAX_BYTES), push-driven hook-config freshness (reload-stampede fix), admission/back-pressure metrics. - packages/shared (19): audit queue byte-budget + supporting transport, storage, and traffic helpers. - packages/compliance-proxy (2), packages/control-plane (2), packages/agent (1): companion updates. Runtime artifacts (logs/, dev-certs/) intentionally excluded. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../internal/sync/status/statusapi_server.go | 29 ++ packages/ai-gateway/ai-gateway.dev.yaml | 16 +- .../ai-gateway/cmd/ai-gateway/wiring/boot.go | 5 + .../cmd/ai-gateway/wiring/observability.go | 11 + .../cmd/ai-gateway/wiring/routes_test.go | 5 +- .../cmd/ai-gateway/wiring/wiring_init_test.go | 2 + .../wiring/wiring_lifecycle_test.go | 1 + .../internal/cache/layer/alias_lookup_test.go | 82 ++++ .../ai-gateway/internal/cache/layer/layer.go | 8 +- .../internal/cache/layer/loaders.go | 29 ++ .../internal/cache/layer/lookups.go | 16 + .../internal/config/config_cov_test.go | 20 + .../ai-gateway/internal/config/config_env.go | 3 + .../internal/config/config_pools.go | 42 +- .../canonicalbridge/stream_encoders.go | 11 +- .../stream_encoders_fastpath.go | 62 +++ .../stream_encoders_fastpath_test.go | 141 +++++++ .../executor/executor_internals_test.go | 28 +- .../proxy/admission_and_alias_review_test.go | 178 +++++++++ .../internal/ingress/proxy/admission_gate.go | 138 +++++++ .../ingress/proxy/admission_gate_test.go | 90 +++++ .../ingress/proxy/chunk_stream_reader_test.go | 6 + .../proxy/embeddings_crossformat_test.go | 8 +- .../internal/ingress/proxy/estimate_test.go | 6 + .../ingress/proxy/integration_test.go | 8 +- .../internal/ingress/proxy/interfaces.go | 4 + .../internal/ingress/proxy/proxy.go | 16 + .../ingress/proxy/proxy_cache_live.go | 13 +- .../ingress/proxy/proxy_cache_live_test.go | 86 +++- .../internal/ingress/proxy/proxy_e2e_test.go | 48 ++- .../ingress/proxy/proxy_fakeexec_test.go | 4 +- .../proxy/proxy_passthrough_fallback_test.go | 7 + .../ingress/proxy/proxy_serveproxy_test.go | 8 +- .../ingress/proxy/stage_accounting.go | 14 + .../internal/ingress/proxy/stage_admission.go | 19 +- .../internal/ingress/proxy/stage_context.go | 4 + .../internal/ingress/proxy/stage_hooks.go | 11 + .../internal/ingress/proxy/stage_routing.go | 7 +- .../internal/platform/audit/adaptive.go | 82 ++-- .../internal/platform/audit/audit.go | 46 +-- .../internal/platform/audit/audit_test.go | 17 +- .../internal/platform/audit/backpressure.go | 191 +++++++-- .../platform/audit/coverage_gaps_test.go | 56 +-- .../internal/platform/audit/loss_mode_test.go | 19 +- .../internal/platform/audit/mem_budget.go | 69 ++++ .../platform/audit/mem_budget_config_test.go | 78 ++++ .../platform/audit/mem_budget_test.go | 373 ++++++++++++++++++ .../audit/publish_failure_bound_test.go | 85 ++++ .../audit/publish_failure_extra_test.go | 175 ++++++++ .../internal/platform/audit/record.go | 20 + .../platform/audit/record_bodypool.go | 8 + .../audit/record_bodypool_release_test.go | 61 +++ .../internal/platform/audit/spill_recovery.go | 18 + .../internal/platform/audit/writer.go | 37 +- .../internal/platform/audit/writer_batch.go | 17 +- .../platform/audit/writer_concurrency_test.go | 15 +- .../platform/audit/writer_lifecycle.go | 40 +- .../internal/platform/audit/writer_metrics.go | 41 +- .../internal/platform/audit/writer_publish.go | 92 ++++- .../audit/writer_quota_backpressure_test.go | 147 +++++++ .../platform/audit/writer_spill_test.go | 21 +- .../internal/platform/audit/writer_splice.go | 2 + .../platform/middleware/connection_stage.go | 9 +- .../internal/platform/store/health.go | 261 +++++++++--- .../internal/platform/store/health_test.go | 199 ++++++++-- .../internal/platform/streaming/live.go | 104 ++++- .../platform/streaming/live_cadence_test.go | 170 ++++++++ .../platform/streaming/live_nohooks_test.go | 100 +++++ .../streaming/live_pipeline_decisions_test.go | 31 ++ .../internal/providers/core/spec.go | 23 ++ .../dispatch/registry_execute_test.go | 14 +- .../providers/dispatch/spec_adapter.go | 35 +- .../spec_adapter_modelrewrite_test.go | 200 +++++++++- .../providers/specs/anthropic/spec.go | 9 + .../routing/core/health_ranker_test.go | 27 +- .../routing/resolver_error_paths_test.go | 27 +- .../cmd/compliance-proxy/main.go | 7 + .../cmd/compliance-proxy/wiring/compliance.go | 6 +- .../authserver/login/displayname_test.go | 12 +- .../store/federated_store_mock_test.go | 4 +- packages/shared/audit/ndjson/ndjson.go | 32 +- .../shared/audit/ndjson/ndjson_quota_test.go | 87 ++++ packages/shared/core/bytebudget/bytebudget.go | 138 +++++++ .../shared/core/bytebudget/bytebudget_test.go | 184 +++++++++ .../validators/rulepack_engine_redact.go | 9 + .../rulepack_engine_redact_fastpath_test.go | 67 ++++ .../shared/policy/pipeline/config_cache.go | 106 +++-- .../pipeline/config_cache_refresh_test.go | 203 ++++++++++ .../policy/pipeline/config_cache_test.go | 21 +- .../shared/policy/rulepack/store_resolve.go | 7 +- .../shared/storage/configcache/metrics.go | 23 ++ .../shared/storage/configcache/snapshot.go | 2 + .../storage/configcache/snapshot_test.go | 35 ++ packages/shared/traffic/tracing.go | 88 +++-- packages/shared/traffic/tracing_ttfb_test.go | 144 +++++++ packages/shared/transport/mq/streams.go | 10 +- packages/shared/transport/mq/streams_test.go | 4 +- packages/shared/transport/streaming/live.go | 38 +- .../transport/streaming/live_cadence_test.go | 77 ++++ 99 files changed, 4910 insertions(+), 499 deletions(-) create mode 100644 packages/ai-gateway/internal/cache/layer/alias_lookup_test.go create mode 100644 packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_fastpath.go create mode 100644 packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_fastpath_test.go create mode 100644 packages/ai-gateway/internal/ingress/proxy/admission_and_alias_review_test.go create mode 100644 packages/ai-gateway/internal/ingress/proxy/admission_gate.go create mode 100644 packages/ai-gateway/internal/ingress/proxy/admission_gate_test.go create mode 100644 packages/ai-gateway/internal/platform/audit/mem_budget.go create mode 100644 packages/ai-gateway/internal/platform/audit/mem_budget_config_test.go create mode 100644 packages/ai-gateway/internal/platform/audit/mem_budget_test.go create mode 100644 packages/ai-gateway/internal/platform/audit/publish_failure_bound_test.go create mode 100644 packages/ai-gateway/internal/platform/audit/publish_failure_extra_test.go create mode 100644 packages/ai-gateway/internal/platform/audit/record_bodypool_release_test.go create mode 100644 packages/ai-gateway/internal/platform/audit/writer_quota_backpressure_test.go create mode 100644 packages/ai-gateway/internal/platform/streaming/live_cadence_test.go create mode 100644 packages/ai-gateway/internal/platform/streaming/live_nohooks_test.go create mode 100644 packages/shared/audit/ndjson/ndjson_quota_test.go create mode 100644 packages/shared/core/bytebudget/bytebudget.go create mode 100644 packages/shared/core/bytebudget/bytebudget_test.go create mode 100644 packages/shared/policy/hooks/validators/rulepack_engine_redact_fastpath_test.go create mode 100644 packages/shared/policy/pipeline/config_cache_refresh_test.go create mode 100644 packages/shared/storage/configcache/metrics.go create mode 100644 packages/shared/traffic/tracing_ttfb_test.go create mode 100644 packages/shared/transport/streaming/live_cadence_test.go diff --git a/packages/agent/internal/sync/status/statusapi_server.go b/packages/agent/internal/sync/status/statusapi_server.go index 9724de98..54e49dd2 100644 --- a/packages/agent/internal/sync/status/statusapi_server.go +++ b/packages/agent/internal/sync/status/statusapi_server.go @@ -280,6 +280,16 @@ type Server struct { wg sync.WaitGroup done chan struct{} stopOnce sync.Once + // closeMu serialises the accept loop's wg.Add(1) against Stop()'s + // wg.Wait(). Go requires every positive-delta Add to happen-before the + // Wait; without this an in-flight Accept() could call wg.Add(1) + // concurrently with Stop()'s wg.Wait(), which the race detector + // (correctly) flags as a data race on the WaitGroup counter. Stop() + // sets `closing` under this lock before waiting, so the accept loop + // either adds strictly before Stop() (mutex happens-before → before + // Wait) or observes closing and never adds at all. + closeMu sync.Mutex + closing bool // sem caps simultaneous live IPC connections. sem chan struct{} } @@ -361,7 +371,19 @@ func (s *Server) Start() error { _ = conn.Close() continue } + // Register the handler under closeMu so wg.Add(1) can never race + // Stop()'s wg.Wait(). If Stop() has already begun, drop this + // just-accepted connection and end the accept loop instead of + // adding to a WaitGroup that is being waited on. + s.closeMu.Lock() + if s.closing { + s.closeMu.Unlock() + <-s.sem + _ = conn.Close() + return nil + } s.wg.Add(1) + s.closeMu.Unlock() go func() { defer s.wg.Done() defer func() { <-s.sem }() @@ -491,6 +513,13 @@ func (s *Server) SetRefreshPoliciesFn(fn RefreshPoliciesFn) { // Stop gracefully stops the server. func (s *Server) Stop() { s.stopOnce.Do(func() { close(s.done) }) + // Mark closing before waiting so the accept loop stops calling + // wg.Add(1). Taken under closeMu, this establishes the happens-before + // the WaitGroup contract requires between any in-flight Add and the + // Wait below (see closeMu's field comment). + s.closeMu.Lock() + s.closing = true + s.closeMu.Unlock() s.listenerMu.Lock() ln := s.listener s.listenerMu.Unlock() diff --git a/packages/ai-gateway/ai-gateway.dev.yaml b/packages/ai-gateway/ai-gateway.dev.yaml index 7cac37aa..e88d3999 100644 --- a/packages/ai-gateway/ai-gateway.dev.yaml +++ b/packages/ai-gateway/ai-gateway.dev.yaml @@ -241,13 +241,17 @@ audit: # constrained one. 0 → 10000. Env: AI_GATEWAY_AUDIT_MAX_QUEUED_RECORDS. maxQueuedRecords: 10000 # Overflow policy. Durable audit is a product promise + compliance requirement, - # so the default "block" is NO-LOSS back-pressure: when the buffer fills, the - # gateway slows the request path until the audit pipeline drains (NATS disk is - # the burst buffer) — never dropping a record. Lossy opt-out for non-compliance - # callers: "spill" (async durable NDJSON spill off the request path; bounded drop + # so the default "spillBlock" is NO-LOSS with the cheap on-disk spool as the + # PRIMARY overflow buffer: when the in-heap buffer fills, records go to the + # durable NDJSON spool first (batched, off the publish worker), and the request + # path is only back-pressured when that large spool is ALSO saturated — never + # dropping a record (the spool is wired via spoolDir below; if no spool is wired + # it falls back to block-style back-pressure). "block" is the stricter variant + # that back-pressures at the small in-heap queue without using the spool. Lossy + # opt-out for non-compliance callers: "spill" (async durable spill; bounded drop # only if the spill is also saturated) or "drop" (counted bounded drop, max - # throughput). Empty/unknown → "block". Env: AI_GATEWAY_AUDIT_LOSS_MODE. - lossMode: block + # throughput). Empty/unknown → "spillBlock". Env: AI_GATEWAY_AUDIT_LOSS_MODE. + lossMode: spillBlock # End-to-end zstd compression of large captured bodies: producer compresses off # the request path, body rides the NATS wire compressed, the Hub persists the # compressed bytes verbatim (no decompress on ingest), and only the CP view diff --git a/packages/ai-gateway/cmd/ai-gateway/wiring/boot.go b/packages/ai-gateway/cmd/ai-gateway/wiring/boot.go index c9c549ea..718b7912 100644 --- a/packages/ai-gateway/cmd/ai-gateway/wiring/boot.go +++ b/packages/ai-gateway/cmd/ai-gateway/wiring/boot.go @@ -245,6 +245,11 @@ func Boot(ctx context.Context, cfg *config.Config, logger *slog.Logger) (*BootDe if d.Tp != nil { _ = d.Tp.Shutdown(context.Background()) } + if d.HealthTracker != nil { + // Stop the health writer goroutine for a clean shutdown (symmetry with + // NewHealthTracker; prod runs a single process-lifetime instance). + d.HealthTracker.Stop() + } } return d, cleanup, nil } diff --git a/packages/ai-gateway/cmd/ai-gateway/wiring/observability.go b/packages/ai-gateway/cmd/ai-gateway/wiring/observability.go index c313d250..06bdbb19 100644 --- a/packages/ai-gateway/cmd/ai-gateway/wiring/observability.go +++ b/packages/ai-gateway/cmd/ai-gateway/wiring/observability.go @@ -86,6 +86,17 @@ func InitAuditWriter( logger.Info("audit overflow policy resolved", "configured", auditCfg.LossMode, "effective", auditWriter.LossMode()) + // In-memory audit byte budget — the memory half of the bounded audit queue. + // Same semantics as NEXUS_EVENTS_MAX_BYTES: empty/"auto" auto-sizes to ~15% of + // available RAM; an explicit size ("8GB") pins it. Log the RESOLVED value so + // operators see what the box defaulted to and how to pin it. + auditWriter.WithMemMaxBytes(auditCfg.MemMaxBytes) + logger.Info("audit memory budget resolved", + "configured", auditCfg.MemMaxBytes, + "effective_bytes", auditWriter.MemBudgetBytes(), + "effective_gib", float64(auditWriter.MemBudgetBytes())/(1<<30), + "override", "set AI_GATEWAY_AUDIT_MEM_MAX_BYTES= (e.g. 8GB) to pin a fixed value") + // End-to-end zstd compression of large captured bodies: the producer marks // large bodies for compression in NewInlineBody and compresses lazily in the // async marshal worker (off the request path); the Hub persists the compressed diff --git a/packages/ai-gateway/cmd/ai-gateway/wiring/routes_test.go b/packages/ai-gateway/cmd/ai-gateway/wiring/routes_test.go index 78f1ba05..67b329ef 100644 --- a/packages/ai-gateway/cmd/ai-gateway/wiring/routes_test.go +++ b/packages/ai-gateway/cmd/ai-gateway/wiring/routes_test.go @@ -41,6 +41,9 @@ func buildMinimalRouteDeps(t *testing.T) RouteDeps { } db := store.NewWithPgxPool(mock) + ht := store.NewHealthTracker() + t.Cleanup(ht.Stop) + return RouteDeps{ Config: &config.Config{ // Auth.InternalServiceToken gates the /internal/* operator routes; @@ -62,7 +65,7 @@ func buildMinimalRouteDeps(t *testing.T) RouteDeps { HookConfigCache: hookCache, GWHookRegistry: reg, ProviderReg: adapterReg, - HealthTracker: store.NewHealthTracker(), + HealthTracker: ht, PayloadCapture: pcs, Allowlist: allowlist, Logger: discardLogger(), diff --git a/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_init_test.go b/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_init_test.go index 1c8f8128..4b59e990 100644 --- a/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_init_test.go +++ b/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_init_test.go @@ -163,6 +163,7 @@ func TestInitExecutor_withNilPtResolverAndStats(t *testing.T) { } adapterReg := InitProviderRegistry(allowlist, discardLogger()) ht := store.NewHealthTracker() + t.Cleanup(ht.Stop) bridge, exec := InitExecutor(adapterReg, nil, ht, nil, discardLogger()) if bridge == nil { @@ -399,6 +400,7 @@ func TestInitRouter_withNonNilCacheLayerAndPtResolver(t *testing.T) { ptResolver := NewResolver(l, mgr, nil) ht := store.NewHealthTracker() + t.Cleanup(ht.Stop) stratReg, healthRanker, resolver, capCache := InitRouter(l, ht, ptResolver, adapterReg, discardLogger()) if stratReg == nil { t.Fatal("expected non-nil strategyReg") diff --git a/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_lifecycle_test.go b/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_lifecycle_test.go index 5f0c7151..275ec63c 100644 --- a/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_lifecycle_test.go +++ b/packages/ai-gateway/cmd/ai-gateway/wiring/wiring_lifecycle_test.go @@ -182,6 +182,7 @@ func TestMountRoutes_nilAiguardNoMountAIGuard(t *testing.T) { hookCache := InitHookConfigCache(nil, hookReg, discardLogger()) pcs := payloadcapture.NewStore(payloadcapture.DefaultConfig()) ht := store.NewHealthTracker() + t.Cleanup(ht.Stop) bd := &BootDeps{ Cfg: &config.Config{ diff --git a/packages/ai-gateway/internal/cache/layer/alias_lookup_test.go b/packages/ai-gateway/internal/cache/layer/alias_lookup_test.go new file mode 100644 index 00000000..e40d2928 --- /dev/null +++ b/packages/ai-gateway/internal/cache/layer/alias_lookup_test.go @@ -0,0 +1,82 @@ +package cachelayer + +import ( + "context" + "testing" + + "github.com/pashagolub/pgxmock/v4" +) + +// aliases is column index 18 in the loadModels SELECT (see makeModelRow). +const aliasesCol = 18 + +// TestGetModelByCodeOrAlias_ResolvesCodeAndAlias asserts the O(1) code-or-alias +// index resolves both a model's code and each of its aliases to the same model, +// while GetModelByCode stays strict (aliases miss). +func TestGetModelByCodeOrAlias_ResolvesCodeAndAlias(t *testing.T) { + mock, l := newMockLayer(t, Config{}) + row := makeModelRow("m1", "opus", "p1", true) + row[aliasesCol] = []string{"fast", "smart"} + mock.ExpectQuery(`FROM "Model" m`). + WillReturnRows(pgxmock.NewRows(modelCols).AddRow(row...)) + if _, err := l.loadModels(context.Background()); err != nil { + t.Fatalf("loadModels: %v", err) + } + + for _, key := range []string{"opus", "fast", "smart"} { + m, err := l.GetModelByCodeOrAlias(context.Background(), key) + if err != nil { + t.Fatalf("key %q must resolve: %v", key, err) + } + if m.Code != "opus" || m.ID != "m1" { + t.Errorf("key %q resolved to wrong model: %+v", key, m) + } + } + // GetModelByCode must NOT resolve an alias — its narrow contract is unchanged. + if _, err := l.GetModelByCode(context.Background(), "fast"); !IsNotFound(err) { + t.Errorf("GetModelByCode must not resolve aliases; got %v", err) + } + if _, err := l.GetModelByCodeOrAlias(context.Background(), "unknown"); !IsNotFound(err) { + t.Errorf("unknown key must be not-found; got %v", err) + } +} + +// TestGetModelByCodeOrAlias_CodeWinsOverAliasCollision pins the priority rule: +// when one model's alias collides with another model's real code, the code wins. +func TestGetModelByCodeOrAlias_CodeWinsOverAliasCollision(t *testing.T) { + mock, l := newMockLayer(t, Config{}) + rowA := makeModelRow("mA", "shared", "p1", true) // real code "shared" + rowB := makeModelRow("mB", "other", "p1", true) + rowB[aliasesCol] = []string{"shared"} // alias collides with A's code + mock.ExpectQuery(`FROM "Model" m`). + WillReturnRows(pgxmock.NewRows(modelCols).AddRow(rowA...).AddRow(rowB...)) + if _, err := l.loadModels(context.Background()); err != nil { + t.Fatalf("loadModels: %v", err) + } + m, err := l.GetModelByCodeOrAlias(context.Background(), "shared") + if err != nil { + t.Fatalf("resolve: %v", err) + } + if m.ID != "mA" { + t.Errorf("code must win over alias collision; got model %s", m.ID) + } +} + +// TestGetModelByCodeOrAlias_DisabledExcluded asserts a disabled model's alias +// is not routable (mirrors the enabled-only byCode filter). +func TestGetModelByCodeOrAlias_DisabledExcluded(t *testing.T) { + mock, l := newMockLayer(t, Config{}) + row := makeModelRow("m1", "opus", "p1", false) // disabled + row[aliasesCol] = []string{"fast"} + mock.ExpectQuery(`FROM "Model" m`). + WillReturnRows(pgxmock.NewRows(modelCols).AddRow(row...)) + if _, err := l.loadModels(context.Background()); err != nil { + t.Fatalf("loadModels: %v", err) + } + if _, err := l.GetModelByCodeOrAlias(context.Background(), "fast"); !IsNotFound(err) { + t.Errorf("disabled model's alias must not resolve; got %v", err) + } + if _, err := l.GetModelByCodeOrAlias(context.Background(), "opus"); !IsNotFound(err) { + t.Errorf("disabled model's code must not resolve; got %v", err) + } +} diff --git a/packages/ai-gateway/internal/cache/layer/layer.go b/packages/ai-gateway/internal/cache/layer/layer.go index 31a9bd80..a08c5bb7 100644 --- a/packages/ai-gateway/internal/cache/layer/layer.go +++ b/packages/ai-gateway/internal/cache/layer/layer.go @@ -65,7 +65,13 @@ type Layer struct { // Secondary indices computed alongside snapshot loads. Replaced // atomically; readers never see a torn state. - modelsByCode atomic.Pointer[map[string]store.Model] + modelsByCode atomic.Pointer[map[string]store.Model] + // modelsByCodeOrAlias resolves a customer-supplied identifier that may be + // a Model.code OR one of its Model.aliases in O(1) — the routing + // passthrough fallback needs alias resolution without an O(n) scan or a + // per-request DB read. Code keys take priority over alias keys, so an + // alias can never shadow another model's real code. + modelsByCodeOrAlias atomic.Pointer[map[string]store.Model] credentialsByProviderFirst atomic.Pointer[map[string]store.Credential] // invalidationCount tracks how many entries the cache evicted on diff --git a/packages/ai-gateway/internal/cache/layer/loaders.go b/packages/ai-gateway/internal/cache/layer/loaders.go index 14533496..9a7b0e5b 100644 --- a/packages/ai-gateway/internal/cache/layer/loaders.go +++ b/packages/ai-gateway/internal/cache/layer/loaders.go @@ -54,6 +54,7 @@ func (l *Layer) loadModels(ctx context.Context) (map[string]store.Model, error) byID := map[string]store.Model{} byCode := map[string]store.Model{} + enabled := make([]store.Model, 0) for rows.Next() { var m store.Model var inPrice, outPrice, cachedReadPrice, cachedWritePrice *string @@ -92,8 +93,36 @@ func (l *Layer) loadModels(ctx context.Context) (map[string]store.Model, error) if m.Enabled && m.Code != "" { byCode[m.Code] = m } + if m.Enabled { + enabled = append(enabled, m) + } + } + // Build the code-or-alias index. Pass 1 seeds every enabled code so a real + // code always wins over any alias — this is order-independent and is the + // only load-bearing guarantee. Pass 2 adds aliases only where the key is + // still free. Alias uniqueness is NOT enforced by the schema or admin API, + // so if two enabled models share the same alias the winner is unspecified + // (whichever row the catalog query returns first); an ambiguous alias is a + // misconfiguration, not a case this index promises to disambiguate. + byCodeOrAlias := make(map[string]store.Model, len(byCode)) + for _, m := range enabled { + if m.Code != "" { + byCodeOrAlias[m.Code] = m + } + } + for _, m := range enabled { + for _, a := range m.Aliases { + if a == "" { + continue + } + if _, taken := byCodeOrAlias[a]; taken { + continue + } + byCodeOrAlias[a] = m + } } l.modelsByCode.Store(&byCode) + l.modelsByCodeOrAlias.Store(&byCodeOrAlias) return byID, nil } diff --git a/packages/ai-gateway/internal/cache/layer/lookups.go b/packages/ai-gateway/internal/cache/layer/lookups.go index afa8f706..bf509f1f 100644 --- a/packages/ai-gateway/internal/cache/layer/lookups.go +++ b/packages/ai-gateway/internal/cache/layer/lookups.go @@ -48,6 +48,22 @@ func (l *Layer) GetModelByCode(ctx context.Context, code string) (*store.Model, return nil, fmt.Errorf("cachelayer: model code %q: %w", code, errNotFound) } +// GetModelByCodeOrAlias returns the enabled Model row whose customer-facing +// code OR one of its aliases matches key, resolved in O(1) from the in-memory +// code-or-alias index (no per-request scan, no DB read). Codes take priority +// over aliases. Used by the routing passthrough fallback so an aliased model +// routes without an explicit routing rule; GetModelByCode stays strict-code. +func (l *Layer) GetModelByCodeOrAlias(ctx context.Context, key string) (*store.Model, error) { + idx := l.modelsByCodeOrAlias.Load() + if idx != nil { + if m, ok := (*idx)[key]; ok { + v := m + return &v, nil + } + } + return nil, fmt.Errorf("cachelayer: model code/alias %q: %w", key, errNotFound) +} + // AllModels returns a copy of every Model row in the current snapshot. // Used by the capability cache rebuild hook in configdispatch after a // models reload so the pre-filter stays in sync without a second DB query. diff --git a/packages/ai-gateway/internal/config/config_cov_test.go b/packages/ai-gateway/internal/config/config_cov_test.go index cee66665..0ea8b52d 100644 --- a/packages/ai-gateway/internal/config/config_cov_test.go +++ b/packages/ai-gateway/internal/config/config_cov_test.go @@ -26,6 +26,26 @@ func TestLoad_AuditSpoolMB_EnvOverride(t *testing.T) { } } +// AI_GATEWAY_AUDIT_MEM_MAX_BYTES flows verbatim into Audit.MemMaxBytes (the +// audit writer parses it and keeps its auto default on a bad value), and an +// unset env leaves the default empty string (= auto). +func TestLoad_AuditMemMaxBytes_EnvOverride(t *testing.T) { + p := writeYAML(t, "server:\n port: 3050\n") + setRequiredEnvBaseline(t) + t.Setenv("AI_GATEWAY_AUDIT_MEM_MAX_BYTES", "8GB") + + cfg, err := Load(p) + if err != nil { + t.Fatalf("Load: %v", err) + } + if cfg.Audit.MemMaxBytes != "8GB" { + t.Errorf("Audit.MemMaxBytes = %q, want %q (from env)", cfg.Audit.MemMaxBytes, "8GB") + } + if def := defaults().Audit.MemMaxBytes; def != "" { + t.Errorf("default Audit.MemMaxBytes = %q, want empty (= auto)", def) + } +} + // TestLoad_AuditSpoolMB_MalformedKeepsDefault asserts the best-effort parse // contract: a non-numeric value is ignored and the default survives. We capture // the default from a clean Load (env unset) and prove the malformed Load yields diff --git a/packages/ai-gateway/internal/config/config_env.go b/packages/ai-gateway/internal/config/config_env.go index c8f343ab..6d8f7861 100644 --- a/packages/ai-gateway/internal/config/config_env.go +++ b/packages/ai-gateway/internal/config/config_env.go @@ -24,6 +24,9 @@ func applyEnvOverrides(cfg *Config) { // best-effort: a malformed value leaves the default (10000) in place. _, _ = fmt.Sscanf(v, "%d", &cfg.Audit.MaxQueuedRecords) } + if v := os.Getenv("AI_GATEWAY_AUDIT_MEM_MAX_BYTES"); v != "" { + cfg.Audit.MemMaxBytes = v + } if v := os.Getenv("AI_GATEWAY_AUDIT_LOSS_MODE"); v != "" { cfg.Audit.LossMode = v } diff --git a/packages/ai-gateway/internal/config/config_pools.go b/packages/ai-gateway/internal/config/config_pools.go index 423053f8..187924ab 100644 --- a/packages/ai-gateway/internal/config/config_pools.go +++ b/packages/ai-gateway/internal/config/config_pools.go @@ -70,22 +70,40 @@ type AuditConfig struct { // one. 0 falls back to the 10000 default. Env: AI_GATEWAY_AUDIT_MAX_QUEUED_RECORDS. MaxQueuedRecords int `yaml:"maxQueuedRecords"` + // MemMaxBytes bounds the BYTES the in-memory audit queue may pin (captured + // request/response bodies), the memory half of the bounded audit queue: on a + // full budget the no-loss modes back-pressure the request path, the lossy + // modes shed with a counted drop. Semantics mirror NEXUS_EVENTS_MAX_BYTES on + // the Hub side exactly: empty or "auto" (default) auto-sizes to ~15% of the + // box's available RAM (2 GiB fallback where /proc/meminfo is unreadable); an + // explicit human size ("4GB", "2048MB", raw bytes) pins it — raise it on a + // memory-rich box to absorb bigger bursts before throttling, lower it on a + // constrained one. Env: AI_GATEWAY_AUDIT_MEM_MAX_BYTES. + MemMaxBytes string `yaml:"memMaxBytes"` + // LossMode is the audit overflow policy. Durable audit is a product promise + - // a compliance requirement, so the DEFAULT is "block" — no-loss back-pressure: - // when the in-heap buffer is full the gateway slows the request path until the - // audit pipeline drains, never dropping a record (NATS disk storage is the - // burst buffer). The lossy modes are an explicit opt-out for callers that do - // NOT need compliance audit and prefer raw throughput: - // - "block" (default): back-pressure, no loss. + // a compliance requirement, so the DEFAULT is "spillblock" — no-loss with the + // cheap on-disk spool as the PRIMARY overflow buffer: when the in-heap buffer + // fills, records go to the durable NDJSON spool first (batched, off the publish + // worker), and the request path is back-pressured only when that large spool is + // ALSO saturated — never dropping a record. The lossy modes are an explicit + // opt-out for callers that do NOT need compliance audit and prefer raw throughput: + // - "spillblock" (default): spool-first, back-pressure the request path only when + // the spool ALSO saturates — and then it PARKS (never drops): the + // single spill worker holds the batch and retries while the recovery + // sweeper drains the spool and frees space, so ingest self-throttles + // to the recovery-drain rate. Genuinely no-loss whenever a spool is + // wired (the only drops are a real non-quota disk error or shutdown); + // if no spool is wired (empty SpoolDir or a spool-dir creation failure) + // it downgrades to "block" at startup so it never silently drops. + // - "block": back-pressure at the in-heap queue (never touches the spool), no + // loss and no spool dependency — the stricter no-loss variant. // - "spill": async durable NDJSON spill off the request path; bounded drop // only if the spill is also saturated. - // - "spillblock": like "spill", but when the spill channel is also saturated - // it back-pressures the request path (bounded by backpressureMaxWait - // → durable spill) instead of dropping — lossless up to disk-write - // success, throttling ingest to the disk rate under overload. // - "drop": counted bounded drop on overflow; max throughput, lossy. - // Empty / unrecognised → "block" (audit must not silently turn lossy from a - // config typo). Env: AI_GATEWAY_AUDIT_LOSS_MODE. + // Empty / unrecognised → "spillblock" (audit must not silently turn lossy from a + // config typo; spillblock is no-loss whenever a spool is wired, else downgrades to + // block). Env: AI_GATEWAY_AUDIT_LOSS_MODE. LossMode string `yaml:"lossMode"` // Compress enables end-to-end zstd compression of large captured bodies on diff --git a/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders.go b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders.go index 7c00a6aa..2b3f35a7 100644 --- a/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders.go +++ b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders.go @@ -32,6 +32,11 @@ type openAIStreamEncoder struct { // it (the io.Writer "must not retain p" contract; every call site writes it // out synchronously before the next Write). scratch []byte + // contentSuffix is the precomputed tail of a content-delta frame — everything + // after the (variable) content string. See stream_encoders_fastpath.go + // (emitContentDelta) for how the per-token hot path uses it to avoid + // marshalling the whole envelope struct graph. + contentSuffix []byte } // NewChatCompletionsStreamEncoder returns an encoder that converts canonical @@ -127,8 +132,10 @@ func (e *openAIStreamEncoder) Write(_ context.Context, chunk provcore.Chunk) ([] // finishReason, and usageMetadata into a single SSE frame, so chunk.Delta // can be non-empty even when chunk.Done is also true. if chunk.Delta != "" { - d := chunk.Delta - e.emit(oaiStreamChoice{Delta: oaiStreamDelta{Content: &d}}, nil) + // Per-token hot path: byte-identical to + // emit(oaiStreamChoice{Delta:{Content:&d}}, nil) but skips the envelope + // struct reflection (the dominant streaming encode cost). + e.emitContentDelta(chunk.Delta) } if len(chunk.ToolCallDeltas) > 0 { e.emit(oaiStreamChoice{Delta: oaiStreamDelta{ToolCalls: buildOAIToolCalls(chunk.ToolCallDeltas)}}, nil) diff --git a/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_fastpath.go b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_fastpath.go new file mode 100644 index 00000000..acd02d36 --- /dev/null +++ b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_fastpath.go @@ -0,0 +1,62 @@ +package canonicalbridge + +// stream_encoders_fastpath.go — the per-token content-delta fast path for +// openAIStreamEncoder, split from stream_encoders.go to keep that file under the +// size ratchet. The output is byte-identical to the struct-marshalled envelope +// (pinned by the differential + fuzz tests in stream_encoders_fastpath_test.go); +// it exists only to skip go-json's reflection over the whole envelope struct for +// the dominant streaming frame. + +import ( + "fmt" + + json "github.com/goccy/go-json" +) + +// contentFramePrefix is the invariant head of a content-delta frame up to (and +// including) the opening of the content value. The envelope's fields are encoded +// in alphabetical order (see stream_encoders_openai_types.go), so `choices` is +// first and nothing stream-specific precedes the content — this prefix is a true +// constant shared by every encoder. +var contentFramePrefix = []byte(`data: {"choices":[{"delta":{"content":`) + +// buildContentSuffix assembles the fixed tail of a content-delta frame: +// everything after the content value — the null finish_reason, index, and the +// per-stream-constant created/id/model/object fields, plus the closing framing +// and trailing blank line. created/id/model are emitted via json.Marshal so +// their escaping and numeric formatting are byte-identical to the +// struct-marshalled envelope. +func buildContentSuffix(created int64, id, model string) []byte { + idJSON, _ := json.Marshal(id) + modelJSON, _ := json.Marshal(model) + var s []byte + s = append(s, `},"finish_reason":null,"index":0}],"created":`...) + s = append(s, fmt.Appendf(nil, "%d", created)...) + s = append(s, `,"id":`...) + s = append(s, idJSON...) + s = append(s, `,"model":`...) + s = append(s, modelJSON...) + s = append(s, `,"object":"chat.completion.chunk"}`...) + s = append(s, '\n', '\n') + return s +} + +// emitContentDelta appends a content-delta frame for the per-token hot path. +// It is byte-identical to emit(oaiStreamChoice{Delta:{Content:&content}}, nil) +// but marshals ONLY the content string (go-json's cheap string fast path) plus a +// precomputed constant prefix/suffix, avoiding reflection over the whole envelope +// struct graph. The content string uses json.Marshal so its HTML-safe escaping +// matches the struct path exactly. Pinned by TestOpenAIStreamEncoder_ContentDeltaFastPath. +func (e *openAIStreamEncoder) emitContentDelta(content string) { + // Lazily assemble the per-stream suffix on first use so the fast path is + // robust to any construction (the constructor, a struct literal in tests, or + // a post-construction id/created/model set) — id/created/model are immutable + // once the stream is producing content. + if e.contentSuffix == nil { + e.contentSuffix = buildContentSuffix(e.created, e.id, e.model) + } + contentJSON, _ := json.Marshal(content) + e.scratch = append(e.scratch, contentFramePrefix...) + e.scratch = append(e.scratch, contentJSON...) + e.scratch = append(e.scratch, e.contentSuffix...) +} diff --git a/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_fastpath_test.go b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_fastpath_test.go new file mode 100644 index 00000000..ae51d33b --- /dev/null +++ b/packages/ai-gateway/internal/execution/canonicalbridge/stream_encoders_fastpath_test.go @@ -0,0 +1,141 @@ +package canonicalbridge + +import ( + "context" + "testing" + + json "github.com/goccy/go-json" + + provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" +) + +// referenceContentFrame produces the content-delta frame the way the struct +// encoder would (emit(oaiStreamChoice{Delta:{Content:&c}}, nil)) — the byte +// baseline the fast path must match. +func referenceContentFrame(e *openAIStreamEncoder, content string) string { + cc := content + data, _ := json.Marshal(oaiStreamEnvelope{ + Choices: []oaiStreamChoice{{Delta: oaiStreamDelta{Content: &cc}}}, + Created: e.created, + ID: e.id, + Model: e.model, + Object: "chat.completion.chunk", + }) + return "data: " + string(data) + "\n\n" +} + +// TestOpenAIStreamEncoder_ContentDeltaFastPath pins the per-token fast path to be +// byte-identical to the struct-marshalled envelope for a range of content that +// exercises JSON string escaping (quotes, backslash, control chars, HTML-unsafe +// chars, line/paragraph separators, multibyte, and would-be JSON-injection). +func TestOpenAIStreamEncoder_ContentDeltaFastPath(t *testing.T) { + cases := []string{ + "", + "hello world", + `a "quoted" b`, + `back\slash`, + "line\nbreak\r\ttab", + "html tag & ampersand", + "café 日本語 😀 é", + "line
sep
para", + "ctrl\x00\x01\x1f end", + `"},"finish_reason":"injected"}]}`, // must be escaped, not break the frame + "trailing backslash \\", + } + for _, c := range cases { + e := newOpenAIStreamEncoder("gpt-4o") + e.scratch = e.scratch[:0] + e.emitContentDelta(c) + got := string(e.scratch) + want := referenceContentFrame(e, c) + if got != want { + t.Errorf("content %q:\n got=%q\nwant=%q", c, got, want) + } + // The fast-path output must parse as valid JSON with the right content. + var env struct { + Choices []struct { + Delta struct { + Content string `json:"content"` + } `json:"delta"` + } `json:"choices"` + } + payload := got[len("data: ") : len(got)-2] + if err := json.Unmarshal([]byte(payload), &env); err != nil { + t.Fatalf("content %q: fast-path frame is not valid JSON: %v (%q)", c, err, payload) + } + if len(env.Choices) != 1 || env.Choices[0].Delta.Content != c { + t.Errorf("content %q: round-trip mismatch, got %q", c, env.Choices[0].Delta.Content) + } + } +} + +// FuzzOpenAIStreamEncoder_ContentDelta asserts byte-identity of the fast path vs +// the struct encoder for arbitrary content, so no escaping edge case can diverge. +func FuzzOpenAIStreamEncoder_ContentDelta(f *testing.F) { + for _, s := range []string{"", "x", `"`, "\\", "<&>", "😀", "
", "\x00", `"}]}`} { + f.Add(s) + } + f.Fuzz(func(t *testing.T, content string) { + e := newOpenAIStreamEncoder("model-x") + e.scratch = e.scratch[:0] + e.emitContentDelta(content) + got := string(e.scratch) + want := referenceContentFrame(e, content) + if got != want { + t.Fatalf("content %q:\n got=%q\nwant=%q", content, got, want) + } + }) +} + +// TestOpenAIStreamEncoder_WriteContentUnchanged verifies the full Write path still +// produces the exact same bytes for a content chunk after routing it through the +// fast path (the role-header frame precedes it on the first Write, and stays on +// the struct-marshal path). +func TestOpenAIStreamEncoder_WriteContentUnchanged(t *testing.T) { + e := newOpenAIStreamEncoder("gpt-4o") + c := `hi "there" ` + out, err := e.Write(context.Background(), provcore.Chunk{Delta: c}) + if err != nil { + t.Fatal(err) + } + // First Write emits the role-header frame (struct path) then the content + // frame (fast path); both reference the same encoder id/created, unchanged by + // Write. + empty := "" + roleData, _ := json.Marshal(oaiStreamEnvelope{ + Choices: []oaiStreamChoice{{Delta: oaiStreamDelta{Content: &empty, Role: "assistant"}}}, + Created: e.created, + ID: e.id, + Model: e.model, + Object: "chat.completion.chunk", + }) + want := "data: " + string(roleData) + "\n\n" + referenceContentFrame(e, c) + if string(out) != want { + t.Errorf("Write content frame changed:\n got=%q\nwant=%q", string(out), want) + } +} + +// BenchmarkOpenAIStreamEncoder_ContentDelta is the before/after comparison: the +// per-token fast path vs the struct-marshalled envelope path. +func BenchmarkOpenAIStreamEncoder_ContentDelta(b *testing.B) { + content := "The quick brown fox jumps over the lazy dog, then keeps on going." + b.Run("fast", func(b *testing.B) { + e := newOpenAIStreamEncoder("gpt-4o") + b.ReportAllocs() + b.ResetTimer() + for range b.N { + e.scratch = e.scratch[:0] + e.emitContentDelta(content) + } + }) + b.Run("marshal", func(b *testing.B) { + e := newOpenAIStreamEncoder("gpt-4o") + cc := content + b.ReportAllocs() + b.ResetTimer() + for range b.N { + e.scratch = e.scratch[:0] + e.emit(oaiStreamChoice{Delta: oaiStreamDelta{Content: &cc}}, nil) + } + }) +} diff --git a/packages/ai-gateway/internal/execution/executor/executor_internals_test.go b/packages/ai-gateway/internal/execution/executor/executor_internals_test.go index a4f0b9f1..3eee2c31 100644 --- a/packages/ai-gateway/internal/execution/executor/executor_internals_test.go +++ b/packages/ai-gateway/internal/execution/executor/executor_internals_test.go @@ -481,6 +481,25 @@ func TestExecute_ContextCanceledDuringBackoff(t *testing.T) { // classifyAttempt — recordHealth happy paths (tracker non-nil). +// awaitHealth polls GetHealth until at least one sample is visible or a short +// deadline elapses. The HealthTracker applies records asynchronously (single +// writer), so a read immediately after Execute must wait for the sample to be +// published rather than assuming synchronous visibility. +func awaitHealth(t *testing.T, health *store.HealthTracker, providerID string) store.HealthState { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for { + hs := health.GetHealth(providerID) + if hs.SampleCount > 0 { + return hs + } + if time.Now().After(deadline) { + return hs + } + time.Sleep(2 * time.Millisecond) + } +} + // TestExecute_RecordsHealth_SuccessAndFailure covers executor.go:454-458 // — when a HealthTracker is wired, recordHealth invokes RecordSuccess // (success path) AND RecordFailure (retryable failure + 4xx terminal). @@ -493,6 +512,7 @@ func TestExecute_RecordsHealth_SuccessAndFailure(t *testing.T) { reg := newRegistry(t, adapter) res := okResolver() health := store.NewHealthTracker() + t.Cleanup(health.Stop) exec := New(reg, res, health, nil) tgt := target(providerSlug) @@ -504,7 +524,7 @@ func TestExecute_RecordsHealth_SuccessAndFailure(t *testing.T) { if result.Error != nil { t.Fatalf("unexpected error: %v", result.Error) } - hs := health.GetHealth(tgt.ProviderID) + hs := awaitHealth(t, health, tgt.ProviderID) // SampleCount > 0 proves we reached RecordSuccess. The // HealthTracker contract guarantees a successful call appends // to the sample window. @@ -523,6 +543,7 @@ func TestExecute_RecordsHealth_SuccessAndFailure(t *testing.T) { reg := newRegistry(t, adapter) res := okResolver() health := store.NewHealthTracker() + t.Cleanup(health.Stop) exec := New(reg, res, health, nil) tgt := target(providerSlug) @@ -538,7 +559,7 @@ func TestExecute_RecordsHealth_SuccessAndFailure(t *testing.T) { if !errors.Is(result.Error, ErrAllTargetsExhausted) { t.Fatalf("expected ErrAllTargetsExhausted, got %v", result.Error) } - hs := health.GetHealth(tgt.ProviderID) + hs := awaitHealth(t, health, tgt.ProviderID) if hs.SampleCount == 0 { t.Fatalf("RecordFailure was not called; GetHealth=%+v", hs) } @@ -557,6 +578,7 @@ func TestExecute_RecordsHealth_SuccessAndFailure(t *testing.T) { reg := newRegistry(t, adapter) res := okResolver() health := store.NewHealthTracker() + t.Cleanup(health.Stop) exec := New(reg, res, health, nil) tgt := target(providerSlug) @@ -571,7 +593,7 @@ func TestExecute_RecordsHealth_SuccessAndFailure(t *testing.T) { if result.StatusCode != http.StatusUnauthorized { t.Fatalf("expected 401, got %d", result.StatusCode) } - hs := health.GetHealth(tgt.ProviderID) + hs := awaitHealth(t, health, tgt.ProviderID) if hs.SampleCount == 0 { t.Fatalf("RecordFailure was not called on 4xx terminal; GetHealth=%+v", hs) } diff --git a/packages/ai-gateway/internal/ingress/proxy/admission_and_alias_review_test.go b/packages/ai-gateway/internal/ingress/proxy/admission_and_alias_review_test.go new file mode 100644 index 00000000..88d87e24 --- /dev/null +++ b/packages/ai-gateway/internal/ingress/proxy/admission_and_alias_review_test.go @@ -0,0 +1,178 @@ +// admission_and_alias_review_test.go — regression tests demanded by the +// adversarial review of the in-flight admission gate and the pooled +// request-body release: env resolution semantics, gate wiring through +// ServeProxy (per-ingress 429 shape + shed counter), and the zero-write +// rewrite alias pin (RequestBodyRedacted must never share backing with the +// pooled request buffer). +package proxy + +import ( + "context" + "log/slog" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/prometheus/client_golang/prometheus/testutil" + + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/platform/audit" + compliance "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/pipeline" + + provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" + routingcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/routing/core" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/builtins" + goHooks "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/traffic/adapters/api/openai" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/typology" +) + +// TestResolveAdmissionMax pins the env contract: auto default, explicit +// value, explicit disable, and — the review finding — a typo must fall back +// to the protective default rather than silently disabling the gate. +func TestResolveAdmissionMax(t *testing.T) { + auto, fb := resolveAdmissionMax("") + if auto <= 0 || fb { + t.Fatalf("unset must yield the positive auto default without fallback, got %d %v", auto, fb) + } + if v, fb := resolveAdmissionMax("auto"); v != auto || fb { + t.Fatalf(`"auto" must equal the unset default, got %d %v`, v, fb) + } + if v, fb := resolveAdmissionMax("123"); v != 123 || fb { + t.Fatalf("explicit value must be honored, got %d %v", v, fb) + } + if v, fb := resolveAdmissionMax("0"); v != 0 || fb { + t.Fatalf("0 must disable without fallback, got %d %v", v, fb) + } + if v, fb := resolveAdmissionMax("-5"); v != 0 || fb { + t.Fatalf("negative must disable, got %d %v", v, fb) + } + if v, fb := resolveAdmissionMax("4O96"); v != auto || !fb { + t.Fatalf("garbage must FALL BACK to the auto default (never silently disable), got %d fellBack=%v", v, fb) + } +} + +// TestServeProxy_GateSheds_PerIngressShape verifies the gate is wired into +// ServeProxy and the reject honors the cross-ingress error contract: an +// anthropic /v1/messages caller gets the anthropic error envelope, an +// OpenAI caller the OpenAI shape — both with Retry-After and a shed-counter +// increment. A max=0 gate sheds immediately, so no other deps are touched. +func TestServeProxy_GateSheds_PerIngressShape(t *testing.T) { + h := &Handler{gate: &admissionGate{max: 0}} + + cases := []struct { + name string + format provcore.Format + marker string + absent string + }{ + // Both shapes carry rate_limit semantics; they differ in envelope: + // OpenAI = {"error":{...,"code":"gateway_overloaded"}}, Anthropic = + // {"type":"error","error":{"type":"rate_limit_error",...}}. + {"openai-shape", provcore.FormatOpenAI, `"code":"gateway_overloaded"`, `"type":"error"`}, + {"anthropic-shape", provcore.FormatAnthropic, `"type":"error"`, `"gateway_overloaded"`}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + before := testutil.ToFloat64(admissionShedTotal) + w := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodPost, "/v1/x", strings.NewReader(`{}`)) + h.ServeProxy(Ingress{WireShape: typology.WireShapeOpenAIChat, BodyFormat: tc.format})(w, req) + + if w.Code != http.StatusTooManyRequests { + t.Fatalf("status=%d want 429", w.Code) + } + if w.Header().Get("Retry-After") != "1" { + t.Fatal("Retry-After missing on shed response") + } + if body := w.Body.String(); !strings.Contains(body, tc.marker) || strings.Contains(body, tc.absent) { + t.Fatalf("shed body not in the caller's ingress shape: %s", body) + } + if after := testutil.ToFloat64(admissionShedTotal); after != before+1 { + t.Fatalf("shed counter did not increment: %v -> %v", before, after) + } + }) + } +} + +// zeroWriteModifyHook mimics a webhook-forward style hook: Decision=Modify +// with NO ModifiedContent and NO TransformSpans — the aggregate carries a +// redaction demand but the adapter rewrite performs zero writes and returns +// the input slice unchanged. +type zeroWriteModifyHook struct { + goHooks.AnyEndpointAnyModality +} + +func (zeroWriteModifyHook) Execute(_ context.Context, _ *goHooks.HookInput) (*goHooks.HookResult, error) { + return &goHooks.HookResult{Decision: goHooks.Modify}, nil +} + +// newZeroWriteModifyRequestHookCache serves one request-stage hook whose +// Modify carries no content — driving the zero-write rewrite branch. +func newZeroWriteModifyRequestHookCache(t *testing.T) *compliance.HookConfigCache { + t.Helper() + reg := builtins.Registry.Clone() + reg.Register("zero-write-modify", func(_ *goHooks.HookConfig) (goHooks.Hook, error) { + return zeroWriteModifyHook{}, nil + }) + reg.Freeze() + loader := func(_ context.Context) ([]goHooks.HookConfig, error) { + return []goHooks.HookConfig{{ + ID: "zw-1", ImplementationID: "zero-write-modify", Name: "zw", + Priority: 1, Enabled: true, Stage: "request", FailBehavior: "fail-closed", + TimeoutMs: 1000, ApplicableIngress: []string{"ALL"}, + Config: map[string]any{}, + }}, nil + } + hc := compliance.NewHookConfigCache(loader, reg, 0, slog.Default()) + if err := hc.Start(context.Background()); err != nil { + t.Fatalf("hookCache.Start: %v", err) + } + time.Sleep(50 * time.Millisecond) + return hc +} + +// TestRunRequestHooks_ZeroWriteModify_RedactedCopyNotAliased pins the +// review's MAJOR: when a Modify decision produces ZERO rewrites, the +// adapter returns the input slice — which is the pooled request buffer that +// finalizeAudit releases for reuse under bodies-off. The stamped +// rec.RequestBodyRedacted must NOT share that backing array: after the +// original buffer is overwritten (simulating the next request's +// AcquireRequestBody copy), the record's redacted copy must be unchanged. +func TestRunRequestHooks_ZeroWriteModify_RedactedCopyNotAliased(t *testing.T) { + h := &Handler{deps: &Deps{ + HookConfigCache: newZeroWriteModifyRequestHookCache(t), + TrafficAdapter: &openai.Adapter{}, + Logger: slog.Default(), + }} + + src := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"tenant-A-secret"}]}`) + body, handle := audit.AcquireRequestBody(src) + if handle == nil { + t.Fatal("expected pooled handle") + } + + req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(string(src))) + w := httptest.NewRecorder() + rec := &audit.Record{RequestID: "req-alias"} + + rewritten, _, rejected := h.runRequestHooks(req, w, rec, "req-alias", body, routingcore.RoutingTarget{}, openAIIngress, nil, slog.Default()) + if rejected { + t.Fatalf("zero-write modify must not reject; response=%s", w.Body.String()) + } + if len(rec.RequestBodyRedacted) == 0 { + t.Fatal("precondition: zero-write modify must stamp RequestBodyRedacted") + } + want := string(rec.RequestBodyRedacted) + + // Simulate the released buffer being recycled by the NEXT request. + for i := range body { + body[i] = 'X' + } + if got := string(rec.RequestBodyRedacted); got != want { + t.Fatalf("RequestBodyRedacted aliases the pooled buffer — next request's bytes bled into this record's redacted copy:\n got=%q\nwant=%q", got, want) + } + _ = rewritten + audit.ReleaseRequestBuffer(handle) +} diff --git a/packages/ai-gateway/internal/ingress/proxy/admission_gate.go b/packages/ai-gateway/internal/ingress/proxy/admission_gate.go new file mode 100644 index 00000000..6a4eaa94 --- /dev/null +++ b/packages/ai-gateway/internal/ingress/proxy/admission_gate.go @@ -0,0 +1,138 @@ +package proxy + +import ( + "log/slog" + "net/http" + "os" + "runtime" + "strconv" + "sync" + "sync/atomic" + + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" + + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/ingress/envelope" + provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" +) + +// admissionGate bounds the number of in-flight proxy requests so overload +// degrades into fast, retryable 429s instead of unbounded in-heap queueing. +// Without it, an arrival rate above the box's service rate grows the +// in-flight set (goroutine stacks + request bodies + response buffers) +// without limit until the heap reaches GOMEMLIMIT and the collector's +// assist tax collapses throughput — or the kernel OOM-kills the process. +// A bounded in-flight count keeps that memory footprint bounded. +// +// The admitted path costs two atomic ops (one add on entry, one on exit). +// A shed request still traverses the outer middleware (RequestID, OTel +// span, Logger — which logs 4xx at Warn) but skips all per-request proxy +// work and gets no audit record: the shed path's observability is the +// admission_shed_total counter. Health, metrics, and admin endpoints are +// outside ServeProxy and therefore never gated — a saturated data plane +// must still answer its liveness probes. +type admissionGate struct { + max int64 + inflight atomic.Int64 +} + +// resolveAdmissionMax parses an AI_GATEWAY_MAX_INFLIGHT value: +// +// unset / "auto" → 1024 × GOMAXPROCS. Scaled by core count so the cap +// tracks the box; sized generously because an SSE stream +// holds its slot for the stream's full duration, so a +// tight cap would false-shed a healthy streaming-heavy +// deployment. The bound exists to cap heap growth under +// overload, not to throttle normal traffic. +// 0 / negative → gate disabled (explicit operator choice) +// garbage → the auto default (a typo must never silently disable +// overload protection); the caller logs the fallback +func resolveAdmissionMax(v string) (max int64, fellBack bool) { + auto := int64(1024 * runtime.GOMAXPROCS(0)) + if v == "" || v == "auto" { + return auto, false + } + n, err := strconv.ParseInt(v, 10, 64) + if err != nil { + return auto, true + } + if n <= 0 { + return 0, false // explicit disable + } + return n, false +} + +var ( + admissionOnce sync.Once + admissionMax int64 +) + +// newAdmissionGate builds the gate from AI_GATEWAY_MAX_INFLIGHT (resolved +// once per process); nil when disabled. +func newAdmissionGate() *admissionGate { + admissionOnce.Do(func() { + v := os.Getenv("AI_GATEWAY_MAX_INFLIGHT") + m, fellBack := resolveAdmissionMax(v) + if fellBack { + slog.Warn("AI_GATEWAY_MAX_INFLIGHT is not a number; using the auto default", + "value", v, "max", m) + } + admissionMax = m + }) + if admissionMax > 0 { + return &admissionGate{max: admissionMax} + } + return nil +} + +// acquire admits the request unless the in-flight cap is reached. The +// increment-then-check shape keeps the hot path a single atomic add; the +// decrement on rejection cannot strand capacity because every path that +// observes n > max gives its slot back before returning. +func (g *admissionGate) acquire() bool { + if g.inflight.Add(1) > g.max { + g.inflight.Add(-1) + admissionShedTotal.Inc() + return false + } + return true +} + +// release returns the request's slot. Deferred by the ServeProxy wrapper so +// it covers every exit path, including streaming completions and panics +// unwound through the handler's defers before Recovery catches them. +func (g *admissionGate) release() { g.inflight.Add(-1) } + +// admissionShedTotal counts requests rejected by the in-flight gate. The +// counter is the shed path's observability — rejected requests get no audit +// record on purpose (auditing a shed storm would itself be load). +var admissionShedTotal = promauto.NewCounter(prometheus.CounterOpts{ + Name: "nexus_ai_gateway_admission_shed_total", + Help: "Requests rejected with 429 by the in-flight admission gate.", +}) + +// writeOverloaded writes the fast-reject response in the CALLER's ingress +// wire shape (same cross-ingress error contract as writeIngressError — +// anthropic /v1/messages gets {"type":"error",...}, gemini gets its +// envelope; OpenAI-family and unknown get the OpenAI error shape), plus +// Retry-After so generic clients back off. The shape is derived from the +// route's static BodyFormat: the gate runs before per-request state, so the +// x-nexus-aigw-body-format override (OpenAI-family niche) is not consulted. +func writeOverloaded(w http.ResponseWriter, ingress provcore.Format) { + const msg = "gateway is at capacity, retry shortly" + var body []byte + if ingress != "" && !ingress.IsOpenAIFamily() { + // CodeRateLimited maps to each envelope's semantically-correct type + // (anthropic "rate_limit_error", gemini RESOURCE_EXHAUSTED) so SDK + // classification matches the OpenAI-shape fallback below. + body = envelope.EncodeErrorEnvelopeForIngress(ingress, ingress, &provcore.ProviderError{ + Status: http.StatusTooManyRequests, Code: provcore.CodeRateLimited, Message: msg, + }) + } else { + body = []byte(`{"error":{"message":"` + msg + `","type":"rate_limit_error","code":"gateway_overloaded"}}`) + } + w.Header().Set("Content-Type", "application/json") + w.Header().Set("Retry-After", "1") + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write(body) +} diff --git a/packages/ai-gateway/internal/ingress/proxy/admission_gate_test.go b/packages/ai-gateway/internal/ingress/proxy/admission_gate_test.go new file mode 100644 index 00000000..0d971360 --- /dev/null +++ b/packages/ai-gateway/internal/ingress/proxy/admission_gate_test.go @@ -0,0 +1,90 @@ +package proxy + +import ( + "net/http/httptest" + "strings" + "sync" + "sync/atomic" + "testing" + + provcore "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/providers/core" +) + +// TestAdmissionGate_AdmitsUpToCapAndSheds asserts the business contract: +// exactly max requests may be in flight; the (max+1)th is shed; a released +// slot re-admits. +func TestAdmissionGate_AdmitsUpToCapAndSheds(t *testing.T) { + g := &admissionGate{max: 3} + for i := range 3 { + if !g.acquire() { + t.Fatalf("request %d within cap must be admitted", i+1) + } + } + if g.acquire() { + t.Fatal("request beyond cap must be shed") + } + g.release() + if !g.acquire() { + t.Fatal("a released slot must re-admit") + } +} + +// TestAdmissionGate_ConcurrentNeverExceedsCap hammers the gate from many +// goroutines and asserts the observed concurrent in-flight count never +// exceeds max — the property that bounds heap growth under overload. +func TestAdmissionGate_ConcurrentNeverExceedsCap(t *testing.T) { + const cap64, workers, iters = 8, 64, 200 + g := &admissionGate{max: cap64} + var cur, peak, shed atomic.Int64 + var wg sync.WaitGroup + for range workers { + wg.Add(1) + go func() { + defer wg.Done() + for range iters { + if !g.acquire() { + shed.Add(1) + continue + } + n := cur.Add(1) + for { + p := peak.Load() + if n <= p || peak.CompareAndSwap(p, n) { + break + } + } + cur.Add(-1) + g.release() + } + }() + } + wg.Wait() + if p := peak.Load(); p > cap64 { + t.Fatalf("in-flight peaked at %d, cap is %d — gate leaked capacity", p, cap64) + } + if g.inflight.Load() != 0 { + t.Fatalf("in-flight must drain to zero, got %d", g.inflight.Load()) + } + if shed.Load() == 0 { + t.Log("note: no sheds observed (timing-dependent); cap property still verified") + } +} + +// TestWriteOverloaded_ResponseShape asserts the reject is a retryable, +// SDK-parseable 429: Retry-After present, OpenAI-shaped error envelope. +func TestWriteOverloaded_ResponseShape(t *testing.T) { + rr := httptest.NewRecorder() + writeOverloaded(rr, provcore.FormatOpenAI) + if rr.Code != 429 { + t.Fatalf("want 429, got %d", rr.Code) + } + if rr.Header().Get("Retry-After") != "1" { + t.Fatalf("Retry-After header missing") + } + body := rr.Body.String() + for _, want := range []string{`"rate_limit_error"`, `"gateway_overloaded"`} { + if !strings.Contains(body, want) { + t.Fatalf("429 body missing %s: %s", want, body) + } + } +} diff --git a/packages/ai-gateway/internal/ingress/proxy/chunk_stream_reader_test.go b/packages/ai-gateway/internal/ingress/proxy/chunk_stream_reader_test.go index d331e308..b341fbd7 100644 --- a/packages/ai-gateway/internal/ingress/proxy/chunk_stream_reader_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/chunk_stream_reader_test.go @@ -636,6 +636,12 @@ func (f *fakeModelsAndPricing) GetModelByCode(_ context.Context, code string) (* } return nil, errors.New("not found") } +func (f *fakeModelsAndPricing) GetModelByCodeOrAlias(_ context.Context, key string) (*store.Model, error) { + if m, ok := f.byCode[key]; ok { + return m, nil + } + return nil, errors.New("not found") +} func (f *fakeModelsAndPricing) ListEnabledModels(_ context.Context) ([]store.Model, error) { return nil, nil } diff --git a/packages/ai-gateway/internal/ingress/proxy/embeddings_crossformat_test.go b/packages/ai-gateway/internal/ingress/proxy/embeddings_crossformat_test.go index 3ce4e1e8..3469230b 100644 --- a/packages/ai-gateway/internal/ingress/proxy/embeddings_crossformat_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/embeddings_crossformat_test.go @@ -88,7 +88,11 @@ func TestProxy_Embeddings_OpenAIIngress_GeminiOnlyTarget_NoCompatibleProvider(t provReg := provcore.NewRegistry() provbuiltins.Register(provReg, nil, logger) bridge := canonicalbridge.New(provbuiltins.SchemaCodecs(logger)) - exec := executor.New(provReg, embeddingsStubProvResolve{}, store.NewHealthTracker(), bridge) + execHT := store.NewHealthTracker() + t.Cleanup(execHT.Stop) + exec := executor.New(provReg, embeddingsStubProvResolve{}, execHT, bridge) + depsHT := store.NewHealthTracker() + t.Cleanup(depsHT.Stop) deps := &Deps{ Models: nil, @@ -108,7 +112,7 @@ func TestProxy_Embeddings_OpenAIIngress_GeminiOnlyTarget_NoCompatibleProvider(t Executor: exec, HookConfigCache: hookCache, ProviderReg: provReg, - HealthTracker: store.NewHealthTracker(), + HealthTracker: depsHT, AuditWriter: audit.NewWriter(nil, "nexus.event.ai-traffic", nil, logger), CanonicalBridge: bridge, Logger: logger, diff --git a/packages/ai-gateway/internal/ingress/proxy/estimate_test.go b/packages/ai-gateway/internal/ingress/proxy/estimate_test.go index 771d99f7..a71c4a39 100644 --- a/packages/ai-gateway/internal/ingress/proxy/estimate_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/estimate_test.go @@ -53,6 +53,12 @@ func (s *stubModels) GetModelByCode(_ context.Context, code string) (*store.Mode } return nil, errors.New("not found") } +func (s *stubModels) GetModelByCodeOrAlias(_ context.Context, key string) (*store.Model, error) { + if m, ok := s.byCode[key]; ok { + return m, nil + } + return nil, errors.New("not found") +} func (s *stubModels) ListEnabledModels(_ context.Context) ([]store.Model, error) { return nil, nil } diff --git a/packages/ai-gateway/internal/ingress/proxy/integration_test.go b/packages/ai-gateway/internal/ingress/proxy/integration_test.go index 8b1bf98f..af790795 100644 --- a/packages/ai-gateway/internal/ingress/proxy/integration_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/integration_test.go @@ -42,7 +42,8 @@ func (stubResolver) Resolve(_ context.Context, providerID, modelID string, _ pro // testDeps creates handler.Deps with no DB (stubs VK auth and routing). // The upstream URL is injected via the provider registry mock. -func testDeps(upstreamURL string) *proxy.Deps { +func testDeps(t *testing.T, upstreamURL string) *proxy.Deps { + t.Helper() logger := slog.Default() // Provider adapter registry populated with the nine built-in stub @@ -63,6 +64,7 @@ func testDeps(upstreamURL string) *proxy.Deps { credMgr := credmanager.NewManager(nil, nil) healthTracker := store.NewHealthTracker() + t.Cleanup(healthTracker.Stop) return &proxy.Deps{ Models: nil, // no DB in test @@ -92,7 +94,7 @@ func (okVKAuth) Authenticate(_ context.Context, _ *http.Request) (*vkauth.VKMeta func TestProxyHandler_MissingModel(t *testing.T) { // auth runs BEFORE the body read, so authenticate with a // passing VKAuth stub to reach the body-level model validation. - deps := testDeps("") + deps := testDeps(t, "") deps.VKAuth = okVKAuth{} h := proxy.NewHandler(deps).ServeProxy(proxy.Ingress{ WireShape: typology.WireShapeOpenAIChat, @@ -121,7 +123,7 @@ func TestProxyHandler_MissingModel(t *testing.T) { // payload capture pre-auth. A missing-model body that would 400 once // authenticated must instead surface the 401. func TestProxyHandler_AuthBeforeBody(t *testing.T) { - deps := testDeps("") + deps := testDeps(t, "") deps.VKAuth = stubAuthErr{err: vkauth.ErrMissing} h := proxy.NewHandler(deps).ServeProxy(proxy.Ingress{ WireShape: typology.WireShapeOpenAIChat, diff --git a/packages/ai-gateway/internal/ingress/proxy/interfaces.go b/packages/ai-gateway/internal/ingress/proxy/interfaces.go index 6f27bc90..c579795a 100644 --- a/packages/ai-gateway/internal/ingress/proxy/interfaces.go +++ b/packages/ai-gateway/internal/ingress/proxy/interfaces.go @@ -37,6 +37,10 @@ type CredentialLookup interface { type ModelLookup interface { GetModel(ctx context.Context, id string) (*store.Model, error) GetModelByCode(ctx context.Context, idOrName string) (*store.Model, error) + // GetModelByCodeOrAlias resolves a code OR an alias (code wins on + // collision). Used by the routing passthrough fallback so an aliased + // model routes without an explicit routing rule. + GetModelByCodeOrAlias(ctx context.Context, key string) (*store.Model, error) ListEnabledModels(ctx context.Context) ([]store.Model, error) // FetchModelPricing returns pricing rows for the requested IDs. // Production wires *cachelayer.Layer which reads from the in-memory diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy.go b/packages/ai-gateway/internal/ingress/proxy/proxy.go index c907c1be..1af5cf7b 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy.go @@ -209,6 +209,10 @@ type Handler struct { // skips the eager request-body Normalize entirely (~29% of request-path CPU // on a 50 KB body). Set NEXUS_LAZY_CANONICAL=0 to force always-compute. lazyCanonical bool + // gate bounds in-flight proxy requests (nil = disabled). See admissionGate: + // overload degrades into fast retryable 429s instead of unbounded in-heap + // queueing toward the GOMEMLIMIT collapse. + gate *admissionGate } // NewHandler creates a Handler with the given dependencies. @@ -216,6 +220,7 @@ func NewHandler(deps *Deps) *Handler { return &Handler{ deps: deps, lazyCanonical: os.Getenv("NEXUS_LAZY_CANONICAL") != "0", + gate: newAdmissionGate(), } } @@ -262,6 +267,17 @@ func (h *Handler) streamCaptureHardCap() int64 { // record on every exit path. func (h *Handler) ServeProxy(in Ingress) http.HandlerFunc { return func(w http.ResponseWriter, r *http.Request) { + // In-flight admission gate: reject-fast BEFORE any per-request setup + // (no state, no audit record — the shed path's observability is the + // admission_shed_total counter). Health/metrics/admin routes live + // outside ServeProxy and are never gated. + if h.gate != nil { + if !h.gate.acquire() { + writeOverloaded(w, in.BodyFormat) + return + } + defer h.gate.release() + } s, ok := h.newProxyState(in, w, r) if !ok { return // invalid body-format override; 400 already written diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_live.go b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_live.go index 215f36eb..0113ae09 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_live.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_live.go @@ -24,10 +24,21 @@ func runLiveStream(ctx context.Context, d runStreamDeps) { if d.SSEReader == nil || d.Tee == nil { return } + // When no response-stage rule is bound (HasResponseHooks=false — the default + // deployment shape), pass a nil runner so Process skips both the per-checkpoint + // scan (which would resolve to nil and stamp a synthetic Approve) AND the + // per-chunk ExtractDeltaText accumulation that only feeds it. Mirrors the + // shared LivePipeline's `l.pipeline == nil` guard; the captured body still + // persists via rec's default ActionApprove, matching the non-stream rule-free + // path (both leave response_hook_decision empty). + hookRun := d.HookRunner + if !d.HasResponseHooks { + hookRun = nil + } lp := streaming.NewLivePipeline(streaming.LiveConfig{ EmitOpenAIDone: d.EmitDone, MaxBufferSize: d.MaxBufferBytes, - }, d.HookRunner, nil, d.Logger) + }, hookRun, nil, d.Logger) // The PreHook re-normalizes the cumulative raw SSE through the Registry at // every checkpoint so a response hook sees the same structured payload the diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_live_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_live_test.go index a16e6c3c..d916af88 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_cache_live_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_cache_live_test.go @@ -45,21 +45,87 @@ func TestRunLiveStream_HappyPath_FlowsThroughLivePipeline(t *testing.T) { deps := &Deps{NormalizeRegistry: reg} runLiveStream(context.Background(), runStreamDeps{ - Deps: deps, - AdapterType: "openai", - Path: "/v1/chat/completions", - AcceptHeader: "text/event-stream", - HookRunner: runner, - HookCtx: hookCtx, - SSEReader: body, - Tee: tee, - Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), - EmitDone: true, + Deps: deps, + AdapterType: "openai", + Path: "/v1/chat/completions", + AcceptHeader: "text/event-stream", + HookRunner: runner, + HookCtx: hookCtx, + HasResponseHooks: true, // response rules bound → PreHook installed + checkpoints run + SSEReader: body, + Tee: tee, + Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + EmitDone: true, }) if teeBuf.Len() == 0 { t.Errorf("expected LivePipeline to forward bytes to tee, got 0") } + // With response hooks bound the checkpoint MUST fire (the runner is kept, + // not nil'd) — without this assertion the test went blind to the + // HasResponseHooks gate. + if hookCalls.Load() == 0 { + t.Error("expected OnCheckpoint to fire with HasResponseHooks=true") + } +} + +// TestRunLiveStream_ResponseHooksGate pins the HasResponseHooks gate: with a +// response rule bound the runner is kept and checkpoints fire; with none bound +// the runner is nil'd and checkpoints are skipped — but in BOTH cases the +// persisted (tee) body carries the full transcript, since skipping the audit +// scan must never drop delivered/persisted bytes. +func TestRunLiveStream_ResponseHooksGate(t *testing.T) { + for _, tc := range []struct { + name string + hasResponseHooks bool + wantCheckpoints bool + }{ + {"hooks_bound_runs_checkpoints", true, true}, + {"no_hooks_skips_checkpoints", false, false}, + } { + t.Run(tc.name, func(t *testing.T) { + body := strings.NewReader("data: {\"choices\":[{\"delta\":{\"content\":\"hello world\"}}]}\n\ndata: {\"choices\":[{\"delta\":{\"content\":\" more\"}}]}\n\ndata: [DONE]\n\n") + var teeBuf bytes.Buffer + tee := &testWriter{Buffer: &teeBuf, header: http.Header{}} + + var hookCalls atomic.Int32 + hookCtx := &streaming.StreamHookContext{ + RequestID: "gate-" + tc.name, + IngressType: "AI_GATEWAY", + OnCheckpoint: func(*hookcore.CompliancePipelineResult) { hookCalls.Add(1) }, + } + runner := func(_ context.Context, _ *hookcore.HookInput) *hookcore.CompliancePipelineResult { + return &hookcore.CompliancePipelineResult{Decision: hookcore.Approve} + } + reg := normcore.NewRegistry() + codecs.RegisterDefaultAIBuiltins(reg) + + runLiveStream(context.Background(), runStreamDeps{ + Deps: &Deps{NormalizeRegistry: reg}, + AdapterType: "openai", + Path: "/v1/chat/completions", + HookRunner: runner, + HookCtx: hookCtx, + HasResponseHooks: tc.hasResponseHooks, + SSEReader: body, + Tee: tee, + Logger: slog.New(slog.NewTextHandler(io.Discard, nil)), + EmitDone: true, + }) + + // (a) OnCheckpoint fires ONLY when response hooks are bound. + if got := hookCalls.Load() > 0; got != tc.wantCheckpoints { + t.Errorf("checkpoints fired=%v, want %v (HasResponseHooks=%v)", got, tc.wantCheckpoints, tc.hasResponseHooks) + } + // (b) The persisted (tee) body carries the FULL transcript in BOTH cases. + persisted := teeBuf.String() + for _, want := range []string{"hello world", " more", "[DONE]"} { + if !strings.Contains(persisted, want) { + t.Errorf("persisted tee body missing %q (HasResponseHooks=%v); got %q", want, tc.hasResponseHooks, persisted) + } + } + }) + } } // TestRunLiveStream_NilDeps_NoPreHook_StillRuns — when Deps is nil diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_e2e_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_e2e_test.go index 1e9281d8..619388d0 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_e2e_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_e2e_test.go @@ -86,7 +86,9 @@ func TestServeProxy_NonStreamHappyPath_DriversHandleNonStream(t *testing.T) { provReg.Freeze() bridge := canonicalbridge.New(provbuiltins.SchemaCodecs(logger)) - exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, store.NewHealthTracker(), bridge) + execHT := store.NewHealthTracker() + t.Cleanup(execHT.Stop) + exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, execHT, bridge) hookCache := compliance.NewHookConfigCache( func(_ context.Context) ([]goHooks.HookConfig, error) { return nil, nil }, @@ -99,6 +101,8 @@ func TestServeProxy_NonStreamHappyPath_DriversHandleNonStream(t *testing.T) { prod := &captureProducer{} auditWriter := audit.NewWriter(prod, "nexus.event.ai-traffic", nil, logger) + depsHT := store.NewHealthTracker() + t.Cleanup(depsHT.Stop) deps := &Deps{ VKAuth: &stubVKAuthCacheTest{meta: &vkauth.VKMeta{ @@ -120,7 +124,7 @@ func TestServeProxy_NonStreamHappyPath_DriversHandleNonStream(t *testing.T) { Executor: exec, HookConfigCache: hookCache, ProviderReg: provReg, - HealthTracker: store.NewHealthTracker(), + HealthTracker: depsHT, AuditWriter: auditWriter, CanonicalBridge: bridge, PayloadCapture: payloadcapture.NewStore(payloadcapture.Config{ @@ -172,7 +176,9 @@ func TestServeProxy_NonStream_UpstreamErrorPath(t *testing.T) { provReg.Freeze() bridge := canonicalbridge.New(provbuiltins.SchemaCodecs(logger)) - exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, store.NewHealthTracker(), bridge) + execHT := store.NewHealthTracker() + t.Cleanup(execHT.Stop) + exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, execHT, bridge) hookCache := compliance.NewHookConfigCache( func(_ context.Context) ([]goHooks.HookConfig, error) { return nil, nil }, @@ -185,6 +191,8 @@ func TestServeProxy_NonStream_UpstreamErrorPath(t *testing.T) { prod := &captureProducer{} auditWriter := audit.NewWriter(prod, "nexus.event.ai-traffic", nil, logger) + depsHT := store.NewHealthTracker() + t.Cleanup(depsHT.Stop) deps := &Deps{ VKAuth: &stubVKAuthCacheTest{meta: &vkauth.VKMeta{ @@ -202,7 +210,7 @@ func TestServeProxy_NonStream_UpstreamErrorPath(t *testing.T) { Executor: exec, HookConfigCache: hookCache, ProviderReg: provReg, - HealthTracker: store.NewHealthTracker(), + HealthTracker: depsHT, AuditWriter: auditWriter, CanonicalBridge: bridge, Logger: logger, @@ -261,7 +269,9 @@ func TestServeProxy_NonStream_CacheMISS_DirectPath(t *testing.T) { provReg.Freeze() bridge := canonicalbridge.New(provbuiltins.SchemaCodecs(logger)) - exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, store.NewHealthTracker(), bridge) + execHT := store.NewHealthTracker() + t.Cleanup(execHT.Stop) + exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, execHT, bridge) rcache := cache.New(rdb, cache.Config{Enabled: true}, logger) if rcache == nil { @@ -279,6 +289,8 @@ func TestServeProxy_NonStream_CacheMISS_DirectPath(t *testing.T) { prod := &captureProducer{} auditWriter := audit.NewWriter(prod, "nexus.event.ai-traffic", nil, logger) + depsHT := store.NewHealthTracker() + t.Cleanup(depsHT.Stop) deps := &Deps{ VKAuth: &stubVKAuthCacheTest{meta: &vkauth.VKMeta{ @@ -296,7 +308,7 @@ func TestServeProxy_NonStream_CacheMISS_DirectPath(t *testing.T) { Executor: exec, HookConfigCache: hookCache, ProviderReg: provReg, - HealthTracker: store.NewHealthTracker(), + HealthTracker: depsHT, AuditWriter: auditWriter, CanonicalBridge: bridge, Cache: rcache, @@ -361,7 +373,9 @@ func TestServeProxy_NonStream_BrokerMISS_LeaderWritesCache(t *testing.T) { provReg.Freeze() bridge := canonicalbridge.New(provbuiltins.SchemaCodecs(logger)) - exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, store.NewHealthTracker(), bridge) + execHT := store.NewHealthTracker() + t.Cleanup(execHT.Stop) + exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, execHT, bridge) rcache := cache.New(rdb, cache.Config{Enabled: true}, logger) if rcache == nil { @@ -381,6 +395,8 @@ func TestServeProxy_NonStream_BrokerMISS_LeaderWritesCache(t *testing.T) { prod := &captureProducer{} auditWriter := audit.NewWriter(prod, "nexus.event.ai-traffic", nil, logger) + depsHT := store.NewHealthTracker() + t.Cleanup(depsHT.Stop) deps := &Deps{ VKAuth: &stubVKAuthCacheTest{meta: &vkauth.VKMeta{ @@ -398,7 +414,7 @@ func TestServeProxy_NonStream_BrokerMISS_LeaderWritesCache(t *testing.T) { Executor: exec, HookConfigCache: hookCache, ProviderReg: provReg, - HealthTracker: store.NewHealthTracker(), + HealthTracker: depsHT, AuditWriter: auditWriter, CanonicalBridge: bridge, Cache: rcache, @@ -476,7 +492,9 @@ func TestServeProxy_Stream_DirectPath(t *testing.T) { provReg.Freeze() bridge := canonicalbridge.New(provbuiltins.SchemaCodecs(logger)) - exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, store.NewHealthTracker(), bridge) + execHT := store.NewHealthTracker() + t.Cleanup(execHT.Stop) + exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, execHT, bridge) hookCache := compliance.NewHookConfigCache( func(_ context.Context) ([]goHooks.HookConfig, error) { return nil, nil }, @@ -489,6 +507,8 @@ func TestServeProxy_Stream_DirectPath(t *testing.T) { prod := &captureProducer{} auditWriter := audit.NewWriter(prod, "nexus.event.ai-traffic", nil, logger) + depsHT := store.NewHealthTracker() + t.Cleanup(depsHT.Stop) deps := &Deps{ VKAuth: &stubVKAuthCacheTest{meta: &vkauth.VKMeta{ @@ -506,7 +526,7 @@ func TestServeProxy_Stream_DirectPath(t *testing.T) { Executor: exec, HookConfigCache: hookCache, ProviderReg: provReg, - HealthTracker: store.NewHealthTracker(), + HealthTracker: depsHT, AuditWriter: auditWriter, CanonicalBridge: bridge, Logger: logger, @@ -576,7 +596,9 @@ func TestServeProxy_Stream_BrokerMISS_LeaderPath(t *testing.T) { provReg.Freeze() bridge := canonicalbridge.New(provbuiltins.SchemaCodecs(logger)) - exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, store.NewHealthTracker(), bridge) + execHT := store.NewHealthTracker() + t.Cleanup(execHT.Stop) + exec := executor.New(provReg, &e2eUpstreamResolver{baseURL: upstream.URL}, execHT, bridge) rcache := cache.New(rdb, cache.Config{Enabled: true}, logger) if rcache == nil { @@ -596,6 +618,8 @@ func TestServeProxy_Stream_BrokerMISS_LeaderPath(t *testing.T) { prod := &captureProducer{} auditWriter := audit.NewWriter(prod, "nexus.event.ai-traffic", nil, logger) + depsHT := store.NewHealthTracker() + t.Cleanup(depsHT.Stop) deps := &Deps{ VKAuth: &stubVKAuthCacheTest{meta: &vkauth.VKMeta{ @@ -613,7 +637,7 @@ func TestServeProxy_Stream_BrokerMISS_LeaderPath(t *testing.T) { Executor: exec, HookConfigCache: hookCache, ProviderReg: provReg, - HealthTracker: store.NewHealthTracker(), + HealthTracker: depsHT, AuditWriter: auditWriter, CanonicalBridge: bridge, Cache: rcache, diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_fakeexec_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_fakeexec_test.go index 816c694d..a5e30c74 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_fakeexec_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_fakeexec_test.go @@ -235,6 +235,8 @@ func makeFakeDeps(t *testing.T, fexec *fakeExecutor, fbridge *fakeBridge) *Deps prod := &captureProducer{} auditWriter := audit.NewWriter(prod, "nexus.event.ai-traffic", nil, logger) + ht := store.NewHealthTracker() + t.Cleanup(ht.Stop) deps := &Deps{ VKAuth: &stubVKAuthCacheTest{meta: &vkauth.VKMeta{ @@ -256,7 +258,7 @@ func makeFakeDeps(t *testing.T, fexec *fakeExecutor, fbridge *fakeBridge) *Deps Executor: fexec, HookConfigCache: hookCache, ProviderReg: provReg, - HealthTracker: store.NewHealthTracker(), + HealthTracker: ht, AuditWriter: auditWriter, CanonicalBridge: fbridge, TrafficAdapters: trafficReg, diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_passthrough_fallback_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_passthrough_fallback_test.go index f576626e..6bf9d2c9 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_passthrough_fallback_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_passthrough_fallback_test.go @@ -27,6 +27,13 @@ func (s fallbackModelLookupStub) GetModelByCode(ctx context.Context, idOrName st return s.model, nil } +func (s fallbackModelLookupStub) GetModelByCodeOrAlias(ctx context.Context, key string) (*store.Model, error) { + if s.err != nil { + return nil, s.err + } + return s.model, nil +} + func (s fallbackModelLookupStub) ListEnabledModels(ctx context.Context) ([]store.Model, error) { return nil, errors.New("not used") } diff --git a/packages/ai-gateway/internal/ingress/proxy/proxy_serveproxy_test.go b/packages/ai-gateway/internal/ingress/proxy/proxy_serveproxy_test.go index 4d92b030..5839aa93 100644 --- a/packages/ai-gateway/internal/ingress/proxy/proxy_serveproxy_test.go +++ b/packages/ai-gateway/internal/ingress/proxy/proxy_serveproxy_test.go @@ -213,7 +213,9 @@ func makeOpenAIDeps(t *testing.T, upstreamURL string, hookCache *compliance.Hook provReg.Freeze() bridge := canonicalbridge.New(provbuiltins.SchemaCodecs(logger)) - exec := executor.New(provReg, &coverageUpstreamResolver{baseURL: upstreamURL}, store.NewHealthTracker(), bridge) + execHT := store.NewHealthTracker() + t.Cleanup(execHT.Stop) + exec := executor.New(provReg, &coverageUpstreamResolver{baseURL: upstreamURL}, execHT, bridge) // Wire the traffic adapter registry so response-hook Modify branches // have an extractor to call (without it, RewriteResponseBody panics @@ -225,6 +227,8 @@ func makeOpenAIDeps(t *testing.T, upstreamURL string, hookCache *compliance.Hook prod := &captureProducer{} auditWriter := audit.NewWriter(prod, "nexus.event.ai-traffic", nil, logger) + depsHT := store.NewHealthTracker() + t.Cleanup(depsHT.Stop) deps := &Deps{ VKAuth: &stubVKAuthCacheTest{meta: &vkauth.VKMeta{ @@ -246,7 +250,7 @@ func makeOpenAIDeps(t *testing.T, upstreamURL string, hookCache *compliance.Hook Executor: exec, HookConfigCache: hookCache, ProviderReg: provReg, - HealthTracker: store.NewHealthTracker(), + HealthTracker: depsHT, AuditWriter: auditWriter, CanonicalBridge: bridge, TrafficAdapters: trafficReg, diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_accounting.go b/packages/ai-gateway/internal/ingress/proxy/stage_accounting.go index 9359a9f0..e22b0859 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stage_accounting.go +++ b/packages/ai-gateway/internal/ingress/proxy/stage_accounting.go @@ -6,6 +6,7 @@ package proxy import ( "time" + "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/platform/audit" "github.com/AlphaBitCore/nexus-gateway/packages/shared/traffic" ) @@ -15,6 +16,19 @@ import ( // ServeProxy immediately after the state is built, so it covers every // stage's exit path. func (s *proxyState) finalizeAudit() { + // Return the uncaptured request-body buffer to its pool. Safe here: every + // reader of s.body (upstream forward, hooks, lazy normalize via the + // respond stage) completed before this defer runs, and the record never + // references the buffer — a zero-write hook rewrite clones before stamping + // RequestBodyRedacted for exactly this reason. Nil when the body was + // captured (writer reclaims via the record) or absent. Accepted residual: + // on an errored upstream round-trip the http transport may in principle + // still be draining the body reader for a microsecond-scale window + // (mirrors the captured path's reclaim-at-marshal exposure). + if s.unattachedBodyHandle != nil { + audit.ReleaseRequestBuffer(s.unattachedBodyHandle) + s.unattachedBodyHandle = nil + } // Pure-forward benchmark mode: skip the entire audit tail (no traffic_event). // The record carries pooled request/response body buffers that the async // writer normally returns downstream of Enqueue; since we skip Enqueue, return diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_admission.go b/packages/ai-gateway/internal/ingress/proxy/stage_admission.go index a47e35a2..da358982 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stage_admission.go +++ b/packages/ai-gateway/internal/ingress/proxy/stage_admission.go @@ -155,10 +155,16 @@ func (st admissionStage) run() bool { s.rec.RequestBody = body s.rec.RequestContentType = s.r.Header.Get("Content-Type") // The captured body IS the pooled buffer; hand its handle to the record so - // the audit writer returns it to the pool at terminal resolution. When NOT - // captured (this branch skipped), bodyHandle is simply dropped and the - // buffer is GC'd — still used by s.body for upstream forwarding this request. + // the audit writer returns it to the pool at terminal resolution. s.rec.AttachPooledRequestBody(bodyHandle) + } else { + // Not captured: the record never references the buffer, so no terminal + // reclaim will run. Stash the handle; finalizeAudit returns it to the + // pool after the body's last reader (upstream forward, hooks, lazy + // normalize — all complete before that defer) has finished. Dropping + // it instead would re-allocate the pooled body buffer on every + // request instead of reusing it. + s.unattachedBodyHandle = bodyHandle } // Phase 3.5: Build the canonical request context. One @@ -233,19 +239,24 @@ func (h *Handler) readBody(r *http.Request, in Ingress) (body []byte, bodyHandle // audit writer, which returns the buffer to the pool at the record's terminal // resolution (severing the per-request fresh allocation, the #1 hot-path // allocator). bodyHandle is the pool handle; the caller attaches it to the - // audit Record when the body is captured, else drops it (the buffer GC's). + // audit Record when the body is captured, stashes it for the finalizeAudit + // release when not, and the reject paths below return it directly — a + // malformed-body storm must not bypass the pool and re-allocate per request. body, bodyHandle = audit.AcquireRequestBody(buf.Bytes()) modelID, isStream, err = ExtractIngressModel(in, r, body) if err != nil { + audit.ReleaseRequestBuffer(bodyHandle) return nil, nil, "", false, err } if modelID == "" { + audit.ReleaseRequestBuffer(bodyHandle) return nil, nil, "", false, fmt.Errorf("model is required") } if modelID == "auto" && typology.KindFromWireShape(in.WireShape) == typology.EndpointKindEmbeddings { + audit.ReleaseRequestBuffer(bodyHandle) return nil, nil, "", false, fmt.Errorf("model \"auto\" is not supported for embeddings") } diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_context.go b/packages/ai-gateway/internal/ingress/proxy/stage_context.go index 0632dceb..b9570f14 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stage_context.go +++ b/packages/ai-gateway/internal/ingress/proxy/stage_context.go @@ -68,6 +68,10 @@ type proxyState struct { modelID string isStream bool rctxFull *requestcontext.RequestContext + // unattachedBodyHandle is the pooled request-body handle when payload + // capture did NOT attach it to the record (bodies-off). finalizeAudit + // returns it to the pool at request end; nil when captured or no body. + unattachedBodyHandle *[]byte // Routing outputs. routeResult *routingcore.RouteResult diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_hooks.go b/packages/ai-gateway/internal/ingress/proxy/stage_hooks.go index bb52fa60..be0590f9 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stage_hooks.go +++ b/packages/ai-gateway/internal/ingress/proxy/stage_hooks.go @@ -290,6 +290,17 @@ func (h *Handler) runRequestHooks(r *http.Request, w http.ResponseWriter, rec *a // the pre-hook bytes for normalization and must never reach // raw storage when redaction is demanded — without this stamp // the writer fail-safes the raw copy to NULL). + // + // Zero-write rewrites return the INPUT slice — which is the pooled + // request-body buffer that finalizeAudit releases for reuse when the + // body is not captured. The record outlives that release (the async + // audit writer holds it), so a stamped alias would let the next + // request's bytes bleed into this record's redacted copy. Clone the + // rare zero-write case; real rewrites (n > 0) already produced a + // fresh buffer via sjson. + if n == 0 { + rewritten = bytes.Clone(rewritten) + } rec.RequestBodyRedacted = rewritten return rewritten, hookResult, false } diff --git a/packages/ai-gateway/internal/ingress/proxy/stage_routing.go b/packages/ai-gateway/internal/ingress/proxy/stage_routing.go index b9e66f3b..ad847d66 100644 --- a/packages/ai-gateway/internal/ingress/proxy/stage_routing.go +++ b/packages/ai-gateway/internal/ingress/proxy/stage_routing.go @@ -213,7 +213,12 @@ func (h *Handler) resolveNoMatchPassthrough(ctx context.Context, requestedModel } } - model, err := h.deps.Models.GetModelByCode(ctx, requestedModel) + // Resolve by code OR alias: a client may address a model by an + // admin-configured alias with no routing rule, which must still route to + // the model (the model's ProviderModelID then reaches the wire, and the + // passthrough body rewrite swaps the alias for it). O(1) from the + // in-memory index — no per-request DB read. + model, err := h.deps.Models.GetModelByCodeOrAlias(ctx, requestedModel) if err != nil || model == nil { return nil, &routingFallbackError{ status: http.StatusNotFound, diff --git a/packages/ai-gateway/internal/platform/audit/adaptive.go b/packages/ai-gateway/internal/platform/audit/adaptive.go index 1939b9e8..92f6d23f 100644 --- a/packages/ai-gateway/internal/platform/audit/adaptive.go +++ b/packages/ai-gateway/internal/platform/audit/adaptive.go @@ -52,60 +52,48 @@ func availableMemoryBytes() uint64 { } const ( - // auditMemFraction is the share of AVAILABLE memory the audit side-path's - // in-heap + spill-handoff buffers may occupy. Audit is a side-path, so it takes - // a minority slice; the rest stays for the request path + page cache. - auditMemFraction = 0.25 - - // avgRecordPinBytes is the assumed memory a queued record pins until it is - // marshaled out of the pools — dominated by the captured request/response body - // (the inline cap is 256 KiB, but typical bodies + the small-record mix average - // far lower). Used only to turn a byte budget into a channel capacity; the exact - // value is not load-bearing (it sizes a buffer, not a correctness bound). - avgRecordPinBytes = 48 << 10 // 48 KiB - - // in-heap : spill-handoff split of the buffer budget. The in-heap queue is the - // fast path (workers drain it to NATS); the spill channel is the overflow runway - // to disk. Weight the fast path heavier. - recChBudgetShare = 0.6 - spillChBudgetShare = 0.4 - - // Capacity clamps so a tiny box still gets a usable buffer and a huge box does - // not allocate an absurd channel. These bound the AUTO value, they are not the - // value — the machine picks within the band. - minRecChCap = 2000 - maxRecChCap = 2_000_000 - minSpillChCap = 1000 - maxSpillChCap = 1_000_000 - - // fixedBudgetFallback is used when MemAvailable cannot be read (non-Linux). + // auditMemFraction is the share of AVAILABLE memory the audit in-memory queue may + // pin, as a BYTE budget (see adaptiveMemBudgetBytes + shared/core/bytebudget). ~15% + // of RAM — the same order as the NATS JetStream memory-store limit, a side-path + // minority slice. A producer Acquires its record's REAL byte weight before enqueuing + // and blocks when this budget is exhausted (back-pressure to the request path), so + // the in-memory audit heap is bounded by BYTES (the actual record size), never by a + // slot count derived from an ASSUMED body size (the old sizing OOM-killed the process + // at 18 GiB when real bodies exceeded the assumed average). Configure it below the + // machine's real headroom (e.g. 15% of RAM) so the soft budget's small overshoot + // never matters. + auditMemFraction = 0.15 + + // recChStructuralCap / spillChStructuralCap are FIXED structural channel depths — + // how many *Record POINTERS the queues may hold. NOT the memory bound (the + // byteBudget bounds bytes); a slot is ~8 B, so a deep channel is cheap. Sized + // generously so the byte budget binds first for realistic large-body traffic, while + // the count still caps a pathological flood of tiny records. Body-size-INDEPENDENT + // on purpose: the memory limit is the byte budget, computed from the ACTUAL record. + recChStructuralCap = 200_000 + spillChStructuralCap = 100_000 + + // fixedBudgetFallback is the byte budget used when MemAvailable cannot be read + // (non-Linux / restricted /proc). fixedBudgetFallback = 2 << 30 // 2 GiB ) -func clampInt(v, lo, hi int) int { - if v < lo { - return lo - } - if v > hi { - return hi +// adaptiveMemBudgetBytes is the in-memory audit BYTE budget — a fraction of the live +// MemAvailable (fallback fixedBudgetFallback off-Linux). The byteBudget back-pressures +// Enqueue on this, bounding the audit heap by ACTUAL bytes regardless of body size. +func adaptiveMemBudgetBytes() int64 { + avail := availableMemoryBytes() + if avail == 0 { + return fixedBudgetFallback } - return v + return int64(float64(avail) * auditMemFraction) } -// adaptiveBufferCaps computes the in-heap (recCh) and spill-handoff (spillCh) -// channel capacities from available memory — replacing the former fixed -// maxQueueSize / spillChanCap magic numbers. A bigger box automatically gets -// deeper buffers (more burst absorption before any back-pressure), a smaller box -// gets shallower ones, with no config. +// adaptiveBufferCaps returns the STRUCTURAL channel depths (pointer counts) for +// recCh / spillCh. These are a count ceiling, not the memory bound — the byteBudget +// bounds bytes. Fixed and body-size-independent by design. func adaptiveBufferCaps() (recCh, spillCh int) { - budget := availableMemoryBytes() - if budget == 0 { - budget = fixedBudgetFallback - } - budget = uint64(float64(budget) * auditMemFraction) - recCh = clampInt(int(float64(budget)*recChBudgetShare)/avgRecordPinBytes, minRecChCap, maxRecChCap) - spillCh = clampInt(int(float64(budget)*spillChBudgetShare)/avgRecordPinBytes, minSpillChCap, maxSpillChCap) - return recCh, spillCh + return recChStructuralCap, spillChStructuralCap } // adaptiveSpillFlush adapts the spill worker's flush size to measured disk write diff --git a/packages/ai-gateway/internal/platform/audit/audit.go b/packages/ai-gateway/internal/platform/audit/audit.go index 4d020f85..3430319e 100644 --- a/packages/ai-gateway/internal/platform/audit/audit.go +++ b/packages/ai-gateway/internal/platform/audit/audit.go @@ -18,25 +18,14 @@ import ( ) const ( - // maxQueueSize is the DEFAULT in-memory record-buffer cap (overridable per - // Writer via WithMaxQueuedRecords / AuditConfig.MaxQueuedRecords). On overflow - // Enqueue spills to the durable NDJSON sink (never a silent drop). Each queued - // record PINS its pooled ~50 KB request/response body until it is marshaled, so - // the cap directly bounds the audit body pool's working set: at the old 50000 a - // slow-publish burst pinned ~5 GB of bodies (the dominant gw retained heap); - // 10000 holds the pool near ~1 GB with the same measured spill/drop rate and a - // ~70% lower GC pause. The buffer of last resort is NATS (disk) + the spill - // sink, NOT this in-heap queue — keep it bounded, raise it only where a - // memory-rich box wants extra absorption headroom. + // maxQueueSize is the LAST-RESORT record-buffer cap, used only when a Writer + // was constructed without NewWriter's structural sizing and without a + // WithMaxQueuedRecords override (effectiveMaxQueue's zero fallback). Like the + // structural caps it bounds pointer count only; the audit heap is bounded by + // the byte budget (Writer.memBudget), which accounts each record's REAL body + // bytes at Enqueue admission. maxQueueSize = 10000 - // flushHighWater triggers an immediate flush once the buffer reaches - // this depth, instead of waiting for the next ticker. Without it a - // burst toward maxQueueSize within one flush interval would force - // backpressure/spill even when the pipeline has drain capacity. Set - // well below maxQueueSize so a spike flushes early. - flushHighWater = 1000 - // Audit loss modes (AuditConfig.LossMode). Every default is zero-loss, because // durable audit is a product promise + a compliance requirement — the gateway // must never silently drop an audit record. Two distinct "defaults" exist and @@ -47,11 +36,13 @@ const ( // lossy modes are an explicit opt-out for callers that do NOT need compliance // audit and prefer raw throughput. // - spillblock (CONFIG default): like spill, but on a full spill channel it - // back-pressures the request goroutine until a slot frees instead of - // dropping (bounded by backpressureMaxWait → durable spill). Lossless up to - // disk-write success, throttling ingest to the disk rate; spills to disk - // before it ever blocks. Identical to spill in the normal regime; differs - // only at the extreme where spill would drop. + // back-pressures the request goroutine (parks on the channel) until a slot + // frees instead of dropping. It also NEVER drops on a full spool QUOTA: the + // single spill worker keeps the batch and retries while the recovery sweeper + // drains + frees space, so ingest self-throttles to the recovery-drain rate + // rather than shedding records. Genuinely lossless whenever a spool is wired; + // the only drops are a genuine (non-quota) disk I/O error or shutdown. Identical + // to spill in the normal regime; differs where spill would drop. // - block (empty/unknown fallback): when the in-heap buffer is full, Enqueue // BACK-PRESSURES the request path (bounded wait for the flush to free space; // durable spill only if the pipeline is genuinely wedged past the wait). @@ -105,6 +96,17 @@ const ( // stack-trace Error was itself a top allocator under the overload it // reports — see Writer.Enqueue. dropLogEvery = 2000 + + // maxPublishRetries bounds handlePublishFailure's in-memory re-queue retries per + // record. A publish failure re-queues onto recCh for a fast in-memory re-publish + // (transient NAK recovery); the cap stops that from becoming an unbounded + // busy-spin under a SUSTAINED outage (stream full → every publish 503s, workers + // keep draining recCh so a re-queue always wins a slot → the record circulates + // forever, pinning its marshaled copy + cratering throughput). After this many + // attempts the record is routed to the durable BATCHED spill (spillCh) instead. + // 3 keeps genuine transient blips recovering in-memory (fast, no disk) while + // bounding a wedged pipeline to a handful of attempts before it drains to spool. + maxPublishRetries = 3 ) // EndpointType is the typed-string alias used to classify the API diff --git a/packages/ai-gateway/internal/platform/audit/audit_test.go b/packages/ai-gateway/internal/platform/audit/audit_test.go index 9a0cd0a0..8bccc61d 100644 --- a/packages/ai-gateway/internal/platform/audit/audit_test.go +++ b/packages/ai-gateway/internal/platform/audit/audit_test.go @@ -717,13 +717,24 @@ func TestPublishFailure_OverflowSpillsNotLost(t *testing.T) { for i := range batch { batch[i] = &Record{RequestID: fmt.Sprintf("rec-stuck-%d", i), Timestamp: time.Now().Add(time.Duration(i))} } - w.publishBatchOn(0, batch) // every publish fails → re-queue full → durable spill + w.publishBatchOn(0, batch) // every publish fails → re-queue full → durable batched spill if got := len(prod.msgs()); got != 0 { t.Errorf("alwaysFail producer must not record messages; got %d", got) } - // The stuck records landed durably on the spill sink — no silent loss, no - // infinite retry. + // The stuck records were handed to the batched spill channel (off the publish + // worker) rather than written synchronously; drive the worker's write step for + // each and confirm they land durably — no silent loss, no infinite retry. + for i := range n { + select { + case rec := <-w.spillCh: + if !w.spillRecord(rec) { + t.Fatalf("spillRecord must durably write stuck record %d", i) + } + default: + t.Fatalf("stuck record %d was not handed to the spill channel", i) + } + } if got := len(readSpool(t, dir, "test")); got != n { t.Errorf("sustained-outage records must spill durably; spooled %d, want %d", got, n) } diff --git a/packages/ai-gateway/internal/platform/audit/backpressure.go b/packages/ai-gateway/internal/platform/audit/backpressure.go index 0d74daf7..c9c5a0fe 100644 --- a/packages/ai-gateway/internal/platform/audit/backpressure.go +++ b/packages/ai-gateway/internal/platform/audit/backpressure.go @@ -1,6 +1,7 @@ package audit import ( + "errors" "os" "time" @@ -14,11 +15,11 @@ import ( var perfNoSpill = os.Getenv("NEXUS_PERF_NO_SPILL") == "1" // WithMaxQueuedRecords sets the in-heap record-buffer cap (overflow → durable -// spill). n <= 0 keeps the maxQueueSize default. The cap bounds the audit body -// pool's working set, since every queued record pins its pooled ~50 KB body until -// it is marshaled — so this is the primary control over the audit side-path's gw -// heap footprint. Wired from AuditConfig.MaxQueuedRecords at startup; call before -// any Enqueue. Returns the receiver for chaining. +// spill). n <= 0 keeps the structural default. The cap bounds the queue's +// POINTER count only — the audit side-path's heap footprint is bounded by the +// byte budget (memBudget), which accounts each record's REAL body bytes at +// Enqueue admission. Wired from AuditConfig.MaxQueuedRecords at startup; call +// before any Enqueue. Returns the receiver for chaining. func (w *Writer) WithMaxQueuedRecords(n int) *Writer { if n > 0 { w.maxQueued = n @@ -34,19 +35,24 @@ func (w *Writer) effectiveMaxQueue() int { return maxQueueSize } -// WithLossMode selects the overflow policy: lossModeBlock (no-loss back-pressure, -// the compliance default), lossModeSpill (async durable spill, counted drop only -// when the spill channel is also saturated), lossModeSpillBlock (spill, but -// back-pressure instead of that last-resort drop), or lossModeDrop (counted -// bounded drop). An empty or unrecognised value keeps the block default — audit -// must never silently start dropping because of a config typo. Wired from -// AuditConfig.LossMode at startup; call before any Enqueue. Returns the receiver. +// WithLossMode selects the overflow policy: lossModeSpillBlock (the no-loss +// default — spill to the durable on-disk spool first, and only back-pressure the +// request path when that large spool is ALSO saturated; makes the cheap 50GB +// spool the primary overflow buffer instead of stalling intake at the small +// in-heap queue), lossModeBlock (no-loss back-pressure at the in-heap queue, +// never touches the spool), lossModeSpill (async durable spill, counted drop only +// when the spill channel is also saturated), or lossModeDrop (counted bounded +// drop). An empty or unrecognised value keeps the spillBlock default — audit must +// never silently start dropping because of a config typo, and spillBlock is +// no-loss whenever the spool is wired (prod/rig always set AI_GATEWAY_AUDIT_SPOOL_DIR; +// spillOverflow falls back to block-style behaviour if no spool sink is wired). +// Wired from AuditConfig.LossMode at startup; call before any Enqueue. Returns the receiver. func (w *Writer) WithLossMode(mode string) *Writer { switch mode { - case lossModeSpill, lossModeDrop, lossModeSpillBlock: + case lossModeSpill, lossModeDrop, lossModeBlock: w.lossMode = mode default: - w.lossMode = lossModeBlock + w.lossMode = lossModeSpillBlock } return w } @@ -112,6 +118,40 @@ func (w *Writer) Enqueue(rec *Record) { } w.ensureStarted() + // Byte-bounded admission: reserve the record's REAL captured-body weight + // against the in-memory byte budget before it may pin heap. This is the memory + // bound — the channel caps bound pointer counts only (a slot count derived + // from an ASSUMED body size let ~145K queued ~100 KiB records OOM the process + // at 18 GiB). No-loss modes BLOCK here on a full budget, back-pressuring the + // request path to the audit drain rate; lossy modes never block — a full + // budget is an immediate counted drop, because handing an UNACCOUNTED record + // to spillCh would re-open the unbounded-memory hole this bound closes. + if n := int64(len(rec.RequestBody) + len(rec.ResponseBody)); n > 0 { + switch w.lossMode { + case lossModeBlock, lossModeSpillBlock: + if !w.memBudget.TryAcquire(n) { + w.metrics.incMemBackpressure() + if !w.memBudget.Acquire(n) { + // Shutdown fired while blocked: nothing was reserved — spill + // durably without enqueuing (still no loss). + w.spillRecord(rec) + return + } + } + rec.reserved = n + default: // lossModeSpill, lossModeDrop — never block the request path + if !w.memBudget.TryAcquire(n) { + // Count the budget-exhausted event on the same counter the blocking + // modes use, so an operator can tell "byte budget full" from a plain + // queue-full drop; the drop itself is counted in dropOverflow. + w.metrics.incMemBackpressure() + w.dropOverflow(rec) + return + } + rec.reserved = n + } + } + switch w.lossMode { case lossModeDrop: select { @@ -163,11 +203,14 @@ func (w *Writer) blockEnqueue(rec *Record) { // loss mode: // - lossModeSpill: a counted, throttled-log drop — never blocks the request // goroutine (the `dropped_total` counter makes the rare overload drop loud). -// - lossModeSpillBlock: back-pressure the request goroutine until a spill slot -// frees, bounded by backpressureMaxWait → durable spill, then a stopCh escape -// for shutdown. This keeps spilling to disk as long as the disk can absorb it -// and only throttles ingest (to the disk rate) under genuine overload, so it is -// lossless up to disk-write success rather than shedding records. +// - lossModeSpillBlock: park the request goroutine on the bounded spillCh until +// the SINGLE spill worker frees a slot — pure channel back-pressure, no +// per-goroutine ndjson.Write/dirSize (the worker is the sole spool writer). The +// worker NEVER drops a full spool quota (it holds the batch and retries while +// recovery frees space), and a genuinely error-ing disk makes the worker drop + +// keep draining, so this park always resolves except on a hung disk. Only +// shutdown escapes (stopCh → one-shot durable spill). So it is lossless up to a +// real (non-quota) disk failure, throttling ingest to the recovery-drain rate. // // No-loss vs never-block is a physical trilemma once ingest exceeds disk capacity: // you cannot be lossless AND non-blocking AND memory-bounded. lossModeSpill chooses @@ -188,17 +231,16 @@ func (w *Writer) spillOverflow(rec *Record) { w.dropOverflow(rec) // lossModeSpill: counted drop, never blocks the request path return } - // lossModeSpillBlock: the spill channel is full, so the worker is saturated - // against the disk write rate. Park the request goroutine until a slot frees - // rather than dropping. Bound the wait so a fully STALLED (not merely slow) - // disk falls back to a durable spill instead of parking intake forever, and - // escape on shutdown so Close() can drain. - timer := time.NewTimer(backpressureMaxWait) - defer timer.Stop() + // lossModeSpillBlock: the spill channel is full, so the single spill worker is + // saturated (writing at the disk rate, or holding a batch while it back-pressures + // a full spool quota). Park the request goroutine on the CHANNEL until a slot + // frees — pure channel back-pressure, no per-goroutine ndjson.Write / dirSize, so + // the spool stays single-writer. The worker never drops a full spool quota (it + // retries), and a genuinely error-ing disk makes the worker drop + keep draining, + // so this park always resolves except on a hung disk. Only shutdown escapes, with + // a one-shot durable spill (quota-full at shutdown → counted drop, bounded Close). select { case w.spillCh <- rec: - case <-timer.C: - w.spillRecord(rec) case <-w.stopCh: w.spillRecord(rec) } @@ -207,6 +249,7 @@ func (w *Writer) spillOverflow(rec *Record) { // dropOverflow is the lossy "drop" policy: a counted, throttled-log bounded drop. func (w *Writer) dropOverflow(rec *Record) { w.metrics.incDropped() + w.releaseRecordMem(rec) w.reclaimRecordBody(rec) if n := w.dropLogCount.Add(1); n%dropLogEvery == 1 { w.logger.Warn("audit overflow, dropping records (lossy mode, throttled)", @@ -224,29 +267,89 @@ func (w *Writer) spillLoop() { defer ticker.Stop() var buf []byte // accumulated NDJSON lines; cap retained across flushes for reuse cnt := 0 // records in buf (for the spilled / dropped counters) + // bufReserved aggregates the byte-budget reservations of the records batched + // into buf. The *Record is discarded at add() (only its marshaled bytes live + // on, in buf), so per-record release is impossible at the flush terminal — + // add() transfers each reservation here and flush releases the aggregate once + // the batch resolves (written, or dropped). While a quota-full back-pressure + // retry holds the batch, the reservation is correctly still held: the bytes + // really are pinned in buf. The accounting is knowingly LOW for this buffer — + // buf holds base64-encoded marshaled lines (~1.33× the wire bytes plus + // envelope overhead over the raw bodies the reservation counted) — but the + // error is capped by the adaptive flush ceiling (spillFlushMaxBytes, 256 MiB), + // noise against a memory-scale budget. + var bufReserved int64 flush := func() { if len(buf) == 0 { return } + releaseBatch := func() { + w.memBudget.Release(bufReserved) + bufReserved = 0 + } if w.ndjsonSpill == nil { w.metrics.addDropped(cnt) - } else { + releaseBatch() + buf = buf[:0] + cnt = 0 + return + } + // The spill worker is the SINGLE writer to the spool: request/consume + // goroutines only ever park on the bounded channels, never call ndjson here, + // so the O(files) quota scan runs from just this goroutine at the retry + // cadence and never thrashes the mutex the recovery sweeper needs. + retried := false + retryLoop: + for { start := time.Now() err := w.ndjsonSpill.WriteBatch(buf) - if err != nil { - w.metrics.addDropped(cnt) - if n := w.spillLogCount.Add(1); n%dropLogEvery == 1 { - w.logger.Warn("audit: durable spill batch failed, dropping records (throttled)", - "error", err, "records", cnt) - } - } else { + if err == nil { w.metrics.addSpilled(cnt) - // Feed this write's size + latency back to the adaptive flush sizer so - // the next flush threshold rides just under the disk's real write rate. - w.spillFlush.observe(len(buf), time.Since(start)) + // Only a clean write is a valid disk-write-latency sample for the + // adaptive flush sizer. A batch that waited out a quota back-pressure is + // contention, not device bandwidth — feeding that wait into the EMA would + // collapse the flush threshold to the floor and re-inflate IOPS once + // recovery frees space, so skip observe for a retried batch. + if !retried { + w.spillFlush.observe(len(buf), time.Since(start)) + } + break + } + // spillBlock + spool AT QUOTA (soft, relieved by the recovery drain): keep + // the SAME batch and back-pressure — wait for recovery to free space, then + // retry. Never a drop. Every OTHER error (genuine I/O / ENOSPC) and every + // non-spillBlock mode falls through to the counted drop below. + if w.lossMode == lossModeSpillBlock && errors.Is(err, sharedndjson.ErrSpoolQuotaExceeded) { + if !retried { + retried = true + w.metrics.incSpillBackpressure() + if n := w.spillLogCount.Add(1); n%dropLogEvery == 1 { + w.logger.Warn("audit: spool quota full — back-pressuring request path "+ + "(recovery draining; check NATS/Hub/PG if sustained)", "records", cnt) + } + } + t := time.NewTimer(spillFlushInterval) + select { + case <-t.C: + continue // retry the same batch after a pause + case <-w.stopCh: + t.Stop() + w.metrics.addDropped(cnt) // shutdown one-shot → Close()/wg.Wait() stays bounded + break retryLoop + } + } + w.metrics.addDropped(cnt) + if n := w.spillLogCount.Add(1); n%dropLogEvery == 1 { + w.logger.Warn("audit: durable spill batch failed, dropping records (throttled)", + "error", err, "records", cnt) } + break } + // Terminal for the whole batch — written durably or counted-dropped either + // way; the batched bytes are about to be discarded, so return their + // aggregated byte-budget reservation. + releaseBatch() buf = buf[:0] cnt = 0 } @@ -260,9 +363,16 @@ func (w *Writer) spillLoop() { // body; the returned bytes alias msgBuf, so reclaim it once they are copied. data, msgBuf, ok := w.marshalRecord(rec) if !ok { + // Terminal: marshal failure dropped the record (marshalRecord released + // its byte-budget reservation). w.metrics.incDropped() return } + // The *Record is discarded past this point (only its marshaled bytes live + // on, in buf) — transfer its reservation to the batch aggregate, released + // at the flush terminal. + bufReserved += rec.reserved + rec.reserved = 0 buf = appendSpillLine(buf, data) // base64 line copies data — binary-safe for the spool reclaimMsgBuf(msgBuf) // data captured in buf; release the pooled buffer cnt++ @@ -306,6 +416,9 @@ func (w *Writer) spillLoop() { // quota-full burst the logging itself became the dominant CPU + allocation cost — a // self-DoS on the very overload it reported. The dropped_total metric stays exact. func (w *Writer) spillRecord(rec *Record) bool { + // Every exit is a terminal (written durably, or counted-dropped) — the record + // leaves the in-memory pipeline here, so its byte-budget reservation returns. + defer w.releaseRecordMem(rec) if w.ndjsonSpill == nil { w.metrics.incDropped() w.reclaimRecordBody(rec) diff --git a/packages/ai-gateway/internal/platform/audit/coverage_gaps_test.go b/packages/ai-gateway/internal/platform/audit/coverage_gaps_test.go index febb1810..d1cee4f8 100644 --- a/packages/ai-gateway/internal/platform/audit/coverage_gaps_test.go +++ b/packages/ai-gateway/internal/platform/audit/coverage_gaps_test.go @@ -74,36 +74,40 @@ func TestAvailableMemoryBytes_MissingFileReturnsZero(t *testing.T) { } } -// adaptiveBufferCaps consumes availableMemoryBytes: with a large fixture the recCh -// cap must land within the adaptive clamp band, proving the parsed value flows into -// the sizing rather than the fixed fallback. -func TestAdaptiveBufferCaps_UsesParsedMemory(t *testing.T) { - withMeminfo(t, "MemAvailable: 8388608 kB\n") // 8 GiB available +// adaptiveBufferCaps returns FIXED structural pointer-count depths — deliberately +// body-size-INDEPENDENT (the byte budget is the memory bound; a count derived from +// an assumed body size is exactly the sizing that OOM'd). Any meminfo content must +// not change them. +func TestAdaptiveBufferCaps_FixedStructuralDepths(t *testing.T) { + withMeminfo(t, "MemAvailable: 8388608 kB\n") // 8 GiB available — must not matter recCh, spillCh := adaptiveBufferCaps() - if recCh < minRecChCap || recCh > maxRecChCap { - t.Fatalf("recCh cap %d outside [%d,%d]", recCh, minRecChCap, maxRecChCap) - } - if spillCh < minSpillChCap || spillCh > maxSpillChCap { - t.Fatalf("spillCh cap %d outside [%d,%d]", spillCh, minSpillChCap, maxSpillChCap) + if recCh != recChStructuralCap || spillCh != spillChStructuralCap { + t.Fatalf("adaptiveBufferCaps() = (%d,%d), want fixed (%d,%d)", + recCh, spillCh, recChStructuralCap, spillChStructuralCap) } } // --------------------------------------------------------------------------- -// adaptive.go — clampInt below/above/within band. +// adaptive.go — adaptiveMemBudgetBytes: the in-memory audit BYTE budget is the +// RAM-fraction of parsed MemAvailable, with a fixed fallback when unreadable. // --------------------------------------------------------------------------- -func TestClampInt_Bounds(t *testing.T) { - cases := []struct{ v, lo, hi, want int }{ - {v: -5, lo: 0, hi: 10, want: 0}, // below low - {v: 99, lo: 0, hi: 10, want: 10}, // above high - {v: 7, lo: 0, hi: 10, want: 7}, // within band - {v: 0, lo: 0, hi: 10, want: 0}, // at low edge - {v: 10, lo: 0, hi: 10, want: 10}, // at high edge - } - for _, c := range cases { - if got := clampInt(c.v, c.lo, c.hi); got != c.want { - t.Errorf("clampInt(%d,%d,%d)=%d want %d", c.v, c.lo, c.hi, got, c.want) - } +func TestAdaptiveMemBudgetBytes_FractionOfParsedMemory(t *testing.T) { + withMeminfo(t, "MemAvailable: 8388608 kB\n") // 8 GiB available + availBytes := float64(uint64(8388608) * 1024) + want := int64(availBytes * auditMemFraction) + if got := adaptiveMemBudgetBytes(); got != want { + t.Fatalf("adaptiveMemBudgetBytes() = %d, want %d (%.0f%% of 8 GiB)", + got, want, auditMemFraction*100) + } +} + +func TestAdaptiveMemBudgetBytes_FallbackWhenUnreadable(t *testing.T) { + orig := meminfoPath + meminfoPath = filepath.Join(t.TempDir(), "does-not-exist") + t.Cleanup(func() { meminfoPath = orig }) + if got := adaptiveMemBudgetBytes(); got != int64(fixedBudgetFallback) { + t.Fatalf("adaptiveMemBudgetBytes() fallback = %d, want %d", got, int64(fixedBudgetFallback)) } } @@ -589,8 +593,10 @@ func TestSpillLoop_WriteBatchErrorCountsDrops(t *testing.T) { if err != nil { t.Fatalf("ndjson.New: %v", err) } - // Pre-seed past the quota so the loop's batched WriteBatch is refused. - if err := os.WriteFile(filepath.Join(dir, "gw", "seed.ndjson"), make([]byte, 1100*1024), 0o600); err != nil { + // Pre-seed past the quota so the loop's batched WriteBatch is refused. The seed + // must be a countable audit-*.ndjson file — the reclaimable-quota gate excludes + // non-audit entries (e.g. .poison), so a differently-named filler would not count. + if err := os.WriteFile(filepath.Join(dir, "gw", "audit-20260101-0001.ndjson"), make([]byte, 1100*1024), 0o600); err != nil { t.Fatalf("seed: %v", err) } w := NewWriter(nil, "q", opsmetrics.NewRegistry(prom), slog.New(slog.NewTextHandler(io.Discard, nil))). diff --git a/packages/ai-gateway/internal/platform/audit/loss_mode_test.go b/packages/ai-gateway/internal/platform/audit/loss_mode_test.go index 1cf8c394..cac444be 100644 --- a/packages/ai-gateway/internal/platform/audit/loss_mode_test.go +++ b/packages/ai-gateway/internal/platform/audit/loss_mode_test.go @@ -81,16 +81,18 @@ func quietWriter(reg *registry.Registry) *Writer { } // TestWriter_WithLossMode_Selection: the overflow policy resolves to the named -// mode, and an empty/unknown value falls back to no-loss block (audit must never -// silently turn lossy from a config typo). +// mode, and an empty/unknown value falls back to no-loss spillBlock — the durable +// spool is the primary overflow buffer and the request path is back-pressured only +// when it too saturates (audit must never silently turn lossy from a config typo; +// spillBlock is no-loss whenever a spool sink is wired, which prod/rig always do). func TestWriter_WithLossMode_Selection(t *testing.T) { cases := map[string]string{ lossModeSpill: lossModeSpill, lossModeDrop: lossModeDrop, lossModeSpillBlock: lossModeSpillBlock, "block": lossModeBlock, - "": lossModeBlock, - "garbage": lossModeBlock, + "": lossModeSpillBlock, + "garbage": lossModeSpillBlock, } for in, want := range cases { if got := quietWriter(nil).WithLossMode(in).LossMode(); got != want { @@ -103,11 +105,10 @@ func TestWriter_WithLossMode_Selection(t *testing.T) { // current value (default when never set). func TestWriter_EffectiveMaxQueue(t *testing.T) { w := quietWriter(nil) - // The default cap is sized to available memory (adaptiveBufferCaps), not a fixed - // constant, so assert it lands within the adaptive clamp band rather than a magic - // number — the whole point of the adaptive sizing is that there is no fixed value. - if c := w.effectiveMaxQueue(); c < minRecChCap || c > maxRecChCap { - t.Fatalf("default cap = %d, want within adaptive band [%d,%d]", c, minRecChCap, maxRecChCap) + // The default cap is the FIXED structural pointer-count depth — body-size- + // independent by design (the byte budget is the memory bound, not this count). + if c := w.effectiveMaxQueue(); c != recChStructuralCap { + t.Fatalf("default cap = %d, want structural cap %d", c, recChStructuralCap) } w.WithMaxQueuedRecords(42) if w.effectiveMaxQueue() != 42 { diff --git a/packages/ai-gateway/internal/platform/audit/mem_budget.go b/packages/ai-gateway/internal/platform/audit/mem_budget.go new file mode 100644 index 00000000..579de1f6 --- /dev/null +++ b/packages/ai-gateway/internal/platform/audit/mem_budget.go @@ -0,0 +1,69 @@ +// mem_budget.go — the audit Writer's in-memory byte budget: the config surface +// (WithMemMaxBytes / MemBudgetBytes) and the exactly-once release helper the +// terminal sites call. The admission (reserve) side lives in Enqueue +// (backpressure.go), which owns the whole buffer-admission path. +package audit + +import ( + "strings" + + "github.com/AlphaBitCore/nexus-gateway/packages/shared/core/bytebudget" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/transport/mq" +) + +// memPinWarnFraction is the pinned-budget share of available RAM past which +// WithMemMaxBytes warns. The budget accounts only raw body bytes; the process's +// REAL memory is roughly 2× the pinned value (Go GC headroom at GOGC=100, plus +// marshal/compression/frame copies) on a box that also hosts the NATS memory +// stream (~15% RAM) — a 10 GB pin on a 32 GB box OOM-killed the gateway at +// 15.9 GB RSS in rig validation. 25% keeps 2×pin + NATS + PG under the ceiling. +const memPinWarnFraction = 0.25 + +// WithMemMaxBytes sets the in-memory audit byte budget from config. Semantics +// mirror NEXUS_EVENTS_MAX_BYTES exactly: empty or "auto" keeps the auto default +// (~15% of available RAM, set in NewWriter); an explicit human size ("4GB", +// "2048MB", raw bytes) pins it; an unparseable value keeps the auto default so a +// typo can never disable the OOM bound. A pin above memPinWarnFraction of +// available RAM logs a WARN (never clamps — the operator stays in charge, but an +// OOM-risky value must not pass silently). Call before Start / any Enqueue +// (rebinds the budget, matching the other With* startup-only options). Returns +// the receiver. +func (w *Writer) WithMemMaxBytes(v string) *Writer { + v = strings.TrimSpace(v) + if v == "" || strings.EqualFold(v, "auto") { + return w + } + n := mq.ParseByteSize(v, 0) + if n <= 0 { + return w + } + if avail := availableMemoryBytes(); avail > 0 && n > int64(float64(avail)*memPinWarnFraction) { + w.logger.Warn("audit memory budget pinned above the safe share of available RAM — "+ + "real process memory runs ~2x the budget (GC headroom + marshal copies), OOM risk "+ + "alongside the NATS memory stream; prefer <=20% of RAM", + "pinned_bytes", n, "available_bytes", avail, + "pinned_fraction", float64(n)/float64(avail)) + } + w.memBudget = bytebudget.New(n, w.stopCh) + return w +} + +// MemBudgetBytes reports the resolved in-memory byte budget, so startup wiring +// can log what is actually in effect (auto-sized or pinned). +func (w *Writer) MemBudgetBytes() int64 { return w.memBudget.Budget() } + +// releaseRecordMem returns rec's byte-budget reservation, exactly once +// (idempotent: the first call zeroes reserved, later calls no-op). Call at every +// TERMINAL — the points where a record leaves the in-memory pipeline: published +// OK, durably spilled, or dropped. A missed release leaks budget (permanent +// over-back-pressure → ingest stall); a double release inflates it (re-opens the +// OOM). The zeroing here guards the double; the per-terminal call sites guard +// the miss. Records batched by the spill worker transfer their reservation to +// the batch aggregate instead (spillLoop's bufReserved). +func (w *Writer) releaseRecordMem(rec *Record) { + if rec == nil || rec.reserved == 0 { + return + } + w.memBudget.Release(rec.reserved) + rec.reserved = 0 +} diff --git a/packages/ai-gateway/internal/platform/audit/mem_budget_config_test.go b/packages/ai-gateway/internal/platform/audit/mem_budget_config_test.go new file mode 100644 index 00000000..1396be40 --- /dev/null +++ b/packages/ai-gateway/internal/platform/audit/mem_budget_config_test.go @@ -0,0 +1,78 @@ +package audit + +// mem_budget_config_test.go — WithMemMaxBytes config semantics (mirrors +// NEXUS_EVENTS_MAX_BYTES): empty/"auto" keeps the auto-sized default, an +// explicit human size pins the budget, and an unparseable value keeps the +// default so a typo can never disable the OOM bound. + +import ( + "io" + "log/slog" + "strings" + "testing" +) + +func TestWithMemMaxBytes_Semantics(t *testing.T) { + newW := func() *Writer { + return NewWriter(nil, "q", nil, slog.New(slog.NewTextHandler(io.Discard, nil))) + } + auto := newW().MemBudgetBytes() + if auto <= 0 { + t.Fatalf("auto budget must be positive, got %d", auto) + } + + cases := []struct { + in string + want int64 + }{ + {"", auto}, // unset → auto default kept + {"auto", auto}, // explicit auto → default kept + {"AUTO", auto}, // case-insensitive + {"8GB", 8 << 30}, // human size pins + {"2048MB", 2 << 30}, // MB suffix + {"1073741824", 1 << 30}, // raw bytes + {"garbage", auto}, // unparseable → default kept (never disables the bound) + {"-5GB", auto}, // non-positive → default kept + } + for _, c := range cases { + if got := newW().WithMemMaxBytes(c.in).MemBudgetBytes(); got != c.want { + t.Errorf("WithMemMaxBytes(%q) budget = %d, want %d", c.in, got, c.want) + } + } +} + +// A pinned budget must actually govern admission: a writer pinned to a tiny +// budget back-pressures exactly like a tiny auto budget would. +func TestWithMemMaxBytes_PinnedBudgetGovernsAdmission(t *testing.T) { + w := NewWriter(nil, "q", nil, slog.New(slog.NewTextHandler(io.Discard, nil))). + WithMemMaxBytes("1KB") + if !w.memBudget.TryAcquire(1024) { + t.Fatal("first acquire within the pinned budget must admit") + } + if w.memBudget.TryAcquire(1) { + t.Fatal("pinned 1KB budget must refuse once exhausted") + } +} + +// A pin above the safe share of available RAM must WARN (never clamp): the +// budget still applies verbatim — the operator stays in charge — but the +// OOM-risk signal must land in the log. Rig-validated failure this guards +// against: a 10GB pin on a 32GB box OOM-killed the gateway at 15.9GB RSS. +func TestWithMemMaxBytes_OversizedPinWarnsButApplies(t *testing.T) { + withMeminfo(t, "MemAvailable: 4194304 kB\n") // 4 GiB available + var buf strings.Builder + w := NewWriter(nil, "q", nil, slog.New(slog.NewTextHandler(&buf, nil))). + WithMemMaxBytes("2GB") // 50% of available — far past the 25% warn line + if got := w.MemBudgetBytes(); got != 2<<30 { + t.Fatalf("oversized pin must still apply verbatim, got %d", got) + } + if !strings.Contains(buf.String(), "OOM risk") { + t.Fatalf("oversized pin must log the OOM-risk WARN; log was: %s", buf.String()) + } + // A modest pin (below the warn line) stays silent. + buf.Reset() + NewWriter(nil, "q", nil, slog.New(slog.NewTextHandler(&buf, nil))).WithMemMaxBytes("512MB") + if strings.Contains(buf.String(), "OOM risk") { + t.Fatalf("modest pin must not warn; log was: %s", buf.String()) + } +} diff --git a/packages/ai-gateway/internal/platform/audit/mem_budget_test.go b/packages/ai-gateway/internal/platform/audit/mem_budget_test.go new file mode 100644 index 00000000..100b8356 --- /dev/null +++ b/packages/ai-gateway/internal/platform/audit/mem_budget_test.go @@ -0,0 +1,373 @@ +package audit + +// mem_budget_test.go — the byte-bounded audit queue's business contracts: +// reservations are taken at Enqueue from REAL body sizes and released exactly +// once at every terminal (published / durably spilled / dropped), no-loss modes +// back-pressure the producer on a full budget, lossy modes shed with a counted +// drop instead of blocking, and shutdown always unblocks a parked producer into +// a durable spill. The invariant asserted throughout: after all records reach a +// terminal, memBudget.InUse() == 0 — a leak here means permanent +// over-back-pressure; a double release means the OOM bound is fiction. + +import ( + "context" + "fmt" + "io" + "log/slog" + "sync" + "testing" + "time" + + sharedndjson "github.com/AlphaBitCore/nexus-gateway/packages/shared/audit/ndjson" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/core/bytebudget" + opsmetrics "github.com/AlphaBitCore/nexus-gateway/packages/shared/core/metrics/registry" + "github.com/prometheus/client_golang/prometheus" +) + +// stallProducer implements mq.Producer + batchProducer; every publish parks on +// gate until the test closes it — simulating a broker that is up but not +// draining, so records HOLD their byte-budget reservations in flight. +type stallProducer struct { + gate chan struct{} + mu sync.Mutex + sent int +} + +func (p *stallProducer) Publish(context.Context, string, []byte) error { return nil } +func (p *stallProducer) Enqueue(context.Context, string, []byte) error { + <-p.gate + p.mu.Lock() + p.sent++ + p.mu.Unlock() + return nil +} +func (p *stallProducer) Close() error { return nil } +func (p *stallProducer) EnqueueBatchAsync(_ context.Context, _ string, b [][]byte) ([]error, error) { + <-p.gate + p.mu.Lock() + p.sent += len(b) + p.mu.Unlock() + return make([]error, len(b)), nil +} + +// bodyRecord builds a record whose captured bodies weigh exactly n bytes — the +// value Enqueue must reserve against the byte budget. +func bodyRecord(id string, n int) *Record { + return &Record{ + RequestID: id, + Timestamp: time.Unix(1700000000, 0).UTC(), + RequestBody: make([]byte, n), + } +} + +// tinyBudgetWriter rebinds the writer's byte budget to n bytes (tests must not +// depend on the machine's RAM); call before the first Enqueue/Start. +func tinyBudgetWriter(w *Writer, n int64) *Writer { + w.memBudget = bytebudget.New(n, w.stopCh) + return w +} + +// Every published record must return its reservation: after the pipeline drains, +// InUse is exactly 0 — across BOTH publish paths (per-record fallback and the +// pooled framed batch path). +func TestMemBudget_ReleasedOnPublishOK(t *testing.T) { + t.Run("per-record fallback", func(t *testing.T) { + prod := &memProducer{} + w := NewWriter(prod, "q", nil, slog.New(slog.NewTextHandler(io.Discard, nil))) + for i := range 5 { + w.Enqueue(bodyRecord(fmt.Sprintf("pub-%d", i), 1024)) + } + if got := waitForMsgCount(prod, 5, 2*time.Second); got != 5 { + t.Fatalf("published %d of 5", got) + } + w.Close() + if got := w.memBudget.InUse(); got != 0 { + t.Fatalf("InUse after publish-OK drain = %d, want 0 (reservation leak)", got) + } + }) + t.Run("pooled framed batch", func(t *testing.T) { + p := &pooledFakeProducer{pool: 2} + w := NewWriter(p, "q", nil, slog.New(slog.NewTextHandler(io.Discard, nil))).WithFramePublish(1 << 20) + w.Start() + for i := range 20 { + w.Enqueue(bodyRecord(fmt.Sprintf("fr-%d", i), 2048)) + } + w.Close() + if got := p.recordCount(); got != 20 { + t.Fatalf("framed publish delivered %d of 20", got) + } + if got := w.memBudget.InUse(); got != 0 { + t.Fatalf("InUse after framed drain = %d, want 0 (reservation leak)", got) + } + }) +} + +// No-loss (block) mode: a full byte budget must PARK the producer — the +// back-pressure that bounds the audit heap — and wake it exactly when the +// in-flight record's publish completes and releases its bytes. The distinct +// mem_backpressure signal must record the wait. +func TestMemBudget_BlockModeBackpressuresAtFullBudget(t *testing.T) { + prom := prometheus.NewRegistry() + prod := &stallProducer{gate: make(chan struct{})} + w := NewWriter(prod, "q", opsmetrics.NewRegistry(prom), + slog.New(slog.NewTextHandler(io.Discard, nil))).WithLossMode(lossModeBlock) + tinyBudgetWriter(w, 64) + w.Start() + + w.Enqueue(bodyRecord("holder", 64)) // fills the budget; its publish stalls on gate + + second := make(chan struct{}) + go func() { + w.Enqueue(bodyRecord("parked", 64)) // must park until the holder releases + close(second) + }() + select { + case <-second: + t.Fatal("Enqueue admitted past an exhausted byte budget (no back-pressure)") + case <-time.After(100 * time.Millisecond): + } + if got := counterValue(t, prom, "nexus_audit_mem_backpressure_total"); got < 1 { + t.Fatalf("mem_backpressure_total = %v, want >= 1 while parked", got) + } + + close(prod.gate) // the holder publishes → releases 64 bytes → parked producer wakes + select { + case <-second: + case <-time.After(2 * time.Second): + t.Fatal("parked Enqueue did not wake after the in-flight record released its bytes") + } + w.Close() + if got := w.memBudget.InUse(); got != 0 { + t.Fatalf("InUse after drain = %d, want 0", got) + } + if got := counterValue(t, prom, "nexus_audit_mq_dropped_total"); got != 0 { + t.Fatalf("dropped_total = %v, want 0 (block mode is no-loss)", got) + } +} + +// Lossy modes must NEVER block the request path: a full byte budget is an +// immediate counted drop (handing an unaccounted record onward would re-open +// the unbounded-memory hole the budget closes). +func TestMemBudget_LossyModeShedsInsteadOfBlocking(t *testing.T) { + prom := prometheus.NewRegistry() + prod := &stallProducer{gate: make(chan struct{})} + w := NewWriter(prod, "q", opsmetrics.NewRegistry(prom), + slog.New(slog.NewTextHandler(io.Discard, nil))).WithLossMode(lossModeDrop) + tinyBudgetWriter(w, 64) + w.Start() + + w.Enqueue(bodyRecord("holder", 64)) // fills the budget + + returned := make(chan struct{}) + go func() { + w.Enqueue(bodyRecord("shed", 64)) // budget full → counted drop, no wait + close(returned) + }() + select { + case <-returned: + case <-time.After(2 * time.Second): + t.Fatal("lossy-mode Enqueue blocked on a full byte budget (must shed instead)") + } + if got := counterValue(t, prom, "nexus_audit_mq_dropped_total"); got != 1 { + t.Fatalf("dropped_total = %v, want exactly 1 (the shed record)", got) + } + // The budget-exhausted event is counted even in lossy modes, so an operator + // can tell a budget-full shed from a plain queue-full drop. + if got := counterValue(t, prom, "nexus_audit_mem_backpressure_total"); got != 1 { + t.Fatalf("mem_backpressure_total = %v, want 1 (budget-full shed must be distinguishable)", got) + } + close(prod.gate) + w.Close() + if got := w.memBudget.InUse(); got != 0 { + t.Fatalf("InUse after drain = %d, want 0", got) + } +} + +// The spill worker batches records and discards the *Record after marshal — the +// reservation must transfer to the batch aggregate and be released at the flush +// terminal (spilled durably), never leaked and never double-released. +func TestMemBudget_SpillFlushReleasesBatchAggregate(t *testing.T) { + prom := prometheus.NewRegistry() + spool, err := sharedndjson.New(t.TempDir(), "gw", 64, 1, nil) + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + w := NewWriter(nil, "q", opsmetrics.NewRegistry(prom), + slog.New(slog.NewTextHandler(io.Discard, nil))). + WithNDJSONSpill(spool).WithLossMode(lossModeSpillBlock) + tinyBudgetWriter(w, 1<<20) + w.spillFlush.targetBytes.Store(1) // flush on the first record + w.Start() + defer w.Close() + + // Reserve exactly as Enqueue would, then hand the record straight to the + // spill worker (the overflow hand-off path). + rec := bodyRecord("sp-1", 4096) + if !w.memBudget.Acquire(4096) { + t.Fatal("Acquire on an empty budget must admit") + } + rec.reserved = 4096 + w.spillCh <- rec + + if !waitFor(2*time.Second, func() bool { + return counterValue(t, prom, "nexus_audit_mq_spilled_total") >= 1 + }) { + t.Fatal("spill worker did not durably write the record") + } + if !waitFor(2*time.Second, func() bool { return w.memBudget.InUse() == 0 }) { + t.Fatalf("InUse after spill flush = %d, want 0 (batch aggregate not released)", w.memBudget.InUse()) + } + if got := counterValue(t, prom, "nexus_audit_mq_dropped_total"); got != 0 { + t.Fatalf("dropped_total = %v, want 0", got) + } +} + +// Shutdown must unblock a producer parked on the byte budget and still lose +// nothing: the record that could not be admitted spills durably instead. +func TestMemBudget_ShutdownUnblocksParkedEnqueueIntoDurableSpill(t *testing.T) { + prom := prometheus.NewRegistry() + spool, err := sharedndjson.New(t.TempDir(), "gw", 64, 1, nil) + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + prod := &stallProducer{gate: make(chan struct{})} + w := NewWriter(prod, "q", opsmetrics.NewRegistry(prom), + slog.New(slog.NewTextHandler(io.Discard, nil))). + WithNDJSONSpill(spool).WithLossMode(lossModeBlock) + tinyBudgetWriter(w, 64) + w.Start() + + w.Enqueue(bodyRecord("holder", 64)) // fills the budget; publish stalls + + parked := make(chan struct{}) + go func() { + w.Enqueue(bodyRecord("straggler", 64)) // parks on the full budget + close(parked) + }() + select { + case <-parked: + t.Fatal("Enqueue admitted past an exhausted budget") + case <-time.After(100 * time.Millisecond): + } + + closed := make(chan struct{}) + go func() { w.Close(); close(closed) }() + select { + case <-parked: // stopCh fired → Acquire returns false → durable spill, Enqueue returns + case <-time.After(2 * time.Second): + t.Fatal("shutdown did not unblock the parked Enqueue") + } + close(prod.gate) // let the in-flight holder publish so Close can finish draining + select { + case <-closed: + case <-time.After(5 * time.Second): + t.Fatal("Close did not complete after the gate opened") + } + if got := counterValue(t, prom, "nexus_audit_mq_spilled_total"); got < 1 { + t.Fatalf("spilled_total = %v, want >= 1 (the straggler must spill durably, not vanish)", got) + } + if got := counterValue(t, prom, "nexus_audit_mq_dropped_total"); got != 0 { + t.Fatalf("dropped_total = %v, want 0 (shutdown must not lose the straggler)", got) + } + if got := w.memBudget.InUse(); got != 0 { + t.Fatalf("InUse after shutdown = %d, want 0", got) + } +} + +// handlePublishFailure terminals must release the reservation exactly once — +// each of its four resolutions is driven directly on an UNSTARTED writer (nil +// recCh → the in-memory re-queue is never ready, so the routing is +// deterministic) with a pre-reserved record, then the budget must read 0. +func TestMemBudget_PublishFailureTerminalsRelease(t *testing.T) { + newRec := func(w *Writer) *Record { + rec := bodyRecord("pf", 128) + if !w.memBudget.Acquire(128) { + t.Fatal("Acquire on an empty budget must admit") + } + rec.reserved = 128 + rec.publishRetries = maxPublishRetries // next failure exhausts the in-memory retry + return rec + } + + t.Run("no spool: counted drop releases", func(t *testing.T) { + w := NewWriter(nil, "q", nil, + slog.New(slog.NewTextHandler(io.Discard, nil))).WithLossMode(lossModeDrop) + tinyBudgetWriter(w, 1<<20) + rec := newRec(w) + w.handlePublishFailure([]byte("{}"), rec) + if got := w.memBudget.InUse(); got != 0 { + t.Fatalf("InUse after nil-spool drop = %d, want 0", got) + } + }) + + t.Run("spillCh hand-off keeps the reservation", func(t *testing.T) { + w := NewWriter(nil, "q", nil, + slog.New(slog.NewTextHandler(io.Discard, nil))).WithLossMode(lossModeSpill) + spool, err := sharedndjson.New(t.TempDir(), "gw", 64, 1, nil) + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + w.WithNDJSONSpill(spool) + rec := newRec(w) + w.handlePublishFailure([]byte("{}"), rec) + // The record now sits in spillCh with its reservation intact — the spill + // worker's flush is the terminal, not this hand-off. + if got := w.memBudget.InUse(); got != 128 { + t.Fatalf("InUse after spillCh hand-off = %d, want 128 (still pinned)", got) + } + if rec.reserved != 128 { + t.Fatalf("reserved after hand-off = %d, want 128", rec.reserved) + } + }) + + t.Run("saturated spill worker, lossy: durable spillData releases", func(t *testing.T) { + w := NewWriter(nil, "q", nil, + slog.New(slog.NewTextHandler(io.Discard, nil))).WithLossMode(lossModeSpill) + spool, err := sharedndjson.New(t.TempDir(), "gw", 64, 1, nil) + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + w.WithNDJSONSpill(spool) + w.spillCh = make(chan *Record) // unbuffered + no worker = saturated hand-off + rec := newRec(w) + w.handlePublishFailure([]byte("{}"), rec) + if got := w.memBudget.InUse(); got != 0 { + t.Fatalf("InUse after lossy spillData terminal = %d, want 0", got) + } + }) + + t.Run("saturated spill worker, spillBlock at shutdown: one-shot spill releases", func(t *testing.T) { + w := NewWriter(nil, "q", nil, + slog.New(slog.NewTextHandler(io.Discard, nil))).WithLossMode(lossModeSpillBlock) + spool, err := sharedndjson.New(t.TempDir(), "gw", 64, 1, nil) + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + w.WithNDJSONSpill(spool) + w.spillCh = make(chan *Record) // saturated park target + close(w.stopCh) // shutdown escapes the park + rec := newRec(w) + w.handlePublishFailure([]byte("{}"), rec) + if got := w.memBudget.InUse(); got != 0 { + t.Fatalf("InUse after spillBlock shutdown terminal = %d, want 0", got) + } + }) +} + +// Records with no captured bodies never touch the budget: even a 1-byte budget +// admits them without blocking (the bound is on body bytes, not record count). +func TestMemBudget_BodylessRecordsBypassAccounting(t *testing.T) { + prod := &memProducer{} + w := NewWriter(prod, "q", nil, slog.New(slog.NewTextHandler(io.Discard, nil))).WithLossMode(lossModeBlock) + tinyBudgetWriter(w, 1) + for i := range 10 { + w.Enqueue(&Record{RequestID: fmt.Sprintf("meta-%d", i)}) + } + if got := waitForMsgCount(prod, 10, 2*time.Second); got != 10 { + t.Fatalf("published %d of 10 bodyless records", got) + } + w.Close() + if got := w.memBudget.InUse(); got != 0 { + t.Fatalf("InUse = %d, want 0 (bodyless records must not reserve)", got) + } +} diff --git a/packages/ai-gateway/internal/platform/audit/publish_failure_bound_test.go b/packages/ai-gateway/internal/platform/audit/publish_failure_bound_test.go new file mode 100644 index 00000000..57336d93 --- /dev/null +++ b/packages/ai-gateway/internal/platform/audit/publish_failure_bound_test.go @@ -0,0 +1,85 @@ +package audit + +import ( + "fmt" + "log/slog" + "testing" + "time" + + sharedndjson "github.com/AlphaBitCore/nexus-gateway/packages/shared/audit/ndjson" +) + +// TestPublishFailure_SustainedOutage_BoundedNoCirculation pins the death-spiral +// fix. Under a SUSTAINED publish outage (every publish fails, e.g. a full NATS +// stream returning 503) with the consumer workers actively draining recCh, a naive +// re-queue always wins a queue slot, so the failed record would circulate on recCh +// FOREVER — pinning its marshaled copy and busy-spinning the workers (the observed +// 226MB-heap + CPU-thrash collapse). The fix bounds the in-memory retry to +// maxPublishRetries, then routes the record to the durable BATCHED spill (spillCh). +// This test proves the loop is bounded: every record lands durably on the spool +// within a bounded time (never lost, never circulating), and the failing producer +// accepts nothing. +func TestPublishFailure_SustainedOutage_BoundedNoCirculation(t *testing.T) { + dir := t.TempDir() + // Generous per-file + total quota so the outage records all fit durably. + spill, err := sharedndjson.New(dir, "test", 4096, 64<<20, nil) + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + prod := &memProducer{alwaysFail: true} + w := NewWriter(prod, "nexus.event.ai-traffic", nil, slog.Default()). + WithNDJSONSpill(spill).Start() + defer w.Close() + + const n = 50 + for i := range n { + w.Enqueue(&Record{RequestID: fmt.Sprintf("stuck-%d", i), Timestamp: time.Now().Add(time.Duration(i))}) + } + + // Every record must land durably on the spool within a bounded time — proving the + // bounded retry exits to the batched spill instead of circulating on recCh forever. + deadline := time.Now().Add(5 * time.Second) + for len(readSpool(t, dir, "test")) < n { + if time.Now().After(deadline) { + t.Fatalf("sustained-outage records must all spill durably (bounded circulation); "+ + "got %d, want %d", len(readSpool(t, dir, "test")), n) + } + time.Sleep(20 * time.Millisecond) + } + if got := len(prod.msgs()); got != 0 { + t.Errorf("alwaysFail producer must not record messages; got %d", got) + } +} + +// TestPublishFailure_TransientRetriesInMemory pins that a GENUINE transient failure +// (recovers within the retry cap) is still recovered IN MEMORY via the recCh +// re-queue — the fast-path preserved by the bounded-retry fix, not spilled to disk. +// A producer that fails a bounded number of times then succeeds must surface the +// record on the producer (re-published from recCh), with nothing on the spool. +func TestPublishFailure_TransientRetriesInMemory(t *testing.T) { + dir := t.TempDir() + spill, err := sharedndjson.New(dir, "test", 4096, 64<<20, nil) + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + // Fail the first 2 publishes (< maxPublishRetries=3), then accept — a transient + // blip that the in-memory re-queue must ride out without touching the spool. + prod := &memProducer{failCount: 2} + w := NewWriter(prod, "nexus.event.ai-traffic", nil, slog.Default()). + WithNDJSONSpill(spill).Start() + defer w.Close() + + w.Enqueue(&Record{RequestID: "transient", Timestamp: time.Now()}) + + deadline := time.Now().Add(5 * time.Second) + for len(prod.msgs()) < 1 { + if time.Now().After(deadline) { + t.Fatalf("transient failure must recover in-memory and publish; got %d messages", len(prod.msgs())) + } + time.Sleep(20 * time.Millisecond) + } + // A transient blip within the cap recovers in memory — it must NOT have spilled. + if got := len(readSpool(t, dir, "test")); got != 0 { + t.Errorf("transient (in-cap) failure must recover in-memory, not spill; spooled %d", got) + } +} diff --git a/packages/ai-gateway/internal/platform/audit/publish_failure_extra_test.go b/packages/ai-gateway/internal/platform/audit/publish_failure_extra_test.go new file mode 100644 index 00000000..302f9608 --- /dev/null +++ b/packages/ai-gateway/internal/platform/audit/publish_failure_extra_test.go @@ -0,0 +1,175 @@ +package audit + +import ( + "context" + "errors" + "fmt" + "log/slog" + "testing" + "time" + + sharedndjson "github.com/AlphaBitCore/nexus-gateway/packages/shared/audit/ndjson" +) + +// errFramePublish is the injected publish failure for the framed-path tests. +var errFramePublish = errors.New("frame publish boom") + +// failBatchProducer is a pooled batchProducer whose every frame publish fails — +// exercises the FRAMED failure path (publishFramed routes every record in a failed +// frame through handlePublishFailure), the path the memProducer (per-record) tests +// do not cover. +type failBatchProducer struct{} + +func (failBatchProducer) Publish(context.Context, string, []byte) error { return errFramePublish } +func (failBatchProducer) Enqueue(context.Context, string, []byte) error { return errFramePublish } +func (failBatchProducer) Close() error { return nil } +func (failBatchProducer) PoolSize() int { return 1 } + +func (failBatchProducer) EnqueueBatchAsync(_ context.Context, _ string, b [][]byte) ([]error, error) { + errs := make([]error, len(b)) + for i := range errs { + errs[i] = errFramePublish + } + return errs, nil +} + +func (failBatchProducer) EnqueueBatchAsyncOn(_ context.Context, _ string, b [][]byte, _ int) ([]error, error) { + errs := make([]error, len(b)) + for i := range errs { + errs[i] = errFramePublish + } + return errs, nil +} + +// TestPublishFailure_SpillBlockNoSpool_DowngradesToBlock pins the no-loss fix for +// the block→spillBlock default change: spillBlock uses the durable spool as its +// overflow buffer, so WITHOUT a spool wired (empty spoolDir OR a spool-dir creation +// failure) it must downgrade to block at Start — the stricter no-loss mode that +// back-pressures at the in-heap queue and needs no spool — rather than dropping on +// the first overflow (which spillOverflow's no-sink branch would do). +func TestPublishFailure_SpillBlockNoSpool_DowngradesToBlock(t *testing.T) { + // spillBlock selected, but NO WithNDJSONSpill → ndjsonSpill stays nil. + w := NewWriter(&memProducer{}, "nexus.event.ai-traffic", nil, slog.Default()). + WithLossMode(lossModeSpillBlock) + if got := w.LossMode(); got != lossModeSpillBlock { + t.Fatalf("pre-Start lossMode=%q want %q", got, lossModeSpillBlock) + } + w.Start() + defer w.Close() + if got := w.LossMode(); got != lossModeBlock { + t.Errorf("spillBlock without a spool must downgrade to block at Start (no-loss "+ + "back-pressure, not first-overflow drop); got %q", got) + } +} + +// TestPublishFailure_SpillBlockWithSpool_Kept verifies the downgrade is conditional: +// with a spool wired, spillBlock is kept (the spool is its no-loss overflow buffer). +func TestPublishFailure_SpillBlockWithSpool_Kept(t *testing.T) { + dir := t.TempDir() + spill, err := sharedndjson.New(dir, "test", 4096, 64<<20, nil) + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + w := NewWriter(&memProducer{}, "nexus.event.ai-traffic", nil, slog.Default()). + WithLossMode(lossModeSpillBlock).WithNDJSONSpill(spill).Start() + defer w.Close() + if got := w.LossMode(); got != lossModeSpillBlock { + t.Errorf("spillBlock with a spool wired must be kept; got %q", got) + } +} + +// TestPublishFailure_FramedPath_Bounded covers the FRAMED failure path: a batch +// producer whose every frame publish fails routes every record through +// handlePublishFailure, which must bound the retry and spill durably — no infinite +// circulation, no loss — exactly like the per-record path. +func TestPublishFailure_FramedPath_Bounded(t *testing.T) { + dir := t.TempDir() + spill, err := sharedndjson.New(dir, "test", 4096, 64<<20, nil) + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + w := NewWriter(failBatchProducer{}, "nexus.event.ai-traffic", nil, slog.Default()). + WithFramePublish(256 << 10). // force the framed publish path + WithNDJSONSpill(spill).Start() + defer w.Close() + + const n = 50 + for i := range n { + w.Enqueue(&Record{RequestID: fmt.Sprintf("frame-stuck-%d", i), Timestamp: time.Now().Add(time.Duration(i))}) + } + + deadline := time.Now().Add(5 * time.Second) + for len(readSpool(t, dir, "test")) < n { + if time.Now().After(deadline) { + t.Fatalf("framed sustained-outage records must all spill durably (bounded); got %d, want %d", + len(readSpool(t, dir, "test")), n) + } + time.Sleep(20 * time.Millisecond) + } +} + +// BenchmarkHandlePublishFailure_SpillRouting proves the wedged-outage routing does +// NOT become a new bottleneck (the user's explicit constraint: the fix must not turn +// into a choke-point). It drives handlePublishFailure from N parallel workers with a +// spool PRE-POPULATED with many sealed files (so the per-record spillData path's +// O(spool-files) dirSize cost is realistic), comparing the two routings: +// +// spillch — the fix: cap-exhausted records hand off to the batched async spill +// worker (spillCh → WriteBatch, one dirSize per batch, off the caller). +// spilldata — the pre-fix per-record path: synchronous ndjson.Write under a single +// mutex + a dirSize (readDir + stat every spool file) PER record. +// +// The spillch arm should stay flat as -cpu (contention) rises while the spilldata +// arm degrades — the empirical "not a bottleneck" evidence the design calls for. +// +// go test -run x -bench BenchmarkHandlePublishFailure_SpillRouting -cpu 1,8,32 \ +// -benchmem ./internal/platform/store/... (audit pkg) +func BenchmarkHandlePublishFailure_SpillRouting(b *testing.B) { + mkWriter := func(b *testing.B) (*Writer, string) { + dir := b.TempDir() + // Small per-file cap so a modest volume seals MANY files → realistic dirSize cost. + spill, err := sharedndjson.New(dir, "bench", 4096, 1<<30, nil) + if err != nil { + b.Fatalf("ndjson.New: %v", err) + } + // Pre-seal ~150 files so dirSize (readDir + stat-per-file) is non-trivial. + for i := range 150 * 4096 / 64 { + _ = spill.Write([]byte(fmt.Sprintf("seed-%d\n", i))) + } + w := NewWriter(&memProducer{}, "nexus.event.ai-traffic", nil, + slog.New(slog.NewTextHandler(discardWriter{}, nil))). + WithNDJSONSpill(spill).Start() + return w, dir + } + + b.Run("spillch", func(b *testing.B) { + w, _ := mkWriter(b) + defer w.Close() + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + rec := &Record{RequestID: "x", Timestamp: time.Now(), publishRetries: maxPublishRetries + 1} + rec.marshaled = []byte(`{"requestId":"x"}`) + w.handlePublishFailure(rec.marshaled, rec) // cap already exceeded → spillCh + } + }) + }) + + b.Run("spilldata", func(b *testing.B) { + w, _ := mkWriter(b) + defer w.Close() + data := []byte(`{"requestId":"x"}`) + b.ReportAllocs() + b.ResetTimer() + b.RunParallel(func(pb *testing.PB) { + for pb.Next() { + w.spillData(data) // the pre-fix per-record synchronous path + } + }) + }) +} + +type discardWriter struct{} + +func (discardWriter) Write(p []byte) (int, error) { return len(p), nil } diff --git a/packages/ai-gateway/internal/platform/audit/record.go b/packages/ai-gateway/internal/platform/audit/record.go index ae17f02b..8f540e7f 100644 --- a/packages/ai-gateway/internal/platform/audit/record.go +++ b/packages/ai-gateway/internal/platform/audit/record.go @@ -315,6 +315,26 @@ type Record struct { // In-process only — never serialized. marshaled []byte + // publishRetries counts how many times handlePublishFailure has re-queued this + // record onto recCh for an in-memory re-publish. It bounds the transient retry: + // under a SUSTAINED broker outage (stream full → every publish 503s) the workers + // keep draining recCh, so a naive re-queue always wins a slot and the record + // circulates FOREVER — a busy-spin that pins its marshaled copy in the queue and + // craters throughput. After maxPublishRetries in-memory attempts the record is + // routed to the durable batched spill instead of re-queued. In-process only — + // never serialized. + publishRetries int + + // reserved is the byte weight this record holds against the Writer's in-memory + // byte budget (Writer.memBudget): Acquired at Enqueue admission from the REAL + // captured body sizes, Released exactly once when the record leaves the + // in-memory pipeline — published OK, durably spilled, or dropped — via + // releaseRecordMem (or transferred to the spill worker's batch aggregate, which + // releases at its flush terminal). 0 when the record carried no captured bodies, + // the budget is disabled, or the reservation has already been released. + // In-process only — never serialized. + reserved int64 + // RequestContentType / ResponseContentType travel with the captured // bytes onto traffic_event_payload.{request,response}_content_type. // Empty when not detected; consumers default to inferring from the diff --git a/packages/ai-gateway/internal/platform/audit/record_bodypool.go b/packages/ai-gateway/internal/platform/audit/record_bodypool.go index 8ecca037..78d5cf4e 100644 --- a/packages/ai-gateway/internal/platform/audit/record_bodypool.go +++ b/packages/ai-gateway/internal/platform/audit/record_bodypool.go @@ -50,6 +50,14 @@ func (r *Record) AttachPooledRequestBody(handle *[]byte) { r.reqBodyHandle = handle } +// ReleaseRequestBuffer returns an UNATTACHED request-body buffer to the pool. +// This is the not-captured counterpart of the writer's terminal reclaim: when +// payload capture is off, the handle is never attached to a Record, so no +// terminal reclaim runs — the request owner returns the buffer itself once the +// body's last reader (upstream forward, hooks, lazy normalize) has finished. +// The caller must not read the buffer afterwards. Nil-safe. +func ReleaseRequestBuffer(hp *[]byte) { releaseRequestBody(hp) } + // releaseRequestBody returns a body buffer to the pool unless it ballooned past // the cap (dropped to GC instead). func releaseRequestBody(hp *[]byte) { diff --git a/packages/ai-gateway/internal/platform/audit/record_bodypool_release_test.go b/packages/ai-gateway/internal/platform/audit/record_bodypool_release_test.go new file mode 100644 index 00000000..e2fd1e39 --- /dev/null +++ b/packages/ai-gateway/internal/platform/audit/record_bodypool_release_test.go @@ -0,0 +1,61 @@ +package audit + +import "testing" + +// TestReleaseRequestBuffer_SteadyStateZeroAlloc asserts the business property +// the release path exists for: when the caller returns unattached body +// buffers (bodies-off / not-captured requests), the acquire→release cycle +// reuses pooled memory instead of allocating a fresh buffer per request. +// Without the release, every iteration allocates. +func TestReleaseRequestBuffer_SteadyStateZeroAlloc(t *testing.T) { + src := make([]byte, 2048) + for i := range src { + src[i] = byte(i) + } + // Warm the pool so the measurement window is steady-state. + for range 16 { + _, h := AcquireRequestBody(src) + ReleaseRequestBuffer(h) + } + avg := testing.AllocsPerRun(200, func() { + body, h := AcquireRequestBody(src) + if len(body) != len(src) { + t.Fatalf("body len %d != src len %d", len(body), len(src)) + } + ReleaseRequestBuffer(h) + }) + // Steady state must not allocate a fresh buffer per cycle. Allow <1 to + // absorb a rare GC clearing the pool mid-run; 1+ means the release path + // is broken and every request pays a fresh allocation again. + if avg >= 1 { + t.Fatalf("acquire/release cycle allocates %.1f allocs/op; pool reuse broken", avg) + } +} + +// TestReleaseRequestBuffer_NilSafe asserts the not-acquired path (empty body +// → nil handle) is a no-op, matching how finalizeAudit calls it on requests +// that never read a body. +func TestReleaseRequestBuffer_NilSafe(t *testing.T) { + ReleaseRequestBuffer(nil) // must not panic + body, h := AcquireRequestBody(nil) + if body != nil || h != nil { + t.Fatalf("empty src must yield nil body+handle, got %v %v", body, h) + } + ReleaseRequestBuffer(h) +} + +// TestReleaseRequestBuffer_OversizedDropped asserts a buffer grown past the +// pool cap is dropped to GC rather than poisoning the pool with a huge +// backing array (same contract as the writer-side terminal reclaim). +func TestReleaseRequestBuffer_OversizedDropped(t *testing.T) { + big := make([]byte, requestBodyPoolCap+1) + _, h := AcquireRequestBody(big) + if h == nil { + t.Fatal("expected a handle for a non-empty body") + } + if cap(*h) <= requestBodyPoolCap { + t.Fatalf("test setup: expected regrown buffer > cap, got %d", cap(*h)) + } + ReleaseRequestBuffer(h) // must drop, not pool — nothing to assert beyond no panic; + // the pooling decision is releaseRequestBody's existing tested contract. +} diff --git a/packages/ai-gateway/internal/platform/audit/spill_recovery.go b/packages/ai-gateway/internal/platform/audit/spill_recovery.go index 0828dcf5..fa253715 100644 --- a/packages/ai-gateway/internal/platform/audit/spill_recovery.go +++ b/packages/ai-gateway/internal/platform/audit/spill_recovery.go @@ -87,6 +87,16 @@ type spillRecovery struct { // leaves its file for a future pass, the pre-recovery behaviour). maxRecordBytes int + // maxPayload, when non-nil, re-reads the broker's negotiated max_payload so + // runOnce can populate maxRecordBytes on a LATER pass if the broker was not yet + // connected when the sweeper started (MaxPayload returns 0 pre-connect). Without + // this refresh an oversize record spooled during that window would never be + // dead-lettered and its sealed file would never drain — permanently consuming the + // reclaimable spool quota, which HARD-WEDGES audit intake once the writer + // back-pressures on that quota (spillBlock). Read once per sweep from the sole run + // goroutine, so no lock is needed. + maxPayload func() int64 + // onReingested / onError / onPoisoned are metric hooks; nil-safe via helpers. onReingested func(records int) onError func() @@ -175,6 +185,14 @@ func (r *spillRecovery) run(ctx context.Context, base time.Duration) { // dir exercises the full read→frame→publish→delete path without a goroutine or a // live broker. func (r *spillRecovery) runOnce(ctx context.Context) int { + // Late-bind the broker max_payload: if it was unknown at start (NATS not yet + // connected → MaxPayload()==0), pick it up now so oversize records start being + // dead-lettered to .poison instead of wedging their file (and the quota) forever. + if r.maxRecordBytes == 0 && r.maxPayload != nil { + if mp := r.maxPayload(); mp > maxPayloadMargin { + r.maxRecordBytes = int(mp - maxPayloadMargin) + } + } if err := r.src.Rotate(); err != nil { r.logf("audit: spill recovery rotate failed", "error", err) // Continue: already-sealed files from earlier rotations are still drainable. diff --git a/packages/ai-gateway/internal/platform/audit/writer.go b/packages/ai-gateway/internal/platform/audit/writer.go index a86b8609..5e19a59e 100644 --- a/packages/ai-gateway/internal/platform/audit/writer.go +++ b/packages/ai-gateway/internal/platform/audit/writer.go @@ -10,6 +10,7 @@ import ( "time" sharedndjson "github.com/AlphaBitCore/nexus-gateway/packages/shared/audit/ndjson" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/core/bytebudget" opsmetrics "github.com/AlphaBitCore/nexus-gateway/packages/shared/core/metrics/registry" "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/payloadcapture" "github.com/AlphaBitCore/nexus-gateway/packages/shared/storage/spillstore" @@ -87,8 +88,9 @@ type Writer struct { // the producer; N publishWorkers are the consumers. This IS the standard // bounded producer/multi-consumer pattern: a blocking send (block mode) is the // no-loss back-pressure, a non-blocking send (spill/drop modes) is the lossy - // opt-out. Each queued record pins its pooled ~50 KB body until a worker - // marshals it, so the cap bounds the audit body-pool working set. + // opt-out. The cap bounds POINTER count only (a guard against a flood of tiny + // records); the memory each queued record pins is bounded by memBudget, which + // accounts the REAL body bytes at Enqueue admission. recCh chan *Record // maxQueued is recCh's capacity (the bounded-queue depth). Defaults to @@ -147,6 +149,18 @@ type Writer struct { // spillFlush adapts the spill worker's flush size to measured disk-write // latency (replacing a fixed flush-bytes magic number). Set in NewWriter. spillFlush *adaptiveSpillFlush + + // memBudget bounds the BYTES the in-memory audit pipeline may pin (captured + // request/response bodies and the marshaled copies that supersede them), + // reserved at Enqueue admission and released exactly once per record at its + // terminal (published / durably spilled / dropped). The structural channel + // caps bound POINTER counts only; this is the real memory bound — sized from + // available RAM (adaptiveMemBudgetBytes), soft and lock-free (see + // shared/core/bytebudget). No-loss modes block at Enqueue on a full budget + // (back-pressure to the request path); lossy modes shed with a counted drop + // instead of blocking. This replaced the count-derived sizing that assumed an + // average body size and OOM'd when real bodies exceeded it. + memBudget *bytebudget.Budget } // NewWriter creates an audit writer that publishes to the given MQ producer. @@ -166,11 +180,15 @@ func NewWriter(producer mq.Producer, queue string, reg *opsmetrics.Registry, log // choice. recChCap, spillChCap := adaptiveBufferCaps() w := &Writer{ - producer: producer, - queue: queue, - logger: logger, - metrics: newAuditMetrics(reg), - maxQueued: recChCap, + producer: producer, + queue: queue, + logger: logger, + metrics: newAuditMetrics(reg), + maxQueued: recChCap, + // Pre-config default is the strictest no-loss mode (back-pressure, no spool + // dependency); WithLossMode applies the deployment default (spillBlock — spool + // as primary buffer) from AuditConfig.LossMode, whose empty/unknown fallback is + // spillBlock. A raw constructor without WithLossMode therefore stays safest. lossMode: lossModeBlock, workers: workers, spillCh: make(chan *Record, spillChCap), @@ -182,6 +200,9 @@ func NewWriter(producer mq.Producer, queue string, reg *opsmetrics.Registry, log // out (the legacy path, kept until it is deleted). wireBinary: !strings.EqualFold(strings.TrimSpace(os.Getenv("NEXUS_AUDIT_WIRE")), "json"), } + // The byte budget shares stopCh so an Enqueue blocked on a full budget always + // unblocks at shutdown (it then spills durably instead of enqueuing). + w.memBudget = bytebudget.New(adaptiveMemBudgetBytes(), w.stopCh) return w } @@ -207,7 +228,7 @@ func (w *Writer) WithPayloadCaptureStore(s *payloadcapture.Store) *Writer { // human-readable name onto every TrafficEventMessage. Persisted as // traffic_event.thing_id / thing_name. Returns the receiver for chaining. // -// Must be called during startup before any Enqueue / flushLoop runs; +// Must be called during startup before any Enqueue / consumer worker runs; // mutates w.thingID / w.thingName without a lock, matching the // WithSpillStore / WithPayloadCaptureStore startup-only convention. func (w *Writer) WithThingIdentity(id, name string) *Writer { diff --git a/packages/ai-gateway/internal/platform/audit/writer_batch.go b/packages/ai-gateway/internal/platform/audit/writer_batch.go index ab6fb8ec..ab565e99 100644 --- a/packages/ai-gateway/internal/platform/audit/writer_batch.go +++ b/packages/ai-gateway/internal/platform/audit/writer_batch.go @@ -167,9 +167,8 @@ func (w *Writer) batchPublish(bp batchProducer, connIdx int, ctx context.Context const ( // batchPublishTimeout bounds how long ONE chunk waits for its async batch to // be acked before treating still-unresolved records as failed - // (re-buffer/spill). Kept under closeShutdownDeadline so Close's drainBuffer - // loop (which checks the wall only between flushes) cannot overrun even - // across two chunk flushes. + // (re-buffer/spill), so a wedged broker cannot pin a publish worker — or the + // shutdown drain that waits on it — indefinitely. batchPublishTimeout = 5 * time.Second ) @@ -344,11 +343,12 @@ func (w *Writer) publishFramed(bp batchProducer, connIdx int, datas [][]byte, re } continue } - for range fr { + for _, rec := range fr { w.metrics.incEnqueueTotal() // Terminal: published OK. The pooled body was already reclaimed at - // marshal time (the frame copied the marshaled bytes), so nothing to - // reclaim here. + // marshal time (the frame copied the marshaled bytes); the byte-budget + // reservation returns here, once the bytes are on the wire. + w.releaseRecordMem(rec) } } } @@ -379,7 +379,9 @@ func (w *Writer) publishChunk(bp batchProducer, connIdx int, datas [][]byte, rec } w.metrics.incEnqueueTotal() // Terminal: published OK. The pooled body was already reclaimed at marshal - // time (EnqueueBatchAsync took the bytes), so nothing to reclaim here. + // time (EnqueueBatchAsync took the bytes); the byte-budget reservation + // returns here, once the bytes are on the wire. + w.releaseRecordMem(recs[j]) } } @@ -424,6 +426,7 @@ func (w *Writer) marshalRecord(rec *Record) (data []byte, buf *bytes.Buffer, ok w.logger.Error("audit: marshal failed", "requestId", rec.RequestID, "error", err) reclaimMsgBuf(enc) w.reclaimRecordBody(rec) // dropped record: return its pooled body to the pool + w.releaseRecordMem(rec) // terminal: the record leaves the pipeline here return nil, nil, false } // json.Encoder appends a trailing '\n'; drop it so each record is a single diff --git a/packages/ai-gateway/internal/platform/audit/writer_concurrency_test.go b/packages/ai-gateway/internal/platform/audit/writer_concurrency_test.go index 1c97bee6..5f1d78a2 100644 --- a/packages/ai-gateway/internal/platform/audit/writer_concurrency_test.go +++ b/packages/ai-gateway/internal/platform/audit/writer_concurrency_test.go @@ -36,7 +36,7 @@ func TestWriter_BurstDrainsPromptly(t *testing.T) { w := NewWriter(prod, "q", nil, slog.Default()) defer w.Close() - const burst = flushHighWater + const burst = 1000 for i := range burst { w.Enqueue(&Record{RequestID: fmt.Sprintf("r%d", i)}) } @@ -160,19 +160,18 @@ func TestWriter_RequeueOnTransientFailure(t *testing.T) { // a small N back-pressures/spills at N, not the maxQueueSize default — the knob // that bounds the audit body pool's working set. n<=0 keeps the default. func TestWriter_WithMaxQueuedRecords_BoundsBuffer(t *testing.T) { - // Default: no override → effective cap is sized to available memory - // (adaptiveBufferCaps), so it lands within the adaptive clamp band, not a fixed - // constant. + // Default: no override → the FIXED structural pointer-count depth (the byte + // budget bounds memory; the count cap is body-size-independent by design). wDef := NewWriter(&memProducer{}, "q", nil, slog.Default()) def := wDef.effectiveMaxQueue() - if def < minRecChCap || def > maxRecChCap { - t.Fatalf("default effectiveMaxQueue = %d, want within adaptive band [%d,%d]", def, minRecChCap, maxRecChCap) + if def != recChStructuralCap { + t.Fatalf("default effectiveMaxQueue = %d, want structural cap %d", def, recChStructuralCap) } wDef.Close() - // n<=0 is ignored (keeps the adaptive default). + // n<=0 is ignored (keeps the structural default). if got := NewWriter(&memProducer{}, "q", nil, slog.Default()).WithMaxQueuedRecords(0).effectiveMaxQueue(); got != def { - t.Fatalf("WithMaxQueuedRecords(0) changed the cap to %d, want adaptive default %d", got, def) + t.Fatalf("WithMaxQueuedRecords(0) changed the cap to %d, want structural default %d", got, def) } // A small override sizes the bounded queue at N: Start() sizes recCh from the diff --git a/packages/ai-gateway/internal/platform/audit/writer_lifecycle.go b/packages/ai-gateway/internal/platform/audit/writer_lifecycle.go index 9688f96b..40cdb0ed 100644 --- a/packages/ai-gateway/internal/platform/audit/writer_lifecycle.go +++ b/packages/ai-gateway/internal/platform/audit/writer_lifecycle.go @@ -21,6 +21,20 @@ func (w *Writer) Start() *Writer { func (w *Writer) ensureStarted() { w.startOnce.Do(func() { + // spillBlock uses the durable on-disk spool as its overflow buffer and only + // back-pressures the request path when that spool ALSO saturates. Without a + // spool wired — an explicit empty spoolDir, OR a spool-dir creation failure at + // wiring time (which only logs and leaves ndjsonSpill nil) — spillOverflow has + // no durable sink and would DROP on the first overflow. Downgrade to block, the + // stricter no-loss mode that back-pressures at the in-heap queue and needs no + // spool, so audit is never silently lossy from a missing/failed spool. This is + // the documented "spillBlock without a spool falls back to block-style + // back-pressure" — enforced here rather than assumed. + if w.lossMode == lossModeSpillBlock && w.ndjsonSpill == nil { + w.logger.Warn("audit: lossMode=spillBlock but no durable spool is wired; " + + "downgrading to block (no-loss back-pressure) — wire a spool dir for spillBlock") + w.lossMode = lossModeBlock + } // Binary wire MUST always travel framed: an unframed (per-record) binary // message begins with field-id 1's uvarint (0x01), which is exactly the frame // magic — the Hub would mis-detect it as a multi-record binary frame. The @@ -76,6 +90,10 @@ func (w *Writer) startSpillRecovery() { // (subject + headers). Unknown (producer without the accessor) → 0 = no // proactive dead-letter. if mp, ok := w.producer.(interface{ MaxPayload() int64 }); ok { + // Wire the accessor so runOnce can late-bind maxRecordBytes if the broker is + // not connected yet at wiring time (MaxPayload()==0 pre-connect); also seed it + // now for the common case where NATS is already up. + r.maxPayload = mp.MaxPayload if max := mp.MaxPayload(); max > maxPayloadMargin { r.maxRecordBytes = int(max - maxPayloadMargin) } @@ -93,19 +111,27 @@ func (w *Writer) startSpillRecovery() { // Close stops accepting records and waits for the consumer + spill workers to // drain and publish/spill everything in flight. Nothing is lost: workers drain // recCh on stopCh, then the final sweep spills any straggler that raced in after -// their last drain check. +// their last drain check — from BOTH recCh and spillCh. spillCh must be drained +// too: after stopCh a producer's `spillCh <- rec` (the primary spillBlock +// back-pressure path) can land a record in the channel buffer AFTER the spill +// worker's final drain already returned; without this sweep that straggler is +// neither published, spilled, nor counted — a silent shutdown loss. func (w *Writer) Close() { close(w.stopCh) w.wg.Wait() if w.recCh == nil { return // never started } - for { - select { - case rec := <-w.recCh: - w.spillRecord(rec) - default: - return + drain := func(ch <-chan *Record) { + for { + select { + case rec := <-ch: + w.spillRecord(rec) + default: + return + } } } + drain(w.recCh) + drain(w.spillCh) } diff --git a/packages/ai-gateway/internal/platform/audit/writer_metrics.go b/packages/ai-gateway/internal/platform/audit/writer_metrics.go index 927830ed..696b270e 100644 --- a/packages/ai-gateway/internal/platform/audit/writer_metrics.go +++ b/packages/ai-gateway/internal/platform/audit/writer_metrics.go @@ -25,6 +25,21 @@ type auditMetrics struct { // value means audit data is durably retained but NOT in the queryable store — // an operator signal to fix the inline-body cap / enable out-of-band body spill. recoveryPoisoned *opsmetrics.CounterPin + // spillBackpressure counts how many times the spill worker entered a quota-full + // back-pressure retry (spool at its total-size quota under a no-loss mode). It is + // the DISTINCT "spool full → throttling the request path, recovery cannot keep up" + // signal: a rising value with dropped_total flat at 0 means the gateway is + // deliberately back-pressuring (not dropping, not hung) because NATS/Hub/PG drain + // < audit ingest. On-call uses it to tell an audit-spool wedge from an app hang. + spillBackpressure *opsmetrics.CounterPin + // memBackpressure counts how many times Enqueue found the in-memory byte + // budget exhausted. No-loss modes then BLOCK the request goroutine until drain + // frees bytes; lossy modes shed with a counted drop (this counter is what + // distinguishes a budget-full drop from a plain queue-full drop). Rising with + // dropped_total flat at 0 = the gateway is deliberately throttling ingest + // because in-memory audit bytes hit their RAM share — the designed overload + // behavior (bounded memory, no loss), not a hang. + memBackpressure *opsmetrics.CounterPin } func newAuditMetrics(reg *opsmetrics.Registry) *auditMetrics { @@ -35,13 +50,15 @@ func newAuditMetrics(reg *opsmetrics.Registry) *auditMetrics { // still applies; With() with zero values returns a CounterPin bound to // the empty label set. return &auditMetrics{ - enqueueTotal: reg.NewCounter("audit.mq_enqueue_total", nil).With(), - enqueueErrors: reg.NewCounter("audit.mq_enqueue_errors_total", nil).With(), - dropped: reg.NewCounter("audit.mq_dropped_total", nil).With(), - spilled: reg.NewCounter("audit.mq_spilled_total", nil).With(), - reingested: reg.NewCounter("audit.mq_reingested_total", nil).With(), - recoveryErrors: reg.NewCounter("audit.mq_recovery_errors_total", nil).With(), - recoveryPoisoned: reg.NewCounter("audit.mq_recovery_poisoned_total", nil).With(), + enqueueTotal: reg.NewCounter("audit.mq_enqueue_total", nil).With(), + enqueueErrors: reg.NewCounter("audit.mq_enqueue_errors_total", nil).With(), + dropped: reg.NewCounter("audit.mq_dropped_total", nil).With(), + spilled: reg.NewCounter("audit.mq_spilled_total", nil).With(), + reingested: reg.NewCounter("audit.mq_reingested_total", nil).With(), + recoveryErrors: reg.NewCounter("audit.mq_recovery_errors_total", nil).With(), + recoveryPoisoned: reg.NewCounter("audit.mq_recovery_poisoned_total", nil).With(), + spillBackpressure: reg.NewCounter("audit.mq_spill_backpressure_total", nil).With(), + memBackpressure: reg.NewCounter("audit.mem_backpressure_total", nil).With(), } } @@ -90,3 +107,13 @@ func (m *auditMetrics) addPoisoned(n int) { m.recoveryPoisoned.Add(float64(n)) } } +func (m *auditMetrics) incSpillBackpressure() { + if m != nil { + m.spillBackpressure.Inc() + } +} +func (m *auditMetrics) incMemBackpressure() { + if m != nil { + m.memBackpressure.Inc() + } +} diff --git a/packages/ai-gateway/internal/platform/audit/writer_publish.go b/packages/ai-gateway/internal/platform/audit/writer_publish.go index eda14d25..1bd06935 100644 --- a/packages/ai-gateway/internal/platform/audit/writer_publish.go +++ b/packages/ai-gateway/internal/platform/audit/writer_publish.go @@ -20,7 +20,8 @@ func (w *Writer) consumeLoop(connIdx int) { defer w.wg.Done() w.batchPathLogOnce.Do(func() { w.logger.Info("audit: bounded-queue consumer workers engaged", - "workers", w.workers, "queueCap", cap(w.recCh), "chunk", batchMaxCount) + "workers", w.workers, "queueCap", cap(w.recCh), "chunk", batchMaxCount, + "memBudgetBytes", w.memBudget.Budget()) }) batch := make([]*Record, 0, batchMaxCount) timer := time.NewTimer(time.Hour) @@ -104,6 +105,13 @@ func (w *Writer) drainOnStop(connIdx int, batch []*Record) { // failures route through handlePublishFailure (retry/spill, never silent loss). func (w *Writer) publishBatchOn(connIdx int, batch []*Record) { if w.producer == nil { + // No-op mode (tests / unwired producer): the batch is discarded — a + // terminal. Return the pooled bodies and byte-budget reservations, or the + // no-op Writer would permanently leak budget and stall its own Enqueue. + for _, rec := range batch { + w.reclaimRecordBody(rec) + w.releaseRecordMem(rec) + } return } bp, ok := w.producer.(batchProducer) @@ -115,6 +123,11 @@ func (w *Writer) publishBatchOn(connIdx int, batch []*Record) { } datas, bufs, recs := w.marshalChunkSerial(batch) if perfNoPublish { + // Perf-ablation terminal: records vanish here — release their reservations + // or the ablated pipeline back-pressures itself into a stall. + for _, rec := range recs { + w.releaseRecordMem(rec) + } reclaimMsgBufs(bufs) return } @@ -134,25 +147,82 @@ func (w *Writer) publishBatchOn(connIdx int, batch []*Record) { reclaimMsgBufs(bufs) } -// handlePublishFailure retries a record whose publish failed by re-queuing its -// already-marshaled bytes (non-blocking), or spilling them durably when the queue -// is full. The pooled body was reclaimed at marshal, so the retry unit is the bytes +// handlePublishFailure routes a record whose publish failed so it is never lost. +// The pooled body was reclaimed at marshal, so the retry unit is the bytes // (rec.marshaled), never a re-marshal. data aliases a pooled buffer the caller -// reclaims after this returns, so re-queue takes a copy. Never a silent loss. +// reclaims after this returns, so the re-queue/spill take a copy (rec.marshaled). +// +// Retry policy (BOUNDED — this is the death-spiral fix): +// - Up to maxPublishRetries in-memory re-queues onto recCh for a fast re-publish +// (transient NAK recovery). When recCh is full, a durable synchronous spill is +// the immediate fallback (the sustained-wedge no-loss contract). +// - PAST the retry cap the pipeline is wedged (a sustained stream-full outage: +// every publish 503s, workers keep draining recCh so a naive re-queue would win +// a slot and the record would circulate FOREVER — the busy-spin that pins the +// marshaled copy in the queue and craters throughput). The record is then handed +// to the durable BATCHED spill (spillCh → spillLoop → WriteBatch, off this +// publish worker; spillLoop's marshalRecord replays rec.marshaled verbatim), +// NOT re-queued and NOT per-record spillData (whose single mutex + O(spool-files) +// dirSize would itself bottleneck the overflow main path). lossMode is honoured +// on a saturated spill worker: spillBlock back-pressures this worker (which +// propagates to intake back-pressure via recCh saturation), else a last-resort +// synchronous spillData. Never a silent loss. func (w *Writer) handlePublishFailure(data []byte, rec *Record) { if rec.marshaled == nil { rec.marshaled = append([]byte(nil), data...) rec.RequestBody = nil // body already reclaimed; never re-read it rec.ResponseBody = nil } + rec.publishRetries++ + if rec.publishRetries <= maxPublishRetries { + select { + case w.recCh <- rec: // bounded in-memory retry (a worker re-publishes rec.marshaled) + return + default: + // Queue full: fall through to the durable batched spill hand-off (spillBlock + // back-pressures on a full spool; lossy modes drop), not a per-record spillData. + } + } + // Retry cap exhausted (or recCh full during retry) → hand off to the durable + // BATCHED spill worker instead of circulating on recCh or paying per-record + // spillData (whose single mutex + O(spool-files) dirSize would itself bottleneck). + // With no durable sink wired, the batched hand-off would only async-drop at the + // nil-sink flush; drop synchronously + counted instead (a spillBlock config + // without a spool is already downgraded to block at Start, so this branch is only + // reached in the lossy spill/drop modes or a no-spool block — all lossy by config). + if w.ndjsonSpill == nil { + w.metrics.incDropped() + w.releaseRecordMem(rec) // terminal: counted drop + return + } select { - case w.recCh <- rec: // re-queue for a retry (a worker re-publishes rec.marshaled) + case w.spillCh <- rec: + return default: - // Queue full: spill the marshaled bytes durably rather than block a worker. - if !w.spillData(data) { - w.metrics.incDropped() + } + if w.lossMode == lossModeSpillBlock { + // Spill worker saturated (writing at the disk rate, or holding a batch while it + // back-pressures a full spool quota). Park THIS publish worker on the channel + // until a slot frees — no per-goroutine ndjson.Write, so the spool stays + // single-writer. A parked publish worker stops draining recCh, so intake + // back-pressures via Enqueue. The worker never drops a full spool quota (it + // retries), so this resolves once recovery frees space; only shutdown escapes, + // with a one-shot durable spill (quota-full at shutdown → counted drop). + select { + case w.spillCh <- rec: + return + case <-w.stopCh: + if !w.spillData(data) { + w.metrics.incDropped() + } + w.releaseRecordMem(rec) // terminal: spilled durably or counted drop + return } } + if !w.spillData(data) { + w.metrics.incDropped() + } + w.releaseRecordMem(rec) // terminal: spilled durably or counted drop } // spillData writes already-marshaled wire bytes to the durable NDJSON fallback. @@ -204,5 +274,7 @@ func (w *Writer) publishRecord(rec *Record) { return } w.metrics.incEnqueueTotal() - // Terminal: published OK. The pooled body was already reclaimed at marshal time. + // Terminal: published OK. The pooled body was already reclaimed at marshal + // time; the byte-budget reservation returns here. + w.releaseRecordMem(rec) } diff --git a/packages/ai-gateway/internal/platform/audit/writer_quota_backpressure_test.go b/packages/ai-gateway/internal/platform/audit/writer_quota_backpressure_test.go new file mode 100644 index 00000000..671100cd --- /dev/null +++ b/packages/ai-gateway/internal/platform/audit/writer_quota_backpressure_test.go @@ -0,0 +1,147 @@ +package audit + +import ( + "io" + "log/slog" + "os" + "path/filepath" + "testing" + "time" + + sharedndjson "github.com/AlphaBitCore/nexus-gateway/packages/shared/audit/ndjson" + opsmetrics "github.com/AlphaBitCore/nexus-gateway/packages/shared/core/metrics/registry" + "github.com/prometheus/client_golang/prometheus" +) + +// waitFor polls cond until it is true or the deadline elapses, returning cond's +// final value — used to synchronise on the async spill worker without a fixed sleep. +func waitFor(d time.Duration, cond func() bool) bool { + deadline := time.Now().Add(d) + for time.Now().Before(deadline) { + if cond() { + return true + } + time.Sleep(5 * time.Millisecond) + } + return cond() +} + +// TestSpillBlock_QuotaFull_BackPressuresThenSpills is the core no-loss proof for the +// spool-quota back-pressure fix: under lossModeSpillBlock, a spool AT its total-size +// quota must NOT drop the record (as it did before — the measured 253778-drop bug); +// the spill worker back-pressures (retries the same batch, incrementing the distinct +// spill_backpressure signal) until the recovery drain frees space, then the record +// spills durably. Zero drops throughout. +func TestSpillBlock_QuotaFull_BackPressuresThenSpills(t *testing.T) { + prom := prometheus.NewRegistry() + dir := t.TempDir() + spool, err := sharedndjson.New(dir, "gw", 1, 1, nil) // 1 MB quota + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + // Seed a countable audit-*.ndjson file PAST the quota so the worker's first + // WriteBatch is refused with ErrSpoolQuotaExceeded. + seed := filepath.Join(dir, "gw", "audit-20260101-0001.ndjson") + if err := os.WriteFile(seed, make([]byte, 1100*1024), 0o600); err != nil { + t.Fatalf("seed: %v", err) + } + w := NewWriter(nil, "q", opsmetrics.NewRegistry(prom), slog.New(slog.NewTextHandler(io.Discard, nil))). + WithNDJSONSpill(spool).WithLossMode(lossModeSpillBlock) + w.spillFlush.targetBytes.Store(1) // flush on the first record + w.Start() + defer w.Close() + + w.spillCh <- &Record{RequestID: "bp-1", Timestamp: time.Unix(1700000000, 0).UTC()} + + // The worker hits the quota and back-pressures — a rising spill_backpressure + // signal, and crucially NO drop. + if !waitFor(2*time.Second, func() bool { + return counterValue(t, prom, "nexus_audit_mq_spill_backpressure_total") >= 1 + }) { + t.Fatal("spillBlock quota-full must raise the spill_backpressure signal") + } + if got := counterValue(t, prom, "nexus_audit_mq_dropped_total"); got != 0 { + t.Fatalf("quota-full under spillBlock must back-pressure, not drop; dropped=%v", got) + } + if got := counterValue(t, prom, "nexus_audit_mq_spilled_total"); got != 0 { + t.Fatalf("nothing may spill while the quota is full; spilled=%v", got) + } + + // Free the quota — the recovery drain's real-world effect. The worker's next + // retry now succeeds and the held record spills durably, still zero drops. + if err := os.Remove(seed); err != nil { + t.Fatalf("free quota: %v", err) + } + if !waitFor(3*time.Second, func() bool { + return counterValue(t, prom, "nexus_audit_mq_spilled_total") >= 1 + }) { + t.Fatal("record must spill durably once the quota frees") + } + if got := counterValue(t, prom, "nexus_audit_mq_dropped_total"); got != 0 { + t.Fatalf("no record may be dropped across the whole back-pressure→spill cycle; dropped=%v", got) + } + if got := len(readSpool(t, dir, "gw")); got < 1 { + t.Fatalf("the back-pressured record must be durably spooled after space frees; spooled=%d", got) + } +} + +// TestSpillBlock_QuotaFull_ShutdownDropsBounded pins the Close() escape: with the +// spool permanently at quota and no drain, the worker back-pressures — but Close() +// must still terminate in bounded time, converting the held batch to a single +// counted shutdown drop rather than hanging wg.Wait() forever. +func TestSpillBlock_QuotaFull_ShutdownDropsBounded(t *testing.T) { + prom := prometheus.NewRegistry() + dir := t.TempDir() + spool, err := sharedndjson.New(dir, "gw", 1, 1, nil) + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + seed := filepath.Join(dir, "gw", "audit-20260101-0001.ndjson") + if err := os.WriteFile(seed, make([]byte, 1100*1024), 0o600); err != nil { + t.Fatalf("seed: %v", err) + } + w := NewWriter(nil, "q", opsmetrics.NewRegistry(prom), slog.New(slog.NewTextHandler(io.Discard, nil))). + WithNDJSONSpill(spool).WithLossMode(lossModeSpillBlock) + w.spillFlush.targetBytes.Store(1) + w.Start() + w.spillCh <- &Record{RequestID: "sd-1", Timestamp: time.Unix(1700000000, 0).UTC()} + + // Let the worker enter back-pressure, then Close with the quota still full. + waitFor(2*time.Second, func() bool { + return counterValue(t, prom, "nexus_audit_mq_spill_backpressure_total") >= 1 + }) + done := make(chan struct{}) + go func() { w.Close(); close(done) }() + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("Close() hung with the spool at quota — the shutdown escape is missing") + } + if got := counterValue(t, prom, "nexus_audit_mq_dropped_total"); got < 1 { + t.Fatalf("a batch held at quota through shutdown must be a counted drop; dropped=%v", got) + } +} + +// TestSpillBlock_Close_DrainsSpillCh covers the shutdown-race sweep (N3): a record +// left in the spillCh buffer at Close (the primary spillBlock back-pressure path) +// must be drained to the durable spool, never silently lost. Deterministic: recCh is +// sized manually and the spill worker is NOT started, so wg.Wait() returns at once +// and the ONLY thing that can rescue the buffered record is Close's own spillCh +// drain — exactly the post-worker-exit race the sweep closes. +func TestSpillBlock_Close_DrainsSpillCh(t *testing.T) { + dir := t.TempDir() + spool, err := sharedndjson.New(dir, "gw", 64, 512, nil) + if err != nil { + t.Fatalf("ndjson.New: %v", err) + } + w := NewWriter(nil, "q", nil, slog.New(slog.NewTextHandler(io.Discard, nil))). + WithNDJSONSpill(spool).WithLossMode(lossModeSpillBlock) + w.recCh = make(chan *Record, 1) // non-nil so Close proceeds to the drain; no worker started + w.spillCh <- &Record{RequestID: "drain-me", Timestamp: time.Unix(1700000000, 0).UTC()} + + w.Close() // wg.Wait() returns immediately (no worker); drain(recCh)+drain(spillCh) rescue the record + + if got := len(readSpool(t, dir, "gw")); got != 1 { + t.Fatalf("Close must drain a buffered spillCh record to the spool; spooled lines=%d, want 1", got) + } +} diff --git a/packages/ai-gateway/internal/platform/audit/writer_spill_test.go b/packages/ai-gateway/internal/platform/audit/writer_spill_test.go index d527543b..acd21318 100644 --- a/packages/ai-gateway/internal/platform/audit/writer_spill_test.go +++ b/packages/ai-gateway/internal/platform/audit/writer_spill_test.go @@ -192,8 +192,10 @@ func TestWriter_SpillWriteErrorFallsToLoudDrop(t *testing.T) { if err != nil { t.Fatalf("ndjson.New: %v", err) } - // Pre-seed past the quota so the next spill write is refused. - if err := os.WriteFile(filepath.Join(dir, "test", "seed.ndjson"), make([]byte, 1100*1024), 0o600); err != nil { + // Pre-seed past the quota so the next spill write is refused. The seed must be a + // countable audit-*.ndjson file: the reclaimable-quota gate excludes non-audit + // entries (e.g. .poison), so a differently-named filler would not count. + if err := os.WriteFile(filepath.Join(dir, "test", "audit-20260101-0001.ndjson"), make([]byte, 1100*1024), 0o600); err != nil { t.Fatalf("seed: %v", err) } w := NewWriter(&memProducer{}, "q", nil, slog.Default()).WithNDJSONSpill(spill).WithLossMode(lossModeSpill) @@ -252,7 +254,18 @@ func TestWriter_PublishRecord_SpillsWhenBufferFull(t *testing.T) { w := NewWriter(&memProducer{alwaysFail: true}, "q", nil, slog.Default()).WithNDJSONSpill(spill).WithLossMode(lossModeSpill) saturate(w) - w.publishRecord(&Record{RequestID: "pub-spill"}) // publish fails, buffer full → spill + w.publishRecord(&Record{RequestID: "pub-spill"}) // publish fails, buffer full → async spill hand-off + + // The overflow is handed to the batched spill channel (off the request path), + // not written synchronously; drain it and drive the worker's write step. + select { + case rec := <-w.spillCh: + if !w.spillRecord(rec) { + t.Fatal("spillRecord must durably write the handed-off overflow record") + } + default: + t.Fatal("publishRecord overflow must hand the record to the async spill channel") + } lines := readSpool(t, dir, "test") if len(lines) != 1 || !strings.Contains(lines[0], "pub-spill") { @@ -325,7 +338,7 @@ func TestSpillData_NilSinkAndWriteError(t *testing.T) { if err != nil { t.Fatalf("ndjson.New: %v", err) } - if err := os.WriteFile(filepath.Join(dir, "test", "seed.ndjson"), make([]byte, 1100*1024), 0o600); err != nil { + if err := os.WriteFile(filepath.Join(dir, "test", "audit-20260101-0001.ndjson"), make([]byte, 1100*1024), 0o600); err != nil { t.Fatalf("seed: %v", err) } w.WithNDJSONSpill(spill).WithLossMode(lossModeSpill) diff --git a/packages/ai-gateway/internal/platform/audit/writer_splice.go b/packages/ai-gateway/internal/platform/audit/writer_splice.go index 39044249..00bbf8aa 100644 --- a/packages/ai-gateway/internal/platform/audit/writer_splice.go +++ b/packages/ai-gateway/internal/platform/audit/writer_splice.go @@ -137,6 +137,8 @@ func (w *Writer) marshalRecordPlain(rec *Record) (data []byte, buf *bytes.Buffer if err := json.NewEncoder(buf).Encode(msg); err != nil { w.logger.Error("audit: marshal failed", "requestId", rec.RequestID, "error", err) reclaimMsgBuf(buf) + w.reclaimRecordBody(rec) // dropped record: return its pooled body to the pool + w.releaseRecordMem(rec) // terminal: the record leaves the pipeline here return nil, nil, false } b := buf.Bytes() diff --git a/packages/ai-gateway/internal/platform/middleware/connection_stage.go b/packages/ai-gateway/internal/platform/middleware/connection_stage.go index c88586c6..9a548207 100644 --- a/packages/ai-gateway/internal/platform/middleware/connection_stage.go +++ b/packages/ai-gateway/internal/platform/middleware/connection_stage.go @@ -23,10 +23,11 @@ import ( // wrapper). This is the default in tests / local runs that don't have a // HookConfigCache configured. // -// The supplier callback is invoked per-request to honor the TTL-based reload -// semantics of pipeline.HookConfigCache.Resolver: the cache refreshes when -// stale, so we must not capture a single *PolicyResolver at construction -// time. +// The supplier callback is invoked per-request so the middleware keeps +// working if a future cache implementation ever hands out per-generation +// resolver instances. With pipeline.HookConfigCache the call is a pure +// pointer read (the resolver is swapped in place by push/backstop reloads), +// so the indirection costs nothing on the hot path. func ConnectionStage( resolverSupplier func(ctx context.Context) (*pipeline.PolicyResolver, error), perHookTimeout, totalTimeout time.Duration, diff --git a/packages/ai-gateway/internal/platform/store/health.go b/packages/ai-gateway/internal/platform/store/health.go index 13c5a19b..731c83a0 100644 --- a/packages/ai-gateway/internal/platform/store/health.go +++ b/packages/ai-gateway/internal/platform/store/health.go @@ -2,6 +2,7 @@ package store import ( "sync" + "sync/atomic" "time" ) @@ -11,6 +12,20 @@ const ( healthThresholdDegraded = 0.05 healthThresholdUnavailable = 0.25 + + // healthSampleChanCap bounds the hot-path sample buffer. record() does a + // non-blocking send; a full channel drops the sample (advisory-only, counted + // in dropped). Sized well above the per-instance record rate so steady load + // never drops. + healthSampleChanCap = 8192 + + // healthPublishInterval is how often the single writer republishes the + // immutable snapshot when new samples have been applied. Batching the publish + // (rather than rebuilding the snapshot on every sample) keeps the per-response + // hot path allocation-light: a per-sample whole-map rebuild at high request + // rate would be pure waste, while the health signal feeds only a 5-min-window, + // reorder-only router where sub-100ms advisory staleness is immaterial. + healthPublishInterval = 100 * time.Millisecond ) // HealthStatus represents the health state of a provider. @@ -22,33 +37,77 @@ const ( HealthStatusUnavailable HealthStatus = "unavailable" ) +// healthSample is one recorded outcome. The timestamp is captured in the hot +// path (record), not in the writer, so a queue backlog cannot skew the 5-min +// window. type healthSample struct { timestamp time.Time success bool latencyMs int } +// healthWindow holds one provider's recent samples. It is used both as the +// writer's private mutable window and, with a freshly COPIED samples slice, as +// an immutable snapshot window (see buildSnapshot). +type healthWindow struct { + providerName string + samples []healthSample +} + +// healthSnapshot is an immutable, atomically-published view of every provider's +// window. Once published it is never mutated; the writer builds a brand-new +// snapshot (new map + copied sample slices) on each publish, so concurrent +// GetHealth readers holding an older pointer are always race-free. +type healthSnapshot struct { + windows map[string]*healthWindow +} + +// healthMsg rides the single writer channel. A data message carries a sample +// for a provider; a flush message carries an ack channel the writer closes AFTER +// applying every prior (FIFO) sample and republishing the snapshot. The flush +// marker travels the SAME channel as samples so single-channel FIFO ordering +// guarantees a same-goroutine "record then flush" observes its own write. +type healthMsg struct { + providerID string + providerName string + sample healthSample + isFlush bool + ack chan struct{} +} + // HealthTracker tracks provider health using an in-process sliding window of // samples. It is used exclusively for per-instance routing decisions (avoiding // unhealthy providers). Durable health state for the status page is computed // centrally by the Hub ProviderHealthRollupJob over traffic_event. +// +// Concurrency model: a single background writer owns the mutable windows; the +// per-response hot path (record) only does a non-blocking channel send, so it +// never contends a mutex. GetHealth reads an immutable snapshot published via an +// atomic.Pointer, so routing reads are lock-free. This removes the process-wide +// mutex that every upstream response previously serialized on. type HealthTracker struct { - mu sync.Mutex - windows map[string]*healthWindow // keyed by providerID + ch chan healthMsg + stop chan struct{} + stopped chan struct{} + snap atomic.Pointer[healthSnapshot] + dropped atomic.Int64 + stopOnce sync.Once } -type healthWindow struct { - providerName string - samples []healthSample - lastRequest time.Time - lastError *time.Time -} - -// NewHealthTracker creates a health tracker. +// NewHealthTracker creates a health tracker and starts its writer goroutine. +// Callers must call Stop when done (production uses a process-lifetime singleton; +// tests should defer/Cleanup Stop) to avoid leaking the writer. func NewHealthTracker() *HealthTracker { - return &HealthTracker{ - windows: make(map[string]*healthWindow), + ht := &HealthTracker{ + ch: make(chan healthMsg, healthSampleChanCap), + stop: make(chan struct{}), + stopped: make(chan struct{}), } + // Pre-publish an empty snapshot so a GetHealth that races ahead of the first + // record never dereferences a nil snapshot. + ht.snap.Store(&healthSnapshot{windows: map[string]*healthWindow{}}) + go ht.run(healthPublishInterval) + return ht } // RecordSuccess records a successful request to a provider. @@ -62,33 +121,110 @@ func (ht *HealthTracker) RecordFailure(providerID, providerName string, latencyM } func (ht *HealthTracker) record(providerID, providerName string, success bool, latencyMs int) { - ht.mu.Lock() - defer ht.mu.Unlock() + ht.recordAt(providerID, providerName, success, latencyMs, time.Now()) +} - w := ht.windows[providerID] - if w == nil { - w = &healthWindow{providerName: providerName} - ht.windows[providerID] = w +// recordAt is the timestamp-explicit form of record. Production always passes +// time.Now() (the timestamp is captured on the hot path so a writer backlog +// cannot skew the 5-min window); it is factored out so tests can inject aged +// samples to exercise the read-time window prune. +func (ht *HealthTracker) recordAt(providerID, providerName string, success bool, latencyMs int, ts time.Time) { + msg := healthMsg{ + providerID: providerID, + providerName: providerName, + sample: healthSample{ + timestamp: ts, + success: success, + latencyMs: latencyMs, + }, } - - now := time.Now() - w.lastRequest = now - if !success { - w.lastError = &now + select { + case ht.ch <- msg: + default: + // Channel full: drop the sample rather than block the request path. The + // health signal is advisory (the router only reorders, never drops a + // target), so a dropped sample at most yields a slightly stale ordering. + ht.dropped.Add(1) } +} - w.samples = append(w.samples, healthSample{ - timestamp: now, - success: success, - latencyMs: latencyMs, +// Stop shuts the writer down. Idempotent and safe to call concurrently with +// record: it closes only ht.stop (never ht.ch — a send on a closed channel from +// an in-flight record would panic), then waits for the writer to exit. +func (ht *HealthTracker) Stop() { + ht.stopOnce.Do(func() { + close(ht.stop) + <-ht.stopped }) +} + +// run is the single writer. It owns the mutable windows map exclusively, applies +// each sample on dequeue, and republishes an immutable snapshot on a tick (when +// dirty) or immediately on a flush marker. +func (ht *HealthTracker) run(interval time.Duration) { + defer close(ht.stopped) - // Cap samples. + windows := make(map[string]*healthWindow) + ticker := time.NewTicker(interval) + defer ticker.Stop() + + dirty := false + publish := func() { + ht.snap.Store(buildSnapshot(windows)) + dirty = false + } + + for { + select { + case msg := <-ht.ch: + if msg.isFlush { + // Every sample enqueued before this marker (FIFO) has been applied + // above; publish unconditionally so the flush caller observes them, + // then release the caller. + publish() + close(msg.ack) + continue + } + applySample(windows, msg) + dirty = true + case <-ticker.C: + if dirty { + publish() + } + case <-ht.stop: + return + } + } +} + +// applySample appends to the provider's private window and caps at the most +// recent maxHealthSamples. Mutates only the writer-private windows map. +func applySample(windows map[string]*healthWindow, msg healthMsg) { + w := windows[msg.providerID] + if w == nil { + w = &healthWindow{providerName: msg.providerName} + windows[msg.providerID] = w + } + w.samples = append(w.samples, msg.sample) if len(w.samples) > maxHealthSamples { w.samples = w.samples[len(w.samples)-maxHealthSamples:] } } +// buildSnapshot deep-copies the writer's private windows into a fresh immutable +// snapshot: a new map plus a COPIED samples slice per window. Copying the slice +// (not just aliasing it) is mandatory — a later writer-side append could +// otherwise overwrite array elements a concurrent GetHealth reader is scanning. +func buildSnapshot(windows map[string]*healthWindow) *healthSnapshot { + m := make(map[string]*healthWindow, len(windows)) + for id, w := range windows { + samples := make([]healthSample, len(w.samples)) + copy(samples, w.samples) + m[id] = &healthWindow{providerName: w.providerName, samples: samples} + } + return &healthSnapshot{windows: m} +} + // HealthState holds computed health metrics for a provider. type HealthState struct { Status HealthStatus @@ -97,34 +233,42 @@ type HealthState struct { SampleCount int } -// GetHealth returns the current health state for a provider. +// GetHealth returns the current health state for a provider. It reads the +// immutable snapshot lock-free and applies the 5-min cutoff at READ time — so an +// idle provider whose samples have all aged out is reported Healthy without +// needing a new sample to trigger recovery, matching the prior read-time prune. +// It never writes any field of the shared snapshot (read-only counting). func (ht *HealthTracker) GetHealth(providerID string) HealthState { - ht.mu.Lock() - defer ht.mu.Unlock() - - w := ht.windows[providerID] + // snap is always non-nil: NewHealthTracker pre-publishes an empty snapshot + // before returning, and the tracker is only reachable through it. + s := ht.snap.Load() + w := s.windows[providerID] if w == nil { return HealthState{Status: HealthStatusHealthy} } cutoff := time.Now().Add(-healthWindowDuration) - w.prune(cutoff) - - if len(w.samples) == 0 { - return HealthState{Status: HealthStatusHealthy} - } - failures := 0 + total := 0 totalLatency := 0 - for _, s := range w.samples { - if !s.success { + for i := range w.samples { + sm := &w.samples[i] + if sm.timestamp.Before(cutoff) { + continue + } + total++ + if !sm.success { failures++ } - totalLatency += s.latencyMs + totalLatency += sm.latencyMs } - errorRate := float64(failures) / float64(len(w.samples)) - avgLatency := totalLatency / len(w.samples) + if total == 0 { + return HealthState{Status: HealthStatusHealthy} + } + + errorRate := float64(failures) / float64(total) + avgLatency := totalLatency / total status := HealthStatusHealthy if errorRate > healthThresholdUnavailable { @@ -137,16 +281,31 @@ func (ht *HealthTracker) GetHealth(providerID string) HealthState { Status: status, ErrorRate: errorRate, AvgLatencyMs: avgLatency, - SampleCount: len(w.samples), + SampleCount: total, } } -func (w *healthWindow) prune(cutoff time.Time) { - i := 0 - for i < len(w.samples) && w.samples[i].timestamp.Before(cutoff) { - i++ +// droppedSamples returns the number of samples the hot path dropped because the +// channel was full — an internal diagnostic (a sustained non-zero value means the +// writer cannot keep up, degrading the advisory health signal). Read by the +// drop-path test; not yet wired to a production metric. +func (ht *HealthTracker) droppedSamples() int64 { + return ht.dropped.Load() +} + +// flush blocks until every sample enqueued before this call has been applied and +// republished into the snapshot. It is a deterministic test seam for +// read-your-write assertions; production never needs it (routing reads health +// during resolve, before the same request records its outcome during execute). +func (ht *HealthTracker) flush() { + ack := make(chan struct{}) + select { + case ht.ch <- healthMsg{isFlush: true, ack: ack}: + case <-ht.stopped: + return } - if i > 0 { - w.samples = w.samples[i:] + select { + case <-ack: + case <-ht.stopped: } } diff --git a/packages/ai-gateway/internal/platform/store/health_test.go b/packages/ai-gateway/internal/platform/store/health_test.go index 2b3283fb..221c25ee 100644 --- a/packages/ai-gateway/internal/platform/store/health_test.go +++ b/packages/ai-gateway/internal/platform/store/health_test.go @@ -1,12 +1,22 @@ package store import ( + "sync" "testing" "time" ) -func TestHealthTracker_HealthyByDefault(t *testing.T) { +// newTestTracker builds a tracker and registers Stop for cleanup so the writer +// goroutine never leaks across tests. +func newTestTracker(t *testing.T) *HealthTracker { + t.Helper() ht := NewHealthTracker() + t.Cleanup(ht.Stop) + return ht +} + +func TestHealthTracker_HealthyByDefault(t *testing.T) { + ht := newTestTracker(t) state := ht.GetHealth("unknown-provider") if state.Status != HealthStatusHealthy { t.Errorf("expected healthy, got %s", state.Status) @@ -14,9 +24,10 @@ func TestHealthTracker_HealthyByDefault(t *testing.T) { } func TestHealthTracker_RecordSuccess(t *testing.T) { - ht := NewHealthTracker() + ht := newTestTracker(t) ht.RecordSuccess("p1", "provider1", 100) ht.RecordSuccess("p1", "provider1", 200) + ht.flush() state := ht.GetHealth("p1") if state.Status != HealthStatusHealthy { @@ -34,12 +45,13 @@ func TestHealthTracker_RecordSuccess(t *testing.T) { } func TestHealthTracker_Degraded(t *testing.T) { - ht := NewHealthTracker() + ht := newTestTracker(t) // 10% error rate → degraded (threshold is 5%). for range 9 { ht.RecordSuccess("p1", "provider1", 50) } ht.RecordFailure("p1", "provider1", 500) + ht.flush() state := ht.GetHealth("p1") if state.Status != HealthStatusDegraded { @@ -48,7 +60,7 @@ func TestHealthTracker_Degraded(t *testing.T) { } func TestHealthTracker_Unavailable(t *testing.T) { - ht := NewHealthTracker() + ht := newTestTracker(t) // 50% error rate → unavailable (threshold is 25%). for range 5 { ht.RecordSuccess("p1", "provider1", 50) @@ -56,6 +68,7 @@ func TestHealthTracker_Unavailable(t *testing.T) { for range 5 { ht.RecordFailure("p1", "provider1", 500) } + ht.flush() state := ht.GetHealth("p1") if state.Status != HealthStatusUnavailable { @@ -63,44 +76,89 @@ func TestHealthTracker_Unavailable(t *testing.T) { } } +// TestHealthTracker_ThresholdBoundaries pins the exact strict-`>` threshold +// semantics: 5% is NOT degraded (must exceed), just over 5% is; 25% is NOT +// unavailable, just over is. This is the byte-identity crux vs the prior +// implementation. +func TestHealthTracker_ThresholdBoundaries(t *testing.T) { + // Exactly 5% (1 failure / 20) → healthy (5% is not > 5%). + ht := newTestTracker(t) + for range 19 { + ht.RecordSuccess("p", "prov", 10) + } + ht.RecordFailure("p", "prov", 10) + ht.flush() + if s := ht.GetHealth("p"); s.Status != HealthStatusHealthy { + t.Errorf("errorRate 0.05 must stay healthy (strict >), got %s", s.Status) + } + + // Exactly 25% (5 failure / 20) → degraded (25% is not > 25%, but is > 5%). + ht2 := newTestTracker(t) + for range 15 { + ht2.RecordSuccess("p", "prov", 10) + } + for range 5 { + ht2.RecordFailure("p", "prov", 10) + } + ht2.flush() + if s := ht2.GetHealth("p"); s.Status != HealthStatusDegraded { + t.Errorf("errorRate 0.25 must be degraded not unavailable (strict >), got %s", s.Status) + } +} + +// TestHealthTracker_SampleCap verifies the window keeps only the most recent +// maxHealthSamples, exercised through the public GetHealth (no internals). func TestHealthTracker_SampleCap(t *testing.T) { - ht := NewHealthTracker() + ht := newTestTracker(t) for range maxHealthSamples + 50 { ht.RecordSuccess("p1", "provider1", 10) } + ht.flush() - ht.mu.Lock() - count := len(ht.windows["p1"].samples) - ht.mu.Unlock() - - if count > maxHealthSamples { - t.Errorf("samples should be capped at %d, got %d", maxHealthSamples, count) + if c := ht.GetHealth("p1").SampleCount; c != maxHealthSamples { + t.Errorf("samples should be capped at %d, got %d", maxHealthSamples, c) } } -// TestHealthTracker_PruneAllSamplesExpired forces every sample to fall -// outside the health window so prune() drops them all, returning to a -// healthy zero-sample state. Covers the `i > 0` branch in prune. -func TestHealthTracker_PruneAllSamplesExpired(t *testing.T) { - ht := NewHealthTracker() - // Record one sample, then back-date it past the window. - ht.RecordSuccess("p1", "provider1", 100) - ht.mu.Lock() - for i := range ht.windows["p1"].samples { - ht.windows["p1"].samples[i].timestamp = time.Now().Add(-2 * healthWindowDuration) - } - ht.mu.Unlock() +// TestHealthTracker_IdleRecoveryReadTimePrune injects one aged sample; the +// read-time 5-min cutoff must drop it so the provider reads healthy with zero +// in-window samples WITHOUT a new sample arriving (idle recovery). This is the +// F-R2-1 correctness trap. +func TestHealthTracker_IdleRecoveryReadTimePrune(t *testing.T) { + ht := newTestTracker(t) + // A single failure that already fell outside the window. + ht.recordAt("p1", "provider1", false, 100, time.Now().Add(-2*healthWindowDuration)) + ht.flush() state := ht.GetHealth("p1") if state.Status != HealthStatusHealthy || state.SampleCount != 0 { - t.Errorf("expired samples should prune to healthy zero; got %+v", state) + t.Errorf("aged-out samples must prune to healthy zero at read time; got %+v", state) + } +} + +// TestHealthTracker_MixedInWindowAndExpired confirms only in-window samples +// count: an expired failure is ignored while a recent success keeps the +// provider healthy. +func TestHealthTracker_MixedInWindowAndExpired(t *testing.T) { + ht := newTestTracker(t) + ht.recordAt("p1", "provider1", false, 100, time.Now().Add(-2*healthWindowDuration)) // expired failure + ht.RecordSuccess("p1", "provider1", 20) // recent success + ht.flush() + + state := ht.GetHealth("p1") + if state.SampleCount != 1 { + t.Fatalf("only the in-window sample should count, got %d", state.SampleCount) + } + if state.ErrorRate != 0 || state.Status != HealthStatusHealthy { + t.Errorf("expired failure must not affect error rate; got %+v", state) } } func TestHealthTracker_IndependentProviders(t *testing.T) { - ht := NewHealthTracker() + ht := newTestTracker(t) ht.RecordSuccess("p1", "provider1", 100) ht.RecordFailure("p2", "provider2", 500) + ht.flush() s1 := ht.GetHealth("p1") s2 := ht.GetHealth("p2") @@ -111,3 +169,94 @@ func TestHealthTracker_IndependentProviders(t *testing.T) { t.Error("p2 should have non-zero error rate") } } + +// TestHealthTracker_TickPublish exercises the periodic (ticker-driven) snapshot +// publish path — the production default cadence — by recording WITHOUT calling +// flush and polling until the tick makes the sample visible. +func TestHealthTracker_TickPublish(t *testing.T) { + ht := newTestTracker(t) + ht.RecordFailure("p1", "provider1", 100) + + deadline := time.Now().Add(2 * time.Second) + var state HealthState + for { + state = ht.GetHealth("p1") + if state.SampleCount > 0 || time.Now().After(deadline) { + break + } + time.Sleep(2 * time.Millisecond) + } + if state.SampleCount != 1 { + t.Fatalf("tick publish should make the sample visible, got %+v", state) + } + if state.Status != HealthStatusUnavailable { + t.Errorf("single failure should read unavailable, got %s", state.Status) + } +} + +// TestHealthTracker_StopIdempotent verifies Stop can be called multiple times +// (and concurrently) without panicking. +func TestHealthTracker_StopIdempotent(t *testing.T) { + ht := NewHealthTracker() + var wg sync.WaitGroup + for range 5 { + wg.Add(1) + go func() { defer wg.Done(); ht.Stop() }() + } + wg.Wait() + // A flush after Stop must not block or panic. + ht.flush() +} + +// TestHealthTracker_DropOnFull verifies the hot path drops (never blocks) when +// the channel is saturated: after Stop the writer is gone, so the buffer fills +// and further records are dropped and counted. +func TestHealthTracker_DropOnFull(t *testing.T) { + ht := NewHealthTracker() + ht.Stop() // writer gone; nothing drains ht.ch + for range healthSampleChanCap + 100 { + ht.RecordSuccess("p1", "provider1", 10) + } + if ht.droppedSamples() == 0 { + t.Error("expected drops once the channel saturated with no writer draining") + } +} + +// TestHealthTracker_ConcurrentRecordAndRead hammers record + GetHealth + Stop +// concurrently; run under -race it proves the lock-free path is race-clean. +func TestHealthTracker_ConcurrentRecordAndRead(t *testing.T) { + ht := NewHealthTracker() + defer ht.Stop() + + var wg sync.WaitGroup + for w := range 8 { + wg.Add(1) + go func(id int) { + defer wg.Done() + for range 500 { + ht.RecordSuccess("p1", "provider1", id) + ht.RecordFailure("p2", "provider2", id) + } + }(w) + } + for range 4 { + wg.Add(1) + go func() { + defer wg.Done() + for range 500 { + _ = ht.GetHealth("p1") + _ = ht.GetHealth("p2") + } + }() + } + wg.Wait() + + ht.flush() + // After all records, p1 is all-success (healthy) and p2 all-failure. + if s := ht.GetHealth("p1"); s.Status != HealthStatusHealthy { + t.Errorf("p1 all-success should be healthy, got %s", s.Status) + } + if s := ht.GetHealth("p2"); s.Status != HealthStatusUnavailable { + t.Errorf("p2 all-failure should be unavailable, got %s", s.Status) + } +} diff --git a/packages/ai-gateway/internal/platform/streaming/live.go b/packages/ai-gateway/internal/platform/streaming/live.go index d2b37b9c..affbfe7c 100644 --- a/packages/ai-gateway/internal/platform/streaming/live.go +++ b/packages/ai-gateway/internal/platform/streaming/live.go @@ -228,6 +228,9 @@ func (lp *LivePipeline) Process( // off-path redacted storage copy, but NEVER blocks or rewrites the already- // delivered wire. runCheckpoint := func() { + if lp.hookRun == nil { + return // no response-stage hooks bound: nothing to scan or stamp. + } input := &hookcore.HookInput{ RequestID: hookCtx.RequestID, Stage: "response", @@ -253,12 +256,32 @@ func (lp *LivePipeline) Process( } } - for ch := range eventCh { - delta := format.ExtractDeltaText(ch.data) + // scanForAudit is false when no response-stage hooks are bound (nil runner). + // In that case the per-event canonical-text accumulation (a per-chunk JSON + // parse) and the audit checkpoints are pure no-op work — the checkpoint would + // resolve BuildPipeline("response") to nil and stamp a synthetic Approve every + // window. Skipping mirrors the shared LivePipeline's `l.pipeline == nil` guard + // and the non-stream rule-free path; rec.ResponseAction already defaults to + // ActionApprove upstream so the captured body still persists (only the + // synthetic response_hook_decision="APPROVE" is dropped, matching non-stream). + scanForAudit := lp.hookRun != nil + + // writeEvent delivers ONE chunk: it enforces the buffer cap, writes the SSE + // frame (WITHOUT flushing — the caller coalesces flushes across a burst), + // accumulates the canonical text, and runs the byte-window audit checkpoint. + // It returns true when the buffer cap overflowed, in which case it has already + // emitted the error frame, flushed, cancelled, closed the upstream, and set + // blocked — the caller must stop. + writeEvent := func(ch chunk) (overflow bool) { totalBytes += len(ch.rawData) - if totalBytes > lp.config.MaxBufferSize { lp.logger.Error("stream buffer exceeded", "bytes", totalBytes) + // Audit the content accumulated so far before aborting: the + // mandatory EOF checkpoint below is skipped on this blocked path, so + // without this scan the widening cadence would leave the tail since + // the last intermediate checkpoint unaudited. The buffer is already + // in hand — one O(n) scan, not a per-chunk cost. + runCheckpoint() // best-effort: error notification to client; we cancel below regardless. // Flush BEFORE cancel — without the flush, the error frame stays // in the kernel buffer and the client sees a silent disconnect @@ -276,7 +299,7 @@ func (lp *LivePipeline) Process( // unblock the reader so wg.Wait() can return. sharedstreaming.CloseUpstreamOnExit(upstream) blocked = true - break + return true } // AUDIT-ONLY (B1): deliver every chunk in real time — delivery is never @@ -285,24 +308,79 @@ func (lp *LivePipeline) Process( // traffic). Accumulate the canonical text so the periodic + final // checkpoints scan it for the audit tag and the off-path redacted copy. _ = format.WriteTypedEvent(client, ch.eventType, ch.data) + + // Skip the audit-scan bookkeeping entirely when no response hooks are + // bound: the per-chunk ExtractDeltaText JSON parse + accumulation exist + // only to feed the checkpoint scan. + if scanForAudit { + accBuf.WriteString(format.ExtractDeltaText(ch.data)) + // Observe-only audit checkpoint on a widening byte cadence; it does + // NOT gate delivery. Each checkpoint re-normalizes the FULL accumulated + // stream (the pre-hook parses cumulative wire bytes, not just the new + // delta), so a fixed byte-step makes the total re-normalization work + // grow with the square of the response length. Growing the step + // proportionally to the transcript caps the checkpoint count so the + // total work stays linear; short streams keep the fine ReinspectStepChars + // cadence (the proportional term overtakes only past ~8x the step), and + // the mandatory EOF checkpoint below is always the authoritative full + // scan, so coarser intermediate spacing never changes the final result. + if accBuf.Len() >= nextInspect { + runCheckpoint() + step := lp.config.ReinspectStepChars + if grow := accBuf.Len() / 8; grow > step { + step = grow + } + nextInspect = accBuf.Len() + step + } + } + return false + } + + for ch := range eventCh { + if writeEvent(ch) { + break + } + + // Opportunistic drain-then-flush: write any events ALREADY queued without + // a per-event flush, then issue ONE flush for the whole burst. This + // coalesces the flush syscall under load (the dominant cost on the live + // path) while adding zero latency — a lone event finds the channel empty on + // the first drain iteration and flushes immediately, so TTFT is unchanged. + // The comma-ok receive is load-bearing: once the reader closes eventCh a + // bare receive would spin returning zero-value chunks forever. + draining := true + for draining { + select { + case ch2, ok := <-eventCh: + if !ok { + // Reader closed the channel; stop draining. The outer range + // exits on its next receive. + draining = false + break + } + if writeEvent(ch2) { + draining = false + } + default: + // Channel momentarily empty — flush the burst delivered so far. + draining = false + } + } + if canFlush { flusher.Flush() } - accBuf.WriteString(delta) - - // Byte-window cadence: an observe-only audit checkpoint roughly every - // ReinspectStepChars of new content. It does NOT gate delivery. - if accBuf.Len() >= nextInspect { - runCheckpoint() - nextInspect = accBuf.Len() + lp.config.ReinspectStepChars + if blocked { + break } } // Mandatory final checkpoint — the authoritative audit scan of the FULL // response, always run at EOF (never gated behind a content-length threshold) // so a stream shorter than the first inspect window is still scanned for the - // audit tag and the off-path redacted storage copy. - if !blocked { + // audit tag and the off-path redacted storage copy. Skipped when no response + // hooks are bound (scanForAudit false) — there is nothing to scan or stamp. + if !blocked && scanForAudit { runCheckpoint() } diff --git a/packages/ai-gateway/internal/platform/streaming/live_cadence_test.go b/packages/ai-gateway/internal/platform/streaming/live_cadence_test.go new file mode 100644 index 00000000..4eb657a6 --- /dev/null +++ b/packages/ai-gateway/internal/platform/streaming/live_cadence_test.go @@ -0,0 +1,170 @@ +package streaming + +import ( + "context" + "log/slog" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" + + goHooks "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// countingHook records how many times the checkpoint hook ran (one call per +// runCheckpoint), and the content length seen at the last call. +func countingHook(calls *atomic.Int64, lastLen *atomic.Int64) StreamHookRunner { + return func(_ context.Context, in *goHooks.HookInput) *goHooks.CompliancePipelineResult { + calls.Add(1) + if in != nil && in.Normalized != nil { + n := 0 + for _, s := range in.Normalized.TextProjection() { + n += len(s) + } + lastLen.Store(int64(n)) + } + return &goHooks.CompliancePipelineResult{Decision: goHooks.Approve} + } +} + +// manyChunkStream builds an SSE stream of n content chunks each `each` chars. +func manyChunkStream(n, each int) string { + var b strings.Builder + seg := strings.Repeat("x", each) + for range n { + b.WriteString(`data: {"choices":[{"delta":{"content":"` + seg + `"}}]}` + "\n\n") + } + b.WriteString("data: [DONE]\n\n") + return b.String() +} + +// TestLiveCadence_LongStreamCheckpointsSublinear is the O(n²) fix guard: on a +// long stream the number of intermediate checkpoints must grow far slower than +// the transcript length. With the old fixed ReinspectStepChars cadence the +// count was contentLen/step; the widening cadence must cut it well below that. +func TestLiveCadence_LongStreamCheckpointsSublinear(t *testing.T) { + const chunks, each, step = 400, 40, 128 // ~16000 content chars + var calls, lastLen atomic.Int64 + lp := NewLivePipeline(LiveConfig{ + FirstInspectChars: 400, + ReinspectStepChars: step, + EmitOpenAIDone: true, + }, countingHook(&calls, &lastLen), nil, slog.Default()) + + rec := httptest.NewRecorder() + blocked := lp.Process(context.Background(), strings.NewReader(manyChunkStream(chunks, each)), + rec, &StreamHookContext{IngressType: "AI_GATEWAY", Path: "/v1/chat/completions"}) + if blocked { + t.Fatal("clean stream must not be blocked") + } + + contentLen := chunks * each // 16000 + fixedCadenceCount := int64(contentLen / step) + got := calls.Load() + // The widening cadence must be well under half the fixed-cadence count on a + // stream this long; otherwise the quadratic re-normalization is back. + if got >= fixedCadenceCount/2 { + t.Fatalf("checkpoint count %d not sublinear vs fixed-cadence %d (content %d chars)", got, fixedCadenceCount, contentLen) + } + // The final checkpoint must have scanned the FULL transcript. + if lastLen.Load() < int64(contentLen) { + t.Fatalf("final checkpoint scanned %d chars, want the full %d", lastLen.Load(), contentLen) + } +} + +// TestLiveCadence_ShortStreamUnchanged pins that short streams keep the +// original fine-grained cadence — the proportional term must not kick in +// below ~8× the step, so behavior for normal-length responses is unchanged. +func TestLiveCadence_ShortStreamUnchanged(t *testing.T) { + const step = 128 + // ~900 content chars: below 8*step=1024, so cadence stays fixed at `step`. + var calls, lastLen atomic.Int64 + lp := NewLivePipeline(LiveConfig{ + FirstInspectChars: 400, + ReinspectStepChars: step, + EmitOpenAIDone: true, + }, countingHook(&calls, &lastLen), nil, slog.Default()) + + rec := httptest.NewRecorder() + _ = lp.Process(context.Background(), strings.NewReader(manyChunkStream(30, 30)), // 900 chars + rec, &StreamHookContext{IngressType: "AI_GATEWAY", Path: "/v1/chat/completions"}) + + // First at 400, then +128 each: 400,528,656,784,(900 end) → 4 intermediate + // + 1 mandatory EOF = 5. Assert we are in the fine-grained regime (>=4), + // proving the widening did not coarsen a short stream. + if got := calls.Load(); got < 4 { + t.Fatalf("short stream got only %d checkpoints; fine-grained cadence regressed", got) + } +} + +// TestLiveCadence_DetectionStillFiresLateInStream ensures the widening cadence +// never drops the EOF scan: a match that only appears in the final chunk of a +// long stream must still be seen by a checkpoint (the mandatory EOF one). +func TestLiveCadence_DetectionStillFiresLateInStream(t *testing.T) { + var sawMarker atomic.Bool + hook := func(_ context.Context, in *goHooks.HookInput) *goHooks.CompliancePipelineResult { + if in != nil && in.Normalized != nil { + for _, s := range in.Normalized.TextProjection() { + if strings.Contains(s, "SECRET_MARKER") { + sawMarker.Store(true) + } + } + } + return &goHooks.CompliancePipelineResult{Decision: goHooks.Approve} + } + lp := NewLivePipeline(LiveConfig{ + FirstInspectChars: 400, ReinspectStepChars: 128, EmitOpenAIDone: true, + }, hook, nil, slog.Default()) + + stream := manyChunkStream(300, 40) // long benign prefix + stream = strings.Replace(stream, "data: [DONE]", + `data: {"choices":[{"delta":{"content":"SECRET_MARKER"}}]}`+"\n\ndata: [DONE]", 1) + + rec := httptest.NewRecorder() + _ = lp.Process(context.Background(), strings.NewReader(stream), + rec, &StreamHookContext{IngressType: "AI_GATEWAY", Path: "/v1/chat/completions"}) + + if !sawMarker.Load() { + t.Fatal("a marker in the final chunk must be seen by the mandatory EOF checkpoint") + } +} + +// TestLiveCadence_OverflowStillAudits pins MINOR-1: when a stream exceeds +// MaxBufferSize the mandatory EOF checkpoint is skipped, so the overflow path +// must run one checkpoint over the content accumulated so far — otherwise the +// widened cadence leaves the tail since the last intermediate checkpoint +// unaudited on the audit-only path. +func TestLiveCadence_OverflowStillAudits(t *testing.T) { + var calls, lastLen atomic.Int64 + lp := NewLivePipeline(LiveConfig{ + FirstInspectChars: 400, + ReinspectStepChars: 128, + MaxBufferSize: 131072, // large enough that the step at overflow (~accLen/8) is a material tail + EmitOpenAIDone: true, + }, countingHook(&calls, &lastLen), nil, slog.Default()) + + rec := httptest.NewRecorder() + // ~16000 raw bytes → crosses the 4096 buffer cap mid-stream. + blocked := lp.Process(context.Background(), strings.NewReader(manyChunkStream(4000, 40)), + rec, &StreamHookContext{IngressType: "AI_GATEWAY", Path: "/v1/chat/completions"}) + if !blocked { + t.Fatal("stream exceeding MaxBufferSize must set blocked") + } + // Pin the fix's business behavior (not just "a checkpoint ran"): the + // LAST checkpoint on the overflow path must have scanned ALL delivered + // content. Without the overflow-path runCheckpoint, the last checkpoint + // is the previous INTERMEDIATE one, which trails the delivered tail by up + // to accLen/8 — so lastLen would be strictly less than what was delivered. + delivered := strings.Count(rec.Body.String(), "x") + if delivered == 0 { + t.Fatal("precondition: some content must have been delivered before overflow") + } + // The overflow checkpoint must cover essentially all delivered content — + // the unaudited tail must be far smaller than the widened step (accLen/8, + // ~260 chars here). Without the overflow-path runCheckpoint the last scan + // is the previous intermediate checkpoint, leaving a whole step's worth of + // delivered tail unaudited (gap >> one base step). + if gap := delivered - int(lastLen.Load()); gap >= 128 { + t.Fatalf("overflow left %d delivered chars unaudited (scanned %d of %d) — the tail since the last intermediate checkpoint went unaudited", gap, lastLen.Load(), delivered) + } +} diff --git a/packages/ai-gateway/internal/platform/streaming/live_nohooks_test.go b/packages/ai-gateway/internal/platform/streaming/live_nohooks_test.go new file mode 100644 index 00000000..b26cf465 --- /dev/null +++ b/packages/ai-gateway/internal/platform/streaming/live_nohooks_test.go @@ -0,0 +1,100 @@ +package streaming + +import ( + "context" + "log/slog" + "net/http/httptest" + "strings" + "testing" + + goHooks "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// TestLivePipeline_NilHookRun_SkipsCheckpoints verifies that when no response +// hooks are bound (nil runner), the pipeline skips EVERY audit checkpoint +// (intermediate + final) yet still delivers the full body + [DONE]. The +// FirstInspect/Reinspect chars are set to 1 so a checkpoint would fire after +// every chunk if the scan ran — proving the skip. +func TestLivePipeline_NilHookRun_SkipsCheckpoints(t *testing.T) { + input := makeSSEStream( + `{"choices":[{"delta":{"content":"Hello "}}]}`, + `{"choices":[{"delta":{"content":"world"}}]}`, + ) + lp := NewLivePipeline(LiveConfig{ + FirstInspectChars: 1, + ReinspectStepChars: 1, + EmitOpenAIDone: true, + }, nil, nil, slog.Default()) + rec := httptest.NewRecorder() + calls := 0 + hookCtx := &StreamHookContext{ + IngressType: "AI_GATEWAY", + Path: "/v1/chat/completions", + OnCheckpoint: func(*goHooks.CompliancePipelineResult) { calls++ }, + } + if lp.Process(context.Background(), strings.NewReader(input), rec, hookCtx) { + t.Fatal("nil-hook stream must not block") + } + if calls != 0 { + t.Errorf("nil hookRun must skip ALL checkpoints (intermediate + final); OnCheckpoint fired %d times", calls) + } + body := rec.Body.String() + for _, want := range []string{"Hello ", "world", "[DONE]"} { + if !strings.Contains(body, want) { + t.Errorf("body must contain %q (full delivery); got %q", want, body) + } + } +} + +// TestLivePipeline_WithHookRun_StillCheckpoints is the counterpart: a bound +// (non-nil) runner must STILL run checkpoints — the skip is gated strictly on +// "no response hooks", never on the mode. +func TestLivePipeline_WithHookRun_StillCheckpoints(t *testing.T) { + input := makeSSEStream( + `{"choices":[{"delta":{"content":"aaaaa"}}]}`, + `{"choices":[{"delta":{"content":"bbbbb"}}]}`, + ) + lp := NewLivePipeline(LiveConfig{FirstInspectChars: 1, ReinspectStepChars: 1}, approveStreamHook, nil, slog.Default()) + rec := httptest.NewRecorder() + calls := 0 + hookCtx := &StreamHookContext{ + IngressType: "AI_GATEWAY", + Path: "/v1/chat/completions", + OnCheckpoint: func(*goHooks.CompliancePipelineResult) { calls++ }, + } + lp.Process(context.Background(), strings.NewReader(input), rec, hookCtx) + if calls == 0 { + t.Error("a bound hook runner must still fire checkpoints") + } +} + +// BenchmarkLivePipeline_Process is the before/after: no-response-hooks (nil +// runner, the skip path) vs a bound runner, over a 50-token stream. The nil arm +// should drop the per-checkpoint HookInput + PayloadFromTextSegments allocs and +// the per-chunk ExtractDeltaText JSON parse. +func BenchmarkLivePipeline_Process(b *testing.B) { + frames := make([]string, 50) + for i := range frames { + frames[i] = `{"choices":[{"delta":{"content":"token "}}]}` + } + input := makeSSEStream(frames...) + newCtx := func() *StreamHookContext { + return &StreamHookContext{IngressType: "AI_GATEWAY", Path: "/v1/chat/completions"} + } + b.Run("no_response_hooks_nil", func(b *testing.B) { + lp := NewLivePipeline(LiveConfig{EmitOpenAIDone: true}, nil, nil, slog.Default()) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + lp.Process(context.Background(), strings.NewReader(input), httptest.NewRecorder(), newCtx()) + } + }) + b.Run("with_hooks", func(b *testing.B) { + lp := NewLivePipeline(LiveConfig{EmitOpenAIDone: true}, approveStreamHook, nil, slog.Default()) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + lp.Process(context.Background(), strings.NewReader(input), httptest.NewRecorder(), newCtx()) + } + }) +} diff --git a/packages/ai-gateway/internal/platform/streaming/live_pipeline_decisions_test.go b/packages/ai-gateway/internal/platform/streaming/live_pipeline_decisions_test.go index 066c6643..d001e728 100644 --- a/packages/ai-gateway/internal/platform/streaming/live_pipeline_decisions_test.go +++ b/packages/ai-gateway/internal/platform/streaming/live_pipeline_decisions_test.go @@ -496,6 +496,37 @@ func TestLivePipeline_MaxBufferSizeAborts(t *testing.T) { } } +// TestLivePipeline_MaxBufferSizeOverflowDuringDrain drives the overflow on a +// LATER chunk inside the opportunistic drain-then-flush burst (not the first +// chunk). Every frame individually fits; the cumulative size crosses +// MaxBufferSize a few frames in. Because the in-memory reader fills the event +// channel before the main loop drains, the overflow is detected while draining +// the burst — exercising the drain-path overflow branch (not just the +// first-chunk path). The observable contract is identical: buffer-exceeded error +// frame + blocked=true. +func TestLivePipeline_MaxBufferSizeOverflowDuringDrain(t *testing.T) { + // Each frame's raw JSON is ~70 bytes; frame 1 fits under 180, the cumulative + // crosses it around frame 3. + chunk := strings.Repeat("y", 30) + frame := `{"choices":[{"delta":{"content":"` + chunk + `"}}]}` + input := makeSSEStream(frame, frame, frame, frame, frame) + + lp := NewLivePipeline(LiveConfig{ + FirstInspectChars: 100000, // no compliance checkpoint fires + MaxBufferSize: 180, // frame 1 fits (~70), cumulative overflows mid-burst + }, approveStreamHook, nil, slog.Default()) + rec := httptest.NewRecorder() + hookCtx := &StreamHookContext{IngressType: "AI_GATEWAY", Path: "/v1/chat/completions"} + + blocked := lp.Process(context.Background(), strings.NewReader(input), rec, hookCtx) + if !blocked { + t.Fatal("cumulative overflow must be reported as blocked") + } + if !strings.Contains(rec.Body.String(), "stream buffer exceeded") { + t.Errorf("client must receive buffer-exceeded SSE error frame: %q", rec.Body.String()) + } +} + // When the last accumulated chunk is short enough that no checkpoint // fired during the loop, the post-loop `len(pendingText) > 0` branch // runs the final checkpoint. Pins L314-317 (here with Approve → diff --git a/packages/ai-gateway/internal/providers/core/spec.go b/packages/ai-gateway/internal/providers/core/spec.go index 71e7a151..8650ff30 100644 --- a/packages/ai-gateway/internal/providers/core/spec.go +++ b/packages/ai-gateway/internal/providers/core/spec.go @@ -49,6 +49,29 @@ type AdapterSpec struct { // adapters that do not provide the probe. PassthroughRewriteApplies func(modelID string) bool + // PassthroughModelInBody declares that this adapter's native wire + // carries the model at the JSON body top-level `model` field (as the + // OpenAI shape does), so the passthrough path must rewrite it to the + // resolved CallTarget.ProviderModelID before upstream dispatch — + // otherwise a client-facing alias / a routing-changed model is sent + // verbatim and the upstream 404s the unknown name. + // + // The generic dispatcher already does this rewrite for OpenAI-family + // formats (gated by Format.IsOpenAIFamily). This flag extends the same + // surgical top-level `model` rewrite to a non-OpenAI-family adapter whose + // wire also puts the model in the body root — today only Anthropic + // (`/v1/messages`). Adapters that carry the model elsewhere leave it + // false: Gemini (URL path, rewritten by Transport.BuildURL from + // ProviderModelID), Bedrock (model deleted from body, encoded into the + // URL by its codec). Cross-format routes never rely on this flag — the + // codec's EncodeRequest already stamps ProviderModelID. + // + // Per provider-adapter-architecture.md §3 (Ingress shape preservation) + // and §3a Rule 3: the model-location fact is owned by the adapter that + // talks to that wire, declared here rather than as a format switch in + // the generic dispatcher. + PassthroughModelInBody bool + // RequestShapes lists the typology.WireShape values this adapter // natively serves at the codec boundary (i.e. the WireShape values // the codec's EncodeRequest/DecodeResponse will accept without diff --git a/packages/ai-gateway/internal/providers/dispatch/registry_execute_test.go b/packages/ai-gateway/internal/providers/dispatch/registry_execute_test.go index 1b644f4e..cc085c44 100644 --- a/packages/ai-gateway/internal/providers/dispatch/registry_execute_test.go +++ b/packages/ai-gateway/internal/providers/dispatch/registry_execute_test.go @@ -777,14 +777,14 @@ func TestPrepareBody_NonChatEndpoint_PassesBodyThrough(t *testing.T) { } } -// TestPrepareBody_NonOpenAIWire_NoRewrite asserts that bodies in a -// non-OpenAI wire format (e.g. when adapters share Format but the body -// shape isn't OpenAI-compatible) are not rewritten by the generic path. -// We must reach the same-format passthrough path with a non-OpenAI-wire -// Format so the !IsOpenAIFamily branch fires. +// TestPrepareBody_NonOpenAIWire_NoRewrite asserts that a non-OpenAI-family +// body is NOT rewritten by the generic passthrough when the adapter does not +// declare PassthroughModelInBody — the Gemini/Bedrock contract (model applied +// by the transport/codec, not the body). A spec that DOES declare the +// capability is covered by TestPrepareBody_Anthropic_ModelInBody_Wired. func TestPrepareBody_NonOpenAIWire_NoRewrite(t *testing.T) { - // Anthropic is not an OpenAI-wire shape — same-Format passthrough hits - // the !IsOpenAIFamily branch in rewritePassthroughModel. + // specFrom leaves PassthroughModelInBody false, so the model-in-body gate + // (OpenAI-family wire OR capability) is not satisfied for FormatAnthropic. ad := NewSpecAdapter(specFrom(&fakeTransport{}, &fakeCodec{}, &fakeStreamDecoder{}, &fakeErrorNormalizer{}, FormatAnthropic), slog.Default()) body := []byte(`{"model":"claude-3-5-sonnet","messages":[]}`) got, rw, _, err := ad.PrepareBody(Request{ diff --git a/packages/ai-gateway/internal/providers/dispatch/spec_adapter.go b/packages/ai-gateway/internal/providers/dispatch/spec_adapter.go index 0b29151e..52d81be4 100644 --- a/packages/ai-gateway/internal/providers/dispatch/spec_adapter.go +++ b/packages/ai-gateway/internal/providers/dispatch/spec_adapter.go @@ -404,7 +404,7 @@ func (a *specAdapter) prepareBodyFull(req Request) (body []byte, rewrites []stri // codec EncodeRequest on those adapters is an identity pass that would // leave the original model ID in the body. if req.BodyFormat == a.spec.Format || (req.BodyFormat.IsOpenAIFamily() && a.spec.Format.IsOpenAIFamily()) { - b, rw, e := rewritePassthroughModel(req, a.spec.PassthroughRewrite, a.spec.PassthroughRewriteApplies) + b, rw, e := rewritePassthroughModel(req, a.spec.PassthroughRewrite, a.spec.PassthroughRewriteApplies, a.spec.PassthroughModelInBody) return b, rw, "", e } // Canonical OpenAI input needs codec translation. Codecs may apply @@ -440,7 +440,7 @@ func applyURLOverride(baseURL, override string) string { return override } -func rewritePassthroughModel(req Request, passthroughRewrite func(map[string]any, string) []string, rewriteApplies func(string) bool) ([]byte, []string, error) { +func rewritePassthroughModel(req Request, passthroughRewrite func(map[string]any, string) []string, rewriteApplies func(string) bool, modelInBody bool) ([]byte, []string, error) { // Strip the gateway-internal `nexus` namespace from the body before // any further work The passthrough path forwards req.Body // to upstream verbatim (modulo model rewrite), and the cross-format @@ -455,16 +455,27 @@ func rewritePassthroughModel(req Request, passthroughRewrite func(map[string]any if req.Target.ProviderModelID == "" { return body, nil, nil } - switch req.WireShape { - case typology.WireShapeOpenAIChat, typology.WireShapeOpenAIEmbeddings, typology.WireShapeOpenAICompletionsLegacy: - default: - return body, nil, nil - } - if !req.BodyFormat.IsOpenAIFamily() { - // Non-OpenAI-shape bodies (Anthropic Messages, Gemini generateContent, - // Bedrock, Cohere, Replicate, ...) carry the model field in different - // places or names; their per-format SchemaCodec.EncodeRequest is the - // site that applies ct.ProviderModelID, not this passthrough path. + // The passthrough rewrites the top-level `model` field only for wires that + // carry it there. Two sources qualify: + // 1. OpenAI-family bodies on an OpenAI wire shape — the model sits at the + // JSON root and `payload["model"] = X` substitution is valid. + // 2. A non-OpenAI-family adapter that declares PassthroughModelInBody + // (today only Anthropic `/v1/messages`), whose native wire also puts + // the model at the body root but whose EncodeRequest — the usual + // ProviderModelID-stamping site — is skipped on the same-format + // native passthrough. Without this the client-facing alias / a + // routing-changed model reaches the upstream verbatim and 404s. + // Wires that carry the model elsewhere leave both false and pass through + // untouched: Gemini (URL path, set by Transport.BuildURL from + // ProviderModelID), Bedrock (model deleted from body, encoded into the URL + // by its codec). The stream_options / PassthroughRewrite map path below + // stays OpenAI-chat-only, so an Anthropic body only ever takes the + // surgical top-level `model` rewrite — never OpenAI-specific injection. + openAIWire := req.WireShape == typology.WireShapeOpenAIChat || + req.WireShape == typology.WireShapeOpenAIEmbeddings || + req.WireShape == typology.WireShapeOpenAICompletionsLegacy + rewriteBodyModel := (openAIWire && req.BodyFormat.IsOpenAIFamily()) || modelInBody + if !rewriteBodyModel { return body, nil, nil } if len(body) == 0 { diff --git a/packages/ai-gateway/internal/providers/dispatch/spec_adapter_modelrewrite_test.go b/packages/ai-gateway/internal/providers/dispatch/spec_adapter_modelrewrite_test.go index e9795b04..14887be2 100644 --- a/packages/ai-gateway/internal/providers/dispatch/spec_adapter_modelrewrite_test.go +++ b/packages/ai-gateway/internal/providers/dispatch/spec_adapter_modelrewrite_test.go @@ -44,10 +44,30 @@ func reqFor(body []byte, stream bool) Request { } } +// rpm calls rewritePassthroughModel with modelInBody=false — the OpenAI-family +// default every reqFor-based test exercises (model rewrite driven by the +// OpenAI wire shape, not the PassthroughModelInBody capability). +func rpm(req Request, pr func(map[string]any, string) []string, ap func(string) bool) ([]byte, []string, error) { + return rewritePassthroughModel(req, pr, ap, false) +} + +// anthropicReqFor builds an Anthropic /v1/messages passthrough request: the +// native wire shape, FormatAnthropic body, and a model that differs from the +// resolved ProviderModelID (the alias/routing case). modelInBody must be true +// for the passthrough to apply ProviderModelID to the body `model`. +func anthropicReqFor(body []byte) Request { + return Request{ + WireShape: typology.WireShapeAnthropicMessages, + BodyFormat: FormatAnthropic, + Body: body, + Target: CallTarget{ProviderModelID: "claude-opus-4-8"}, + } +} + // TestA2_FastPath_SetsModelAndStripsNexus: non-stream, no rewrite → fast path. func TestA2_FastPath_SetsModelAndStripsNexus(t *testing.T) { body := chatBody(`"nexus":{"ext":{"x":1}},"temperature":0.5`) - out, rewrites, err := rewritePassthroughModel(reqFor(body, false), nil, nil) + out, rewrites, err := rpm(reqFor(body, false), nil, nil) if err != nil { t.Fatal(err) } @@ -73,13 +93,13 @@ func TestA2_FastPath_SetsModelAndStripsNexus(t *testing.T) { // the old map round-trip for a normal body. func TestA2_FastPath_EqualsMapPath(t *testing.T) { body := chatBody(`"temperature":0.5,"top_p":0.9`) - fast, _, err := rewritePassthroughModel(reqFor(body, false), nil, nil) + fast, _, err := rpm(reqFor(body, false), nil, nil) if err != nil { t.Fatal(err) } // Force the map path with a no-op rewrite callback. noop := func(m map[string]any, id string) []string { return nil } - mapOut, _, err := rewritePassthroughModel(reqFor(body, false), noop, nil) + mapOut, _, err := rpm(reqFor(body, false), noop, nil) if err != nil { t.Fatal(err) } @@ -92,7 +112,7 @@ func TestA2_FastPath_EqualsMapPath(t *testing.T) { // resolve to the provider model with NO stale duplicate (map path collapses). func TestA2_DuplicateModel_FallsBackToMapPath(t *testing.T) { body := []byte(`{"model":"alias-a","messages":[],"model":"alias-b"}`) - out, _, err := rewritePassthroughModel(reqFor(body, false), nil, nil) + out, _, err := rpm(reqFor(body, false), nil, nil) if err != nil { t.Fatal(err) } @@ -114,7 +134,7 @@ func TestA2_DuplicateModel_FallsBackToMapPath(t *testing.T) { // stream_options.include_usage + stream:true, and preserve sibling fields. func TestA2_Streaming_NonConformant_MapPath_AppliesUsageOption(t *testing.T) { body := chatBody(`"temperature":0.5`) - out, rewrites, err := rewritePassthroughModel(reqFor(body, true), nil, nil) + out, rewrites, err := rpm(reqFor(body, true), nil, nil) if err != nil { t.Fatal(err) } @@ -146,7 +166,7 @@ func TestA2_Streaming_NonConformant_MapPath_AppliesUsageOption(t *testing.T) { // the loadtest and real OpenAI-stream clients that send the upstream model name. func TestA2_Streaming_AlreadyConformant_ZeroRewrite(t *testing.T) { body := []byte(`{"model":"provider-real-model","stream":true,"stream_options":{"include_usage":true},"messages":[{"role":"user","content":"hi"}]}`) - out, _, err := rewritePassthroughModel(reqFor(body, true), nil, nil) + out, _, err := rpm(reqFor(body, true), nil, nil) if err != nil { t.Fatal(err) } @@ -162,7 +182,7 @@ func TestA2_Streaming_ConformantButRewriteApplies_TakesMapPath(t *testing.T) { body := []byte(`{"model":"provider-real-model","stream":true,"stream_options":{"include_usage":true},"messages":[]}`) called := false rw := func(m map[string]any, id string) []string { called = true; return []string{"r"} } - out, rewrites, err := rewritePassthroughModel(reqFor(body, true), rw, nil) + out, rewrites, err := rpm(reqFor(body, true), rw, nil) if err != nil { t.Fatal(err) } @@ -192,13 +212,13 @@ func TestA2_EdgeCases(t *testing.T) { // either gets its model rewritten and forwarded (upstream returns the // error) or sjson itself errors — both acceptable; the gateway must not // crash and must produce a deterministic result. - out, _, err := rewritePassthroughModel(reqFor([]byte(`{"model":"x"`), false), nil, nil) + out, _, err := rpm(reqFor([]byte(`{"model":"x"`), false), nil, nil) if err == nil && len(out) == 0 { t.Error("expected either a forwarded body or an error, got neither") } }) t.Run("absent_model_added", func(t *testing.T) { - out, _, err := rewritePassthroughModel(reqFor([]byte(`{"messages":[{"role":"user","content":"hi"}]}`), false), nil, nil) + out, _, err := rpm(reqFor([]byte(`{"messages":[{"role":"user","content":"hi"}]}`), false), nil, nil) if err != nil { t.Fatal(err) } @@ -211,7 +231,7 @@ func TestA2_EdgeCases(t *testing.T) { } }) t.Run("numeric_model_overwritten_with_string", func(t *testing.T) { - out, _, err := rewritePassthroughModel(reqFor([]byte(`{"model":123,"messages":[]}`), false), nil, nil) + out, _, err := rpm(reqFor([]byte(`{"model":123,"messages":[]}`), false), nil, nil) if err != nil { t.Fatal(err) } @@ -230,7 +250,7 @@ func TestA2_EdgeCases(t *testing.T) { Body: []byte(`{"model":"client-alias","input":"hello"}`), Target: CallTarget{ProviderModelID: "provider-real-model"}, } - out, _, err := rewritePassthroughModel(req, nil, nil) + out, _, err := rpm(req, nil, nil) if err != nil { t.Fatal(err) } @@ -266,14 +286,14 @@ func BenchmarkA2ModelRewrite(b *testing.B) { req := reqFor(body, false) b.ReportAllocs() for range b.N { - _, _, _ = rewritePassthroughModel(req, nil, nil) + _, _, _ = rpm(req, nil, nil) } }) b.Run("old_map_roundtrip", func(b *testing.B) { req := reqFor(body, false) b.ReportAllocs() for range b.N { - _, _, _ = rewritePassthroughModel(req, noop, nil) + _, _, _ = rpm(req, noop, nil) } }) } @@ -289,14 +309,14 @@ func BenchmarkA2StreamAllSkip(b *testing.B) { req := reqFor(conformant, true) b.ReportAllocs() for range b.N { - _, _, _ = rewritePassthroughModel(req, nil, nil) + _, _, _ = rpm(req, nil, nil) } }) b.Run("nonconformant_map_path", func(b *testing.B) { req := reqFor(rewrite, true) b.ReportAllocs() for range b.N { - _, _, _ = rewritePassthroughModel(req, nil, nil) + _, _, _ = rpm(req, nil, nil) } }) } @@ -311,7 +331,7 @@ func TestPassthroughRewriteApplies_GatesMapPath(t *testing.T) { // Probe false → surgical sjson path: rewrite NOT invoked, model set, other // fields (temperature) preserved verbatim. - out, _, err := rewritePassthroughModel(reqFor(body, false), rewrite, func(string) bool { return false }) + out, _, err := rpm(reqFor(body, false), rewrite, func(string) bool { return false }) if err != nil { t.Fatal(err) } @@ -327,7 +347,7 @@ func TestPassthroughRewriteApplies_GatesMapPath(t *testing.T) { // Probe true → map path: rewrite invoked. rewriteCalls = 0 - if _, _, err = rewritePassthroughModel(reqFor(body, false), rewrite, func(string) bool { return true }); err != nil { + if _, _, err = rpm(reqFor(body, false), rewrite, func(string) bool { return true }); err != nil { t.Fatal(err) } if rewriteCalls != 1 { @@ -336,10 +356,154 @@ func TestPassthroughRewriteApplies_GatesMapPath(t *testing.T) { // Nil probe ⇒ conservative map path (prior behavior). rewriteCalls = 0 - if _, _, err = rewritePassthroughModel(reqFor(body, false), rewrite, nil); err != nil { + if _, _, err = rpm(reqFor(body, false), rewrite, nil); err != nil { t.Fatal(err) } if rewriteCalls != 1 { t.Fatalf("nil probe must keep the map path; got %d", rewriteCalls) } } + +// TestModelInBody_Anthropic_AliasRewritten is the core fix: an Anthropic +// /v1/messages native passthrough (non-OpenAI wire, non-OpenAI-family body) +// whose adapter declares PassthroughModelInBody MUST rewrite the top-level +// `model` from the client-facing alias to the resolved ProviderModelID — +// otherwise the alias reaches Anthropic verbatim and 404s. +func TestModelInBody_Anthropic_AliasRewritten(t *testing.T) { + body := []byte(`{"model":"my-fast-alias","messages":[{"role":"user","content":"hi"}],"nexus":{"ext":{"anthropic":{"topK":42}}}}`) + out, rewrites, err := rewritePassthroughModel(anthropicReqFor(body), nil, nil, true) + if err != nil { + t.Fatal(err) + } + if rewrites != nil { + t.Errorf("model rewrite is not a coercion; want nil rewrites, got %v", rewrites) + } + var m map[string]any + if err := json.Unmarshal(out, &m); err != nil { + t.Fatal(err) + } + if m["model"] != "claude-opus-4-8" { + t.Errorf("model = %v, want claude-opus-4-8 (alias resolved to ProviderModelID)", m["model"]) + } + if _, ok := m["nexus"]; ok { + t.Error("nexus namespace must be stripped before upstream send") + } + // The messages payload must survive untouched. + if _, ok := m["messages"]; !ok { + t.Error("messages array lost during model rewrite") + } +} + +// TestModelInBody_Anthropic_NoOpWhenEqual: when the client already sent the +// provider model name (no alias), the body must be returned unchanged — the +// GetBytes==ProviderModelID short-circuit avoids a wasted sjson rebuild on the +// hot path (nexus strip aside). +func TestModelInBody_Anthropic_NoOpWhenEqual(t *testing.T) { + body := []byte(`{"model":"claude-opus-4-8","messages":[{"role":"user","content":"hi"}]}`) + out, _, err := rewritePassthroughModel(anthropicReqFor(body), nil, nil, true) + if err != nil { + t.Fatal(err) + } + if string(out) != string(body) { + t.Errorf("no-alias body must be returned byte-identical; got %s", out) + } +} + +// TestModelInBody_False_NonOpenAI_NoRewrite: without the PassthroughModelInBody +// capability, a non-OpenAI-family body is NOT rewritten even when +// ProviderModelID differs — this is the Gemini/Bedrock contract (model lives in +// the URL / is deleted from the body, applied by the transport/codec instead). +func TestModelInBody_False_NonOpenAI_NoRewrite(t *testing.T) { + body := []byte(`{"model":"gemini-1.5-pro","contents":[]}`) + req := Request{ + WireShape: typology.WireShapeGeminiGenerateContent, + BodyFormat: FormatGemini, + Body: body, + Target: CallTarget{ProviderModelID: "gemini-2.0-flash"}, + } + out, rw, err := rewritePassthroughModel(req, nil, nil, false) + if err != nil { + t.Fatal(err) + } + if string(out) != string(body) || rw != nil { + t.Errorf("no-capability non-OpenAI body must pass through unchanged; got %s rewrites=%v", out, rw) + } +} + +// TestModelInBody_Anthropic_Idempotent: PrepareBody runs on both the cache-prep +// and execute paths, so the rewrite must be idempotent — a second run over the +// already-rewritten body is a no-op. +func TestModelInBody_Anthropic_Idempotent(t *testing.T) { + body := []byte(`{"model":"my-fast-alias","messages":[]}`) + first, _, err := rewritePassthroughModel(anthropicReqFor(body), nil, nil, true) + if err != nil { + t.Fatal(err) + } + second, _, err := rewritePassthroughModel(anthropicReqFor(first), nil, nil, true) + if err != nil { + t.Fatal(err) + } + if string(second) != string(first) { + t.Errorf("second rewrite must be a byte-identical no-op; first=%s second=%s", first, second) + } +} + +// TestPrepareBody_Anthropic_ModelInBody_Wired proves the call site threads +// a.spec.PassthroughModelInBody into the rewrite: an adapter whose spec +// declares the capability rewrites the Anthropic body model through the public +// PrepareBody entry point (the executor + cache-prep path), not just the +// internal helper. +func TestPrepareBody_Anthropic_ModelInBody_Wired(t *testing.T) { + spec := specFrom(&fakeTransport{}, &fakeCodec{}, &fakeStreamDecoder{}, &fakeErrorNormalizer{}, FormatAnthropic) + spec.PassthroughModelInBody = true + ad := NewSpecAdapter(spec, nil) + body := []byte(`{"model":"my-fast-alias","messages":[]}`) + got, _, _, err := ad.PrepareBody(Request{ + WireShape: typology.WireShapeAnthropicMessages, + BodyFormat: FormatAnthropic, + Body: body, + Target: CallTarget{ProviderModelID: "claude-opus-4-8"}, + }) + if err != nil { + t.Fatalf("PrepareBody: %v", err) + } + if !strings.Contains(string(got), `"model":"claude-opus-4-8"`) { + t.Fatalf("capability spec must rewrite the Anthropic body model to ProviderModelID; got %s", got) + } + // A spec WITHOUT the capability must NOT rewrite (Gemini/Bedrock contract). + specNoCap := specFrom(&fakeTransport{}, &fakeCodec{}, &fakeStreamDecoder{}, &fakeErrorNormalizer{}, FormatAnthropic) + adNoCap := NewSpecAdapter(specNoCap, nil) + gotNoCap, _, _, err := adNoCap.PrepareBody(Request{ + WireShape: typology.WireShapeAnthropicMessages, + BodyFormat: FormatAnthropic, + Body: body, + Target: CallTarget{ProviderModelID: "claude-opus-4-8"}, + }) + if err != nil { + t.Fatalf("PrepareBody(no-cap): %v", err) + } + if string(gotNoCap) != string(body) { + t.Fatalf("no-capability spec must pass the body through unchanged; got %s", gotNoCap) + } +} + +// TestModelInBody_Anthropic_DuplicateModel_MapPath: a pathological duplicate +// top-level model must still resolve to ProviderModelID with a single key (the +// map path collapses last-wins), matching the OpenAI dup-key guarantee. +func TestModelInBody_Anthropic_DuplicateModel_MapPath(t *testing.T) { + body := []byte(`{"model":"alias-a","messages":[],"model":"alias-b"}`) + out, _, err := rewritePassthroughModel(anthropicReqFor(body), nil, nil, true) + if err != nil { + t.Fatal(err) + } + var m map[string]any + if err := json.Unmarshal(out, &m); err != nil { + t.Fatal(err) + } + if m["model"] != "claude-opus-4-8" { + t.Errorf("model = %v, want claude-opus-4-8 (no stale duplicate)", m["model"]) + } + if strings.Count(string(out), `"model"`) != 1 { + t.Errorf("expected exactly one model key after rewrite, got %q", out) + } +} diff --git a/packages/ai-gateway/internal/providers/specs/anthropic/spec.go b/packages/ai-gateway/internal/providers/specs/anthropic/spec.go index 0862f6d0..bb093940 100644 --- a/packages/ai-gateway/internal/providers/specs/anthropic/spec.go +++ b/packages/ai-gateway/internal/providers/specs/anthropic/spec.go @@ -25,6 +25,15 @@ func NewSpec(log *slog.Logger) provcore.AdapterSpec { SchemaCodec: apcodec.NewCodec(), StreamDecoder: apstream.NewStreamDecoder(log), ErrorNormalizer: specerrors.ErrorNormalizer{}, + // Anthropic `/v1/messages` carries the model at the body top-level, + // like the OpenAI shape. On the same-format native passthrough the + // codec's EncodeRequest (which stamps ProviderModelID) is skipped, so + // the generic dispatcher must apply the resolved ProviderModelID to + // the body `model` itself — otherwise a client-facing alias reaches + // Anthropic verbatim and 404s. Cross-format routes still stamp it via + // the codec. Sampling-param wire quirks stay in the codec (Rule 3) and + // remain transparent on the native passthrough by design. + PassthroughModelInBody: true, } } diff --git a/packages/ai-gateway/internal/routing/core/health_ranker_test.go b/packages/ai-gateway/internal/routing/core/health_ranker_test.go index 5f11e490..f88137f7 100644 --- a/packages/ai-gateway/internal/routing/core/health_ranker_test.go +++ b/packages/ai-gateway/internal/routing/core/health_ranker_test.go @@ -2,17 +2,37 @@ package core import ( "testing" + "time" "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/platform/store" ) +// awaitStatus polls until a provider reaches the expected health status. The +// HealthTracker applies records asynchronously (single writer), so a read +// immediately after recording must wait for the samples to be published. +func awaitStatus(t *testing.T, tracker *store.HealthTracker, providerID string, want store.HealthStatus) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for { + if tracker.GetHealth(providerID).Status == want { + return + } + if time.Now().After(deadline) { + t.Fatalf("provider %s did not reach %s in time; got %+v", providerID, want, tracker.GetHealth(providerID)) + } + time.Sleep(2 * time.Millisecond) + } +} + func TestHealthRanker_Reorder(t *testing.T) { tracker := store.NewHealthTracker() + t.Cleanup(tracker.Stop) // Record failures for provider B to make it unhealthy for range 20 { tracker.RecordFailure("provB", "provider-b", 100) } // Provider A is healthy (no records = healthy) + awaitStatus(t, tracker, "provB", store.HealthStatusUnavailable) ranker := NewHealthRanker(tracker) targets := []RoutingTarget{ @@ -38,7 +58,9 @@ func TestHealthRanker_NilTracker(t *testing.T) { } func TestHealthRanker_SingleTarget(t *testing.T) { - ranker := NewHealthRanker(store.NewHealthTracker()) + tracker := store.NewHealthTracker() + t.Cleanup(tracker.Stop) + ranker := NewHealthRanker(tracker) targets := []RoutingTarget{{ProviderID: "a"}} result := ranker.Reorder(targets) if len(result) != 1 { @@ -50,6 +72,7 @@ func TestHealthRanker_SingleTarget(t *testing.T) { // > 5% but ≤ 25%) is ranked after healthy but before unavailable. func TestHealthRanker_DegradedProvider(t *testing.T) { tracker := store.NewHealthTracker() + t.Cleanup(tracker.Stop) // 10 failures + 90 successes = 10% error rate → degraded (>5%, ≤25%) for range 10 { tracker.RecordFailure("provDegraded", "provider-degraded", 100) @@ -62,6 +85,8 @@ func TestHealthRanker_DegradedProvider(t *testing.T) { for range 20 { tracker.RecordFailure("provUnavailable", "provider-unavailable", 100) } + awaitStatus(t, tracker, "provDegraded", store.HealthStatusDegraded) + awaitStatus(t, tracker, "provUnavailable", store.HealthStatusUnavailable) ranker := NewHealthRanker(tracker) targets := []RoutingTarget{ diff --git a/packages/ai-gateway/internal/routing/resolver_error_paths_test.go b/packages/ai-gateway/internal/routing/resolver_error_paths_test.go index 35b71561..bf3a7b9c 100644 --- a/packages/ai-gateway/internal/routing/resolver_error_paths_test.go +++ b/packages/ai-gateway/internal/routing/resolver_error_paths_test.go @@ -10,6 +10,7 @@ import ( "net/http" "strings" "testing" + "time" "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/platform/store" "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/routing/core" @@ -17,6 +18,23 @@ import ( "github.com/AlphaBitCore/nexus-gateway/packages/ai-gateway/internal/routing/strategies" ) +// awaitStatus polls until a provider reaches the expected health status. The +// HealthTracker applies records asynchronously (single writer), so a read +// immediately after recording must wait for the samples to be published. +func awaitStatus(t *testing.T, tracker *store.HealthTracker, providerID string, want store.HealthStatus) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for { + if tracker.GetHealth(providerID).Status == want { + return + } + if time.Now().After(deadline) { + t.Fatalf("provider %s did not reach %s in time; got %+v", providerID, want, tracker.GetHealth(providerID)) + } + time.Sleep(2 * time.Millisecond) + } +} + // Test helpers for routing package coverage tests. // These mirror the helpers in resolver_test.go to avoid the import cycle that // would occur if this file referenced the routing_test fakeStore directly. @@ -164,6 +182,7 @@ func TestFormatTarget_PartialFields(t *testing.T) { func TestHealthRanker_DegradedAndUnavailable_RankedLast(t *testing.T) { tracker := store.NewHealthTracker() + t.Cleanup(tracker.Stop) // degraded: ~10% failure rate (just above 0.05 threshold). // thresholdDegraded = 0.05, thresholdUnavailable = 0.25. // One failure out of 11 samples = ~0.09 error rate => degraded. @@ -176,6 +195,8 @@ func TestHealthRanker_DegradedAndUnavailable_RankedLast(t *testing.T) { for range 5 { tracker.RecordFailure("provUnavail", "una", 200) } + awaitStatus(t, tracker, "provDegraded", store.HealthStatusDegraded) + awaitStatus(t, tracker, "provUnavail", store.HealthStatusUnavailable) ranker := core.NewHealthRanker(tracker) in := []core.RoutingTarget{ @@ -270,7 +291,9 @@ func TestNarrowingEngine_Apply_SkipsNonStage0AndNonPolicy(t *testing.T) { func TestNewResolver_WiresAllFields(t *testing.T) { fs := &coverageFakeStore{providers: map[string]*store.Provider{}, models: map[string]*store.Model{}} reg := strategies.NewStrategyRegistry() - ranker := core.NewHealthRanker(store.NewHealthTracker()) + ht := store.NewHealthTracker() + t.Cleanup(ht.Stop) + ranker := core.NewHealthRanker(ht) logger := slog.New(slog.NewTextHandler(io.Discard, nil)) r := NewResolver(fs, reg, ranker, logger, nil) if r == nil { @@ -819,9 +842,11 @@ func TestResolveTargets_FlattensAndHealthRanks(t *testing.T) { // Mark openai unhealthy so ranker should sink it to the end. tracker := store.NewHealthTracker() + t.Cleanup(tracker.Stop) for range 20 { tracker.RecordFailure("openai", "openai", 100) } + awaitStatus(t, tracker, "openai", store.HealthStatusUnavailable) f.resolver.healthRanker = core.NewHealthRanker(tracker) res, err := f.resolver.ResolveTargets(context.Background(), &core.RoutingContext{ diff --git a/packages/compliance-proxy/cmd/compliance-proxy/main.go b/packages/compliance-proxy/cmd/compliance-proxy/main.go index 147ece7d..f605d19d 100644 --- a/packages/compliance-proxy/cmd/compliance-proxy/main.go +++ b/packages/compliance-proxy/cmd/compliance-proxy/main.go @@ -157,6 +157,13 @@ func run() int { ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGTERM, syscall.SIGINT) defer stop() + // Arm the hook-config TTL-backstop ticker (wiring already did the + // fail-hard initial load; this re-load is harmless and the ticker + // covers a degraded Hub push channel for the process lifetime). + if err := compRes.HookConfigCache.Start(ctx); err != nil { + logger.Warn("hook config cache start failed", "error", err) + } + configKeyRecorder := runtimeintrospect.NewKeyStateRecorder() readiness := &atomic.Bool{} readiness.Store(true) diff --git a/packages/compliance-proxy/cmd/compliance-proxy/wiring/compliance.go b/packages/compliance-proxy/cmd/compliance-proxy/wiring/compliance.go index caa73ac1..4f64839e 100644 --- a/packages/compliance-proxy/cmd/compliance-proxy/wiring/compliance.go +++ b/packages/compliance-proxy/cmd/compliance-proxy/wiring/compliance.go @@ -107,8 +107,10 @@ func InitCompliance(cfg *config.Config, cacheManager *cache.Manager, auditWriter 2*time.Minute, logger, ) - // Redis and Start are handled after context creation in main.go. - // Perform initial load synchronously here. + // Initial load happens synchronously here and FAILS the boot on error — + // the compliance plane must never start serving with an empty hook + // config. main.go arms the TTL-backstop ticker via HookConfigCache.Start + // once the signal context exists; Hub push stays the primary update path. initCtx, initCancel := context.WithTimeout(context.Background(), 15*time.Second) if err := result.HookConfigCache.Reload(initCtx); err != nil { initCancel() diff --git a/packages/control-plane/internal/identity/authserver/login/displayname_test.go b/packages/control-plane/internal/identity/authserver/login/displayname_test.go index 39f43740..35333d0b 100644 --- a/packages/control-plane/internal/identity/authserver/login/displayname_test.go +++ b/packages/control-plane/internal/identity/authserver/login/displayname_test.go @@ -24,8 +24,8 @@ func TestSAMLNameWithFallback(t *testing.T) { {"given + family composed", []saml.Attribute{attr(givenURI, "Carol"), attr(surnameURI, "King")}, "", "Carol King"}, {"given only", []saml.Attribute{attr("givenName", "Dave")}, "", "Dave"}, {"no name + no email → empty", []saml.Attribute{attr("dept", "eng")}, "", ""}, - {"email-valued name falls to humanized nickname", []saml.Attribute{attr(fullNameURI, "steve.chen@alphabitcore.com"), attr("http://schemas.auth0.com/nickname", "steve.chen")}, "steve.chen@alphabitcore.com", "steve chen"}, - {"email-valued name, no nickname → humanized email local-part", []saml.Attribute{attr(fullNameURI, "steve.chen@alphabitcore.com")}, "steve.chen@alphabitcore.com", "steve chen"}, + {"email-valued name falls to humanized nickname", []saml.Attribute{attr(fullNameURI, "steve.chen@itechchoice.com"), attr("http://schemas.auth0.com/nickname", "steve.chen")}, "steve.chen@itechchoice.com", "steve chen"}, + {"email-valued name, no nickname → humanized email local-part", []saml.Attribute{attr(fullNameURI, "steve.chen@itechchoice.com")}, "steve.chen@itechchoice.com", "steve chen"}, {"nickname handle when no full/given name", []saml.Attribute{attr("nickname", "ada.lovelace")}, "", "ada lovelace"}, {"no name attribute → humanized email local-part", []saml.Attribute{attr("dept", "eng")}, "grace.hopper@navy.mil", "grace hopper"}, } @@ -45,7 +45,7 @@ func TestHumanizeHandle(t *testing.T) { want string }{ {"steve.chen", "steve chen"}, - {"steve.chen@alphabitcore.com", "steve chen"}, + {"steve.chen@itechchoice.com", "steve chen"}, {"steve_chen", "steve chen"}, {"john.doe.42", "john doe"}, {"gracehopper", "gracehopper"}, @@ -80,9 +80,9 @@ func TestOIDCDisplayName(t *testing.T) { {"preferred_username humanized when no name/given", map[string]any{"preferred_username": "grace.hopper"}, "g@x", "sub-3", "grace hopper"}, {"humanized email local-part when no name claims", map[string]any{}, "harry.styles@x.com", "sub-4", "harry styles"}, {"falls back to subject when no name + no email", map[string]any{}, "", "sub-5", "sub-5"}, - {"email-valued name falls to humanized preferred_username", map[string]any{"name": "steve.chen@alphabitcore.com", "preferred_username": "steve.chen"}, "steve.chen@alphabitcore.com", "sub-6", "steve chen"}, - {"email-valued name falls to humanized nickname", map[string]any{"name": "steve.chen@alphabitcore.com", "nickname": "steve.chen"}, "steve.chen@alphabitcore.com", "sub-7", "steve chen"}, - {"email-valued name, no handles → humanized email local-part", map[string]any{"name": "steve.chen@alphabitcore.com"}, "steve.chen@alphabitcore.com", "sub-8", "steve chen"}, + {"email-valued name falls to humanized preferred_username", map[string]any{"name": "steve.chen@itechchoice.com", "preferred_username": "steve.chen"}, "steve.chen@itechchoice.com", "sub-6", "steve chen"}, + {"email-valued name falls to humanized nickname", map[string]any{"name": "steve.chen@itechchoice.com", "nickname": "steve.chen"}, "steve.chen@itechchoice.com", "sub-7", "steve chen"}, + {"email-valued name, no handles → humanized email local-part", map[string]any{"name": "steve.chen@itechchoice.com"}, "steve.chen@itechchoice.com", "sub-8", "steve chen"}, } for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { diff --git a/packages/control-plane/internal/identity/authserver/store/federated_store_mock_test.go b/packages/control-plane/internal/identity/authserver/store/federated_store_mock_test.go index edb15bc6..ac88b925 100644 --- a/packages/control-plane/internal/identity/authserver/store/federated_store_mock_test.go +++ b/packages/control-plane/internal/identity/authserver/store/federated_store_mock_test.go @@ -269,10 +269,10 @@ func TestFederatedStore_RefreshUserProfile_Success(t *testing.T) { ctx := context.Background() mock.ExpectExec(`UPDATE "NexusUser"\s+SET "displayName" = COALESCE\(NULLIF\(\$2, ''\), "displayName"\),\s+email\s+= COALESCE\(NULLIF\(\$3, ''\), email\),\s+"updatedAt"\s+= NOW\(\)\s+WHERE id = \$1`). - WithArgs("user_1", "steve chen", "steve.chen@alphabitcore.com"). + WithArgs("user_1", "steve chen", "steve.chen@itechchoice.com"). WillReturnResult(pgxmock.NewResult("UPDATE", 1)) - if err := s.RefreshUserProfile(ctx, "user_1", "steve chen", "steve.chen@alphabitcore.com"); err != nil { + if err := s.RefreshUserProfile(ctx, "user_1", "steve chen", "steve.chen@itechchoice.com"); err != nil { t.Fatalf("RefreshUserProfile: %v", err) } } diff --git a/packages/shared/audit/ndjson/ndjson.go b/packages/shared/audit/ndjson/ndjson.go index c91c4295..806fd000 100644 --- a/packages/shared/audit/ndjson/ndjson.go +++ b/packages/shared/audit/ndjson/ndjson.go @@ -16,6 +16,7 @@ package ndjson import ( + "errors" "fmt" "io" "os" @@ -27,6 +28,16 @@ import ( "time" ) +// ErrSpoolQuotaExceeded is returned (wrapped) by Write / WriteBatch ONLY when the +// instance spool is already at its total-size quota — a SOFT, transient condition +// that a recovery sweeper relieves by draining + deleting sealed files. It is +// deliberately distinct from a genuine I/O failure (open / write / fsync / ENOSPC), +// which carries a different error: a caller running a no-loss loss mode retries +// (back-pressures) ONLY on this sentinel, and treats every other error as a hard +// drop — so a truly broken disk can never wedge the caller forever. Test with +// errors.Is(err, ErrSpoolQuotaExceeded). +var ErrSpoolQuotaExceeded = errors.New("audit/ndjson: spool quota exceeded") + // spoolFile is the file surface the writer needs. *os.File satisfies it in // production; tests substitute a fake to exercise the write/close/sync error // paths that a real filesystem will not produce on demand. @@ -163,7 +174,7 @@ func (w *Writer) Write(record []byte) error { // losing a record to a transient stat failure is worse than the small // risk of briefly exceeding the soft quota. if total, err := w.dirSize(); err == nil && total >= w.maxTotalSize { - return fmt.Errorf("audit/ndjson: instance spool %d bytes exceeds quota %d", total, w.maxTotalSize) + return fmt.Errorf("audit/ndjson: instance spool %d bytes exceeds quota %d: %w", total, w.maxTotalSize, ErrSpoolQuotaExceeded) } if w.currentFile != nil && w.currentSize+int64(len(line)) > w.maxFileSize { @@ -208,7 +219,7 @@ func (w *Writer) WriteBatch(block []byte) error { defer w.mu.Unlock() if total, err := w.dirSize(); err == nil && total >= w.maxTotalSize { - return fmt.Errorf("audit/ndjson: instance spool %d bytes exceeds quota %d", total, w.maxTotalSize) + return fmt.Errorf("audit/ndjson: instance spool %d bytes exceeds quota %d: %w", total, w.maxTotalSize, ErrSpoolQuotaExceeded) } if w.currentFile != nil && w.currentSize+int64(len(block)) > w.maxFileSize { @@ -339,8 +350,17 @@ func (w *Writer) SealedFiles() ([]string, error) { return out, nil } -// dirSize sums the sizes of the files in the instance spool directory. Must -// hold w.mu. +// dirSize sums the sizes of the RECLAIMABLE spool files in the instance directory +// — every audit-*.ndjson (the active file plus sealed ones a recovery sweeper can +// drain and delete). It deliberately EXCLUDES .poison sidecars (and any other +// non-audit entry): those are dead-lettered records recovery never drains or +// removes, so counting them against the quota would let undrainable content +// monotonically consume it. Under a no-loss loss mode that back-pressures on the +// quota, an undrainable-content-consumed quota would wedge audit intake forever; +// bounding the quota to reclaimable content keeps back-pressure resolvable while +// recovery drains (given a non-zero drain rate). Poison growth is bounded instead +// by the physical disk (a genuine ENOSPC is a hard, counted drop — not this quota) +// and surfaced via the recovery-poisoned counter. Must hold w.mu. func (w *Writer) dirSize() (int64, error) { instanceDir := filepath.Join(w.dir, w.instanceID) entries, err := readDir(instanceDir) @@ -352,6 +372,10 @@ func (w *Writer) dirSize() (int64, error) { if entry.IsDir() { continue } + name := entry.Name() + if !strings.HasPrefix(name, "audit-") || !strings.HasSuffix(name, ".ndjson") { + continue // exclude .poison + any non-audit file from the reclaimable quota + } info, err := entry.Info() if err != nil { continue diff --git a/packages/shared/audit/ndjson/ndjson_quota_test.go b/packages/shared/audit/ndjson/ndjson_quota_test.go new file mode 100644 index 00000000..5099e302 --- /dev/null +++ b/packages/shared/audit/ndjson/ndjson_quota_test.go @@ -0,0 +1,87 @@ +package ndjson + +import ( + "errors" + "os" + "path/filepath" + "testing" +) + +// TestQuotaGate_WrapsSentinel proves the total-quota refusal carries the +// ErrSpoolQuotaExceeded sentinel (so a no-loss caller can tell "spool full, +// retryable" from a genuine I/O failure), on BOTH Write and WriteBatch. +func TestQuotaGate_WrapsSentinel(t *testing.T) { + dir := t.TempDir() + w, err := New(dir, "inst", 1, 1, nil) // 1 MB file cap, 1 MB total quota + if err != nil { + t.Fatalf("New: %v", err) + } + // Seed a countable audit-*.ndjson file past the quota. + if err := os.WriteFile(filepath.Join(dir, "inst", "audit-20260101-0001.ndjson"), + make([]byte, 1100*1024), 0o600); err != nil { + t.Fatalf("seed: %v", err) + } + + errWrite := w.Write([]byte(`{"id":"a"}`)) + if !errors.Is(errWrite, ErrSpoolQuotaExceeded) { + t.Fatalf("Write over quota must wrap ErrSpoolQuotaExceeded; got %v", errWrite) + } + errBatch := w.WriteBatch([]byte("{\"id\":\"a\"}\n")) + if !errors.Is(errBatch, ErrSpoolQuotaExceeded) { + t.Fatalf("WriteBatch over quota must wrap ErrSpoolQuotaExceeded; got %v", errBatch) + } +} + +// TestGenuineWriteError_NotQuotaSentinel proves a real I/O failure (injected via the +// openSpool seam) does NOT match the quota sentinel — so a no-loss caller drops / +// dead-letters it rather than back-pressuring forever on a broken disk. +func TestGenuineWriteError_NotQuotaSentinel(t *testing.T) { + dir := t.TempDir() + w, err := New(dir, "inst", 64, 512, nil) // well under quota + if err != nil { + t.Fatalf("New: %v", err) + } + orig := openSpool + t.Cleanup(func() { openSpool = orig }) + openSpool = func(string) (spoolFile, int64, error) { + return nil, 0, errors.New("audit/ndjson: open file: injected io error") + } + got := w.Write([]byte(`{"id":"x"}`)) + if got == nil { + t.Fatal("Write must surface the injected open error") + } + if errors.Is(got, ErrSpoolQuotaExceeded) { + t.Fatalf("a genuine I/O error must NOT match the quota sentinel; got %v", got) + } +} + +// TestDirSizeQuota_ExcludesPoison proves the reclaimable-quota gate counts only +// audit-*.ndjson and ignores .poison sidecars (and any non-audit entry). Without +// this, undrainable poison would monotonically consume the quota and — under a +// back-pressuring no-loss caller — wedge audit intake forever. +func TestDirSizeQuota_ExcludesPoison(t *testing.T) { + dir := t.TempDir() + w, err := New(dir, "inst", 1, 1, nil) // 1 MB total quota + if err != nil { + t.Fatalf("New: %v", err) + } + // A big .poison sidecar (well over the quota) that must NOT count. + if err := os.WriteFile(filepath.Join(dir, "inst", "audit-20260101-0001.ndjson.poison"), + make([]byte, 2*1024*1024), 0o600); err != nil { + t.Fatalf("seed poison: %v", err) + } + // A write must SUCCEED — the 2 MB poison does not count toward the 1 MB quota. + if err := w.Write([]byte(`{"id":"ok"}`)); err != nil { + t.Fatalf("write must succeed when only .poison exceeds the quota; got %v", err) + } + + // Now a big countable audit-*.ndjson file DOES push the reclaimable total over + // quota → the next write is refused with the sentinel. + if err := os.WriteFile(filepath.Join(dir, "inst", "audit-20260101-0002.ndjson"), + make([]byte, 1100*1024), 0o600); err != nil { + t.Fatalf("seed audit: %v", err) + } + if got := w.Write([]byte(`{"id":"no"}`)); !errors.Is(got, ErrSpoolQuotaExceeded) { + t.Fatalf("a countable audit-*.ndjson over quota must refuse with the sentinel; got %v", got) + } +} diff --git a/packages/shared/core/bytebudget/bytebudget.go b/packages/shared/core/bytebudget/bytebudget.go new file mode 100644 index 00000000..cd3c82b4 --- /dev/null +++ b/packages/shared/core/bytebudget/bytebudget.go @@ -0,0 +1,138 @@ +// Package bytebudget is a lock-free, soft byte semaphore: it bounds the total bytes +// concurrently "in use" to a budget, blocking Acquire while the budget is exhausted +// and waking blocked acquirers when Release frees space. It is a reusable +// back-pressure primitive for any producer/consumer pipeline whose real resource is +// BYTES, not item count — e.g. an in-memory queue that pins variable-size payloads: a +// producer Acquires its item's byte weight before enqueuing and Releases it when the +// item leaves the pipeline. A full budget blocks the producer, which back-pressures +// upstream instead of dropping or growing without bound. +// +// SOFT and LOCK-FREE by design (no mutex): `used` is an atomic and the +// check-then-add in Acquire races, so the budget may be OVERSHOT by up to roughly +// (concurrent acquirers × their weights). That is deliberate — this is a memory +// back-pressure knob, not exact accounting. A zero-overshoot bound is possible +// lock-free too (a CompareAndSwap retry loop), but it contends hardest exactly at +// the full-budget boundary where every producer is hammering the same word; the +// single wait-free Add never retries, and the overshoot it admits is bounded by +// producer concurrency — negligible against a memory-scale budget configured +// below real headroom. +// +// Wake-up uses a cap-1 buffered channel signalled non-blockingly on Release, plus a +// short polling fallback in Acquire, so a signal that races an about-to-wait acquirer +// can never wedge it forever (a missed edge is re-checked within WakePollInterval). +package bytebudget + +import ( + "sync/atomic" + "time" +) + +// WakePollInterval bounds how long a blocked Acquire waits before re-checking the +// budget even with no Release signal — the safety net against a missed wake-up edge +// (the atomic check and the channel signal are not one operation, so a Release can +// slip between an acquirer's over-budget check and its channel wait). Short enough to +// pick up freed budget promptly, long enough that a fully blocked pipeline does not +// spin. +const WakePollInterval = 20 * time.Millisecond + +// Budget is a soft, lock-free byte semaphore. The zero value is not usable; call New. +type Budget struct { + budget int64 + used atomic.Int64 + notFull chan struct{} // cap 1; a non-blocking send on Release nudges one waiter + stop <-chan struct{} +} + +// New makes a budget of at most `budget` bytes. A non-positive budget disables +// back-pressure (Acquire always admits) — a caller can wire it to turn the bound off. +// stop unblocks any waiter when it is closed/fired (shutdown), so a Close that drains +// the pipeline stays bounded; pass nil for a budget that only ever unblocks on Release. +func New(budget int64, stop <-chan struct{}) *Budget { + return &Budget{ + budget: budget, + notFull: make(chan struct{}, 1), + stop: stop, + } +} + +// TryAcquire reserves n bytes only when the budget has room RIGHT NOW — the +// non-blocking admission. Same rule as Acquire (room, or an empty pipeline so an +// oversized n is never wedged), same soft race. Two callers: pipelines that must +// never block (lossy modes shed the item on false), and blocking callers that +// want to observe "about to wait" (TryAcquire first, count, then Acquire). +func (b *Budget) TryAcquire(n int64) bool { + if b == nil || b.budget <= 0 { + return true + } + // Admit if there is room, OR if the pipeline is empty (used<=0) so an oversized + // item (n > budget) is not stuck forever. Soft: the check and the Add race, so + // used may briefly exceed budget — accepted (see package doc). + if u := b.used.Load(); u < b.budget || u <= 0 { + b.used.Add(n) + return true + } + return false +} + +// Acquire reserves n bytes, blocking while the budget is exhausted. It returns true +// once reserved, or false if stop fired first (shutdown) — the caller then handles the +// item without enqueuing. An n larger than the whole budget is admitted alone once +// used drops to 0, so a single item bigger than the budget can never deadlock the +// pipeline. A non-positive budget admits immediately (back-pressure disabled). +func (b *Budget) Acquire(n int64) bool { + if b == nil || b.budget <= 0 { + return true + } + if b.TryAcquire(n) { // fast path: no timer allocated while the budget has room + return true + } + // One reusable timer per blocked Acquire, Reset each round — a time.After per + // poll iteration allocated a fresh runtime timer every ≤20ms for EVERY blocked + // producer, pure garbage churn in exactly the wedged regime. Reset without a + // drain is safe on Go 1.23+ timer semantics (unbuffered channel; Reset discards + // a pending fire). + t := time.NewTimer(WakePollInterval) + defer t.Stop() + for { + select { + case <-b.notFull: + case <-t.C: + case <-b.stop: + return false + } + if b.TryAcquire(n) { + return true + } + t.Reset(WakePollInterval) + } +} + +// Release returns n bytes to the budget and nudges one blocked acquirer. Safe on a +// nil receiver / disabled budget (no-op). n must match the Acquire that reserved it. +func (b *Budget) Release(n int64) { + if b == nil || b.budget <= 0 || n == 0 { + return + } + b.used.Add(-n) + select { + case b.notFull <- struct{}{}: // wake one waiter; cap-1 buffer coalesces bursts + default: + } +} + +// InUse reports the current reserved bytes (metrics / tests). May transiently exceed +// the budget or go negative under the soft race — treat it as approximate. +func (b *Budget) InUse() int64 { + if b == nil { + return 0 + } + return b.used.Load() +} + +// Budget returns the configured ceiling (0 = disabled). +func (b *Budget) Budget() int64 { + if b == nil { + return 0 + } + return b.budget +} diff --git a/packages/shared/core/bytebudget/bytebudget_test.go b/packages/shared/core/bytebudget/bytebudget_test.go new file mode 100644 index 00000000..10a04d96 --- /dev/null +++ b/packages/shared/core/bytebudget/bytebudget_test.go @@ -0,0 +1,184 @@ +package bytebudget + +import ( + "sync" + "testing" + "time" +) + +// A full budget must block Acquire until Release frees space — the core +// back-pressure contract: the producer parks, then proceeds exactly when the +// consumer returns bytes. +func TestAcquire_BlocksAtBudgetAndWakesOnRelease(t *testing.T) { + b := New(100, nil) + if !b.Acquire(100) { // fills the budget exactly + t.Fatal("first Acquire within budget must admit") + } + acquired := make(chan struct{}) + go func() { + b.Acquire(50) // over budget → must park + close(acquired) + }() + select { + case <-acquired: + t.Fatal("Acquire admitted while the budget was exhausted") + case <-time.After(50 * time.Millisecond): + } + b.Release(100) + select { + case <-acquired: + case <-time.After(2 * time.Second): + t.Fatal("Acquire did not wake after Release freed the budget") + } + if got := b.InUse(); got != 50 { + t.Fatalf("InUse after acquire/release cycle = %d, want 50", got) + } +} + +// TryAcquire is the non-blocking admission: it must admit while there is room +// and refuse (without blocking) once the budget is exhausted, then admit again +// after Release. +func TestTryAcquire_NonBlockingAdmission(t *testing.T) { + b := New(100, nil) + if !b.TryAcquire(60) { + t.Fatal("TryAcquire with room must admit") + } + if !b.TryAcquire(60) { // used=60 < budget=100 → admits (soft overshoot to 120) + t.Fatal("TryAcquire below the budget line must admit even if n overshoots") + } + if b.TryAcquire(1) { // used=120 ≥ budget → refuse + t.Fatal("TryAcquire on an exhausted budget must refuse") + } + b.Release(120) + if !b.TryAcquire(1) { + t.Fatal("TryAcquire after Release must admit again") + } +} + +// The soft budget admits whenever used < budget, so the overshoot is bounded by +// one in-flight acquisition per concurrent producer — never unbounded. +func TestAcquire_SoftOvershootBounded(t *testing.T) { + b := New(100, nil) + if !b.Acquire(99) || !b.Acquire(70) { // 99 < 100 → second admits, used=169 + t.Fatal("soft admission below the budget line must admit") + } + if got := b.InUse(); got != 169 { + t.Fatalf("InUse = %d, want 169 (bounded overshoot)", got) + } + if b.TryAcquire(1) { + t.Fatal("no further admission once used ≥ budget") + } +} + +// An item larger than the whole budget must be admitted once the pipeline is +// empty — a single oversized item can never deadlock the queue. +func TestAcquire_OversizedItemAdmittedWhenEmpty(t *testing.T) { + b := New(100, nil) + if !b.Acquire(500) { + t.Fatal("oversized item on an empty pipeline must admit") + } + if b.TryAcquire(1) { + t.Fatal("budget exhausted by the oversized item must refuse the next") + } + b.Release(500) + if !b.Acquire(500) { + t.Fatal("oversized item must admit again once drained to empty") + } +} + +// A closed stop channel must unblock a parked Acquire with false so shutdown +// never hangs a producer; the caller then handles the item without enqueuing. +func TestAcquire_StopEscapesWithFalse(t *testing.T) { + stop := make(chan struct{}) + b := New(100, stop) + b.Acquire(100) + res := make(chan bool, 1) + go func() { res <- b.Acquire(1) }() + select { + case <-res: + t.Fatal("Acquire returned before stop fired on a full budget") + case <-time.After(50 * time.Millisecond): + } + close(stop) + select { + case ok := <-res: + if ok { + t.Fatal("Acquire after stop must return false") + } + case <-time.After(2 * time.Second): + t.Fatal("Acquire did not unblock on stop") + } +} + +// Even when the Release signal races an about-to-park acquirer (the notFull +// nudge can be consumed before the acquirer waits), the WakePollInterval +// re-check must recover — no missed-wakeup wedge. Hammer the handoff to give a +// real race a chance to bite under -race. +func TestAcquireRelease_NoMissedWakeupUnderChurn(t *testing.T) { + b := New(64, nil) + const producers = 8 + const rounds = 200 + var wg sync.WaitGroup + wg.Add(producers) + for range producers { + go func() { + defer wg.Done() + for range rounds { + if !b.Acquire(16) { + t.Error("Acquire returned false with no stop channel") + return + } + b.Release(16) + } + }() + } + done := make(chan struct{}) + go func() { wg.Wait(); close(done) }() + select { + case <-done: + case <-time.After(30 * time.Second): + t.Fatal("producers wedged — missed wakeup") + } + if got := b.InUse(); got != 0 { + t.Fatalf("InUse after balanced acquire/release churn = %d, want 0", got) + } +} + +// A non-positive budget disables back-pressure entirely: every Acquire admits, +// Release is a no-op — the off switch a caller wires by config. +func TestDisabledBudget_AlwaysAdmits(t *testing.T) { + b := New(0, nil) + if !b.Acquire(1<<40) || !b.TryAcquire(1<<40) { + t.Fatal("disabled budget must always admit") + } + b.Release(1 << 40) // must not underflow anything + if got := b.InUse(); got != 0 { + t.Fatalf("disabled budget InUse = %d, want 0 (accounting off)", got) + } +} + +// Nil receivers are safe no-ops so an unwired consumer (tests constructing the +// owner struct directly) never panics or blocks. +func TestNilBudget_SafeNoOps(t *testing.T) { + var b *Budget + if !b.Acquire(10) || !b.TryAcquire(10) { + t.Fatal("nil budget must admit") + } + b.Release(10) + if b.InUse() != 0 || b.Budget() != 0 { + t.Fatal("nil budget must report zero usage and zero ceiling") + } +} + +// Budget reports the configured ceiling; Release(0) must not emit a wake signal +// (n==0 short-circuits before the atomic). +func TestBudgetAccessorsAndZeroRelease(t *testing.T) { + b := New(4096, nil) + if b.Budget() != 4096 { + t.Fatalf("Budget() = %d, want 4096", b.Budget()) + } + b.Release(0) + if got := b.InUse(); got != 0 { + t.Fatalf("InUse after Release(0) = %d, want 0", got) + } +} diff --git a/packages/shared/policy/hooks/validators/rulepack_engine_redact.go b/packages/shared/policy/hooks/validators/rulepack_engine_redact.go index 250da95c..f77aa040 100644 --- a/packages/shared/policy/hooks/validators/rulepack_engine_redact.go +++ b/packages/shared/policy/hooks/validators/rulepack_engine_redact.go @@ -87,6 +87,15 @@ func (e *RulePackEngine) addressedSegments(input *core.HookInput) []addressedTex // the few rules that hit. Severity still gates enforcement: warn/info matches // only tag (observe), they do not mask. func (e *RulePackEngine) executeRedact(input *core.HookInput, result *core.HookResult, matched map[[2]int]struct{}, complete bool, start time.Time) (*core.HookResult, error) { + // Zero matches on a COMPLETE scan: nothing can tag or mask, so skip the + // re-localization machinery (segment addressing, per-rule loop) — benign + // traffic is the overwhelmingly common case and paid it for nothing. An + // INCOMPLETE scan must fall through: an empty hit set proves nothing + // there, and the fail-safe below re-confirms every rule with RE2. + if complete && len(matched) == 0 { + result.LatencyMs = int(time.Since(start).Milliseconds()) + return result, nil + } var matchedRules map[int]struct{} if complete { matchedRules = make(map[int]struct{}, len(matched)) diff --git a/packages/shared/policy/hooks/validators/rulepack_engine_redact_fastpath_test.go b/packages/shared/policy/hooks/validators/rulepack_engine_redact_fastpath_test.go new file mode 100644 index 00000000..f93cb428 --- /dev/null +++ b/packages/shared/policy/hooks/validators/rulepack_engine_redact_fastpath_test.go @@ -0,0 +1,67 @@ +package validators + +import ( + "testing" + "time" + + core "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// TestRulePackEngine_Redact_BenignNoMatchContract locks the zero-match +// contract of the redact path on a COMPLETE scan: benign content yields an +// untouched Approve result — no tags, no spans, no modified content. The +// zero-match early return must be behaviorally identical to walking the +// full rule loop with an empty hit set. +func TestRulePackEngine_Redact_BenignNoMatchContract(t *testing.T) { + cfg := buildEngineConfig([]rulePackInstall{{ + InstallID: "i", PackName: "pii", PackVersion: "v", Enabled: true, + Rules: []rulePackRule{{RuleID: "ssn", Category: "pii.ssn", Severity: "hard", Pattern: `\b\d{3}-\d{2}-\d{4}\b`}}, + }}) + cfg.Config["onMatch"] = map[string]any{"action": "redact"} + h, err := NewRulePackEngine(cfg) + if err != nil { + t.Fatalf("NewRulePackEngine: %v", err) + } + res, err := h.Execute(t.Context(), &HookInput{Normalized: PayloadFromTextSegments([]string{"a perfectly benign sentence"})}) + if err != nil { + t.Fatalf("Execute: %v", err) + } + if res.Decision != Approve { + t.Errorf("benign content must Approve, got %s", res.Decision) + } + if len(res.Tags) != 0 || len(res.TransformSpans) != 0 || len(res.ModifiedContent) != 0 { + t.Errorf("benign content must leave the result untouched; tags=%v spans=%v modified=%v", + res.Tags, res.TransformSpans, res.ModifiedContent) + } +} + +// TestRulePackEngine_Redact_IncompleteScanFailSafe guards the fail-safe the +// zero-match early return must NOT swallow: when the matcher scan was +// INCOMPLETE, an empty hit set proves nothing — executeRedact must still +// re-confirm every rule with RE2 and mask the PII the truncated scan missed. +func TestRulePackEngine_Redact_IncompleteScanFailSafe(t *testing.T) { + cfg := buildEngineConfig([]rulePackInstall{{ + InstallID: "i", PackName: "pii", PackVersion: "v", Enabled: true, + Rules: []rulePackRule{{RuleID: "ssn", Category: "pii.ssn", Severity: "hard", Pattern: `\b\d{3}-\d{2}-\d{4}\b`}}, + }}) + cfg.Config["onMatch"] = map[string]any{"action": "redact"} + h, err := NewRulePackEngine(cfg) + if err != nil { + t.Fatalf("NewRulePackEngine: %v", err) + } + e, ok := h.(*RulePackEngine) + if !ok { + t.Fatalf("expected *RulePackEngine, got %T", h) + } + input := &HookInput{Normalized: PayloadFromTextSegments([]string{"ssn 123-45-6789 leaked"})} + result := &core.HookResult{Decision: Approve} + // Empty matched + complete=false = the truncated-scan worst case. + res, err := e.executeRedact(input, result, map[[2]int]struct{}{}, false, time.Now()) + if err != nil { + t.Fatalf("executeRedact: %v", err) + } + if res.Decision != core.Modify || len(res.TransformSpans) == 0 { + t.Fatalf("incomplete scan with empty hit set must still fail-safe redact; decision=%s spans=%d", + res.Decision, len(res.TransformSpans)) + } +} diff --git a/packages/shared/policy/pipeline/config_cache.go b/packages/shared/policy/pipeline/config_cache.go index 534dbccb..99b4630c 100644 --- a/packages/shared/policy/pipeline/config_cache.go +++ b/packages/shared/policy/pipeline/config_cache.go @@ -2,6 +2,7 @@ package pipeline import ( "context" + "errors" "log/slog" "sync" "time" @@ -18,25 +19,41 @@ type HookConfigLoader func(ctx context.Context) ([]core.HookConfig, error) // shared/configcache.SnapshotCache for atomic-pointer swap and metrics; // on top of the snapshot a PolicyResolver holds the compiled pipeline state. // +// Config changes are low-frequency, push-driven events: the Hub thingclient +// OnConfigChanged callback calls Reload when an admin edits hook config. The +// request path NEVER loads — Resolver is a pure snapshot getter, so a slow +// database can never stall or stampede request goroutines (thousands of +// concurrent requests each observing a stale TTL and issuing their own full +// reload is exactly the failure mode this design forbids). +// // Two operating modes via the ttl argument to NewHookConfigCache: // // - ttl > 0 (AI Gateway, Compliance Proxy): push-driven invalidation via -// Hub thingclient OnConfigChanged calling Reload, plus a TTL backstop -// on Resolver() that closes any gap if the push channel is degraded. +// Reload, plus a background ticker started by Start that refreshes every +// ttl as a backstop if the push channel is degraded. The backstop runs +// off the request path and holds at most one load in flight. // -// - ttl = 0 (Agent): pure push mode. Resolver() never auto-reloads; the -// configsync layer pulls hook configs from Hub and calls Reload. The -// agent has no direct DB access, so no TTL backstop is meaningful. +// - ttl = 0 (Agent): pure push mode. No ticker; the configsync layer +// pulls hook configs from Hub and calls Reload. The agent has no direct +// DB access, so no TTL backstop is meaningful. type HookConfigCache struct { snap *configcache.SnapshotCache[core.HookConfig] resolver *PolicyResolver ttl time.Duration logger *slog.Logger - mu sync.Mutex - lastLoad time.Time + // loadMu serializes snapshot loads (backstop tick, Hub push Reload, + // startup). Loads commit in start order, so the last committer read the + // freshest rows — a slow backstop load can never overwrite a newer push + // result with stale content. Never touched on the request path. + loadMu sync.Mutex } +// ttlRefreshTimeout bounds a single background backstop load so a wedged +// database cannot pin the refresh goroutine's context forever. Loads are +// serial (one ticker loop), so a slow load also cannot stack. +const ttlRefreshTimeout = 30 * time.Second + // NewHookConfigCache creates a cache. func NewHookConfigCache(loader HookConfigLoader, registry *core.HookRegistry, ttl time.Duration, logger *slog.Logger) *HookConfigCache { if logger == nil { @@ -79,49 +96,78 @@ func NewHookConfigCache(loader HookConfigLoader, registry *core.HookRegistry, tt // generation under load. Real changes (startup, Hub push delta) differ // and swap normally. c.resolver.SwapIfContentChanged(cfgs) - c.mu.Lock() - c.lastLoad = time.Now() - c.mu.Unlock() // Build the (potentially expensive) engines off the request path. // Background and best-effort: it must never block this callback, - // which runs both at startup and on a TTL-staleness reload that a - // request triggers. The lazy build in resolve() is the fallback if a - // request arrives before prewarm finishes. swapGen guards staleness. + // which runs at startup, on Hub push, and on the backstop ticker. + // The lazy build in resolve() is the fallback if a request arrives + // before prewarm finishes. swapGen guards staleness. go c.resolver.Prewarm() }), ) return c } -// Start performs the initial load. The TTL backstop and any external -// invalidation source (currently the Hub thingclient OnConfigChanged -// callback) drive subsequent reloads. +// Start performs the initial load and, in TTL mode, starts the background +// backstop ticker that refreshes the snapshot every ttl. The ticker stops +// when ctx is canceled, so callers must pass a process-lifetime context. +// Push invalidation (Hub thingclient OnConfigChanged → Reload) remains the +// primary update path; the ticker only closes the gap when push is degraded. func (c *HookConfigCache) Start(ctx context.Context) error { - if err := c.snap.Load(ctx); err != nil { + if err := c.load(ctx); err != nil { c.logger.Warn("initial hook config load failed, continuing with empty config", "error", err) } + if c.ttl > 0 { + go c.refreshLoop(ctx) + } return nil } -// Resolver returns the PolicyResolver with the current config snapshot. -// In TTL mode (ttl > 0) triggers a reload if the cache is stale. In -// pure push mode (ttl == 0) returns the resolver directly without any -// TTL bookkeeping. -func (c *HookConfigCache) Resolver(ctx context.Context) *PolicyResolver { - if c.ttl > 0 { - c.mu.Lock() - stale := time.Since(c.lastLoad) > c.ttl - c.mu.Unlock() - if stale { - _ = c.Reload(ctx) +// refreshLoop is the TTL backstop: one serial background loop, so at most +// one backstop load is ever in flight regardless of load or DB latency. +func (c *HookConfigCache) refreshLoop(ctx context.Context) { + ticker := time.NewTicker(c.ttl) + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + if err := c.load(ctx); err != nil { + c.logger.Warn("hook config backstop refresh failed; previous snapshot stays active", "error", err) + } } } +} + +// load runs one serialized, deadline-bounded snapshot load. Every load — +// startup, backstop tick, Hub push — goes through here, so a wedged +// database can pin at most one goroutine for at most ttlRefreshTimeout, +// and loads commit in start order (see loadMu). +func (c *HookConfigCache) load(ctx context.Context) error { + // A nil context must fail gracefully (as snap.Load does) rather than panic + // in WithTimeout — the agent's shadow-apply path surfaces a nil-context + // reload as a normal error and keeps the prior policy. + if ctx == nil { + return errors.New("configcache: nil context") + } + c.loadMu.Lock() + defer c.loadMu.Unlock() + rctx, cancel := context.WithTimeout(ctx, ttlRefreshTimeout) + defer cancel() + return c.snap.Load(rctx) +} + +// Resolver returns the PolicyResolver holding the current config snapshot. +// Pure getter on the request hot path: freshness is owned by push (Reload) +// and the Start backstop ticker, never by request goroutines. +func (c *HookConfigCache) Resolver(_ context.Context) *PolicyResolver { return c.resolver } -// Reload forces an immediate reload from the database. +// Reload forces an immediate reload from the database. Called by the Hub +// push-invalidation path (thingclient OnConfigChanged) and boot wiring. func (c *HookConfigCache) Reload(ctx context.Context) error { - return c.snap.Load(ctx) + return c.load(ctx) } // HookSnapshotEntry is the redacted view of a HookConfig used by runtime diff --git a/packages/shared/policy/pipeline/config_cache_refresh_test.go b/packages/shared/policy/pipeline/config_cache_refresh_test.go new file mode 100644 index 00000000..61ce25e4 --- /dev/null +++ b/packages/shared/policy/pipeline/config_cache_refresh_test.go @@ -0,0 +1,203 @@ +package pipeline + +import ( + "context" + "errors" + "log/slog" + "strings" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/builtins" + "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// waitFor polls cond until it returns true or the timeout elapses. +func waitFor(t *testing.T, timeout time.Duration, cond func() bool) bool { + t.Helper() + deadline := time.Now().Add(timeout) + for time.Now().Before(deadline) { + if cond() { + return true + } + time.Sleep(2 * time.Millisecond) + } + return cond() +} + +// TestHookConfigCache_StaleResolverDoesNotBlock asserts the request path +// never waits on a TTL refresh: with a slow loader and an expired TTL, +// Resolver must return the current snapshot immediately instead of +// blocking behind the database load. +func TestHookConfigCache_StaleResolverDoesNotBlock(t *testing.T) { + var calls atomic.Int32 + loader := func(ctx context.Context) ([]core.HookConfig, error) { + if calls.Add(1) > 1 { // every call after the initial Start load is slow + time.Sleep(200 * time.Millisecond) + } + return nil, nil + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cache := NewHookConfigCache(loader, builtins.Registry, 5*time.Millisecond, slog.Default()) + _ = cache.Start(ctx) + time.Sleep(10 * time.Millisecond) // let the TTL expire + + start := time.Now() + resolver := cache.Resolver(context.Background()) + elapsed := time.Since(start) + + if resolver == nil { + t.Fatal("resolver nil") + } + if elapsed > 50*time.Millisecond { + t.Fatalf("Resolver blocked %v on a stale refresh; the request path must serve the current snapshot immediately", elapsed) + } + // The background refresh must still happen. + if !waitFor(t, time.Second, func() bool { return calls.Load() >= 2 }) { + t.Fatalf("expected a background TTL refresh, loader calls=%d", calls.Load()) + } +} + +// TestHookConfigCache_StaleStampedeSingleFlight asserts that concurrent +// requests arriving while the snapshot is stale trigger AT MOST ONE +// reload — not one per request. The request path never loads; the serial +// backstop ticker holds the only in-flight load. +func TestHookConfigCache_StaleStampedeSingleFlight(t *testing.T) { + var calls atomic.Int32 + release := make(chan struct{}) + loader := func(ctx context.Context) ([]core.HookConfig, error) { + if calls.Add(1) > 1 { // calls after Start block until released + <-release + } + return nil, nil + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cache := NewHookConfigCache(loader, builtins.Registry, 5*time.Millisecond, slog.Default()) + _ = cache.Start(ctx) + time.Sleep(10 * time.Millisecond) // expire the TTL + + const n = 50 + var wg sync.WaitGroup + for range n { + wg.Add(1) + go func() { + defer wg.Done() + _ = cache.Resolver(context.Background()) + }() + } + wg.Wait() + + // All 50 concurrent requests must have triggered zero loads of their own; + // only the backstop ticker's single in-flight load may have run. + if got := calls.Load(); got > 2 { + t.Fatalf("stampede: %d loader calls for %d concurrent stale requests (want <= 2 incl. Start)", got, n) + } + close(release) +} + +// TestHookConfigCache_AsyncRefreshAppliesNewConfig asserts a background +// TTL refresh actually lands: after the refresh completes, the resolver +// serves the newly loaded config. +func TestHookConfigCache_AsyncRefreshAppliesNewConfig(t *testing.T) { + var calls atomic.Int32 + loader := func(ctx context.Context) ([]core.HookConfig, error) { + if calls.Add(1) == 1 { + return nil, nil // initial: no hooks + } + return []core.HookConfig{ + {ID: "h1", ImplementationID: "keyword-filter", Name: "kw", Enabled: true, Stage: "request", + Config: map[string]any{"patterns": []any{map[string]any{"pattern": "bad", "category": "test", "severity": "hard"}}}}, + }, nil + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cache := NewHookConfigCache(loader, builtins.Registry, 5*time.Millisecond, slog.Default()) + _ = cache.Start(ctx) + if cache.Resolver(context.Background()).HasHooks("request") { + t.Fatal("precondition: no hooks after initial empty load") + } + + time.Sleep(10 * time.Millisecond) // expire TTL + _ = cache.Resolver(context.Background()) + + if !waitFor(t, time.Second, func() bool { + return cache.Resolver(context.Background()).HasHooks("request") + }) { + t.Fatal("background refresh did not apply the new config") + } +} + +// TestHookConfigCache_FailedRefreshRetriesNextStale asserts a failed +// backstop refresh does not wedge the loop: the next tick tries again +// and the previous snapshot stays active in between. +func TestHookConfigCache_FailedRefreshRetriesNextStale(t *testing.T) { + var calls atomic.Int32 + loader := func(ctx context.Context) ([]core.HookConfig, error) { + n := calls.Add(1) + if n == 2 { + return nil, errors.New("db down") + } + return nil, nil + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cache := NewHookConfigCache(loader, builtins.Registry, 5*time.Millisecond, slog.Default()) + _ = cache.Start(ctx) + + // Tick #1 fails (call 2). The loop must survive it… + if !waitFor(t, time.Second, func() bool { return calls.Load() >= 2 }) { + t.Fatalf("expected failing refresh attempt, calls=%d", calls.Load()) + } + // …and tick #2 must retry (call 3). + if !waitFor(t, time.Second, func() bool { return calls.Load() >= 3 }) { + t.Fatalf("failed refresh wedged the backstop loop; calls=%d", calls.Load()) + } +} + +// TestHookConfigCache_PushModeNeverAutoReloads asserts ttl=0 (Agent push +// mode) never triggers a load from the request path. +func TestHookConfigCache_PushModeNeverAutoReloads(t *testing.T) { + var calls atomic.Int32 + loader := func(ctx context.Context) ([]core.HookConfig, error) { + calls.Add(1) + return nil, nil + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + cache := NewHookConfigCache(loader, builtins.Registry, 0, slog.Default()) + _ = cache.Start(ctx) + time.Sleep(10 * time.Millisecond) + + for range 10 { + _ = cache.Resolver(context.Background()) + } + if got := calls.Load(); got != 1 { + t.Fatalf("push mode must not auto-reload; loader calls=%d (want 1)", got) + } +} + +// TestHookConfigCache_ReloadNilContext_ErrorsNotPanics pins that a nil-context +// Reload fails gracefully (as the agent shadow-apply path relies on to keep the +// prior policy) instead of panicking inside context.WithTimeout. +func TestHookConfigCache_ReloadNilContext_ErrorsNotPanics(t *testing.T) { + loader := func(ctx context.Context) ([]core.HookConfig, error) { return nil, nil } + cache := NewHookConfigCache(loader, builtins.Registry, 0, slog.Default()) + + var nilCtx context.Context + err := cache.Reload(nilCtx) + if err == nil { + t.Fatal("nil-context Reload must return an error, not succeed") + } + if !strings.Contains(err.Error(), "nil context") { + t.Fatalf("error must identify the nil context; got %q", err.Error()) + } +} diff --git a/packages/shared/policy/pipeline/config_cache_test.go b/packages/shared/policy/pipeline/config_cache_test.go index 2655f1ce..2059ae28 100644 --- a/packages/shared/policy/pipeline/config_cache_test.go +++ b/packages/shared/policy/pipeline/config_cache_test.go @@ -3,6 +3,7 @@ package pipeline import ( "context" "log/slog" + "sync/atomic" "testing" "time" @@ -38,20 +39,24 @@ func TestHookConfigCache_StartAndResolve(t *testing.T) { } func TestHookConfigCache_TTLRefresh(t *testing.T) { - loadCount := 0 + var loadCount atomic.Int32 loader := func(ctx context.Context) ([]core.HookConfig, error) { - loadCount++ + loadCount.Add(1) return nil, nil } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() cache := NewHookConfigCache(loader, builtins.Registry, 10*time.Millisecond, slog.Default()) - _ = cache.Start(context.Background()) + _ = cache.Start(ctx) - time.Sleep(20 * time.Millisecond) - _ = cache.Resolver(context.Background()) - - if loadCount < 2 { - t.Fatalf("expected TTL refresh, got %d loads", loadCount) + // The backstop ticker must refresh on its own — no request traffic needed. + deadline := time.Now().Add(time.Second) + for loadCount.Load() < 2 && time.Now().Before(deadline) { + time.Sleep(5 * time.Millisecond) + } + if loadCount.Load() < 2 { + t.Fatalf("expected TTL backstop refresh, got %d loads", loadCount.Load()) } } diff --git a/packages/shared/policy/rulepack/store_resolve.go b/packages/shared/policy/rulepack/store_resolve.go index 311c61ec..021b7266 100644 --- a/packages/shared/policy/rulepack/store_resolve.go +++ b/packages/shared/policy/rulepack/store_resolve.go @@ -104,13 +104,18 @@ func (s *Store) UpgradeInstallToLatest(ctx context.Context, installID string) (* // including the pack name for convenience. Disabled installs are included // (callers decide whether to evaluate them; the admin UI needs visibility). // Ordered by installedAt ASC so the rule-pack engine scans in install order. +// The id tiebreaker keeps the order deterministic when several installs share +// one timestamp (batch-seeded in a single transaction): unstable tie order +// would reorder the enriched `_rulePackInstalls` payload between loads, making +// a no-change config reload look like a change and churning the compiled +// matcher for nothing. func (s *Store) ListInstallsForHook(ctx context.Context, hookID string) ([]Install, error) { rows, err := s.pool.Query(ctx, ` SELECT i.id, i."packId", p.name, i."pinVersion", i."boundHookId", i.enabled, i."installedAt" FROM "rule_pack_install" i JOIN "rule_pack" p ON p.id = i."packId" WHERE i."boundHookId" = $1 - ORDER BY i."installedAt" ASC`, hookID) + ORDER BY i."installedAt" ASC, i.id ASC`, hookID) if err != nil { return nil, fmt.Errorf("rulepack.ListInstallsForHook: %w", err) } diff --git a/packages/shared/storage/configcache/metrics.go b/packages/shared/storage/configcache/metrics.go new file mode 100644 index 00000000..106a8d2e --- /dev/null +++ b/packages/shared/storage/configcache/metrics.go @@ -0,0 +1,23 @@ +package configcache + +import ( + "github.com/prometheus/client_golang/prometheus" + "github.com/prometheus/client_golang/prometheus/promauto" +) + +// Load-health metrics. A degraded push channel plus a failing database means +// the config plane is frozen on its previous snapshot — before these series +// existed that state was logs-only and invisible to alerting. Operators alert +// on loadFailures increasing, or on now() − lastLoadSuccess exceeding twice +// the service's TTL backstop. +var ( + loadFailures = promauto.NewCounterVec(prometheus.CounterOpts{ + Name: "nexus_configcache_load_failures_total", + Help: "Snapshot loads that failed; the previous snapshot stays active.", + }, []string{"cache"}) + + lastLoadSuccess = promauto.NewGaugeVec(prometheus.GaugeOpts{ + Name: "nexus_configcache_last_success_timestamp_seconds", + Help: "Unix time of the last successful snapshot load.", + }, []string{"cache"}) +) diff --git a/packages/shared/storage/configcache/snapshot.go b/packages/shared/storage/configcache/snapshot.go index f1b93d09..9a744e49 100644 --- a/packages/shared/storage/configcache/snapshot.go +++ b/packages/shared/storage/configcache/snapshot.go @@ -79,12 +79,14 @@ func (c *SnapshotCache[T]) Load(ctx context.Context) error { items, err := c.loader(ctx) if err != nil { c.log.Error("snapshot load failed", "cache", c.name, "error", err) + loadFailures.WithLabelValues(c.name).Inc() return err } if items == nil { items = map[string]T{} } c.snap.Store(&items) + lastLoadSuccess.WithLabelValues(c.name).SetToCurrentTime() c.log.Info("snapshot loaded", "cache", c.name, "size", len(items)) if c.onLoad != nil { c.onLoad(c.name, len(items)) diff --git a/packages/shared/storage/configcache/snapshot_test.go b/packages/shared/storage/configcache/snapshot_test.go index bc0b9c8d..07ee7936 100644 --- a/packages/shared/storage/configcache/snapshot_test.go +++ b/packages/shared/storage/configcache/snapshot_test.go @@ -8,6 +8,8 @@ import ( "sync" "sync/atomic" "testing" + + "github.com/prometheus/client_golang/prometheus/testutil" ) func newDiscardLoggerForTest(t *testing.T) *slog.Logger { @@ -226,3 +228,36 @@ func TestSnapshotCache_NilLoaderPanics(t *testing.T) { func key(i int) string { return string(rune('a'+i%26)) + string(rune('0'+i/26)) } + +// TestSnapshotCache_LoadHealthMetrics asserts the operator-facing load-health +// series move: a failed load increments the failure counter (and keeps the +// previous snapshot); a successful load stamps the last-success gauge. +func TestSnapshotCache_LoadHealthMetrics(t *testing.T) { + fail := true + c := NewSnapshotCache(func(ctx context.Context) (map[string]int, error) { + if fail { + return nil, errors.New("db down") + } + return map[string]int{"a": 1}, nil + }, WithSnapshotName("metrics_test_cache")) + + before := testutil.ToFloat64(loadFailures.WithLabelValues("metrics_test_cache")) + if err := c.Load(context.Background()); err == nil { + t.Fatal("expected load error") + } + after := testutil.ToFloat64(loadFailures.WithLabelValues("metrics_test_cache")) + if after != before+1 { + t.Fatalf("failure counter did not increment: %v -> %v", before, after) + } + + fail = false + if err := c.Load(context.Background()); err != nil { + t.Fatalf("load: %v", err) + } + if ts := testutil.ToFloat64(lastLoadSuccess.WithLabelValues("metrics_test_cache")); ts <= 0 { + t.Fatalf("last-success gauge not stamped: %v", ts) + } + if c.Size() == 0 { + t.Fatal("successful load must install the snapshot") + } +} diff --git a/packages/shared/traffic/tracing.go b/packages/shared/traffic/tracing.go index 40520f8f..7bcceea5 100644 --- a/packages/shared/traffic/tracing.go +++ b/packages/shared/traffic/tracing.go @@ -4,7 +4,6 @@ import ( "context" "io" "net/http" - "net/http/httptrace" "sync" "sync/atomic" "time" @@ -14,8 +13,8 @@ import ( // plus an optional long-tail phase breakdown stamped by deeper layers // (e.g. the spec_adapter codec) where the request handler's PhaseTimer // isn't reachable. Producers (the TracingRoundTripper) write `ttfbMs` -// from httptrace.ClientTrace.GotFirstResponseByte and `totalMs` from a -// response-body Close hook. Long-tail consumers call Breakdown(phase, ms) +// from the response body's first Read and `totalMs` from the body's +// Read/Close stamp. Long-tail consumers call Breakdown(phase, ms) // to add a named entry; the handler reads Breakdown() at finalize time // and merges into the audit Record's latency_breakdown JSONB column. // @@ -151,16 +150,6 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) return t.base.RoundTrip(req) } ps.sendStart = time.Now() - trace := &httptrace.ClientTrace{ - GotFirstResponseByte: func() { - ms := time.Since(ps.sendStart).Milliseconds() - if ms <= 0 { - ms = 1 - } - ps.ttfbMs.Store(ms) - }, - } - req = req.WithContext(httptrace.WithClientTrace(req.Context(), trace)) resp, err := t.base.RoundTrip(req) if err != nil { ms := time.Since(ps.sendStart).Milliseconds() @@ -171,17 +160,10 @@ func (t *tracingTransport) RoundTrip(req *http.Request) (*http.Response, error) return resp, err } if resp != nil && resp.Body != nil { - sendStart := ps.sendStart - stamp := func() { - ms := time.Since(sendStart).Milliseconds() - if ms <= 0 { - ms = 1 - } - ps.totalMs.Store(ms) - } resp.Body = &phaseTrackedBody{ ReadCloser: resp.Body, - stamp: stamp, + ps: ps, + sendStart: ps.sendStart, // copied on this goroutine before resp returns } } return resp, nil @@ -195,35 +177,61 @@ func (t *tracingTransport) Unwrap() http.RoundTripper { return t.base } -// phaseTrackedBody wraps an io.ReadCloser so the upstream-total instant -// is refreshed on EVERY successful Read (last-write-wins on the atomic), -// in addition to a once-only stamp at Close. Read-side stamping is -// required because the streaming broker pump runs Close() in its own -// goroutine (`defer session.Close()` in broker.pump) — the request -// handler reads PhaseSink.TotalMs() during its own defer, which can -// fire before the pump goroutine has reached its defer. Without -// Read-side stamping, the audit row sees totalMs=0 (and writes NULL) -// even though the upstream stream finished successfully. Non-streaming -// reads also benefit: totalMs ends up reflecting the last-byte time -// rather than the Body.Close call site, which is closer to "upstream -// done". Close keeps sync.Once semantics so repeat-close (idempotent -// callers like SSEScanner.Close + defer Close stacks) does not -// re-stamp with a later timestamp. +// phaseTrackedBody wraps an io.ReadCloser to capture two upstream timings +// off the response body, avoiding a per-RoundTrip httptrace.ClientTrace + +// request-context copy: +// +// - ttfbMs is stamped ONCE, on the first Read that returns content +// (n>0). The first body byte reaching the caller is the first SSE +// chunk for a streaming provider — matching the audit Record's +// documented "TTFB = first SSE chunk arrival" intent (the prior +// httptrace hook fired on the first response-HEADER byte, which for a +// slow-thinking model wrongly folded think-time into TTFB). +// - totalMs is refreshed on EVERY successful Read (last-write-wins on +// the atomic) plus a once-only stamp at Close. Read-side stamping is +// required because the streaming broker pump runs Close() in its own +// goroutine (`defer session.Close()` in broker.pump) — the request +// handler reads PhaseSink.TotalMs() during its own defer, which can +// fire before the pump goroutine has reached its defer. Without +// Read-side stamping, the audit row sees totalMs=0 (and writes NULL) +// even though the upstream stream finished successfully. Close keeps +// sync.Once semantics so repeat-close (idempotent callers like +// SSEScanner.Close + defer Close stacks) does not re-stamp with a +// later timestamp. +// +// A response body is not safe for concurrent Read, so ttfbStamped needs no +// synchronisation; cross-goroutine visibility of the timings is via the +// PhaseSink atomics. type phaseTrackedBody struct { io.ReadCloser - stamp func() - closeOnce sync.Once + ps *PhaseSink + sendStart time.Time + ttfbStamped bool + closeOnce sync.Once +} + +func (b *phaseTrackedBody) elapsedMs() int64 { + ms := time.Since(b.sendStart).Milliseconds() + if ms <= 0 { + ms = 1 + } + return ms } func (b *phaseTrackedBody) Read(p []byte) (int, error) { n, err := b.ReadCloser.Read(p) if n > 0 || err != nil { - b.stamp() + ms := b.elapsedMs() + if n > 0 && !b.ttfbStamped { + b.ttfbStamped = true + b.ps.ttfbMs.Store(ms) + } + b.ps.totalMs.Store(ms) } return n, err } func (b *phaseTrackedBody) Close() error { - b.closeOnce.Do(b.stamp) + b.closeOnce.Do(func() { b.ps.totalMs.Store(b.elapsedMs()) }) return b.ReadCloser.Close() } diff --git a/packages/shared/traffic/tracing_ttfb_test.go b/packages/shared/traffic/tracing_ttfb_test.go new file mode 100644 index 00000000..b37c07b3 --- /dev/null +++ b/packages/shared/traffic/tracing_ttfb_test.go @@ -0,0 +1,144 @@ +package traffic + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" +) + +// TestTracingTransport_TtfbOnFirstBodyByte_Streaming verifies TTFB anchors on +// the first body byte the caller reads (≈ first SSE chunk) and is NOT moved by +// later frames. +func TestTracingTransport_TtfbOnFirstBodyByte_Streaming(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusOK) + fl, _ := w.(http.Flusher) + if fl != nil { + fl.Flush() // headers out immediately; body byte comes later + } + time.Sleep(25 * time.Millisecond) + _, _ = w.Write([]byte("frame1")) + if fl != nil { + fl.Flush() + } + time.Sleep(25 * time.Millisecond) + _, _ = w.Write([]byte("frame2")) + })) + defer srv.Close() + + client := &http.Client{Transport: NewTracingTransport(http.DefaultTransport)} + ps := NewPhaseSink() + req, _ := http.NewRequestWithContext(WithPhaseSink(context.Background(), ps), http.MethodGet, srv.URL, nil) + resp, err := client.Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + defer resp.Body.Close() + + buf := make([]byte, 6) + if _, err := io.ReadFull(resp.Body, buf); err != nil { + t.Fatalf("read frame1: %v", err) + } + ttfb1 := ps.TtfbMs() + if ttfb1 == nil { + t.Fatal("TtfbMs must be populated after the first body byte") + } + // TTFB anchors on the body byte, not the (immediate) header flush. + if *ttfb1 < 20 { + t.Errorf("TTFB (%d ms) should reflect the pre-frame delay (~25ms), not the header flush", *ttfb1) + } + + if _, err := io.ReadFull(resp.Body, buf); err != nil { + t.Fatalf("read frame2: %v", err) + } + if ttfb2 := ps.TtfbMs(); ttfb2 == nil || *ttfb2 != *ttfb1 { + t.Errorf("TTFB must not move on later frames: was %v, now %v", ttfb1, ttfb2) + } +} + +// TestTracingTransport_TtfbOnFirstBodyByte_NonStreaming verifies a single-write +// body populates both TTFB and total, with total >= ttfb. +func TestTracingTransport_TtfbOnFirstBodyByte_NonStreaming(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + time.Sleep(15 * time.Millisecond) + _, _ = w.Write([]byte("hello world")) + })) + defer srv.Close() + + client := &http.Client{Transport: NewTracingTransport(http.DefaultTransport)} + ps := NewPhaseSink() + req, _ := http.NewRequestWithContext(WithPhaseSink(context.Background(), ps), http.MethodGet, srv.URL, nil) + resp, err := client.Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + + ttfb, total := ps.TtfbMs(), ps.TotalMs() + if ttfb == nil || *ttfb <= 0 { + t.Fatalf("TtfbMs must populate on a non-streaming body; got %v", ttfb) + } + if total == nil || *total < *ttfb { + t.Fatalf("TotalMs (%v) must be >= TtfbMs (%v)", total, ttfb) + } +} + +// TestTracingTransport_EmptyBody_TtfbNil verifies an empty body leaves TTFB nil +// (no first content byte was observed) while total still populates. +func TestTracingTransport_EmptyBody_TtfbNil(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusNoContent) // 204, no body + })) + defer srv.Close() + + client := &http.Client{Transport: NewTracingTransport(http.DefaultTransport)} + ps := NewPhaseSink() + req, _ := http.NewRequestWithContext(WithPhaseSink(context.Background(), ps), http.MethodGet, srv.URL, nil) + resp, err := client.Do(req) + if err != nil { + t.Fatalf("Do: %v", err) + } + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + + if ttfb := ps.TtfbMs(); ttfb != nil { + t.Errorf("empty body must leave TTFB nil (no content byte); got %d", *ttfb) + } + if total := ps.TotalMs(); total == nil { + t.Error("TotalMs must still populate for an empty body (stamped at Close)") + } +} + +// benchBase is an in-memory RoundTripper returning a fixed body, so the +// benchmark isolates the tracing wrapper's per-RoundTrip cost. +type benchBase struct{ body string } + +func (b benchBase) RoundTrip(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Body: io.NopCloser(strings.NewReader(b.body)), + Request: req, + }, nil +} + +// BenchmarkTracingRoundTrip_WithSink measures the per-RoundTrip allocation of +// the tracing wrapper (RoundTrip + drain + close). After folding TTFB into the +// body wrapper, RoundTrip allocates only &phaseTrackedBody{} — the prior +// httptrace.ClientTrace + request-context copy + two closures are gone. +func BenchmarkTracingRoundTrip_WithSink(b *testing.B) { + tr := NewTracingTransport(benchBase{body: "the quick brown fox jumps over the lazy dog"}) + ctx := WithPhaseSink(context.Background(), NewPhaseSink()) + b.ReportAllocs() + b.ResetTimer() + for range b.N { + req, _ := http.NewRequestWithContext(ctx, http.MethodGet, "http://x", nil) + resp, _ := tr.RoundTrip(req) + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + } +} diff --git a/packages/shared/transport/mq/streams.go b/packages/shared/transport/mq/streams.go index 4e0801fd..bacb82a3 100644 --- a/packages/shared/transport/mq/streams.go +++ b/packages/shared/transport/mq/streams.go @@ -30,11 +30,13 @@ const eventsMaxBytesRAMFraction = 0.15 // on a non-Linux host. Production never reassigns it. var meminfoPath = "/proc/meminfo" -// parseByteSize parses a human byte size ("8GB", "512MB", "1073741824") into +// ParseByteSize parses a human byte size ("8GB", "512MB", "1073741824") into // bytes. Returns fallback for an empty or unparseable value. Recognises KB/MB/GB // (powers of 1024) and a bare integer (bytes); case-insensitive, trailing "B" -// optional. -func parseByteSize(s string, fallback int64) int64 { +// optional. Exported so every byte-size env knob across services (the NATS +// stream cap here, the ai-gateway audit memory budget) parses with ONE set of +// semantics instead of drifting copies. +func ParseByteSize(s string, fallback int64) int64 { s = strings.TrimSpace(strings.ToUpper(s)) if s == "" { return fallback @@ -68,7 +70,7 @@ func eventsMaxBytes() int64 { v = strings.TrimSpace(os.Getenv("NEXUS_STREAM_MAX_BYTES")) } if v != "" && !strings.EqualFold(v, "auto") { - return parseByteSize(v, eventsMaxBytesAuto()) + return ParseByteSize(v, eventsMaxBytesAuto()) } n := eventsMaxBytesAuto() slog.Warn("NEXUS_EVENTS audit-stream cap auto-sized from total RAM", diff --git a/packages/shared/transport/mq/streams_test.go b/packages/shared/transport/mq/streams_test.go index 4bd0c64d..68a542be 100644 --- a/packages/shared/transport/mq/streams_test.go +++ b/packages/shared/transport/mq/streams_test.go @@ -50,8 +50,8 @@ func TestParseByteSize(t *testing.T) { {" 2GB ", 2 * 1024 * 1024 * 1024}, } for _, c := range cases { - if got := parseByteSize(c.in, fb); got != c.want { - t.Errorf("parseByteSize(%q) = %d, want %d", c.in, got, c.want) + if got := ParseByteSize(c.in, fb); got != c.want { + t.Errorf("ParseByteSize(%q) = %d, want %d", c.in, got, c.want) } } } diff --git a/packages/shared/transport/streaming/live.go b/packages/shared/transport/streaming/live.go index a2da0929..508ce652 100644 --- a/packages/shared/transport/streaming/live.go +++ b/packages/shared/transport/streaming/live.go @@ -23,9 +23,14 @@ const ( // LiveConfig configures the live streaming compliance pipeline. type LiveConfig struct { - CheckpointChars int // chars between checkpoints (default 500) - MinCheckpointChars int // adaptive lower bound (default 200) - MaxCheckpointChars int // adaptive upper bound (default 2000) + CheckpointChars int // base chars between checkpoints (default 500); the cadence widens with the transcript + // MinCheckpointChars / MaxCheckpointChars are RESERVED and currently + // unused — no caller sets them and the checkpoint cadence is + // deliberately uncapped (a fixed upper bound would re-introduce + // quadratic re-normalization on very long streams). Kept for wire/field + // compatibility; do not treat MaxCheckpointChars as a cadence ceiling. + MinCheckpointChars int + MaxCheckpointChars int MaxBufferSize int // max total buffer (default 8MB) ChannelSize int // internal channel buffer (default 64) } @@ -100,13 +105,11 @@ func (l *LivePipeline) WithUsageAccumulator(acc UsageAccumulator) *LivePipeline // pipelines see structured chat content rather than the flat-text // fallback buildCheckpointInput would otherwise produce. // -// Cost: each checkpoint re-runs normalize on the cumulative body — for -// long streams that's quadratic in body size if normalize itself is -// O(n). Acceptable in practice because (a) checkpoints fire every -// ~500 chars by default so the per-call body is bounded by -// MaxBufferSize anyway, (b) normalize.Registry is structured around -// O(n) parse + adapter dispatch, (c) Registry tier 1 adapters (claude- -// web, chatgpt-web, anthropic, openai-chat) are byte-parser-fast. +// Cost: each checkpoint re-runs normalize on the cumulative body. To keep +// the total normalize work linear in the response length rather than +// quadratic, the checkpoint cadence widens with the accumulated transcript +// (see the step calc in Process) so the checkpoint count stays sublinear; +// the mandatory final checkpoint is always the authoritative full scan. func (l *LivePipeline) WithPreHook(fn PreHookCallback) *LivePipeline { l.preHook = fn return l @@ -312,7 +315,20 @@ func (l *LivePipeline) Process( auditCapped = true continue } - if pendingLen >= l.config.CheckpointChars { + // Widen the checkpoint cadence with the transcript: each checkpoint + // re-normalizes the FULL accumulated body (the pre-hook parses cumulative + // wire bytes), so a fixed step makes the total re-normalization work grow + // with the square of the response length — a long stream saturates the + // box on parsing alone. Growing the step proportionally to the accumulated + // length keeps the checkpoint count sublinear and the total work linear. + // Short streams keep the fixed CheckpointChars cadence; the mandatory + // final checkpoint below always covers the trailing content, so coarser + // intermediate spacing never changes the final audit result. + step := l.config.CheckpointChars + if grow := accumulatedAll.Len() / 8; grow > step { + step = grow + } + if pendingLen >= step { runCheckpoint() pendingLen = 0 } diff --git a/packages/shared/transport/streaming/live_cadence_test.go b/packages/shared/transport/streaming/live_cadence_test.go new file mode 100644 index 00000000..a5536e0e --- /dev/null +++ b/packages/shared/transport/streaming/live_cadence_test.go @@ -0,0 +1,77 @@ +package streaming + +import ( + "bytes" + "context" + "log/slog" + "strings" + "testing" + + core "github.com/AlphaBitCore/nexus-gateway/packages/shared/policy/hooks/core" +) + +// longSSE builds an SSE stream of n content chunks of `each` chars each. +func longSSE(n, each int) string { + var b strings.Builder + seg := strings.Repeat("x", each) + for range n { + b.WriteString(`data: {"choices":[{"delta":{"content":"` + seg + `"}}]}` + "\n\n") + } + b.WriteString("data: [DONE]\n\n") + return b.String() +} + +// TestLiveCadence_LongStreamSublinear is the O(n²) fix guard for the shared +// live pipeline: on a long stream the checkpoint count must grow far slower +// than transcript/CheckpointChars, and the last checkpoint must still see the +// full transcript. +func TestLiveCadence_LongStreamSublinear(t *testing.T) { + const chunks, each, cp = 400, 40, 100 // 16000 content chars + mp := &mockPipeline{} + lp := NewLivePipeline(LiveConfig{CheckpointChars: cp}, mp, slog.Default()) + + var out bytes.Buffer + _, err := lp.Process(context.Background(), strings.NewReader(longSSE(chunks, each)), &out, + &core.HookInput{Stage: "response", IngressType: "COMPLIANCE_PROXY"}) + if err != nil { + t.Fatalf("process: %v", err) + } + + contentLen := chunks * each + fixed := contentLen / cp // old fixed-cadence checkpoint count + got := len(mp.calls) + if got >= fixed/2 { + t.Fatalf("checkpoint count %d not sublinear vs fixed %d (%d chars)", got, fixed, contentLen) + } + // Last checkpoint must have scanned the whole transcript. + last := mp.calls[len(mp.calls)-1] + total := 0 + for _, s := range last { + total += len(s) + } + if total < contentLen { + t.Fatalf("final checkpoint scanned %d chars, want %d", total, contentLen) + } +} + +// TestLiveCadence_ShortStreamUnchanged pins that a short stream (below 8× +// CheckpointChars) keeps the fixed cadence — the widening term must not +// coarsen normal-length responses. +func TestLiveCadence_ShortStreamUnchanged(t *testing.T) { + const cp = 100 // widening kicks in only past 8*cp=800 accumulated chars + mp := &mockPipeline{} + lp := NewLivePipeline(LiveConfig{CheckpointChars: cp}, mp, slog.Default()) + + var out bytes.Buffer + // ~600 content chars: 20 chunks × 30 → below 800, cadence stays at cp. + _, err := lp.Process(context.Background(), strings.NewReader(longSSE(20, 30)), &out, + &core.HookInput{Stage: "response", IngressType: "COMPLIANCE_PROXY"}) + if err != nil { + t.Fatalf("process: %v", err) + } + // 600 chars / 100 cadence → ~5-6 intermediate checkpoints; assert the + // fine-grained regime survived (>=4). + if got := len(mp.calls); got < 4 { + t.Fatalf("short stream got only %d checkpoints; fine cadence regressed", got) + } +}