feat: stream ensemble judge synthesis - #606

Merged
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming
Jun 15, 2026
Merged

feat: stream ensemble judge synthesis#606
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming

Conversation

@moonming

@moonmingmoonming commented Jun 15, 2026

Copy link
Copy Markdown
Member

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_streamunmodified — 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_kindpanel/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-core259 pass.
  • cargo test -p aisix-proxy440 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.

Support stream:true on ensemble models. The panel still buffers
non-streaming (needed to synthesize), then the judge's tokens stream to
the client by reusing 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 untouched.
- ensemble.rs: split run_ensemble into run_ensemble_panel (phases 1-2 +
judge_request) + the judge call, so the judge leg can be streamed.
- chat.rs: streaming arm in dispatch_ensemble — 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)
and commits panel+judge tokens via add_tokens_post_stream. The
"successful panel members are always billed" invariant extends to the
judge-connect-failure exit (the panel already round-tripped).
- Client sees the ensemble model name on every chunk; judge model names
never leak; judge-resolution errors are redacted.
- Panel-phase keep-alive is deferred (TODO): the panel runs before the
Sse is constructed, no worse than the non-streaming ensemble path.
Tests: +2 e2e (judge synthesis streams + [DONE] + per-sub-call usage;
judge-connect-failure still bills the panel); replaces the obsolete
streaming-rejection test. 259 core + 440 proxy green; clippy -D clean;
schema-drift gate clean.
Refs #601.
@coderabbitai

coderabbitaiBot commented Jun 15, 2026

Copy link
Copy Markdown

Review Change Stack

ℹ️ 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

Copy link
Copy Markdown
MemberAuthor

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).
  • L3stream_options.include_usage:true e2e (shared code path with single-upstream).
  • L4 — panel-phase socket keep-alive (already documented TODO).

@moonming
moonming merged commit 0132ac5 into mainJun 15, 2026
8 checks passed
@moonming
moonming deleted the feat/ensemble-streaming branch June 15, 2026 04:19
moonming added a commit that referenced this pull request Jun 15, 2026
)
Unbreaks main: #606 merged onto a main with #607's MultiReservation API change (semantic conflict, no textual conflict). Ports the streaming-ensemble reservation code to the new API and fixes the latent un-awaited commit_tokens (panel tokens were never billed on streaming error exits). Verified green locally + CI.
Sign up for freeto 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.

1 participant

@moonming
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '
Skip to content

feat: stream ensemble judge synthesis - #606

Merged
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming
Jun 15, 2026
Merged

feat: stream ensemble judge synthesis#606
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming

Conversation

@moonming

@moonmingmoonming commented Jun 15, 2026

Copy link
Copy Markdown
Member

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_streamunmodified — 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_kindpanel/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-core259 pass.
  • cargo test -p aisix-proxy440 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.

Support stream:true on ensemble models. The panel still buffers
non-streaming (needed to synthesize), then the judge's tokens stream to
the client by reusing 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 untouched.
- ensemble.rs: split run_ensemble into run_ensemble_panel (phases 1-2 +
judge_request) + the judge call, so the judge leg can be streamed.
- chat.rs: streaming arm in dispatch_ensemble — 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)
and commits panel+judge tokens via add_tokens_post_stream. The
"successful panel members are always billed" invariant extends to the
judge-connect-failure exit (the panel already round-tripped).
- Client sees the ensemble model name on every chunk; judge model names
never leak; judge-resolution errors are redacted.
- Panel-phase keep-alive is deferred (TODO): the panel runs before the
Sse is constructed, no worse than the non-streaming ensemble path.
Tests: +2 e2e (judge synthesis streams + [DONE] + per-sub-call usage;
judge-connect-failure still bills the panel); replaces the obsolete
streaming-rejection test. 259 core + 440 proxy green; clippy -D clean;
schema-drift gate clean.
Refs #601.
@coderabbitai

coderabbitaiBot commented Jun 15, 2026

