fix: reject incomplete and malformed Messages streams

.rust-file-sizes.json

@@ -4,7 +4,6 @@
     "crates/agentic-server-core/src/executor/accumulator/mod.rs": 666,
     "crates/agentic-server-core/src/executor/engine.rs": 651,
     "crates/agentic-server-core/src/executor/gateway.rs": 653,
-    "crates/agentic-server-core/src/executor/messages_stream.rs": 501,
     "crates/agentic-server-core/src/executor/session.rs": 508,
     "crates/agentic-server-core/src/storage/schema.rs": 552,
     "crates/agentic-server-core/src/tool/registry.rs": 515,

crates/agentic-server-core/src/executor/messages_stream.rs

@@ -16,7 +16,9 @@
 //! (#119/#132); kept deliberately parallel for a future consolidation. Reuses
 //! only the neutral tool layer via [`crate::types::messages::tool_seam`].
 
+mod block;
 mod wire;
+use block::BufferedBlock;
 use wire::{error_sse, executor_error_sse, sse};
 
 use std::collections::{BTreeMap, HashMap, HashSet};
@@ -116,7 +118,7 @@ pub async fn run_messages_stream(
                 for out in acc.push(&line) {
                     yield out;
                 }
-                if acc.has_upstream_error() {
+                if acc.has_error() {
                     return;
                 }
                 if acc.has_completed_round() {
@@ -169,65 +171,12 @@ struct StreamedCall {
     input_json: String,
 }
 
-/// One assistant content block buffered across a round, so the full turn
-/// (`thinking`/`text`/`signature`/`tool_use`, in order) can be reconstructed for
-/// the next round's history — F3. The client-facing SSE is still forwarded live;
-/// this is a parallel record for the fed-back conversation state.
-struct BufferedBlock {
-    /// The `content_block` skeleton from `content_block_start`, mutated by deltas.
-    block: Value,
-    /// Accumulated `input_json_delta` fragments for a `tool_use` block.
-    input_json: String,
-    /// Gateway-owned `tool_use` (drives the loop; suppressed from the client).
-    is_gateway_tool: bool,
-}
-
 #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
 enum RoundState {
     #[default]
     Active,
     Completed,
-    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);
-                }
-            }
-            _ => {}
-        }
-    }
-
-    /// The finished assistant content block. For `tool_use`, parse the
-    /// accumulated arguments (best-effort — a malformed fragment falls back to
-    /// `{}`; the paired error `tool_result` records the failure).
-    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
-    }
-}
-
-/// Append a streamed string fragment onto a string field of `block`, creating it
-/// if absent.
-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);
+    Failed,
 }
 
 /// State machine that turns per-round Anthropic SSE into one client-visible
