fix: normalize multi-round gateway streaming

[harivilasp]

Summary

Revives and updates the GatewayAccumulator work for #119 on top of current main.

The gateway streaming path currently forwards each upstream round as if it were an independent client-visible response. In a multi-round gateway/tool loop that leaks upstream lifecycle and indexing details to clients: duplicate response lifecycle events, sequence numbers that reset between rounds, and output indexes that collide across rounds.

This PR introduces a single stream accumulation layer for gateway responses so the public SSE stream is owned and normalized at one boundary before it reaches the client.

Closes #119.

What Changed

• Added gateway stream accumulation for public SSE events across gateway rounds.
• Deduplicated client-visible lifecycle events so a multi-round stream emits one logical response lifecycle.
• Assigned monotonic public sequence numbers across forwarded upstream frames, synthetic gateway frames, and the terminal event.
• Rebased output indexes so final client-visible output items remain contiguous across rounds instead of restarting at each upstream call.
• Routed forwarded upstream frames, synthetic tool/search frames, and terminal response events through the same normalization path.
• Added a multi-round streaming regression test covering a gateway tool round followed by a final text round.

Why

Issue #119 documents three symptoms with the same root cause: no stage owns the final client-visible SSE sequence for multi-round gateway execution.

Without this, clients that initialize on response.created can reset more than once, clients that order or dedupe by sequence_number see the sequence reset mid-stream, and clients keyed by output_index can overwrite prior round output when later rounds reuse the same upstream indexes.

Validation

Ran the following locally from agent/issue-119-gateway-accumulator-verified:

• cargo fmt --check
• cargo test -p agentic-server-core --test web_search_tool_test multi_round_stream_has_single_lifecycle_and_monotonic_public_sequence
• cargo test -p agentic-server-core --test web_search_tool_test
• cargo clippy -p agentic-server-core --all-targets -- -D warnings
• cargo test -p agentic-server-core
• cargo clippy --all-targets -- -D warnings
• cargo test
• uvx pre-commit run --all-files

Results:

• Targeted regression: 1 passed, 17 filtered out
• Full web_search_tool_test: 18 passed
• agentic-server-core: 337 passed, 1 ignored
• Workspace tests: 376 passed, 1 ignored
• Clippy: no issues found
• Pre-commit: all hooks passed

Notes

This was previously explored in draft PR #129 and cancelled because it was considered too narrow relative to #121. This PR keeps the #119 behavior fix focused, but includes the regression coverage needed to prove the multi-round public stream invariants. It also updates the revived test for the current RequestPayload.cache_salt field added on main.

[ashwing]

The per-round lifecycle leak and colliding output_index are the sharp edges from #119, and one normalization boundary is the right shape. Traced it against the three acceptance criteria — single lifecycle ✅, contiguous output_index ✅, monotonic sequence_number holds on the success path. One gap and two minor notes inline.

[ashwing]

#119's criterion 2 was "monotonic sequence_number on every emitted frame." These error_sse_chunk frames (this arm and the Ok(Err) one) bypass the accumulator, so if a round already emitted frames (seq 0..k) and a later round fails, the client sees an error frame with no sequence_number after a numbered sequence. The accumulator is moved into the spawned task and only handed back on Ok(Ok(...)), so these arms don't have a handle to it today — worth restructuring so error frames get stamped too, or calling it out as a known follow-up before this closes #119.

[ashwing]

stream_events and stream_accumulator are separate Options but always travel together (both Some on the stream path, both None on blocking). This guard — and the matching one in emit_gateway_completed_events — silently drops every synthetic gateway event if a caller ever passes stream_events: Some with stream_accumulator: None. Bundling them into one type (the way StreamEmitContext already does on the upstream path) would make "both or neither" unrepresentable rather than convention.

[ashwing]

Two multi-round edges worth covering since this PR owns the invariant now: (a) a round that emits multiple public output items before the final round — the current test only ever has one item per round, so it never checks output_offset advancing by item count vs. +1; (b) the response.incomplete/max-rounds terminal path — confirms the terminal still gets a monotonic sequence_number after N rounds. Not blocking.

[maralbahari]

@harivilasp Thanks for the fix. Four structural requests, each building on the previous:

