Skip to main content

agentic_core/executor/
messages_stream.rs

1//! Streaming Messages-native gateway tool loop.
2//!
3//! Consumes vLLM's per-round Anthropic SSE and presents the client **one**
4//! logical message across all gateway rounds:
5//!   * `message_start` emitted once (first round only);
6//!   * surfaced `content_block_*` forwarded with client-visible indices rebased
7//!     contiguously across rounds;
8//!   * gateway-owned `tool_use` blocks suppressed (and their `input_json_delta`
9//!     buffered to reconstruct the call for dispatch);
10//!   * intermediate `message_delta`/`message_stop` (the per-round terminals)
11//!     suppressed; the final round's terminal is forwarded once.
12//!
13//! Structurally the Anthropic-native analogue of the Responses `GatewayStreamAccumulator`
14//! (#119/#132); kept deliberately parallel for a future consolidation. Reuses
15//! only the neutral tool layer via [`crate::types::messages::tool_seam`].
16
17use std::collections::{BTreeMap, HashMap, HashSet};
18use std::sync::Arc;
19use std::time::Duration;
20
21use async_stream::stream;
22use futures::StreamExt;
23use serde_json::{Value, json};
24
25use crate::executor::error::{ExecutorError, ExecutorResult};
26use crate::executor::inference::{BoxStream, response_lines, send_request};
27use crate::executor::messages_request::{normalize_native_web_search, web_search_budget_exhausted_result};
28use crate::executor::request::ExecutionContext;
29use crate::proxy::processed_response_headers;
30use crate::tool::ToolRegistry;
31use crate::types::messages::tool_seam;
32use crate::utils::common::{deserialize_from_str, serialize_to_string};
33
34// Shared with the non-streaming loop so the two Messages loops can't drift.
35use crate::executor::messages_loop::{
36    GATEWAY_TOOL_TIMEOUT, MAX_GATEWAY_TOOL_ROUNDS, MessagesResponse, MessagesUpstream,
37};
38/// vLLM streaming chunk timeout (per line). Generous — the loop's own budget is
39/// the round cap, not this.
40const CHUNK_TIMEOUT: Duration = Duration::from_secs(120);
41
42/// Drive the streaming Messages-native loop, yielding Anthropic SSE lines for
43/// the client. Owns the multi-round → single-message accumulation.
44///
45/// # Errors
46///
47/// Returns an executor error when the initial request cannot be serialized or
48/// when the upstream rejects it before streaming begins.
49pub async fn run_messages_stream(
50    mut request: Value,
51    registry: Arc<ToolRegistry>,
52    exec_ctx: Arc<ExecutionContext>,
53    upstream: MessagesUpstream,
54) -> ExecutorResult<MessagesResponse<BoxStream>> {
55    let mut web_search_budget = normalize_native_web_search(&mut request)?;
56    request["stream"] = Value::Bool(true);
57
58    // Prime the first upstream request before the handler commits an HTTP 200.
59    // This lets initial vLLM errors retain their original status and body.
60    let first_body = serialize_to_string(&request)?;
61    let first_response = send_request(
62        &exec_ctx.client,
63        upstream.url(),
64        first_body,
65        None,
66        Some(upstream.headers()),
67    )
68    .await?;
69    let response_headers = processed_response_headers(first_response.headers());
70
71    let body: BoxStream = Box::pin(stream! {
72        let mut acc = MessagesStreamAccumulator::new(exec_ctx.messages_gateway_tools.clone());
73        let mut prepared_response = Some(first_response);
74
75        for _round in 0..MAX_GATEWAY_TOOL_ROUNDS {
76            let response = if let Some(response) = prepared_response.take() {
77                response
78            } else {
79                let body = match serialize_to_string(&request) {
80                    Ok(b) => b,
81                    Err(e) => { yield error_sse(&e.to_string()); return; }
82                };
83                match send_request(
84                    &exec_ctx.client,
85                    upstream.url(),
86                    body,
87                    None,
88                    Some(upstream.headers()),
89                )
90                .await
91                {
92                    Ok(response) => response,
93                    Err(e) => { yield executor_error_sse(&e); return; }
94                }
95            };
96            let mut response_stream = Box::pin(response_lines(response, CHUNK_TIMEOUT));
97
98            acc.begin_round();
99            while let Some(line) = response_stream.next().await {
100                let line = match line {
101                    Ok(l) => l,
102                    Err(e) => { yield error_sse(&e.to_string()); return; }
103                };
104                for out in acc.push(&line) {
105                    yield out;
106                }
107                if acc.has_upstream_error() {
108                    return;
109                }
110            }
111
112            // Round finished. Continue only for a pure gateway-tool round; a
113            // client-owned tool_use (or any non-tool_use stop) is terminal.
114            if !acc.should_continue_loop() {
115                for out in acc.finish() {
116                    yield out;
117                }
118                return;
119            }
120            // Reconstruct the FULL assistant turn (thinking/text/signature +
121            // gateway tool_use, in order) for the next round's history — not just
122            // the gateway tool_use (F3, streaming half). The gateway calls are
123            // derived from the same buffered blocks for dispatch.
124            let (assistant_content, calls) = acc.take_round();
125            let allowed_searches = web_search_budget.reserve(calls.len());
126            let resolved = execute_gateway_calls(
127                &calls,
128                &registry,
129                &exec_ctx.messages_gateway_tools,
130                allowed_searches,
131            ).await;
132            append_round_to_history(&mut request, &assistant_content, &resolved);
133        }
134
135        // Round budget exhausted.
136        yield error_sse(&format!("gateway tool loop exceeded {MAX_GATEWAY_TOOL_ROUNDS} rounds"));
137    });
138    Ok(MessagesResponse {
139        body,
140        headers: response_headers,
141    })
142}
143
144/// A gateway `tool_use` reconstructed from the stream, ready to dispatch.
145struct StreamedCall {
146    id: String,
147    name: String,
148    input_json: String,
149}
150
151/// One assistant content block buffered across a round, so the full turn
152/// (`thinking`/`text`/`signature`/`tool_use`, in order) can be reconstructed for
153/// the next round's history — F3. The client-facing SSE is still forwarded live;
154/// this is a parallel record for the fed-back conversation state.
155struct BufferedBlock {
156    /// The `content_block` skeleton from `content_block_start`, mutated by deltas.
157    block: Value,
158    /// Accumulated `input_json_delta` fragments for a `tool_use` block.
159    input_json: String,
160    /// Gateway-owned `tool_use` (drives the loop; suppressed from the client).
161    is_gateway_tool: bool,
162}
163
164#[derive(Clone, Copy, Debug, Eq, PartialEq)]
165enum RoundState {
166    Active,
167    UpstreamError,
168}
169
170impl BufferedBlock {
171    fn apply_delta(&mut self, delta: &Value) {
172        match delta.get("type").and_then(Value::as_str) {
173            Some("text_delta") => append_str(&mut self.block, "text", delta.get("text")),
174            Some("thinking_delta") => append_str(&mut self.block, "thinking", delta.get("thinking")),
175            Some("signature_delta") => append_str(&mut self.block, "signature", delta.get("signature")),
176            Some("input_json_delta") => {
177                if let Some(partial) = delta.get("partial_json").and_then(Value::as_str) {
178                    self.input_json.push_str(partial);
179                }
180            }
181            _ => {}
182        }
183    }
184
185    /// The finished assistant content block. For `tool_use`, parse the
186    /// accumulated arguments (best-effort — a malformed fragment falls back to
187    /// `{}`; the paired error `tool_result` records the failure).
188    fn to_block(&self) -> Value {
189        let mut block = self.block.clone();
190        if block.get("type").and_then(Value::as_str) == Some("tool_use") {
191            block["input"] = tool_seam::parse_tool_input(&self.input_json).unwrap_or_else(|_| json!({}));
192        }
193        block
194    }
195}
196
197/// Append a streamed string fragment onto a string field of `block`, creating it
198/// if absent.
199fn append_str(block: &mut Value, field: &str, fragment: Option<&Value>) {
200    let Some(fragment) = fragment.and_then(Value::as_str) else {
201        return;
202    };
203    let combined = match block.get(field).and_then(Value::as_str) {
204        Some(existing) => format!("{existing}{fragment}"),
205        None => fragment.to_owned(),
206    };
207    block[field] = Value::from(combined);
208}
209
210/// State machine that turns per-round Anthropic SSE into one client-visible
211/// message. Fed line-by-line via [`Self::push`].
212struct MessagesStreamAccumulator {
213    message_started: bool,
214    /// Next client-visible block index (contiguous across rounds).
215    next_index: u32,
216    /// Map upstream (per-round) block index → client index, for the blocks we
217    /// forward this round. Cleared each round.
218    index_map: HashMap<u64, u32>,
219    /// Upstream indices belonging to a suppressed gateway `tool_use` this round.
220    suppressed_indices: HashSet<u64>,
221    /// Every assistant block this round, keyed by upstream index (ordered), so
222    /// the full turn — `thinking`/`text`/`signature` + gateway `tool_use` — can
223    /// be reconstructed for the next round's history (F3). Cleared each round.
224    blocks: BTreeMap<u64, BufferedBlock>,
225    /// Did this round end with `stop_reason: tool_use`?
226    ended_on_tool_use: bool,
227    /// Did this round surface a client-owned `tool_use`? If so the loop cannot
228    /// continue server-side (the client must supply that tool's result), so it
229    /// is terminal — matching the non-streaming path's E7 handling.
230    has_client_tool_use: bool,
231    /// Buffered terminal `message_delta` from the final round (emitted by `finish`).
232    final_message_delta: Option<Value>,
233    /// Whether this round is still active or terminated with an upstream error.
234    round_state: RoundState,
235    /// Operator-configured client-tool → gateway-executor aliases, so a client
236    /// tool like Claude Code's `WebSearch` is classified gateway-owned (and
237    /// suppressed) the same way the built-in `web_search` is.
238    gateway_map: tool_seam::GatewayToolMap,
239}
240
241impl MessagesStreamAccumulator {
242    fn new(gateway_map: tool_seam::GatewayToolMap) -> Self {
243        Self {
244            message_started: false,
245            next_index: 0,
246            index_map: HashMap::new(),
247            suppressed_indices: HashSet::new(),
248            blocks: BTreeMap::new(),
249            ended_on_tool_use: false,
250            has_client_tool_use: false,
251            final_message_delta: None,
252            round_state: RoundState::Active,
253            gateway_map,
254        }
255    }
256
257    fn begin_round(&mut self) {
258        self.index_map.clear();
259        self.suppressed_indices.clear();
260        self.blocks.clear();
261        self.ended_on_tool_use = false;
262        self.has_client_tool_use = false;
263        self.round_state = RoundState::Active;
264        // F6: clear the previous round's terminal so a clean-EOF round can't
265        // re-emit a stale stop_reason.
266        self.final_message_delta = None;
267    }
268
269    /// Number of gateway `tool_use` blocks buffered this round.
270    fn gateway_call_count(&self) -> usize {
271        self.blocks.values().filter(|b| b.is_gateway_tool).count()
272    }
273
274    /// Consume this round's buffered blocks, returning (full assistant content in
275    /// order, gateway calls to dispatch). The assistant content preserves
276    /// `thinking`/`text`/`signature` and the gateway `tool_use` blocks (F3); the
277    /// calls are the gateway `tool_use` blocks reconstructed for dispatch.
278    fn take_round(&mut self) -> (Vec<Value>, Vec<StreamedCall>) {
279        let blocks = std::mem::take(&mut self.blocks);
280        let mut assistant_content = Vec::with_capacity(blocks.len());
281        let mut calls = Vec::new();
282        for buffered in blocks.values() {
283            assistant_content.push(buffered.to_block());
284            if buffered.is_gateway_tool {
285                calls.push(StreamedCall {
286                    id: buffered.block["id"].as_str().unwrap_or_default().to_owned(),
287                    name: buffered.block["name"].as_str().unwrap_or_default().to_owned(),
288                    input_json: buffered.input_json.clone(),
289                });
290            }
291        }
292        (assistant_content, calls)
293    }
294
295    /// The loop should continue only when the round asked for a gateway tool AND
296    /// did not also surface a client-owned tool (which the client must handle,
297    /// making the round terminal — E7).
298    fn should_continue_loop(&self) -> bool {
299        self.ended_on_tool_use && self.gateway_call_count() > 0 && !self.has_client_tool_use
300    }
301
302    fn has_upstream_error(&self) -> bool {
303        self.round_state == RoundState::UpstreamError
304    }
305
306    /// Translate one upstream SSE line into zero or more client SSE lines.
307    fn push(&mut self, line: &str) -> Vec<String> {
308        let Some(data) = line.strip_prefix("data: ") else {
309            return Vec::new();
310        };
311        let data = data.trim();
312        if data == "[DONE]" {
313            return Vec::new();
314        }
315        let Ok(mut event) = serde_json::from_str::<Value>(data) else {
316            return Vec::new();
317        };
318        match event.get("type").and_then(Value::as_str) {
319            Some("message_start") => self.on_message_start(&event),
320            Some("content_block_start") => self.on_block_start(&mut event),
321            Some("content_block_delta") => self.on_block_delta(&mut event),
322            Some("content_block_stop") => self.on_block_stop(&mut event),
323            Some("message_delta") => {
324                // Buffer as the (possibly) final terminal; suppress mid-loop.
325                self.ended_on_tool_use = event["delta"]["stop_reason"].as_str() == Some("tool_use");
326                self.final_message_delta = Some(event);
327                Vec::new()
328            }
329            Some("error") => {
330                self.round_state = RoundState::UpstreamError;
331                vec![sse("error", &event)]
332            }
333            // `message_stop` (per-round terminal) is suppressed; `finish` emits
334            // the single client-visible terminal. Everything else is dropped.
335            _ => Vec::new(),
336        }
337    }
338
339    fn on_message_start(&mut self, event: &Value) -> Vec<String> {
340        if self.message_started {
341            return Vec::new();
342        }
343        self.message_started = true;
344        vec![sse("message_start", event)]
345    }
346
347    fn on_block_start(&mut self, event: &mut Value) -> Vec<String> {
348        let up_index = event.get("index").and_then(Value::as_u64).unwrap_or(0);
349        let block_type = event["content_block"]["type"].as_str().unwrap_or_default();
350        let name = event["content_block"]["name"].as_str().unwrap_or_default();
351
352        // Buffer every block for history reconstruction (F3), preserving order.
353        let is_gateway_tool = block_type == "tool_use" && self.gateway_map.is_gateway_owned(name);
354        self.blocks.insert(
355            up_index,
356            BufferedBlock {
357                block: event["content_block"].clone(),
358                input_json: String::new(),
359                is_gateway_tool,
360            },
361        );
362
363        if block_type == "tool_use" {
364            if is_gateway_tool {
365                // Suppress gateway-owned tool_use from the client; it stays in the
366                // buffered history only and drives the loop.
367                self.suppressed_indices.insert(up_index);
368                return Vec::new();
369            }
370            // A client-owned tool_use: the client must execute it, so this round
371            // is terminal (E7). Forward it (below) and stop the loop.
372            self.has_client_tool_use = true;
373        }
374
375        // Forward with a rebased contiguous client index.
376        let client_index = self.next_index;
377        self.next_index += 1;
378        self.index_map.insert(up_index, client_index);
379        event["index"] = Value::from(client_index);
380        vec![sse("content_block_start", event)]
381    }
382
383    fn on_block_delta(&mut self, event: &mut Value) -> Vec<String> {
384        let up_index = event.get("index").and_then(Value::as_u64).unwrap_or(0);
385        // Accumulate the delta into the buffered block (for history — F3),
386        // regardless of whether it is forwarded to the client.
387        if let Some(buffered) = self.blocks.get_mut(&up_index) {
388            buffered.apply_delta(&event["delta"]);
389        }
390        // A suppressed gateway tool_use is not forwarded to the client (its
391        // input_json_delta was just buffered above).
392        if self.suppressed_indices.contains(&up_index) {
393            return Vec::new();
394        }
395        let Some(&client_index) = self.index_map.get(&up_index) else {
396            return Vec::new();
397        };
398        event["index"] = Value::from(client_index);
399        vec![sse("content_block_delta", event)]
400    }
401
402    fn on_block_stop(&mut self, event: &mut Value) -> Vec<String> {
403        let up_index = event.get("index").and_then(Value::as_u64).unwrap_or(0);
404        if self.suppressed_indices.contains(&up_index) {
405            return Vec::new();
406        }
407        let Some(&client_index) = self.index_map.get(&up_index) else {
408            return Vec::new();
409        };
410        event["index"] = Value::from(client_index);
411        vec![sse("content_block_stop", event)]
412    }
413
414    /// Emit the terminal `message_delta` + `message_stop` once, at loop end.
415    fn finish(&mut self) -> Vec<String> {
416        let mut out = Vec::new();
417        if let Some(delta) = self.final_message_delta.take() {
418            out.push(sse("message_delta", &delta));
419        }
420        out.push(sse("message_stop", &json!({"type": "message_stop"})));
421        out
422    }
423}
424
425fn sse(event: &str, value: &Value) -> String {
426    let json = serialize_to_string(value).unwrap_or_default();
427    format!("event: {event}\ndata: {json}\n\n")
428}
429
430fn error_sse(message: &str) -> String {
431    let event = json!({"type": "error", "error": {"type": "api_error", "message": message}});
432    let json = serialize_to_string(&event).unwrap_or_default();
433    format!("event: error\ndata: {json}\n\n")
434}
435
436fn executor_error_sse(error: &ExecutorError) -> String {
437    if let ExecutorError::LLMRequest { body, .. } = error
438        && let Ok(value) = deserialize_from_str::<Value>(body)
439        && value.get("type").and_then(Value::as_str) == Some("error")
440    {
441        let data = if body.contains(['\r', '\n']) {
442            serialize_to_string(&value).unwrap_or_else(|_| body.clone())
443        } else {
444            body.clone()
445        };
446        return format!("event: error\ndata: {data}\n\n");
447    }
448    error_sse(&error.to_string())
449}
450
451/// Execute reconstructed gateway calls (concurrent, per-call timeout). Errors
452/// become error `tool_result`s (E5).
453async fn execute_gateway_calls(
454    calls: &[StreamedCall],
455    registry: &ToolRegistry,
456    gateway_map: &tool_seam::GatewayToolMap,
457    allowed_searches: usize,
458) -> Vec<ResolvedStreamCall> {
459    let futures = calls.iter().enumerate().map(|(index, c)| async move {
460        if index >= allowed_searches {
461            return ResolvedStreamCall {
462                tool_result_block: web_search_budget_exhausted_result(&c.id),
463            };
464        }
465        // F4: reject a malformed/incomplete reconstructed input rather than
466        // coercing to {} and dispatching the tool with args the model never sent.
467        let (output, is_error) = match tool_seam::parse_tool_input(&c.input_json) {
468            Ok(input) => {
469                let call = tool_seam::tool_use_to_call(&c.id, &c.name, &input, gateway_map);
470                match tokio::time::timeout(GATEWAY_TOOL_TIMEOUT, registry.dispatch(&call)).await {
471                    Ok(Some(result)) => match result.output {
472                        Ok(o) => (o.output, false),
473                        Err(e) => (format!("tool execution failed: {e}"), true),
474                    },
475                    Ok(None) => (format!("no handler for tool '{}'", c.name), true),
476                    Err(_) => (
477                        format!("gateway tool '{}' timed out after {GATEWAY_TOOL_TIMEOUT:?}", c.name),
478                        true,
479                    ),
480                }
481            }
482            Err(reason) => (format!("{reason}; tool was not run"), true),
483        };
484        ResolvedStreamCall {
485            tool_result_block: tool_seam::tool_result_block(&c.id, &output, is_error),
486        }
487    });
488    futures::future::join_all(futures).await
489}
490
491/// The `tool_result` block for one executed gateway call, fed back next round.
492/// (The assistant turn — including this call's `tool_use` block — is reconstructed
493/// from the accumulator's buffered blocks in [`MessagesStreamAccumulator::take_round`].)
494struct ResolvedStreamCall {
495    tool_result_block: Value,
496}
497
498/// Append the model's full assistant turn (`thinking`/`text`/`signature` +
499/// gateway `tool_use`, order preserved — F3) and a following user turn of `tool_result`s,
500/// so the next upstream round sees the complete conversation state. These stay
501/// internal — the client never sees the gateway call (hide-the-call).
502fn append_round_to_history(request: &mut Value, assistant_content: &[Value], resolved: &[ResolvedStreamCall]) {
503    let assistant = json!({ "role": "assistant", "content": assistant_content });
504    let user = json!({
505        "role": "user",
506        "content": resolved.iter().map(|r| r.tool_result_block.clone()).collect::<Vec<_>>()
507    });
508    if let Some(messages) = request.get_mut("messages").and_then(Value::as_array_mut) {
509        messages.push(assistant);
510        messages.push(user);
511    }
512}
513
514#[cfg(test)]
515mod tests {
516    use super::*;
517
518    fn line(v: &Value) -> String {
519        format!("data: {v}")
520    }
521
522    /// Accumulator with the default gateway map (built-in `web_search` only).
523    fn acc() -> MessagesStreamAccumulator {
524        MessagesStreamAccumulator::new(tool_seam::GatewayToolMap::default())
525    }
526
527    // A single non-tool round: message_start forwarded once, blocks pass through
528    // with contiguous indices, terminal emitted by finish().
529    #[test]
530    fn single_round_text_passes_through() {
531        let mut acc = acc();
532        acc.begin_round();
533        let mut out = Vec::new();
534        out.extend(acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}}))));
535        out.extend(acc.push(&line(
536            &json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
537        )));
538        out.extend(acc.push(&line(
539            &json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}),
540        )));
541        out.extend(acc.push(&line(&json!({"type": "content_block_stop", "index": 0}))));
542        out.extend(acc.push(&line(
543            &json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}),
544        )));
545        out.extend(acc.push(&line(&json!({"type": "message_stop"}))));
546        assert!(!acc.should_continue_loop(), "text-only round is terminal");
547        out.extend(acc.finish());
548        let s = out.join("");
549        assert_eq!(s.matches("event: message_start").count(), 1);
550        assert_eq!(s.matches("event: message_stop").count(), 1);
551        assert!(s.contains("text_delta"));
552        assert!(s.contains("end_turn"));
553    }
554
555    // A gateway tool round: the tool_use block (start/delta/stop) is suppressed,
556    // its input reconstructed, thinking/text forwarded, and no terminal leaks.
557    #[test]
558    fn gateway_tool_round_suppresses_tool_use_and_reconstructs_call() {
559        let mut acc = acc();
560        acc.begin_round();
561        let mut out = Vec::new();
562        out.extend(acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}}))));
563        // thinking idx0 (forward)
564        out.extend(acc.push(&line(
565            &json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}),
566        )));
567        out.extend(acc.push(&line(&json!({"type": "content_block_stop", "index": 0}))));
568        // gateway tool_use idx1 (suppress + reconstruct)
569        out.extend(acc.push(&line(&json!({"type": "content_block_start", "index": 1, "content_block": {"type": "tool_use", "id": "tid", "name": "web_search", "input": {}}}))));
570        out.extend(acc.push(&line(&json!({"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": "{\"query\":"}}))));
571        out.extend(acc.push(&line(&json!({"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": "\"rust\"}"}}))));
572        out.extend(acc.push(&line(&json!({"type": "content_block_stop", "index": 1}))));
573        out.extend(acc.push(&line(
574            &json!({"type": "message_delta", "delta": {"stop_reason": "tool_use"}}),
575        )));
576        out.extend(acc.push(&line(&json!({"type": "message_stop"}))));
577
578        let s = out.join("");
579        assert!(acc.should_continue_loop(), "pure gateway-tool round continues the loop");
580        assert!(!s.contains("tool_use"), "gateway tool_use must not surface: {s}");
581        assert!(!s.contains("message_stop"), "intermediate terminal suppressed");
582        assert!(s.contains("thinking"), "thinking forwarded");
583        let (_assistant, calls) = acc.take_round();
584        assert_eq!(calls.len(), 1);
585        assert_eq!(calls[0].name, "web_search");
586        assert_eq!(calls[0].input_json, "{\"query\":\"rust\"}");
587    }
588
589    // Across two rounds, client-visible block indices stay contiguous (round 1
590    // thinking=0, round 2 text=1) — no reset/collision.
591    #[test]
592    fn indices_are_contiguous_across_rounds() {
593        let mut acc = acc();
594        // round 1: thinking (idx0) + suppressed tool_use (idx1)
595        acc.begin_round();
596        acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}})));
597        acc.push(&line(
598            &json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking"}}),
599        ));
600        acc.push(&line(&json!({"type": "content_block_stop", "index": 0})));
601        acc.push(&line(&json!({"type": "content_block_start", "index": 1, "content_block": {"type": "tool_use", "name": "web_search", "id": "t"}})));
602        acc.push(&line(&json!({"type": "content_block_stop", "index": 1})));
603        // round 2: text (upstream idx0) must map to client idx1
604        acc.begin_round();
605        let out = acc.push(&line(
606            &json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text"}}),
607        ));
608        let started: Value =
609            serde_json::from_str(out[0].lines().nth(1).unwrap().strip_prefix("data: ").unwrap()).unwrap();
610        assert_eq!(started["index"], 1, "round-2 text rebased to contiguous client index 1");
611    }
612
613    // E7 (streaming): a round with a gateway tool_use AND a client-owned tool_use
614    // is terminal — the loop must NOT continue (the client owns the second tool).
615    // The client-owned tool_use is forwarded; the gateway one is suppressed.
616    #[test]
617    fn mixed_client_and_gateway_tool_use_stops_the_loop() {
618        let mut acc = acc();
619        acc.begin_round();
620        acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}})));
621        // gateway tool_use (idx0) — suppressed
622        acc.push(&line(&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "name": "web_search", "id": "g"}})));
623        acc.push(&line(&json!({"type": "content_block_stop", "index": 0})));
624        // client tool_use (idx1) — forwarded
625        let out = acc.push(&line(&json!({"type": "content_block_start", "index": 1, "content_block": {"type": "tool_use", "name": "get_weather", "id": "c"}})));
626        acc.push(&line(&json!({"type": "content_block_stop", "index": 1})));
627        acc.push(&line(
628            &json!({"type": "message_delta", "delta": {"stop_reason": "tool_use"}}),
629        ));
630
631        // Client tool_use surfaces; gateway one does not.
632        let started: Value =
633            serde_json::from_str(out[0].lines().nth(1).unwrap().strip_prefix("data: ").unwrap()).unwrap();
634        assert_eq!(
635            started["content_block"]["name"], "get_weather",
636            "client tool_use forwarded"
637        );
638        // The loop must terminate despite a gateway call being present.
639        assert!(
640            !acc.should_continue_loop(),
641            "mixed round is terminal — loop must not continue"
642        );
643    }
644
645    // F6 (repro): begin_round() must reset final_message_delta. Round 1 ends on a
646    // tool_use terminal; round 2 ends WITHOUT a message_delta (clean EOF). finish()
647    // must NOT emit round 1's stale stop_reason: tool_use.
648    #[test]
649    fn repro_f6_begin_round_resets_stale_terminal() {
650        let mut acc = acc();
651        // Round 1: a gateway tool round → sets final_message_delta = tool_use terminal.
652        acc.begin_round();
653        acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}})));
654        acc.push(&line(&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "name": "web_search", "id": "t"}})));
655        acc.push(&line(&json!({"type": "content_block_stop", "index": 0})));
656        acc.push(&line(
657            &json!({"type": "message_delta", "delta": {"stop_reason": "tool_use"}}),
658        ));
659        // Round 2: text, but upstream ends with NO message_delta (cut short).
660        acc.begin_round();
661        acc.push(&line(
662            &json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text"}}),
663        ));
664        acc.push(&line(
665            &json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}),
666        ));
667        acc.push(&line(&json!({"type": "content_block_stop", "index": 0})));
668        let out = acc.finish().join("");
669        assert!(
670            !out.contains(r#""stop_reason":"tool_use""#),
671            "must not emit round 1's stale tool_use terminal: {out}"
672        );
673    }
674
675    // F3 (repro): the assistant turn fed into the next round's history must
676    // preserve the model's thinking/text/signature blocks in order, not just the
677    // gateway tool_use. (This is the streaming half of Maral's F3 — "also repeated
678    // in messages_stream.rs".)
679    #[test]
680    fn repro_f3_stream_history_preserves_thinking_text_and_signature() {
681        let mut acc = acc();
682        acc.begin_round();
683        acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}})));
684        // thinking idx0 (with a signature delta)
685        acc.push(&line(
686            &json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}),
687        ));
688        acc.push(&line(&json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "let me search"}})));
689        acc.push(&line(&json!({"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": "SIG=="}})));
690        acc.push(&line(&json!({"type": "content_block_stop", "index": 0})));
691        // text idx1
692        acc.push(&line(
693            &json!({"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}),
694        ));
695        acc.push(&line(&json!({"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "Searching..."}})));
696        acc.push(&line(&json!({"type": "content_block_stop", "index": 1})));
697        // gateway tool_use idx2 (suppressed from client, but must appear in history)
698        acc.push(&line(&json!({"type": "content_block_start", "index": 2, "content_block": {"type": "tool_use", "id": "tid", "name": "web_search", "input": {}}})));
699        acc.push(&line(&json!({"type": "content_block_delta", "index": 2, "delta": {"type": "input_json_delta", "partial_json": "{\"query\":\"rust\"}"}})));
700        acc.push(&line(&json!({"type": "content_block_stop", "index": 2})));
701        acc.push(&line(
702            &json!({"type": "message_delta", "delta": {"stop_reason": "tool_use"}}),
703        ));
704
705        let (assistant, _calls) = acc.take_round();
706        let types: Vec<&str> = assistant.iter().filter_map(|b| b["type"].as_str()).collect();
707        assert_eq!(
708            types,
709            vec!["thinking", "text", "tool_use"],
710            "full assistant turn preserved in order, not just the gateway tool_use: {assistant:?}"
711        );
712        assert_eq!(assistant[0]["thinking"], "let me search", "thinking text reconstructed");
713        assert_eq!(
714            assistant[0]["signature"], "SIG==",
715            "signature preserved for the next round"
716        );
717        assert_eq!(assistant[1]["text"], "Searching...", "text reconstructed");
718        assert_eq!(
719            assistant[2]["input"]["query"], "rust",
720            "gateway call input reconstructed"
721        );
722    }
723
724    // F4 (repro): a malformed/incomplete input_json for a gateway call must NOT
725    // silently become `{}` and dispatch the tool with args the model never sent.
726    #[tokio::test]
727    async fn repro_f4_malformed_partial_json_is_not_dispatched_with_empty_args() {
728        let mut acc = acc();
729        acc.begin_round();
730        acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}})));
731        acc.push(&line(&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "name": "web_search", "id": "t"}})));
732        // Incomplete partial_json (stream cut mid-arguments).
733        acc.push(&line(&json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"query\":"}})));
734        let (_assistant, calls) = acc.take_round();
735        assert_eq!(calls.len(), 1);
736        // The reconstructed input is invalid JSON.
737        assert!(
738            serde_json::from_str::<serde_json::Value>(&calls[0].input_json).is_err(),
739            "incomplete partial_json is invalid JSON"
740        );
741        // After the fix, execute_gateway_calls must NOT coerce invalid input to
742        // {} and dispatch — it must produce an error tool_result. Assert the
743        // reconstructed call is flagged invalid rather than silently dispatchable.
744        let resolved = execute_gateway_calls(
745            &calls,
746            &no_op_registry().await,
747            &tool_seam::GatewayToolMap::default(),
748            calls.len(),
749        )
750        .await;
751        let content = resolved[0].tool_result_block["content"].as_str().unwrap_or_default();
752        assert!(
753            content.contains("invalid") || content.contains("malformed") || content.contains("could not"),
754            "malformed args must yield an error tool_result, not an empty-arg dispatch: {content:?}"
755        );
756    }
757
758    /// Registry with no gateway executors — dispatch of any call returns None, so
759    /// the ONLY way `execute_gateway_calls` can produce a non-"no handler" result
760    /// for a malformed input is by rejecting the args before dispatch (the fix).
761    async fn no_op_registry() -> ToolRegistry {
762        let mut tools = [];
763        let mut executors = crate::tool::GatewayExecutors::from_env(std::sync::Arc::new(reqwest::Client::new()));
764        ToolRegistry::build_with_handlers(&mut tools, &mut executors)
765            .await
766            .unwrap()
767    }
768}