use std::collections::{BTreeMap, HashMap, HashSet};
use std::sync::Arc;
use std::time::Duration;
use async_stream::stream;
use futures::StreamExt;
use serde_json::{Value, json};
use crate::executor::error::{ExecutorError, ExecutorResult};
use crate::executor::inference::{BoxStream, response_lines, send_request};
use crate::executor::messages_request::{normalize_native_web_search, web_search_budget_exhausted_result};
use crate::executor::request::ExecutionContext;
use crate::proxy::processed_response_headers;
use crate::tool::ToolRegistry;
use crate::types::messages::tool_seam;
use crate::utils::common::{deserialize_from_str, serialize_to_string};
use crate::executor::messages_loop::{
GATEWAY_TOOL_TIMEOUT, MAX_GATEWAY_TOOL_ROUNDS, MessagesResponse, MessagesUpstream,
};
const CHUNK_TIMEOUT: Duration = Duration::from_secs(120);
pub async fn run_messages_stream(
mut request: Value,
registry: Arc<ToolRegistry>,
exec_ctx: Arc<ExecutionContext>,
upstream: MessagesUpstream,
) -> ExecutorResult<MessagesResponse<BoxStream>> {
let mut web_search_budget = normalize_native_web_search(&mut request)?;
request["stream"] = Value::Bool(true);
let first_body = serialize_to_string(&request)?;
let first_response = send_request(
&exec_ctx.client,
upstream.url(),
first_body,
None,
Some(upstream.headers()),
)
.await?;
let response_headers = processed_response_headers(first_response.headers());
let body: BoxStream = Box::pin(stream! {
let mut acc = MessagesStreamAccumulator::new(exec_ctx.messages_gateway_tools.clone());
let mut prepared_response = Some(first_response);
for _round in 0..MAX_GATEWAY_TOOL_ROUNDS {
let response = if let Some(response) = prepared_response.take() {
response
} else {
let body = match serialize_to_string(&request) {
Ok(b) => b,
Err(e) => { yield error_sse(&e.to_string()); return; }
};
match send_request(
&exec_ctx.client,
upstream.url(),
body,
None,
Some(upstream.headers()),
)
.await
{
Ok(response) => response,
Err(e) => { yield executor_error_sse(&e); return; }
}
};
let mut response_stream = Box::pin(response_lines(response, CHUNK_TIMEOUT));
acc.begin_round();
while let Some(line) = response_stream.next().await {
let line = match line {
Ok(l) => l,
Err(e) => { yield error_sse(&e.to_string()); return; }
};
for out in acc.push(&line) {
yield out;
}
if acc.has_upstream_error() {
return;
}
}
if !acc.should_continue_loop() {
for out in acc.finish() {
yield out;
}
return;
}
let (assistant_content, calls) = acc.take_round();
let allowed_searches = web_search_budget.reserve(calls.len());
let resolved = execute_gateway_calls(
&calls,
®istry,
&exec_ctx.messages_gateway_tools,
allowed_searches,
).await;
append_round_to_history(&mut request, &assistant_content, &resolved);
}
yield error_sse(&format!("gateway tool loop exceeded {MAX_GATEWAY_TOOL_ROUNDS} rounds"));
});
Ok(MessagesResponse {
body,
headers: response_headers,
})
}
struct StreamedCall {
id: String,
name: String,
input_json: String,
}
struct BufferedBlock {
block: Value,
input_json: String,
is_gateway_tool: bool,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum RoundState {
Active,
UpstreamError,
}
impl BufferedBlock {
fn apply_delta(&mut self, delta: &Value) {
match delta.get("type").and_then(Value::as_str) {
Some("text_delta") => append_str(&mut self.block, "text", delta.get("text")),
Some("thinking_delta") => append_str(&mut self.block, "thinking", delta.get("thinking")),
Some("signature_delta") => append_str(&mut self.block, "signature", delta.get("signature")),
Some("input_json_delta") => {
if let Some(partial) = delta.get("partial_json").and_then(Value::as_str) {
self.input_json.push_str(partial);
}
}
_ => {}
}
}
fn to_block(&self) -> Value {
let mut block = self.block.clone();
if block.get("type").and_then(Value::as_str) == Some("tool_use") {
block["input"] = tool_seam::parse_tool_input(&self.input_json).unwrap_or_else(|_| json!({}));
}
block
}
}
fn append_str(block: &mut Value, field: &str, fragment: Option<&Value>) {
let Some(fragment) = fragment.and_then(Value::as_str) else {
return;
};
let combined = match block.get(field).and_then(Value::as_str) {
Some(existing) => format!("{existing}{fragment}"),
None => fragment.to_owned(),
};
block[field] = Value::from(combined);
}
struct MessagesStreamAccumulator {
message_started: bool,
next_index: u32,
index_map: HashMap<u64, u32>,
suppressed_indices: HashSet<u64>,
blocks: BTreeMap<u64, BufferedBlock>,
ended_on_tool_use: bool,
has_client_tool_use: bool,
final_message_delta: Option<Value>,
round_state: RoundState,
gateway_map: tool_seam::GatewayToolMap,
}
impl MessagesStreamAccumulator {
fn new(gateway_map: tool_seam::GatewayToolMap) -> Self {
Self {
message_started: false,
next_index: 0,
index_map: HashMap::new(),
suppressed_indices: HashSet::new(),
blocks: BTreeMap::new(),
ended_on_tool_use: false,
has_client_tool_use: false,
final_message_delta: None,
round_state: RoundState::Active,
gateway_map,
}
}
fn begin_round(&mut self) {
self.index_map.clear();
self.suppressed_indices.clear();
self.blocks.clear();
self.ended_on_tool_use = false;
self.has_client_tool_use = false;
self.round_state = RoundState::Active;
self.final_message_delta = None;
}
fn gateway_call_count(&self) -> usize {
self.blocks.values().filter(|b| b.is_gateway_tool).count()
}
fn take_round(&mut self) -> (Vec<Value>, Vec<StreamedCall>) {
let blocks = std::mem::take(&mut self.blocks);
let mut assistant_content = Vec::with_capacity(blocks.len());
let mut calls = Vec::new();
for buffered in blocks.values() {
assistant_content.push(buffered.to_block());
if buffered.is_gateway_tool {
calls.push(StreamedCall {
id: buffered.block["id"].as_str().unwrap_or_default().to_owned(),
name: buffered.block["name"].as_str().unwrap_or_default().to_owned(),
input_json: buffered.input_json.clone(),
});
}
}
(assistant_content, calls)
}
fn should_continue_loop(&self) -> bool {
self.ended_on_tool_use && self.gateway_call_count() > 0 && !self.has_client_tool_use
}
fn has_upstream_error(&self) -> bool {
self.round_state == RoundState::UpstreamError
}
fn push(&mut self, line: &str) -> Vec<String> {
let Some(data) = line.strip_prefix("data: ") else {
return Vec::new();
};
let data = data.trim();
if data == "[DONE]" {
return Vec::new();
}
let Ok(mut event) = serde_json::from_str::<Value>(data) else {
return Vec::new();
};
match event.get("type").and_then(Value::as_str) {
Some("message_start") => self.on_message_start(&event),
Some("content_block_start") => self.on_block_start(&mut event),
Some("content_block_delta") => self.on_block_delta(&mut event),
Some("content_block_stop") => self.on_block_stop(&mut event),
Some("message_delta") => {
self.ended_on_tool_use = event["delta"]["stop_reason"].as_str() == Some("tool_use");
self.final_message_delta = Some(event);
Vec::new()
}
Some("error") => {
self.round_state = RoundState::UpstreamError;
vec![sse("error", &event)]
}
_ => Vec::new(),
}
}
fn on_message_start(&mut self, event: &Value) -> Vec<String> {
if self.message_started {
return Vec::new();
}
self.message_started = true;
vec![sse("message_start", event)]
}
fn on_block_start(&mut self, event: &mut Value) -> Vec<String> {
let up_index = event.get("index").and_then(Value::as_u64).unwrap_or(0);
let block_type = event["content_block"]["type"].as_str().unwrap_or_default();
let name = event["content_block"]["name"].as_str().unwrap_or_default();
let is_gateway_tool = block_type == "tool_use" && self.gateway_map.is_gateway_owned(name);
self.blocks.insert(
up_index,
BufferedBlock {
block: event["content_block"].clone(),
input_json: String::new(),
is_gateway_tool,
},
);
if block_type == "tool_use" {
if is_gateway_tool {
self.suppressed_indices.insert(up_index);
return Vec::new();
}
self.has_client_tool_use = true;
}
let client_index = self.next_index;
self.next_index += 1;
self.index_map.insert(up_index, client_index);
event["index"] = Value::from(client_index);
vec![sse("content_block_start", event)]
}
fn on_block_delta(&mut self, event: &mut Value) -> Vec<String> {
let up_index = event.get("index").and_then(Value::as_u64).unwrap_or(0);
if let Some(buffered) = self.blocks.get_mut(&up_index) {
buffered.apply_delta(&event["delta"]);
}
if self.suppressed_indices.contains(&up_index) {
return Vec::new();
}
let Some(&client_index) = self.index_map.get(&up_index) else {
return Vec::new();
};
event["index"] = Value::from(client_index);
vec![sse("content_block_delta", event)]
}
fn on_block_stop(&mut self, event: &mut Value) -> Vec<String> {
let up_index = event.get("index").and_then(Value::as_u64).unwrap_or(0);
if self.suppressed_indices.contains(&up_index) {
return Vec::new();
}
let Some(&client_index) = self.index_map.get(&up_index) else {
return Vec::new();
};
event["index"] = Value::from(client_index);
vec![sse("content_block_stop", event)]
}
fn finish(&mut self) -> Vec<String> {
let mut out = Vec::new();
if let Some(delta) = self.final_message_delta.take() {
out.push(sse("message_delta", &delta));
}
out.push(sse("message_stop", &json!({"type": "message_stop"})));
out
}
}
fn sse(event: &str, value: &Value) -> String {
let json = serialize_to_string(value).unwrap_or_default();
format!("event: {event}\ndata: {json}\n\n")
}
fn error_sse(message: &str) -> String {
let event = json!({"type": "error", "error": {"type": "api_error", "message": message}});
let json = serialize_to_string(&event).unwrap_or_default();
format!("event: error\ndata: {json}\n\n")
}
fn executor_error_sse(error: &ExecutorError) -> String {
if let ExecutorError::LLMRequest { body, .. } = error
&& let Ok(value) = deserialize_from_str::<Value>(body)
&& value.get("type").and_then(Value::as_str) == Some("error")
{
let data = if body.contains(['\r', '\n']) {
serialize_to_string(&value).unwrap_or_else(|_| body.clone())
} else {
body.clone()
};
return format!("event: error\ndata: {data}\n\n");
}
error_sse(&error.to_string())
}
async fn execute_gateway_calls(
calls: &[StreamedCall],
registry: &ToolRegistry,
gateway_map: &tool_seam::GatewayToolMap,
allowed_searches: usize,
) -> Vec<ResolvedStreamCall> {
let futures = calls.iter().enumerate().map(|(index, c)| async move {
if index >= allowed_searches {
return ResolvedStreamCall {
tool_result_block: web_search_budget_exhausted_result(&c.id),
};
}
let (output, is_error) = match tool_seam::parse_tool_input(&c.input_json) {
Ok(input) => {
let call = tool_seam::tool_use_to_call(&c.id, &c.name, &input, gateway_map);
match tokio::time::timeout(GATEWAY_TOOL_TIMEOUT, registry.dispatch(&call)).await {
Ok(Some(result)) => match result.output {
Ok(o) => (o.output, false),
Err(e) => (format!("tool execution failed: {e}"), true),
},
Ok(None) => (format!("no handler for tool '{}'", c.name), true),
Err(_) => (
format!("gateway tool '{}' timed out after {GATEWAY_TOOL_TIMEOUT:?}", c.name),
true,
),
}
}
Err(reason) => (format!("{reason}; tool was not run"), true),
};
ResolvedStreamCall {
tool_result_block: tool_seam::tool_result_block(&c.id, &output, is_error),
}
});
futures::future::join_all(futures).await
}
struct ResolvedStreamCall {
tool_result_block: Value,
}
fn append_round_to_history(request: &mut Value, assistant_content: &[Value], resolved: &[ResolvedStreamCall]) {
let assistant = json!({ "role": "assistant", "content": assistant_content });
let user = json!({
"role": "user",
"content": resolved.iter().map(|r| r.tool_result_block.clone()).collect::<Vec<_>>()
});
if let Some(messages) = request.get_mut("messages").and_then(Value::as_array_mut) {
messages.push(assistant);
messages.push(user);
}
}
#[cfg(test)]
mod tests {
use super::*;
fn line(v: &Value) -> String {
format!("data: {v}")
}
fn acc() -> MessagesStreamAccumulator {
MessagesStreamAccumulator::new(tool_seam::GatewayToolMap::default())
}
#[test]
fn single_round_text_passes_through() {
let mut acc = acc();
acc.begin_round();
let mut out = Vec::new();
out.extend(acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}}))));
out.extend(acc.push(&line(
&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
)));
out.extend(acc.push(&line(
&json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}),
)));
out.extend(acc.push(&line(&json!({"type": "content_block_stop", "index": 0}))));
out.extend(acc.push(&line(
&json!({"type": "message_delta", "delta": {"stop_reason": "end_turn"}}),
)));
out.extend(acc.push(&line(&json!({"type": "message_stop"}))));
assert!(!acc.should_continue_loop(), "text-only round is terminal");
out.extend(acc.finish());
let s = out.join("");
assert_eq!(s.matches("event: message_start").count(), 1);
assert_eq!(s.matches("event: message_stop").count(), 1);
assert!(s.contains("text_delta"));
assert!(s.contains("end_turn"));
}
#[test]
fn gateway_tool_round_suppresses_tool_use_and_reconstructs_call() {
let mut acc = acc();
acc.begin_round();
let mut out = Vec::new();
out.extend(acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}}))));
out.extend(acc.push(&line(
&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}),
)));
out.extend(acc.push(&line(&json!({"type": "content_block_stop", "index": 0}))));
out.extend(acc.push(&line(&json!({"type": "content_block_start", "index": 1, "content_block": {"type": "tool_use", "id": "tid", "name": "web_search", "input": {}}}))));
out.extend(acc.push(&line(&json!({"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": "{\"query\":"}}))));
out.extend(acc.push(&line(&json!({"type": "content_block_delta", "index": 1, "delta": {"type": "input_json_delta", "partial_json": "\"rust\"}"}}))));
out.extend(acc.push(&line(&json!({"type": "content_block_stop", "index": 1}))));
out.extend(acc.push(&line(
&json!({"type": "message_delta", "delta": {"stop_reason": "tool_use"}}),
)));
out.extend(acc.push(&line(&json!({"type": "message_stop"}))));
let s = out.join("");
assert!(acc.should_continue_loop(), "pure gateway-tool round continues the loop");
assert!(!s.contains("tool_use"), "gateway tool_use must not surface: {s}");
assert!(!s.contains("message_stop"), "intermediate terminal suppressed");
assert!(s.contains("thinking"), "thinking forwarded");
let (_assistant, calls) = acc.take_round();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].name, "web_search");
assert_eq!(calls[0].input_json, "{\"query\":\"rust\"}");
}
#[test]
fn indices_are_contiguous_across_rounds() {
let mut acc = acc();
acc.begin_round();
acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}})));
acc.push(&line(
&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking"}}),
));
acc.push(&line(&json!({"type": "content_block_stop", "index": 0})));
acc.push(&line(&json!({"type": "content_block_start", "index": 1, "content_block": {"type": "tool_use", "name": "web_search", "id": "t"}})));
acc.push(&line(&json!({"type": "content_block_stop", "index": 1})));
acc.begin_round();
let out = acc.push(&line(
&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text"}}),
));
let started: Value =
serde_json::from_str(out[0].lines().nth(1).unwrap().strip_prefix("data: ").unwrap()).unwrap();
assert_eq!(started["index"], 1, "round-2 text rebased to contiguous client index 1");
}
#[test]
fn mixed_client_and_gateway_tool_use_stops_the_loop() {
let mut acc = acc();
acc.begin_round();
acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}})));
acc.push(&line(&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "name": "web_search", "id": "g"}})));
acc.push(&line(&json!({"type": "content_block_stop", "index": 0})));
let out = acc.push(&line(&json!({"type": "content_block_start", "index": 1, "content_block": {"type": "tool_use", "name": "get_weather", "id": "c"}})));
acc.push(&line(&json!({"type": "content_block_stop", "index": 1})));
acc.push(&line(
&json!({"type": "message_delta", "delta": {"stop_reason": "tool_use"}}),
));
let started: Value =
serde_json::from_str(out[0].lines().nth(1).unwrap().strip_prefix("data: ").unwrap()).unwrap();
assert_eq!(
started["content_block"]["name"], "get_weather",
"client tool_use forwarded"
);
assert!(
!acc.should_continue_loop(),
"mixed round is terminal — loop must not continue"
);
}
#[test]
fn repro_f6_begin_round_resets_stale_terminal() {
let mut acc = acc();
acc.begin_round();
acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}})));
acc.push(&line(&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "name": "web_search", "id": "t"}})));
acc.push(&line(&json!({"type": "content_block_stop", "index": 0})));
acc.push(&line(
&json!({"type": "message_delta", "delta": {"stop_reason": "tool_use"}}),
));
acc.begin_round();
acc.push(&line(
&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "text"}}),
));
acc.push(&line(
&json!({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}),
));
acc.push(&line(&json!({"type": "content_block_stop", "index": 0})));
let out = acc.finish().join("");
assert!(
!out.contains(r#""stop_reason":"tool_use""#),
"must not emit round 1's stale tool_use terminal: {out}"
);
}
#[test]
fn repro_f3_stream_history_preserves_thinking_text_and_signature() {
let mut acc = acc();
acc.begin_round();
acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}})));
acc.push(&line(
&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "thinking", "thinking": ""}}),
));
acc.push(&line(&json!({"type": "content_block_delta", "index": 0, "delta": {"type": "thinking_delta", "thinking": "let me search"}})));
acc.push(&line(&json!({"type": "content_block_delta", "index": 0, "delta": {"type": "signature_delta", "signature": "SIG=="}})));
acc.push(&line(&json!({"type": "content_block_stop", "index": 0})));
acc.push(&line(
&json!({"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}),
));
acc.push(&line(&json!({"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "Searching..."}})));
acc.push(&line(&json!({"type": "content_block_stop", "index": 1})));
acc.push(&line(&json!({"type": "content_block_start", "index": 2, "content_block": {"type": "tool_use", "id": "tid", "name": "web_search", "input": {}}})));
acc.push(&line(&json!({"type": "content_block_delta", "index": 2, "delta": {"type": "input_json_delta", "partial_json": "{\"query\":\"rust\"}"}})));
acc.push(&line(&json!({"type": "content_block_stop", "index": 2})));
acc.push(&line(
&json!({"type": "message_delta", "delta": {"stop_reason": "tool_use"}}),
));
let (assistant, _calls) = acc.take_round();
let types: Vec<&str> = assistant.iter().filter_map(|b| b["type"].as_str()).collect();
assert_eq!(
types,
vec!["thinking", "text", "tool_use"],
"full assistant turn preserved in order, not just the gateway tool_use: {assistant:?}"
);
assert_eq!(assistant[0]["thinking"], "let me search", "thinking text reconstructed");
assert_eq!(
assistant[0]["signature"], "SIG==",
"signature preserved for the next round"
);
assert_eq!(assistant[1]["text"], "Searching...", "text reconstructed");
assert_eq!(
assistant[2]["input"]["query"], "rust",
"gateway call input reconstructed"
);
}
#[tokio::test]
async fn repro_f4_malformed_partial_json_is_not_dispatched_with_empty_args() {
let mut acc = acc();
acc.begin_round();
acc.push(&line(&json!({"type": "message_start", "message": {"id": "m"}})));
acc.push(&line(&json!({"type": "content_block_start", "index": 0, "content_block": {"type": "tool_use", "name": "web_search", "id": "t"}})));
acc.push(&line(&json!({"type": "content_block_delta", "index": 0, "delta": {"type": "input_json_delta", "partial_json": "{\"query\":"}})));
let (_assistant, calls) = acc.take_round();
assert_eq!(calls.len(), 1);
assert!(
serde_json::from_str::<serde_json::Value>(&calls[0].input_json).is_err(),
"incomplete partial_json is invalid JSON"
);
let resolved = execute_gateway_calls(
&calls,
&no_op_registry().await,
&tool_seam::GatewayToolMap::default(),
calls.len(),
)
.await;
let content = resolved[0].tool_result_block["content"].as_str().unwrap_or_default();
assert!(
content.contains("invalid") || content.contains("malformed") || content.contains("could not"),
"malformed args must yield an error tool_result, not an empty-arg dispatch: {content:?}"
);
}
async fn no_op_registry() -> ToolRegistry {
let mut tools = [];
let mut executors = crate::tool::GatewayExecutors::from_env(std::sync::Arc::new(reqwest::Client::new()));
ToolRegistry::build_with_handlers(&mut tools, &mut executors)
.await
.unwrap()
}
}