Skip to content

feat(gateway): Hub/Bridge abstractions + SSE decoder - #6

Merged
moonming merged 1 commit into
mainfrom
feat/gateway-hub-bridge
Apr 17, 2026
Merged

feat(gateway): Hub/Bridge abstractions + SSE decoder#6
moonming merged 1 commit into
mainfrom
feat/gateway-hub-bridge

Conversation

@moonming

Copy link
Copy Markdown
Collaborator

Summary

Provider-agnostic core that every `aisix-provider-*` crate implements
against and that the proxy layer dispatches through.

  • chat.rs — `ChatFormat` (normalised OpenAI-compatible request) +
    `ChatMessage`, `Role`, `ChatResponse`, streaming `ChatChunk`/`ChatDelta`,
    `UsageStats`, `FinishReason`. Unknown request fields land in
    `extra` via `serde(flatten)` so Bridges can forward/ignore.
  • bridge.rs — `Bridge` trait (`chat` + `chat_stream`),
    `BridgeContext` (request id, `Arc`, optional deadline), typed
    `BridgeError` with stable `http_status()` + `error_type()` mapping. 4xx
    upstream passes through; 5xx collapses to 502 to hide infrastructure
    bleed-through.
  • hub.rs — `Provider → Arc` registry backed by
    `DashMap` so runtime swaps don't lock readers.
  • sse.rs — byte-stream-in / `SseEvent`-out SSE decoder with state
    that survives partial feeds. Not built on `eventsource-stream` so the
    Bridge trait stays HTTP-client-agnostic.

Also derives `Hash` on `aisix_core::Provider` so it can be a `DashMap` key.

Test plan

  • `cargo test --workspace` — 103 tests pass (28 new, 75 existing)
  • `cargo clippy --all-targets -- -D warnings` clean
  • `cargo fmt --check` clean
  • CI green across all 6 jobs

The provider-agnostic core that every aisix-provider-* crate implements
against, and that the proxy layer dispatches through.

- chat.rs: ChatFormat (normalised OpenAI-compatible request), ChatMessage,
  Role, ChatResponse, ChatChunk/ChatDelta for streaming, UsageStats,
  FinishReason. Unknown top-level request fields flow through
  `serde(flatten)` into an `extra` map so Bridges can forward or ignore
  per their upstream's tolerance.
- bridge.rs: Bridge trait (async_trait) with `chat` + `chat_stream`,
  BridgeContext carrying request_id / Model / deadline, typed
  BridgeError with stable http_status() and error_type() mapping. 4xx
  upstream statuses pass through; 5xx collapses to 502 so clients never
  see bleed-through from upstream infrastructure.
- hub.rs: Provider → Arc<dyn Bridge> registry. DashMap-backed so bridges
  can be swapped at runtime when a future etcd-driven reconfigure ships.
- sse.rs: provider-agnostic SSE line decoder. Byte-stream in, SseEvent
  (Data / Done) out, with state that survives partial feeds.
  Deliberately not built on eventsource-stream so the Bridge trait stays
  independent of any specific HTTP client.

Also: derive Hash on aisix_core::Provider so it can be used as a DashMap
key in the Hub registry.

28 new unit tests — round-tripping JSON shapes, BridgeError → HTTP
status mapping for every variant, SSE decoder edge cases (split feeds,
multi-line data, CRLF, invalid UTF-8, finish flush, [DONE] sentinel),
Hub register/get/overwrite semantics. 103 total across the workspace.
Copilot AI review requested due to automatic review settings April 17, 2026 06:03

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Introduces a new provider-agnostic “gateway core” (aisix-gateway) that standardizes chat request/response types, defines a Bridge trait for provider implementations, provides a Hub for dispatching by provider, and includes an SSE byte-stream decoder for streaming responses.