Copy link
Copy Markdown

Review Change Stack

ℹ️ 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

Copy link
Copy Markdown
MemberAuthor

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).
  • L3stream_options.include_usage:true e2e (shared code path with single-upstream).
  • L4 — panel-phase socket keep-alive (already documented TODO).

@moonming
moonming merged commit 0132ac5 into mainJun 15, 2026
8 checks passed
@moonming
moonming deleted the feat/ensemble-streaming branch June 15, 2026 04:19
moonming added a commit that referenced this pull request Jun 15, 2026
)
Unbreaks main: #606 merged onto a main with #607's MultiReservation API change (semantic conflict, no textual conflict). Ports the streaming-ensemble reservation code to the new API and fixes the latent un-awaited commit_tokens (panel tokens were never billed on streaming error exits). Verified green locally + CI.
Sign up for freeto 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.

1 participant

@moonming
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix"); } } catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

feat: stream ensemble judge synthesis - #606

Merged
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming
Jun 15, 2026
Merged

feat: stream ensemble judge synthesis#606
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming

Conversation

@moonming

@moonmingmoonming commented Jun 15, 2026

Copy link
Copy Markdown
Member

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_streamunmodified — 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_kindpanel/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-core259 pass.
  • cargo test -p aisix-proxy440 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.

Support stream:true on ensemble models. The panel still buffers
non-streaming (needed to synthesize), then the judge's tokens stream to
the client by reusing 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 untouched.
- ensemble.rs: split run_ensemble into run_ensemble_panel (phases 1-2 +
judge_request) + the judge call, so the judge leg can be streamed.
- chat.rs: streaming arm in dispatch_ensemble — 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)
and commits panel+judge tokens via add_tokens_post_stream. The
"successful panel members are always billed" invariant extends to the
judge-connect-failure exit (the panel already round-tripped).
- Client sees the ensemble model name on every chunk; judge model names
never leak; judge-resolution errors are redacted.
- Panel-phase keep-alive is deferred (TODO): the panel runs before the
Sse is constructed, no worse than the non-streaming ensemble path.
Tests: +2 e2e (judge synthesis streams + [DONE] + per-sub-call usage;
judge-connect-failure still bills the panel); replaces the obsolete
streaming-rejection test. 259 core + 440 proxy green; clippy -D clean;
schema-drift gate clean.
Refs #601.
@coderabbitai

coderabbitaiBot commented Jun 15, 2026

Copy link
Copy Markdown

Review Change Stack

ℹ️ 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

Copy link
Copy Markdown
MemberAuthor

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).
  • L3stream_options.include_usage:true e2e (shared code path with single-upstream).
  • L4 — panel-phase socket keep-alive (already documented TODO).

@moonming
moonming merged commit 0132ac5 into mainJun 15, 2026
8 checks passed
@moonming
moonming deleted the feat/ensemble-streaming branch June 15, 2026 04:19
moonming added a commit that referenced this pull request Jun 15, 2026
)
Unbreaks main: #606 merged onto a main with #607's MultiReservation API change (semantic conflict, no textual conflict). Ports the streaming-ensemble reservation code to the new API and fixes the latent un-awaited commit_tokens (panel tokens were never billed on streaming error exits). Verified green locally + CI.
Sign up for freeto 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.

1 participant

@moonming
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms"); } } catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

feat: stream ensemble judge synthesis - #606

Merged
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming
Jun 15, 2026
Merged

feat: stream ensemble judge synthesis#606
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming

Conversation

@moonming

@moonmingmoonming commented Jun 15, 2026

Copy link
Copy Markdown
Member

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_streamunmodified — 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_kindpanel/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-core259 pass.
  • cargo test -p aisix-proxy440 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.