@@ -252,7 +201,7 @@ struct MessagesStreamAccumulator {
     has_client_tool_use: bool,
     /// Buffered terminal `message_delta` from the final round (emitted by `finish`).
     final_message_delta: Option<Value>,
-    /// Whether this round is active, complete, or terminated with an upstream error.
+    /// Whether this round is active, complete, or terminated with an error.
     round_state: RoundState,
     /// Every consumed round's terminal `usage`, reported once in the final `message_delta`.
     usage: MessagesUsageTotals,
@@ -310,9 +259,7 @@ impl MessagesStreamAccumulator {
             .final_message_delta
             .as_ref()
             .and_then(|event| event["delta"]["stop_reason"].as_str());
-        // Preserve clean-EOF behavior for ordinary tool_use rounds. The named
-        // end_turn compatibility case requires an explicit message_stop.
-        (self.has_completed_round() || stop_reason == Some("tool_use"))
+        self.has_completed_round()
             && self.gateway_call_count() > 0
             && !self.has_client_tool_use
             && ctx.is_tool_call_stop(
@@ -324,8 +271,8 @@ impl MessagesStreamAccumulator {
             )
     }
 
-    fn has_upstream_error(&self) -> bool {
-        self.round_state == RoundState::UpstreamError
+    fn has_error(&self) -> bool {
+        self.round_state == RoundState::Failed
     }
 
     fn has_completed_round(&self) -> bool {
@@ -338,7 +285,8 @@ impl MessagesStreamAccumulator {
             return Vec::new();
         };
         let Ok(mut event) = deserialize_from_str::<Value>(data.as_str()) else {
-            return Vec::new();
+            self.round_state = RoundState::Failed;
+            return vec![error_sse("invalid JSON in upstream Messages stream")];
         };
         match event.get("type").and_then(Value::as_str) {
             Some("message_start") => self.on_message_start(&event),
@@ -352,7 +300,7 @@ impl MessagesStreamAccumulator {
                 Vec::new()
             }
             Some("error") => {
-                self.round_state = RoundState::UpstreamError;
+                self.round_state = RoundState::Failed;
                 vec![sse("error", &event)]
             }
             Some("message_stop") => {
@@ -442,13 +390,15 @@ impl MessagesStreamAccumulator {
         vec![sse("content_block_stop", event)]
     }
 
-    /// Emit the terminal `message_delta` + `message_stop` once, at loop end.
+    /// Emit completion only for an explicitly completed upstream round.
     fn finish(&mut self) -> Vec<String> {
+        if !self.has_completed_round() {
+            return vec![error_sse("upstream Messages stream ended before message_stop")];
+        }
         let mut out = Vec::new();
         if let Some(mut delta) = self.final_message_delta.take() {
-            // A completed client call requires client action, even when vLLM
-            // labels a named call end_turn. Do not hide truncation or clean EOF.
-            if self.has_client_tool_use && self.has_completed_round() && delta["delta"]["stop_reason"] == "end_turn" {
+            // A completed client call requires client action even when vLLM labels it end_turn.
+            if self.has_client_tool_use && delta["delta"]["stop_reason"] == "end_turn" {
                 delta["delta"]["stop_reason"] = json!("tool_use");
             }
             self.usage.finish(&mut delta);
@@ -638,7 +588,14 @@ mod tests {
                     acc.push(&line(&json!({"type":"message_stop"})));
                 }
                 assert!(!acc.should_continue_loop(&context()));
-                if completed && reason == "end_turn" {
+                if !completed {
+                    assert_eq!(
+                        acc.finish(),
+                        vec![error_sse("upstream Messages stream ended before message_stop")]
+                    );
+                    continue;
+                }
+                if reason == "end_turn" {
                     terminal["delta"]["stop_reason"] = json!("tool_use");
                 }
                 assert_eq!(
@@ -877,6 +834,8 @@ mod tests {
         ));
         acc.push(&line(&json!({"type": "content_block_stop", "index": 0})));
         let out = acc.finish().join("");
+        assert!(out.contains("event: error"), "incomplete later round must fail: {out}");
+        assert!(!out.contains("event: message_stop"), "no synthetic completion: {out}");
         assert!(
             !out.contains(r#""stop_reason":"tool_use""#),
             "must not emit round 1's stale tool_use terminal: {out}"

crates/agentic-server-core/src/executor/messages_stream/block.rs

@@ -0,0 +1,57 @@
+//! Reconstruct one streamed Messages content block for tool dispatch and history.
+
+use crate::types::messages::tool_seam;
+use serde_json::{Value, json};
+
+/// One assistant content block buffered across a round, so the full turn
+/// (`thinking`/`text`/`signature`/`tool_use`, in order) can be reconstructed for
+/// the next round's history — F3. The client-facing SSE is still forwarded live;
+/// this is a parallel record for the fed-back conversation state.
+pub(super) struct BufferedBlock {
+    /// The `content_block` skeleton from `content_block_start`, mutated by deltas.
+    pub(super) block: Value,
+    /// Accumulated `input_json_delta` fragments for a `tool_use` block.
+    pub(super) input_json: String,
+    /// Gateway-owned `tool_use` (drives the loop; suppressed from the client).
+    pub(super) is_gateway_tool: bool,
+}
+
+impl BufferedBlock {
+    pub(super) 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);
+                }
+            }
+            _ => {}
+        }
+    }
+
+    /// The finished assistant content block. For `tool_use`, parse the
+    /// accumulated arguments (best-effort — a malformed fragment falls back to
+    /// `{}`; the paired error `tool_result` records the failure).
+    pub(super) 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
+    }
+}
+
+/// Append a streamed string fragment onto a string field of `block`, creating it
+/// if absent.
+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);
+}

crates/agentic-server-core/tests/messages_stream_test.rs

@@ -8,6 +8,7 @@
 //! block indices stay contiguous across rounds, and no raw per-round terminal
 //! leaks.
 
+use std::fmt::Write as _;
 use std::sync::Arc;
 use std::sync::atomic::{AtomicUsize, Ordering};
 
@@ -621,3 +622,130 @@ async fn messages_stream_preserves_claude_code_cache_control_across_rounds() {
     );
     assert!(requests[1]["tools"][0].get("cache_control").is_none());
 }
+
+/// A terminal `stop_reason` or transport `[DONE]` cannot substitute for `message_stop`.
+#[tokio::test]
+async fn incomplete_rounds_never_dispatch_tools_or_report_success() {
+    for tool in [false, true] {
+        for terminal_delta in [false, true] {
+            for done_marker in [false, true] {
+                let block = if tool {
+                    serde_json::json!({"type":"tool_use", "id":"search", "name":"web_search", "input":{}})
+                } else {
+                    serde_json::json!({"type":"text", "text":""})
+                };
+                let mut events = vec![
+                    serde_json::json!({"type":"message_start", "message":{"id":"m"}}),
+                    serde_json::json!({"type":"content_block_start", "index":0, "content_block":block}),
+                ];
+                if tool {
+                    events.push(serde_json::json!({"type":"content_block_delta", "index":0,
+                        "delta":{"type":"input_json_delta", "partial_json":"{\"query\":\"rust\"}"}}));
+                }
+                events.push(serde_json::json!({"type":"content_block_stop", "index":0}));
+                if terminal_delta {
+                    events.push(serde_json::json!({"type":"message_delta", "delta":{
+                        "stop_reason":if tool { "tool_use" } else { "end_turn" }
+                    }}));
+                }
+                let mut body = String::new();
+                for event in events {
+                    write!(body, "data: {event}\n\n").unwrap();
+                }
+                if done_marker {
+                    body.push_str("data: [DONE]\n\n");
+                }
+                assert_failed_stream(vec![body], 0, "message_stop").await;
+            }
+        }
+    }
+    assert_failed_stream(vec![String::new()], 0, "message_stop").await;
+    // A completed tool round must not make an incomplete later round successful.
+    assert_failed_stream(
+        vec![cassette_turn_streams().remove(0), String::new()],
+        1,
+        "message_stop",
+    )
+    .await;
+}
+
+async fn assert_failed_stream(streams: Vec<String>, expected_searches: usize, expected_error: &str) {
+    let (vllm_url, upstream, vllm) = spawn_mock_vllm_stream(streams).await;
+    let searches = Arc::new(AtomicUsize::new(0));
+    let search_calls = Arc::clone(&searches);
+    let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
+    let search_url = format!("http://{}", listener.local_addr().unwrap());
+    let app = Router::new().route(
+        "/v1/search",
+        get(move || {
+            search_calls.fetch_add(1, Ordering::SeqCst);
+            async { Json(serde_json::json!({"results":{"web":[],"news":[]}})) }
+        }),
+    );
+    let search = tokio::spawn(async move {
+        axum::serve(listener, app).await.unwrap();
+    });
+    let exec_ctx = build_exec_ctx(&vllm_url, &search_url).await;
+    let request = serde_json::json!({
+        "model":"test", "max_tokens":64, "stream":true,
+        "messages":[{"role":"user", "content":"Search"}],
+        "tools":[{"name":"web_search", "input_schema":{"type":"object"}}]
+    });
+    let tools: Vec<ToolParam> = serde_json::from_value(request["tools"].clone()).unwrap();
+    let mut params = registry_tools(Some(&tools), &GatewayToolMap::default());
+    let mut handlers = exec_ctx.gateway_executors.clone();
+    let registry = Arc::new(
+        ToolRegistry::build_with_handlers(&mut params, &mut handlers)
+            .await
+            .unwrap(),
+    );
+    let stream = run_test_messages_stream(request, registry, exec_ctx).await;
+    let sse = tokio::time::timeout(std::time::Duration::from_secs(5), stream.collect::<Vec<_>>())
+        .await
+        .expect("finite incomplete stream")
+        .join("");
+    vllm.abort();
+    search.abort();
+    let _ = tokio::join!(vllm, search);
+
+    assert_eq!(
+        searches.load(Ordering::SeqCst),
+        expected_searches,
+        "no tool execution from incomplete round"
+    );
+    assert_eq!(
+        upstream.calls.load(Ordering::SeqCst),
+        expected_searches + 1,
+        "no continuation after EOF"
+    );
+    assert!(!sse.contains("DO_NOT_ECHO"), "error must not expose upstream content");
+    assert_eq!(sse.matches("event: error").count(), 1, "one error: {sse}");
+    assert!(sse.contains(expected_error), "expected error missing: {sse}");
+    assert!(!sse.contains("event: message_stop"), "no successful completion: {sse}");
+    assert!(
+        !sse.contains("event: message_delta"),
+        "no successful terminal metadata: {sse}"
+    );
+}
+
+#[tokio::test]
+async fn messages_stream_rejects_malformed_json_before_tools_or_completion() {
+    let invalid = "data: {\"private\":\"DO_NOT_ECHO\",\n\n";
+    let streams = cassette_turn_streams();
+    for (index, body) in streams.iter().enumerate() {
+        let damaged = body.replacen("event: message_stop", &format!("{invalid}event: message_stop"), 1);
+        assert_ne!(&damaged, body, "fixture contains a terminal");
+        let mut rounds = streams[..index].to_vec();
+        rounds.push(damaged);
+        assert_failed_stream(rounds, index, "invalid JSON in upstream Messages stream").await;
+    }
+}
+
+#[tokio::test]
+async fn messages_stream_ignores_well_formed_unknown_events() {
+    let streams = cassette_turn_streams()
+        .into_iter()
+        .map(|body| format!("data: {{\"type\":\"future_event\",\"extension\":1}}\n\n: heartbeat\n\n{body}"))
+        .collect();
+    assert_messages_stream_presents_one_message(streams).await;
+}