Changes:

  • Add Bridge/BridgeContext/BridgeError abstractions and normalized chat types (ChatFormat, streaming chunks, usage/finish reasons).
  • Add Hub registry backed by DashMap and an SSE decoder (SseDecoder) with tests.
  • Derive Hash for aisix_core::models::Provider to support usage as a DashMap key.

Reviewed changes

Copilot reviewed 7 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
crates/aisix-gateway/src/lib.rs Documents and re-exports the new gateway core modules.
crates/aisix-gateway/src/bridge.rs Defines the provider-facing Bridge trait and typed error/status mapping.
crates/aisix-gateway/src/chat.rs Adds normalized chat request/response and streaming delta types.
crates/aisix-gateway/src/hub.rs Implements provider→bridge registry via DashMap with tests.
crates/aisix-gateway/src/sse.rs Adds a feed-driven SSE event decoder with unit tests.
crates/aisix-gateway/Cargo.toml Adds dashmap and dev-dependency tokio for tests.
crates/aisix-core/src/models/model.rs Adds Hash derive to Provider.
Cargo.lock Records dashmap dependency addition.
Comments suppressed due to low confidence (1)

crates/aisix-gateway/Cargo.toml:31

  • aisix-gateway’s code in this PR doesn’t reference several dependencies currently listed in [dependencies] (e.g. reqwest, eventsource-stream, bytes, tokio-stream, futures-util, anyhow, tracing, http). Keeping unused deps increases compile time and muddies the crate’s “transport-agnostic core” intent; consider removing any that are no longer needed and re-adding when a module actually uses them.
anyhow.workspace = true
tracing.workspace = true
http.workspace = true
dashmap.workspace = true

[dev-dependencies]
tokio = { workspace = true, features = ["macros", "rt"] }


💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

}

