fix(streaming): rehydrate every ordered channel, not just the first choice's content

[ninadphalak]

What this fixes

The proxy's own SSE path gated all rehydration on a string delta.content and read choices[0] without consulting index, against one retention window shared by the whole stream. Three failures fall out of that. All three were reproduced before the fix; all three have tests that fail if it is reverted.

This is the Shield's own pipeline, not the LiteLLM adapter (#31), and not the JSON-RPC path — api/main.py routes jsonrpc: "2.0" to the stateless v3 engine, whose StatelessStreamingLexer regex-scans the raw stream and is field-agnostic, so v3 never had this gap.

1. Streamed tool-call arguments were never rehydrated

A tool-call event carries no string content, so the whole event was forwarded verbatim. The application invoked its tool with [EMAIL_1] as the argument value.

Not a leak — it fails toward the placeholder, nothing sensitive escapes — but it is the reversibility guarantee failing in the channel agentic traffic actually uses. Non-streaming tool calls were already covered by _rehydrate_json_response; only the stream was not.

2. Parallel choices corrupted each other

With n > 1 the choices interleave on one wire and are told apart only by index. One shared window released the characters held back for one choice into whichever event arrived next:

in :  choice 0 "A:[EMA"  |  choice 1 "B:hello"  |  choice 0 "IL_1]"
out:  choice 0 "A:"      |  choice 1 "[EMAB:hello"  |  choice 0 "IL_1]"

Choice 0 reassembles to A:IL_1] — the token destroyed rather than restored. Choice 1 carries a fragment of a vault token it never produced. Both answers wrong, no error raised.

3. Batched choices were skipped

choices is a list and only [0] was walked, so a provider that batches n choices into one event had every choice but the first forwarded unrehydrated.

How

Windows are keyed (choice index, tool-call index or None) — content and each tool call are separate ordered token streams. Bounded by MAX_STREAM_WINDOWS and fail-closed past it (invariant 1) rather than allocating on whatever indices an upstream invents; reusing another channel's window to stay under a cap would reintroduce exactly the splicing the keying prevents.

Tool arguments flush differently from content, and that asymmetry is the point. A client reads content as it arrives, so a tail delivered after finish_reason merely lands late. It accumulates arguments and parses them as JSON the moment the choice reports finished — so a tail delivered after that is parsed against a truncated document. Held tool text therefore rides inside the finishing event, synthesised as an index-only continuation entry where that event names no tool call. The end-of-stream drain remains the net for streams that stop without a finish_reason at all.

arguments joins the buffer-owned key set so the sibling redactor skips it: the buffer already redacts model-originated text before rehydrating, and re-scanning after would redact the caller's own restored values back out. The non-streaming walk passes no skip_keys and still scans the field.

Anthropic keeps one window deliberately — its index counts content blocks, sequential within a message rather than parallel answers, so keying by block would strand each block's held tail until the stream ended.

Declared limits, not hidden

• arguments is JSON text, so a restored value containing a quote, backslash or newline can leave the document unparseable. _rehydrate_json_response has the same property today, and escaping blindly is unsafe because neither site knows whether it sits inside a string.
• Anthropic's streamed input_json_delta tool input is still not restored. It needs its own translation and is deliberately out of this change.

Testing

7 new tests in tests/test_streaming.py. All 7 verified to fail against the unpatched file — they are evidence, not decoration.

Full suite in a clean worktree off main: 928 passed, versus 921 on the same tree without the patch. The two conformance failures in that run (test_every_declared_capability_has_a_proof, test_the_two_ordered_streams_do_not_splice_into_each_other) are pre-existing on clean main — identical without this patch, and both pass in isolation, so they are order-dependent flakes unrelated to this change.

ruff check clean. ruff format --check was already failing on this file before this change and CI does not run it.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge.

Summary

• Tracks parallel and batched choices by their index.
• Gives every tool call its own argument buffer.
• Flushes held arguments into the finishing event and caps streams at 256 windows.
• Keeps Anthropic text on one sequential buffer.

Diagram

sequenceDiagram
    participant Provider
    participant Proxy as rehydrate_sse_stream
    participant Window as Channel window
    participant Vault as Vault view
    participant Client
    Provider->>Proxy: SSE choice delta
    Proxy->>Window: Select by choice and tool index
    Window->>Vault: Restore content or JSON-safe arguments
    Vault-->>Window: Restored text
    Window-->>Proxy: Safe fragment
    alt choice finishes
        Proxy->>Window: Drain tool tail
        Proxy->>Client: Finishing event with tool tail
    else stream continues
        Proxy->>Client: Rehydrated SSE event
    end
    Provider->>Proxy: [DONE]
    Proxy->>Window: Drain every open channel
    Proxy->>Client: Tagged tail events, then [DONE]

Loading

_{Reviews (2) · Last reviewed commit: "fix(streaming): escape restored values i..."}

[ninadphalak]

Both findings fixed in 61e99bd, with a test each that fails against a83124b.

Restored arguments break JSON — fixed rather than declared. Tool-call windows now restore against a vault view whose values are the JSON string body of the original (json.dumps minus its quotes). Only values are rewritten, so token matching, max_token_length retention and the model-originated scan are untouched, and a value needing no escaping passes through byte-identical. One view per stream, rebuilt only when the vault gains tokens, so it stays off the per-delta path. Test round-trips a value carrying a quote, backslash, newline and tab through json.loads.

Window cap rejects 256 — correct, and it was worse than an unused allocation: at exactly the cap it aborted a stream that was completing normally. _flush_tool_windows_into now returns a window it drained and the sibling scan borrows that, so a finishing event never opens a channel. Test drives a choice to the cap with two tool channels and no content window, and asserts the stream completes with both tails intact.

Not fixed here: _rehydrate_json_response on the non-streaming path has the same unescaped-value property. Same class of bug, different path, and this PR does not touch it — it wants its own change rather than a drive-by.

@greptileai