1. Move GatewayStreamAccumulator to its own file. It mirrors ResponseAccumulator, which lives in executor/accumulator.rs, so give it symmetric placement (e.g. executor/gateway_accumulator.rs) and take rebase_output_index (gateway.rs:165) and emit_sse_json (gateway.rs:349) with it. The duplicated Some/None branches in execute_and_emit_output_calls should clean up in the process.

2. Parse each SSE line once. Today every line is parsed up to three times: normalize_sse_line at upstream.rs:115, again in emit_stream_line (upstream.rs:143), and a third time via acc.process_sse_line (upstream.rs:89). normalize_sse_line already has the full JSON in hand and drops it; keep it on the frame, as a typed envelope rather than a bare Value:

#[derive(Serialize, Deserialize)]
pub struct WireEvent {
    #[serde(rename = "type")]
    pub event_type: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub sequence_number: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub output_index: Option<u64>,
    #[serde(flatten)]
    pub rest: serde_json::Map<String, Value>,
}

Renumbering becomes wire.sequence_number = Some(n) instead of value["sequence_number"] = ..., and the flatten catch-all keeps the relay lossless. Don't deserialize relayed events into the existing item types (OutputItem etc.): they have no flatten catch-alls and #[serde(other)] Unknown (types/io/output.rs:358) discards data, so re-emitting through them would silently strip unknown fields from the client stream. The namespace restorer is the only non-trivial adaptation, and it's close to mechanical: restore_response_value_with_map (tool/codex.rs:447) already goes through as_object_mut() at every step, so it can take the rest map directly. The WireEvent envelope is part of this PR; what can wait for a follow-up PR is extending the typing beyond it, promoting fields out of rest into dedicated structs per event and item type (with flatten catch-alls so the relay stays lossless).

3. Mirror ResponseAccumulator's two-level API. Give the gateway accumulator a line-based entry point process_sse_line that normalizes internally and delegates to a frame-based core process_event(&mut EventFrame), the same shape as accumulator.rs:211-229. In the fetch_stream_payload loop, normalize once and call the frame-based process_event of BOTH accumulators; calling the two line-based wrappers there would reintroduce the double parse.

4. Same frame currency for synthetic and deferred events. The gateway-built events (emit_gateway_start_events gateway.rs:360, emit_gateway_completed_events gateway.rs:411, terminal_response_chunk gateway.rs:143) are bare json!, yet every type they emit already has an SSEEventType variant. Add an EventFrame::synthetic(event_type, wire) constructor in events and route them through the same process_event, which also shares sequence numbering automatically. Likewise, pending_unnamed_function_events (upstream.rs:105) should park EventFrames instead of raw Strings that flush_pending_function_events (upstream.rs:211) re-parses on release.

[harivilasp]

will try to address by tomrrow

[harivilasp]

Addressed in aa1675b:

1. Moved the gateway stream accumulator, output-index rebasing, and SSE serialization into gateway_accumulator.rs, with a single stream context.
2. Added a lossless WireEvent envelope on EventFrame, so upstream SSE JSON is parsed once and unknown fields are retained.
3. Added the gateway accumulator line/frame two-level API; the stream loop normalizes once and passes the same frame to both accumulators.
4. Routed synthetic, terminal, and deferred events through EventFrame/process_event, preserving shared sequence numbering.

Also covered lossless wire preservation and namespace restoration. cargo fmt --check, Clippy, and cargo test -p agentic-server-core pass (340 passed, 1 ignored).

[maralbahari]

@harivilasp Thanks for addressing the first round. A few concerns mostly on the new events surface:

1. normalize_sse_line traded the redundant parse for a redundant deep copy. At events/normalize.rs:19, wire.to_value() re-serializes the entire event back into a Value tree on every line, just so extract_payload can read &Value. That is a full copy of every delta and item in the hot path. Reorder instead: parse the line to Value, run extract_payload(&json) (borrows), then serde_json::from_value::<WireEvent>(json) (consumes and moves, no copy). Three lines, extract_payload untouched.

2. WireEvent::into_value is dead code (zero callers), and to_value's three callers should all disappear. Besides normalize.rs:19 above, EventFrame::synthetic (events/types.rs:288) and error_frame (executor/gateway_accumulator.rs:138) snapshot the wire into EventPayload::Raw(wire.to_value()). That payload is never read (synthetic frames flow only through GatewayStreamContext::process_event into serialize_sse_frame, which serializes wire), and it goes stale the moment the accumulator stamps the sequence number. Use EventPayload::None, which exists for exactly this. Then both conversion methods can be deleted; if they were to stay, direct serde_json::to_value with a silent fallback needs to go through utils::common per AGENTS.md.

3. sequence_number is duplicated inside EventFrame. The frame stores it twice: once as frame.sequence_number and once inside frame.wire.sequence_number, and four sites keep the copies in sync by hand (process_event at gateway_accumulator.rs:56-57, normalize_sse_line at normalize.rs:28, EventFrame::synthetic at types.rs:292, error_frame at gateway_accumulator.rs:139). Only the wire copy matters, since serialize_sse_frame serializes frame.wire; nothing in production reads the other one (the sole read in the repo is a test assert at gateway_accumulator.rs:175, whose next line already checks the same value via wire). Remove the duplicated sequence_number field from EventFrame, keep wire.sequence_number as the single source, and expose fn sequence_number(&self) -> Option<u64> { self.wire.sequence_number } for any caller that wants it.

4. Minor: the type string is written twice on the synthetic path. terminal_response_frame and synthetic_event build WireEvent::new(event_type.as_str()), then EventFrame::synthetic overwrites wire.event_type with the same string (types.rs:287). Pick one owner: have EventFrame::synthetic take the SSEEventType and set the wire name itself (via the TryFrom conversion from the inline comment), so callers stop passing the string entirely.

[maralbahari]

Maintaining two parallel 29-arm tables, classify_event_type (normalize.rs:34) and SSEEventType::as_str (types.rs:115), is troublesome; they live in different files and will drift. Implement conversion traits on SSEEventType instead: From<&str> for parsing and TryFrom for &'static str for the wire name (fallible because Other has no wire form and must not serialize as a made-up "unknown"). Then classify_event_type and as_str are both no longer needed, normalize_sse_line just calls SSEEventType::from(wire.event_type.as_str()), and the two impls sit adjacent in types.rs with a round-trip test to keep them in sync.

[maralbahari]

@harivilasp FYI we have some rust best practices mentioned in our repo AGENTS.md apart from those for best agent guidance it is better to follow skills like https://github.com/leonardomso/rust-skills. we internally use this to make sure best practices are followed by agents.

[maralbahari]

Borrowing both fields here forces every downstream signature to carry lifetimes (GatewayStreamContext<'a>, StreamEmitContext<'a, 'stream>) and, worse, forces the Arc<Mutex<GatewayStreamAccumulator>> workaround in engine.rs:217-228, where the spawned task takes the lock and holds the guard across run_until_gateway_tools_complete(...).await, i.e. for the entire streaming run. AGENTS.md forbids holding a guard across .await, and the mutex is not synchronizing anything: the generator only locks after the join handle resolves (engine.rs:242/247/262), so it is never contended. It exists purely to satisfy the borrow checker and plants a deadlock trap for anyone who later locks while the task is running.

Suggestion: delete GatewayStreamContext and keep the accumulator a pure state machine, exactly like ResponseAccumulator: it consumes frames and updates state, and never touches transport. Its API already wants this shape: process_event returns bool without sending, and terminal_response_chunk / error_chunk return Strings without sending. The full target API:

/// Shapes the multi-round gateway stream into one client-visible stream:
/// dedups response.created / response.in_progress across rounds, assigns
/// gateway-global sequence numbers, rebases output_index per round.
///
/// Mirrors `ResponseAccumulator`: consumes frames, updates state, never
/// touches transport.
pub struct GatewayStreamAccumulator {
    next_sequence_number: u64,
    emitted_created: bool,
    emitted_in_progress: bool,
}

impl GatewayStreamAccumulator {
    pub fn new() -> Self;

    /// Line-based entry point, mirroring ResponseAccumulator::process_sse_line:
    /// normalizes internally, then delegates to process_event.
    /// Returns the stamped frame if it should be emitted.
    pub fn process_sse_line(&mut self, line: &str, output_offset: usize) -> Option<EventFrame>;

    /// Core state machine. Stamps wire.sequence_number and rebases
    /// wire.output_index. Returns false when the frame is suppressed
    /// (duplicate lifecycle event); suppressed frames consume no number.
    #[must_use]
    pub fn process_event(&mut self, frame: &mut EventFrame, output_offset: usize) -> bool;

    /// Terminal response.{completed,failed,incomplete} event: stamped,
    /// serialized, returned as a String for the generator to yield.
    pub fn terminal_response_chunk(&mut self, payload: &ResponsePayload) -> ExecutorResult<String>;

    /// Stamped error event, same yield-not-send contract.
    pub fn error_chunk(&mut self, message: &str) -> String;
}

// Module-level transport helpers, used by callers, never by the accumulator:
fn serialize_sse_frame(frame: &EventFrame) -> ExecutorResult<String>;      // from frame.wire only
fn emit_sse_frame(sender: &UnboundedSender<String>, frame: &EventFrame) -> ExecutorResult<()>;
fn synthetic_event(event_type: SSEEventType, rest: ...) -> EventFrame;     // EventPayload::None

Emission belongs to the callers, which already have the sender in scope (a field on StreamEmitContext in upstream.rs, a parameter in gateway.rs):

if acc.process_event(&mut frame, output_offset) {
    emit_sse_frame(sender, &frame)?;
}

Then run_stream mirrors how fetch_stream_payload uses ResponseAccumulator: own it locally, use it, hand it back when done. The spawned task owns the accumulator and returns it with the result in both cases, so no shared ownership is needed:

let mut run_handle = AbortOnDrop::new(tokio::spawn(async move {
    let mut acc = GatewayStreamAccumulator::new();
    let result = run_until_gateway_tools_complete(ctx, exec_ctx_for_run.as_ref(),
        auth.as_deref(), true, Some(&mut acc), &event_tx).await;
    (result, acc)
}));

The join branches use the returned accumulator for terminal_response_chunk / error_chunk; only the JoinError (panic) path lacks one, where a plain unnumbered error chunk is an acceptable fallback. Downstream, Option<&mut GatewayStreamContext<'_>> becomes Option<&mut GatewayStreamAccumulator> plus the sender, and StreamEmitContext drops its second lifetime.

Invariants this preserves: every emitted event (relayed, synthetic, terminal) passes through process_event exactly once, so it stays the sole sequence-numbering authority; wire output is always serialized from frame.wire; and the type does no I/O, so its tests need no channels or runtime. Net: the Arc, the Mutex, the guard-across-await, both lifetime parameters, and the GatewayStreamContext type all disappear.

Further clarification to keep in mind coding in rust:
Why a borrowed sender in a struct is a bad practice specifically:

1. Channel senders are owned handles, not data. tokio::sync::mpsc::UnboundedSender is Clone by design and documented as a cheap, reference-counted handle meant to be cloned once per producer/task; send() takes &self, so a borrowed sender has zero capability an owned clone lacks. The clone is an Arc refcount bump. The borrow-over-clone guideline (own-borrow-over-clone) targets data with meaningful clone cost; for handles, the ecosystem convention (tokio, std::sync::mpsc::Sender) is one owned clone per owner.

2. A struct with &'a fields is a temporary view type and cannot cross a task boundary. tokio::spawn requires 'static futures ("the spawned task must not contain references to data owned outside it"), so a reference-holding context can never be stored, returned from, or moved into a task. That is the direct cause of the Arc<Mutex> here: the state the context borrows had to live somewhere shareable instead of simply being owned by the task. Reference-holding structs are appropriate for short-lived views (iterators, guards, per-call contexts); long-lived state should own its fields (own-lifetime-elision: explicit lifetimes only when required).

3. The workaround it forces violates stronger rules than the one it follows. Keeping the borrow led to Arc<Mutex> that is never contended (own-arc-shared: Arc is for genuinely shared state) and a guard held across .await for the whole run (async-no-lock-await / anti-lock-across-await, and the same rule in AGENTS.md). async-clone-before-await states the intended pattern directly: clone Arc-backed handles before suspension points instead of holding references across them.

[harivilasp]

Sure thank you Maral will address it.

[harivilasp]

GitHub automatically closed this PR when its cross-fork source branch was renamed during the rebase/push update. The implementation and review work were preserved unchanged on the replacement PR: #136.

#136 contains the reviewed fixes for lossless single-parse streaming frames, one gateway-global sequence authority across relayed/generated events, lifecycle/output-index handling across rounds, and the requested regression coverage. It was validated locally with cargo fmt -- --check, cargo clippy --all-targets -- -D warnings, cargo test, and pre-commit run --all-files.

Please continue review on #136; apologies for the PR migration.

[harivilasp]

it was not allowing to restore as well. very sorry for that.