pub fn get(&self, provider: Provider) -> Option<Arc<dyn Bridge>> {
self.bridges.get(&provider).map(|r| r.clone())

Copilot AI Apr 17, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hub::get() uses self.bridges.get(&provider).map(|r| r.clone()). This relies on method resolution through Deref to clone the inner Arc, which is a bit subtle to readers. Consider making the intent explicit (e.g., cloning the value via r.value() / Arc::clone(...)) to avoid confusion about whether the DashMap ref/guard is being cloned.

Suggested change
self.bridges.get(&provider).map(|r| r.clone())
self.bridges.get(&provider).map(|r| Arc::clone(r.value()))

Copilot uses AI. Check for mistakes.
Comment on lines +52 to +58
pub fn feed<'a>(&mut self, bytes: impl Into<Cow<'a, [u8]>>) -> Vec<SseEvent> {
let bytes = bytes.into();
// Non-UTF-8 bytes are replaced rather than erroring — upstreams
// that break encoding still surface a best-effort event so a
// single bad byte doesn't kill the whole stream.
let chunk = String::from_utf8_lossy(&bytes);
self.buffer.push_str(&chunk);

Copilot AI Apr 17, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

feed() decodes each incoming chunk with String::from_utf8_lossy(&bytes) and appends it to a String buffer. If a valid multi-byte UTF-8 codepoint is split across HTTP chunks, decoding per-chunk will emit U+FFFD replacement chars even though the overall stream is valid UTF-8, corrupting JSON payloads. Consider buffering raw bytes (e.g., Vec<u8>/BytesMut) and only UTF-8 decoding once you’ve identified complete \n\n-terminated event frames, so split codepoints are handled correctly.

Copilot uses AI. Check for mistakes.
Comment on lines +29 to +38
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ChatMessage {
pub role: Role,
pub content: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub name: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
}

Copilot AI Apr 17, 2026

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ChatMessage is marked with #[serde(deny_unknown_fields)], which will reject OpenAI-compatible message objects that include additional fields (e.g. tool_calls, legacy function_call, or future extensions). This conflicts with the surrounding goal of being permissive/forward-compatible (like ChatFormat.extra). Consider either removing deny_unknown_fields here or adding an extra field with #[serde(flatten)] on ChatMessage as well.

Copilot uses AI. Check for mistakes.
@moonming
moonming merged commit 0fbf5c2 into main Apr 17, 2026
10 checks passed
@moonming
moonming deleted the feat/gateway-hub-bridge branch April 17, 2026 06:07
moonming added a commit that referenced this pull request Apr 17, 2026
First concrete Bridge implementation against the aisix-gateway trait.

- wire.rs: OpenAI /chat/completions request and response wire types,
  plus the two mappers that round-trip between our ChatFormat /
  ChatResponse / ChatChunk and the upstream shape. Request extras flow
  through `#[serde(flatten)]` so seed/presence_penalty/etc. forward
  without the gateway having to know about them.
- bridge.rs: OpenAiBridge owns a shared reqwest::Client. chat() does
  POST /chat/completions, parses the typed response, and applies the
  BridgeContext deadline via tokio::time::timeout. chat_stream() pipes
  bytes_stream() through the gateway's SseDecoder and yields ChatChunks
  via async_stream, terminating cleanly on the [DONE] sentinel.
- Error mapping matches the BridgeError contract from PR #6:
  transport → Transport, non-2xx → UpstreamStatus (4xx passes through,
  5xx collapses to 502 via http_status()), malformed JSON →
  UpstreamDecode, elapsed deadline → Timeout { elapsed_ms }.
- `with_name()` lets OpenAI-compatible providers (DeepSeek today,
  Gemini-OAI later) reuse this transport with a distinct metrics label.

15 new unit tests across wire and bridge, 10 using wiremock: happy path
(streaming + non-streaming), 429 pass-through, 500 pre-stream, malformed
body → decode error, deadline → timeout, missing api_key → config error,
SSE with role/content/finish_reason/[DONE], and resolve_base trailing-
slash handling. 118 tests pass workspace-wide.
moonming added a commit that referenced this pull request Apr 17, 2026
…at (#7)

First concrete Bridge implementation against the aisix-gateway trait.

- wire.rs: OpenAI /chat/completions request and response wire types,
  plus the two mappers that round-trip between our ChatFormat /
  ChatResponse / ChatChunk and the upstream shape. Request extras flow
  through `#[serde(flatten)]` so seed/presence_penalty/etc. forward
  without the gateway having to know about them.
- bridge.rs: OpenAiBridge owns a shared reqwest::Client. chat() does
  POST /chat/completions, parses the typed response, and applies the
  BridgeContext deadline via tokio::time::timeout. chat_stream() pipes
  bytes_stream() through the gateway's SseDecoder and yields ChatChunks
  via async_stream, terminating cleanly on the [DONE] sentinel.
- Error mapping matches the BridgeError contract from PR #6:
  transport → Transport, non-2xx → UpstreamStatus (4xx passes through,
  5xx collapses to 502 via http_status()), malformed JSON →
  UpstreamDecode, elapsed deadline → Timeout { elapsed_ms }.
- `with_name()` lets OpenAI-compatible providers (DeepSeek today,
  Gemini-OAI later) reuse this transport with a distinct metrics label.

15 new unit tests across wire and bridge, 10 using wiremock: happy path
(streaming + non-streaming), 429 pass-through, 500 pre-stream, malformed
body → decode error, deadline → timeout, missing api_key → config error,
SSE with role/content/finish_reason/[DONE], and resolve_base trailing-
slash handling. 118 tests pass workspace-wide.
moonming added a commit that referenced this pull request May 18, 2026
…ped reads, redact 5xx message, Vertex content-type guard

Five concrete fixes from the Copilot inline review on PR #323. Two
stale comments (#3, #4 — already fixed in commit 3) are skipped.

**#1+#7 — Azure OpenAI-compatible code preservation.**
Azure's envelope omits `error.type` and carries only `error.code`.
The bridge previously put the upstream code into `view.kind` and
left `view.code` as `None`. For OpenAI-compat tokens Azure inherits
unchanged (e.g. `rate_limit_exceeded`), this meant downstream OpenAI
clients received `error.type=rate_limit_exceeded` but
`error.code=null` — exactly the SDK-retry break issue #322 is about.

Fix:
- Azure parser populates BOTH `view.kind` AND `view.code` from the
  upstream `error.code` field.
- `render_openai_envelope`'s AzureOpenAI branch now prefers the
  translation-table-derived code (so explicit Azure tokens like
  `DeploymentNotFound` → `model_not_found` still win), falling back
  to `view.code` for OpenAI-compat pass-through.

**#2 — Drain the response stream after hitting the cap.**
`read_body_capped` previously broke out of the read loop the moment
`limit` bytes were buffered. With reqwest/hyper that leaves unread
bytes in the response and prevents connection reuse — during a burst
of upstream errors the gateway would churn TCP connections instead
of recycling the keep-alive pool. Fix: keep iterating the stream,
discarding chunks past the cap. Memory stays bounded by `limit`.

**#5 — Redact upstream `error.message` on 5xx.**
The 5xx branch of `render_bridge_upstream_envelope` was forwarding
`BridgeError::UpstreamStatus.message` verbatim — which for OpenAI /
Anthropic comes from the parsed upstream `error.message`. Upstream
5xx bodies routinely embed operator-internal detail (engine names,
shard ids, queue depth). Fix: on 5xx, emit a canned
`"upstream returned {status}"` message; the full upstream body
remains in operator logs via tracing.

**#6 — Stale "follow-up" comment.**
The docstring on `render_bridge_upstream_envelope` claimed cross-wire
translation would ship in a follow-up, but it already shipped in
commit 2. Rewrite the comment to describe current behaviour
(4xx → `error_translate`; 5xx → canned envelope; `Unknown` wire →
legacy generic envelope).

**#8 — Content-type guard on Vertex (and Azure, while at it).**
`capture_upstream_error_http` already gates serde parsing on
`Content-Type: application/json` so a 64 KB HTML error page from a
fronting WAF doesn't waste CPU on a doomed JSON parse. The Vertex
and Azure bridges call serde directly because they need a custom
parse path (canned message for redaction) — same guard now applies.
Promoted `content_type_is_json` and added a `response_is_json`
helper to the gateway's public surface; both bridges call it before
`parse_*_error_*`.

New tests:
- `upstream_openai_5xx_with_json_envelope_collapses_and_redacts_message`
  pins the 5xx redaction (asserts `engine offline` / `shard 47` /
  `engine_overloaded` don't reach the customer envelope).
- `chat_429_preserves_openai_compatible_code_for_sdk_retry` (Azure)
  pins that `parsed.code` carries the OpenAI-compat upstream code.
- `chat_400_non_json_body_skips_envelope_parse` (Azure) and
  `chat_gemini_non_json_body_skips_envelope_parse` (Vertex) pin the
  new content-type guard.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
jarvis9443 added a commit that referenced this pull request Jun 2, 2026
…eaming)

Thread the resolved guardrail chain (as Arc) through the /v1/messages
dispatch paths and run output guardrails on the response:

- Non-streaming: cross-provider checks the bridge ChatResponse;
  passthrough extracts response text (content blocks + raw content array
  for tool_use) into a synthetic ChatResponse.
- Streaming: both the cross-provider SSE encoder path and the verbatim
  Anthropic byte-passthrough accumulate assistant text and run the
  guardrail at end-of-stream. Bytes are forwarded live (matching
  /v1/chat/completions and LiteLLM's streaming guardrail), so a block is
  signalled with a terminal Anthropic `error` (content_filter) event.

Completes the output side of #448 #22; with this and the earlier input +
budget work, /v1/messages no longer bypasses the guardrail/quota
pipeline. The remaining findings (#6 count_tokens, #2/#13
reasoning_content, #24 guardrail-vs-rate-limit ordering) are accepted as
standard behavior (LiteLLM has the same gap).

Fixes #448
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants