fix: store provider signatures before terminal streams

[hanakannzashi]

Summary

• hold the raw upstream terminal SSE suffix until the provider signature is persisted
• release the exact held bytes after finalization, preserving provider-byte exactness
• add a frame-by-frame regression test for an immediate signature lookup at [DONE]

This is stacked on #998, which prevents successful signatures for errored streams.

Fixes #999

Testing

• cargo fmt --check
• cargo test -p services completions:: --lib
• cargo test -p api --lib
• cargo test -p api --test e2e_all --no-run

Focused E2E execution is blocked locally by PostgreSQL authentication during test bootstrap.

[ironloopai]

Review · Status

🟩 Completed

IronLoop completed the review and posted it to GitHub.

Result

Open submitted review →

Run details

• Run: 9d22a71a-1f1d-43b7-8e17-9f4e4d3ca62c
• Base: codex/no-signature-after-stream-error-997 at 4384a9c
• Head: codex/provider-signature-before-done-999 at 6230341
• Created: 2026-09-01 06:37 UTC
• Updated: 2026-09-01 06:42 UTC

_{Automatic trigger · attempt 1 of 3 · completed in 4m 51s}

[claude]

Review: hold provider terminal suffix until signature is persisted

Approach is correct. InterceptStream::poll_next only returns Poll::Ready(None) after StreamState::Finalizing resolves the signature future (crates/services/src/completions/mod.rs:557-578), so the route's .chain(stream::once(...)) tail is guaranteed to run after the provider signature is persisted. Releasing the held bytes there is the right hook. Byte-exactness is preserved (the same raw_bytes are re-emitted verbatim), and since upstream_done.store(true) immediately precedes the latch with no early return between them, synthesized_done and the release branch are mutually exclusive — no double [DONE], no silently dropped suffix.

Three issues before merge:

1. Post-[DONE] buffer is unbounded (production safety)

crates/api/src/routes/completions.rs:1697-1705 — once latched, every subsequent Ok event is accumulated into held_provider_terminal_bytes with no cap. Previously these bytes streamed straight through to the client; now they are held in gateway memory until EOF.

In a well-behaved stream this is \n (the SSE parser emits the blank separator line as its own control event). But a broken or hostile upstream that emits data: [DONE] and then keeps streaming turns this into unbounded per-request heap growth in the TEE. The codebase already guards this class elsewhere — MAX_SSE_LINE_BYTES in crates/inference_providers/src/attested/chutes/e2ee_stream.rs:49, and MAX_LEADING_CONTROL_EVENTS in this same handler.

Suggest a cap that fails open (forwards instead of holding) once exceeded:

if holding_provider_terminal_suffix.load(Ordering::Relaxed) {
    let mut held = held_provider_terminal_bytes.lock().await;
    if held.len() + event.raw_bytes.len() > MAX_HELD_TERMINAL_SUFFIX_BYTES {
        tracing::warn!(%organization_id, "Upstream exceeded held terminal suffix cap; forwarding");
        holding_provider_terminal_suffix.store(false, Ordering::Relaxed);
        let flush = std::mem::take(&mut *held);
        drop(held);
        return Some(Ok::<Bytes, Infallible>(Bytes::from(
            [flush.as_slice(), &event.raw_bytes].concat(),
        )));
    }
    held.extend_from_slice(&event.raw_bytes);
    return None;
}

A few KiB is ample for a legitimate terminal suffix.

2. Err events after [DONE] bypass the hold and invert wire order

The latch check lives inside the Ok(event) arm only; the Err(e) arm at completions.rs:1950-1961 still emits sse_error_frame(&e) immediately. So an error arriving after the terminator is written to the socket before the held data: [DONE] that the tail later appends — the terminator is no longer last on the wire, which breaks SSE clients that treat [DONE] as the stop condition (they parse the error frame as a data chunk instead). Pre-fix the ordering was the other way round.

Either route the error frame through the hold buffer too, or drop the held suffix and let the error be terminal. Worth an explicit decision plus a comment either way.

3. hold_provider_terminal_suffix is computed true but inert for alias-served streams

hold_provider_terminal_suffix (completions.rs:1634-1635) is derived only from provider_signature_enabled && !gateway_signature_enabled, but the latch at completions.rs:1750-1760 sits inside the passthrough branch guarded by !alias_served. For an alias-served attested model the [DONE] falls to the re-serialization path (completions.rs:1783-1796), which sets upstream_done and forwards the terminator immediately — the race this PR fixes is still open there, while the flag reads as if it were handled.

Narrow in practice (needs alias + attestation_supported + continuous_usage_stats/E2EE, since otherwise strip_intermediate_usage forces the gateway-signature path), but the flag should reflect reality. Fold !alias_served into the flag and note why, or apply the hold in the control-event branch too.

Minor: the new test could assert the stronger property

test_raw_provider_signature_is_available_when_done_is_emitted breaks out of the read loop at [DONE] and never drains the body. The sibling test test_raw_stream_without_upstream_done_uses_gateway_signature (lines 122-131) asserts no bytes follow the terminator — worth mirroring here, since "nothing escapes after the released suffix" is precisely the invariant the hold introduces.

---

Flagging as an intentional tradeoff rather than a defect: client-visible [DONE] for attested raw streams is now delayed by the provider signature fetch, bounded by FINALIZE_TIMEOUT_SECS = 5. That is the point of the fix, but it is a real TTLB change on the attested path and worth a line in the PR description.

Note: I could not compile in this environment, so the above is from static reading; CI should confirm.

⚠️

[ironloopai]

Review · Summary

Found two medium-severity defects in the terminal-suffix buffering change.

Findings: 🟠 Medium 2

_{Code-specific findings are attached to the diff.}

Validation

• ✅ Static lifecycle inspection — The normal raw provider-signature path was traced through completion-stream finalization before releasing its held terminal bytes.

Review details

• Run: 9d22a71a-1f1d-43b7-8e17-9f4e4d3ca62c
• Attempts: 1

[github-actions]

✅ OpenCodeReview: Review failed: 0 finding(s); 2 of 2 selected item(s) failed.

[think-in-universe]

No actionable findings. Terminal buffering is bounded to the terminal event/separator and preserves signature availability before [DONE].

[think-in-universe]

@ironloopai review

[ironloopai]

Review · Status

🟩 Completed

IronLoop completed the review and posted it to GitHub.

Result

Open submitted review →

Run details

• Run: edda806a-edf6-4f12-b06b-b2d160f479cb
• Base: codex/no-signature-after-stream-error-997 at 49dbdc1
• Head: codex/provider-signature-before-done-999 at db4e6b2
• Created: 2026-09-03 06:41 UTC
• Updated: 2026-09-03 06:58 UTC

_{Manual command by think-in-universe · attempt 1 of 3 · completed in 16m 17s}

[ironloopai]

Review · Summary

Found 1 medium-severity issue.

Findings: 🟠 Medium 1

_{Code-specific findings are attached to the diff.}

Validation

• ✅ Rust formatting — Changed sources conform to repository formatting.
• ✅ Terminal-stream state tests — Both complete and incomplete terminal-marker state-machine tests passed.

Review details

• Run: edda806a-edf6-4f12-b06b-b2d160f479cb
• Attempts: 1

[think-in-universe]

Approved.

Please resolve review comments before merge