Support stream:true on ensemble models. The panel still buffers
non-streaming (needed to synthesize), then the judge's tokens stream to
the client by reusing 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 untouched.
- ensemble.rs: split run_ensemble into run_ensemble_panel (phases 1-2 +
judge_request) + the judge call, so the judge leg can be streamed.
- chat.rs: streaming arm in dispatch_ensemble — 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)
and commits panel+judge tokens via add_tokens_post_stream. The
"successful panel members are always billed" invariant extends to the
judge-connect-failure exit (the panel already round-tripped).
- Client sees the ensemble model name on every chunk; judge model names
never leak; judge-resolution errors are redacted.
- Panel-phase keep-alive is deferred (TODO): the panel runs before the
Sse is constructed, no worse than the non-streaming ensemble path.
Tests: +2 e2e (judge synthesis streams + [DONE] + per-sub-call usage;
judge-connect-failure still bills the panel); replaces the obsolete
streaming-rejection test. 259 core + 440 proxy green; clippy -D clean;
schema-drift gate clean.
Refs #601.
@coderabbitai

coderabbitaiBot commented Jun 15, 2026

Copy link
Copy Markdown

Review Change Stack

ℹ️ 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

Copy link
Copy Markdown
MemberAuthor

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).
  • L3stream_options.include_usage:true e2e (shared code path with single-upstream).
  • L4 — panel-phase socket keep-alive (already documented TODO).

@moonming
moonming merged commit 0132ac5 into mainJun 15, 2026
8 checks passed
@moonming
moonming deleted the feat/ensemble-streaming branch June 15, 2026 04:19
moonming added a commit that referenced this pull request Jun 15, 2026
)
Unbreaks main: #606 merged onto a main with #607's MultiReservation API change (semantic conflict, no textual conflict). Ports the streaming-ensemble reservation code to the new API and fixes the latent un-awaited commit_tokens (panel tokens were never billed on streaming error exits). Verified green locally + CI.
Sign up for freeto 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.

1 participant

@moonming
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links"); } } catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); } })(); (function(){ try { var __m = "youtube.com"; var __re = new RegExp('^' + "youtube\\.com" + '
Skip to content

feat: stream ensemble judge synthesis - #606

Merged
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming
Jun 15, 2026
Merged

feat: stream ensemble judge synthesis#606
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming

Conversation

@moonming

@moonmingmoonming commented Jun 15, 2026

Copy link
Copy Markdown
Member

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_streamunmodified — 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_kindpanel/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-core259 pass.
  • cargo test -p aisix-proxy440 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.

Support stream:true on ensemble models. The panel still buffers
non-streaming (needed to synthesize), then the judge's tokens stream to
the client by reusing 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 untouched.
- ensemble.rs: split run_ensemble into run_ensemble_panel (phases 1-2 +
judge_request) + the judge call, so the judge leg can be streamed.
- chat.rs: streaming arm in dispatch_ensemble — 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)
and commits panel+judge tokens via add_tokens_post_stream. The
"successful panel members are always billed" invariant extends to the
judge-connect-failure exit (the panel already round-tripped).
- Client sees the ensemble model name on every chunk; judge model names
never leak; judge-resolution errors are redacted.
- Panel-phase keep-alive is deferred (TODO): the panel runs before the
Sse is constructed, no worse than the non-streaming ensemble path.
Tests: +2 e2e (judge synthesis streams + [DONE] + per-sub-call usage;
judge-connect-failure still bills the panel); replaces the obsolete
streaming-rejection test. 259 core + 440 proxy green; clippy -D clean;
schema-drift gate clean.
Refs #601.
@coderabbitai

coderabbitaiBot commented Jun 15, 2026

Copy link
Copy Markdown

Review Change Stack

ℹ️ 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

Copy link
Copy Markdown
MemberAuthor

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).
  • L3stream_options.include_usage:true e2e (shared code path with single-upstream).
  • L4 — panel-phase socket keep-alive (already documented TODO).

@moonming
moonming merged commit 0132ac5 into mainJun 15, 2026
8 checks passed
@moonming
moonming deleted the feat/ensemble-streaming branch June 15, 2026 04:19
moonming added a commit that referenced this pull request Jun 15, 2026
)
Unbreaks main: #606 merged onto a main with #607's MultiReservation API change (semantic conflict, no textual conflict). Ports the streaming-ensemble reservation code to the new API and fixes the latent un-awaited commit_tokens (panel tokens were never billed on streaming error exits). Verified green locally + CI.
Sign up for freeto 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.

1 participant

@moonming
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default"); } } catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

feat: stream ensemble judge synthesis - #606

Merged
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming
Jun 15, 2026
Merged

feat: stream ensemble judge synthesis#606
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming

Conversation

@moonming

@moonmingmoonming commented Jun 15, 2026

Copy link
Copy Markdown
Member

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_streamunmodified — 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_kindpanel/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-core259 pass.
  • cargo test -p aisix-proxy440 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.

Support stream:true on ensemble models. The panel still buffers
non-streaming (needed to synthesize), then the judge's tokens stream to
the client by reusing 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 untouched.
- ensemble.rs: split run_ensemble into run_ensemble_panel (phases 1-2 +
judge_request) + the judge call, so the judge leg can be streamed.
- chat.rs: streaming arm in dispatch_ensemble — 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)
and commits panel+judge tokens via add_tokens_post_stream. The
"successful panel members are always billed" invariant extends to the
judge-connect-failure exit (the panel already round-tripped).
- Client sees the ensemble model name on every chunk; judge model names
never leak; judge-resolution errors are redacted.
- Panel-phase keep-alive is deferred (TODO): the panel runs before the
Sse is constructed, no worse than the non-streaming ensemble path.
Tests: +2 e2e (judge synthesis streams + [DONE] + per-sub-call usage;
judge-connect-failure still bills the panel); replaces the obsolete
streaming-rejection test. 259 core + 440 proxy green; clippy -D clean;
schema-drift gate clean.
Refs #601.
@coderabbitai

coderabbitaiBot commented Jun 15, 2026

Copy link
Copy Markdown

Review Change Stack

ℹ️ 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

Copy link
Copy Markdown
MemberAuthor

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).
  • L3stream_options.include_usage:true e2e (shared code path with single-upstream).
  • L4 — panel-phase socket keep-alive (already documented TODO).

@moonming
moonming merged commit 0132ac5 into mainJun 15, 2026
8 checks passed
@moonming
moonming deleted the feat/ensemble-streaming branch June 15, 2026 04:19
moonming added a commit that referenced this pull request Jun 15, 2026
)
Unbreaks main: #606 merged onto a main with #607's MultiReservation API change (semantic conflict, no textual conflict). Ports the streaming-ensemble reservation code to the new API and fixes the latent un-awaited commit_tokens (panel tokens were never billed on streaming error exits). Verified green locally + CI.
Sign up for freeto 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.

1 participant

@moonming
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers"); } } catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); } })(); (function(){ try { var __m = "*"; var __re = new RegExp('^' + ".*" + '
Skip to content

feat: stream ensemble judge synthesis - #606

Merged
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming
Jun 15, 2026
Merged

feat: stream ensemble judge synthesis#606
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming

Conversation

@moonming

@moonmingmoonming commented Jun 15, 2026

Copy link
Copy Markdown
Member

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_streamunmodified — 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_kindpanel/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-core259 pass.
  • cargo test -p aisix-proxy440 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.

Support stream:true on ensemble models. The panel still buffers
non-streaming (needed to synthesize), then the judge's tokens stream to
the client by reusing 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 untouched.
- ensemble.rs: split run_ensemble into run_ensemble_panel (phases 1-2 +
judge_request) + the judge call, so the judge leg can be streamed.
- chat.rs: streaming arm in dispatch_ensemble — 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)
and commits panel+judge tokens via add_tokens_post_stream. The
"successful panel members are always billed" invariant extends to the
judge-connect-failure exit (the panel already round-tripped).
- Client sees the ensemble model name on every chunk; judge model names
never leak; judge-resolution errors are redacted.
- Panel-phase keep-alive is deferred (TODO): the panel runs before the
Sse is constructed, no worse than the non-streaming ensemble path.
Tests: +2 e2e (judge synthesis streams + [DONE] + per-sub-call usage;
judge-connect-failure still bills the panel); replaces the obsolete
streaming-rejection test. 259 core + 440 proxy green; clippy -D clean;
schema-drift gate clean.
Refs #601.
@coderabbitai

