feat: stream ensemble judge synthesis

[moonming]

Summary

Adds streaming support for ensemble models (DP; follow-up to #604, tracked in #601). stream:true to an ensemble model now streams the judge's synthesized answer token-by-token, instead of the previous 400.

Design: Option A — judge-leg streaming, shared builder reused unmodified

The panel must be buffered to synthesize, so phases 1-2 run non-streaming; only the judge's tokens stream. This reuses the existing build_sse_stream unmodified — all change is isolated to the ensemble dispatch branch; the shared SSE builder and the single-upstream streaming path are byte-for-byte untouched. (Safety property: a bug in the shared builder would hit every streaming request, so this PR deliberately doesn't touch it.)

Deferred — panel-phase keep-alive. The panel fan-out runs before the Sse is constructed, so the 15s SSE keep-alive doesn't cover it. This is no worse than the already-merged non-streaming ensemble path, which likewise holds the socket bytes-free across panel + judge; TTFB actually improves (judge-first-token vs full answer). True panel keep-alive needs an early SSE handshake before the panel completes — a self-contained follow-up kept out of this PR to avoid refactoring the shared builder. Marked with a TODO.

What's in it

• ensemble.rs: split run_ensemble into run_ensemble_panel (phases 1-2 + judge_request) + the judge call, so the judge leg can be streamed. No behavior change to the non-streaming path (all existing executor + handler tests pass unchanged).
• chat.rs::dispatch_ensemble streaming arm: run the panel → resolve the judge bridge → judge.chat_stream() → feed it to build_sse_stream with an on_complete that emits one usage event per sub-call (panel + judge, attempt_kind panel/judge, sharing request_id) and commits the panel+judge tokens via add_tokens_post_stream (the streaming token-commit API). telemetry_handled_by_stream: true suppresses the entry-level emit — no double-emit.
• The "successful panel members are always billed" invariant now also covers the judge-connect-failure and judge-resolution-failure exits (the panel already round-tripped upstream).
• Client sees the ensemble model name on every chunk (re-stamped); judge model names never leak; judge-resolution errors are redacted to generic messages (detail to server logs).

Test plan

• cargo test -p aisix-core — 259 pass.
• cargo test -p aisix-proxy — 440 pass, incl. 2 new e2e through real wiremock bridges:
  • ensemble_streams_judge_synthesis (replaces the now-obsolete ensemble_rejects_streaming_with_400): stream:true → SSE body contains the judge content + [DONE], every chunk model == "council" (not the upstream id), and the usage sink shows 2 panel + 1 judge events sharing one request_id.
  • ensemble_streaming_judge_connect_failure_still_bills_panel: judge 500 on the streaming connect → 502, panel usage events still fire.
• cargo clippy --all-targets -- -D warnings clean; schema-drift gate clean.

An independent audit will follow (merge-gate); I'll post findings + resolutions before merge.

Summary by CodeRabbit

• New Features
  • Ensemble models now support streaming API responses, enabling real-time output delivery.
  • Proper usage tracking and billing for streamed ensemble responses.

[coderabbitai]

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Free

Run ID: 0a6df273-cd6b-43de-aaad-85aee3b9dd72

📥 Commits

Reviewing files that changed from the base of the PR and between ca2542e and b693158.

📒 Files selected for processing (3)

• crates/aisix-proxy/src/chat.rs
• crates/aisix-proxy/src/ensemble.rs
• crates/aisix-proxy/src/lib.rs

---

📝 Walkthrough

Walkthrough

Streaming support is added for ensemble models using an "Option A" approach: the panel phase runs buffered (non-streaming), the judge request is flipped to streaming, and its SSE output is forwarded to the client. Post-stream telemetry emits per-panel and judge usage events. run_ensemble_panel is extracted as a crate-visible helper to enable this split.

Changes

Ensemble Streaming (Option A)

Layer / File(s)	Summary
Extract run_ensemble_panel helper crates/aisix-proxy/src/ensemble.rs	Factors the panel fan-out, min-response enforcement, and judge ChatFormat construction into a new pub(crate) run_ensemble_panel function. Refactors run_ensemble to delegate panel work to it and reuse the returned prebuilt judge request.
Streaming ensemble implementation crates/aisix-proxy/src/chat.rs	Removes the early streaming-denial branch from dispatch_ensemble. Adds the full Option A streaming path: calls run_ensemble_panel for buffered panel execution, flips judge request to streaming, resolves judge model/provider/bridge, builds an SSE stream via build_sse_stream, and wires an on_complete callback that emits per-panel (attempt_kind="panel") and judge (attempt_kind="judge") usage events, returning telemetry_handled_by_stream: true.
Integration tests crates/aisix-proxy/src/lib.rs	Replaces the old 400-rejection test with two new scenarios: ensemble_streams_judge_synthesis validates SSE output shape, per-chunk model-name re-stamping, and two panel plus one judge telemetry events; ensemble_streaming_judge_connect_failure_still_bills_panel validates 502 response and panel-only billing when the judge streaming leg fails.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

---

Note

🎁 Summarized by CodeRabbit Free

Your organization has reached its limit of developer seats under the Pro Plan. For new users, CodeRabbit will generate a high-level summary and a walkthrough for each pull request. For a comprehensive line-by-line review, please add seats to your subscription by visiting https://app.coderabbit.ai/login.If you believe this is a mistake and have available seats, please assign one to the pull request author through the subscription management page using the link above.

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[moonming]

Independent audit — safe to merge

A cold-review agent audited all 8 risk points; verdict: safe to merge (zero HIGH, zero MEDIUM).

Verified:

• build_sse_stream + the single-upstream streaming path are byte-for-byte unchanged — the diff only touches dispatch_ensemble (+ a #[cfg(test)] block in lib.rs).
• No double-emit / double-commit: the borrow checker forces the reservation consumed once (commit_tokens on the error exits XOR into_stream_hold on success); telemetry_handled_by_stream:true gates the entry-level emit.
• on_complete fires exactly once via CompleteOnDrop on both normal completion and client disconnect; disconnect-before-first-chunk bills panel-full + judge-prompt-only (matches the feat(messages): route /v1/messages/count_tokens to Anthropic upstream (#418) #419 cost-leak gate).
• The "successful panel members are always billed" invariant holds on all five error exits (insufficient panel, judge connect fail, unknown judge model, unresolved PK, no bridge).
• No judge identity leaks: chunk model re-stamped to the ensemble name, judge upstream 5xx message suppressed, resolution errors are static literals.
• Both new e2e drive the real handler + build_sse_stream + OpenAiBridge + wiremock (no mock caller).

LOW follow-ups (non-blocking; tracked for the hardening pass under #601):

• L1 — add a streaming output-guardrail-block e2e (the non-streaming path has one; the streaming comp.guardrail_blocked→panel-event wiring is audit-verified-correct but not test-locked).
• L2 — client-disconnect-mid-stream billing (correct by inspection; hard to make deterministic).
• L3 — stream_options.include_usage:true e2e (shared code path with single-upstream).
• L4 — panel-phase socket keep-alive (already documented TODO).