coderabbitaiBot commented Jun 15, 2026

Copy link
Copy Markdown

Review Change Stack

ℹ️ 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

Copy link
Copy Markdown
MemberAuthor

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).
  • L3stream_options.include_usage:true e2e (shared code path with single-upstream).
  • L4 — panel-phase socket keep-alive (already documented TODO).

@moonming
moonming merged commit 0132ac5 into mainJun 15, 2026
8 checks passed
@moonming
moonming deleted the feat/ensemble-streaming branch June 15, 2026 04:19
moonming added a commit that referenced this pull request Jun 15, 2026
)
Unbreaks main: #606 merged onto a main with #607's MultiReservation API change (semantic conflict, no textual conflict). Ports the streaming-ensemble reservation code to the new API and fixes the latent un-awaited commit_tokens (panel tokens were never billed on streaming error exits). Verified green locally + CI.
Sign up for freeto 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.

1 participant

@moonming
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode"); } } catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); } })(); })();
Skip to content

feat: stream ensemble judge synthesis - #606

Merged
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming
Jun 15, 2026
Merged

feat: stream ensemble judge synthesis#606
moonming merged 1 commit into
mainfrom
feat/ensemble-streaming

Conversation

@moonming

@moonmingmoonming commented Jun 15, 2026

Copy link
Copy Markdown
Member

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_streamunmodified — 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_kindpanel/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-core259 pass.
  • cargo test -p aisix-proxy440 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.

Support stream:true on ensemble models. The panel still buffers
non-streaming (needed to synthesize), then the judge's tokens stream to
the client by reusing 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 untouched.
- ensemble.rs: split run_ensemble into run_ensemble_panel (phases 1-2 +
judge_request) + the judge call, so the judge leg can be streamed.
- chat.rs: streaming arm in dispatch_ensemble — 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)
and commits panel+judge tokens via add_tokens_post_stream. The
"successful panel members are always billed" invariant extends to the
judge-connect-failure exit (the panel already round-tripped).
- Client sees the ensemble model name on every chunk; judge model names
never leak; judge-resolution errors are redacted.
- Panel-phase keep-alive is deferred (TODO): the panel runs before the
Sse is constructed, no worse than the non-streaming ensemble path.
Tests: +2 e2e (judge synthesis streams + [DONE] + per-sub-call usage;
judge-connect-failure still bills the panel); replaces the obsolete
streaming-rejection test. 259 core + 440 proxy green; clippy -D clean;
schema-drift gate clean.
Refs #601.
@coderabbitai

coderabbitaiBot commented Jun 15, 2026

Copy link
Copy Markdown

Review Change Stack

ℹ️ 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

Copy link
Copy Markdown
MemberAuthor

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).
  • L3stream_options.include_usage:true e2e (shared code path with single-upstream).
  • L4 — panel-phase socket keep-alive (already documented TODO).

@moonming
moonming merged commit 0132ac5 into mainJun 15, 2026
8 checks passed
@moonming
moonming deleted the feat/ensemble-streaming branch June 15, 2026 04:19
moonming added a commit that referenced this pull request Jun 15, 2026
)
Unbreaks main: #606 merged onto a main with #607's MultiReservation API change (semantic conflict, no textual conflict). Ports the streaming-ensemble reservation code to the new API and fixes the latent un-awaited commit_tokens (panel tokens were never billed on streaming error exits). Verified green locally + CI.
Sign up for freeto 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.

1 participant

@moonming