fix(llm): accept complete streams without done

Author: claude

CHANGELOG.md

@@ -21,6 +21,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Preserved bounded fallback notices and failure metadata when web search runs
   inside a batch workflow.
 
+### Fixed
+
+- Accepted a complete schema-valid streamed object after a short terminal-event
+  grace period, so an OpenAI-compatible endpoint that omits or delays its final
+  `Done` event cannot turn an already generated result into a timeout.
+
 ## [6.3.1] - 2026-07-22
 
 ### Fixed

Cargo.lock

@@ -17,7 +17,7 @@ a3s-common = { version = "0.1.1", path = "../../common" }
 a3s-memory = { version = "0.1.2", path = "../../memory" }
 a3s-lane = { version = "0.5", path = "../../lane" }
 a3s-search = { version = "2.1.0", path = "../../search", default-features = false, features = ["lightpanda"] }
-a3s-flow = { version = "0.4.2", path = "../../flow" }
+a3s-flow = { version = "0.4.3", path = "../../flow" }
 
 # Async runtime
 tokio = { version = "1.35", features = [

core/Cargo.toml

@@ -303,18 +303,65 @@ pub async fn generate_streaming(
 
     let cancel_token = CancellationToken::new();
     let mut rx = client
-        .complete_streaming_structured(&messages, Some(&system), &tools, &directive, cancel_token)
+        .complete_streaming_structured(
+            &messages,
+            Some(&system),
+            &tools,
+            &directive,
+            cancel_token.clone(),
+        )
         .await
         .context("LLM streaming call failed during structured generation")?;
 
     let mut json_buffer = String::new();
     let mut last_valid_partial: Option<Value> = None;
     let mut final_response: Option<super::LlmResponse> = None;
     let mut last_parse_len: usize = 0;
+    let mut complete_candidate: Option<(Value, String, tokio::time::Instant)> = None;
     // Minimum bytes of new data before attempting a partial parse (reduces CPU)
     const PARSE_THRESHOLD: usize = 8;
-
-    while let Some(event) = rx.recv().await {
+    // Well-behaved providers send Done immediately after the complete object.
+    // A short grace preserves their final usage metadata while preventing an
+    // otherwise valid result from hanging on a compatible endpoint that never
+    // terminates its stream.
+    const DONE_GRACE: std::time::Duration = std::time::Duration::from_millis(250);
+    loop {
+        let event = if let Some((_, _, deadline)) = complete_candidate.as_ref() {
+            tokio::select! {
+                event = rx.recv() => event,
+                _ = tokio::time::sleep_until(*deadline) => {
+                    let candidate = complete_candidate
+                        .take()
+                        .expect("complete streamed candidate exists");
+                    let (value, raw_text, _) = candidate;
+                    cancel_token.cancel();
+                    on_partial(&value);
+                    return Ok(StructuredResult {
+                        object: value,
+                        raw_text: Some(raw_text),
+                        usage: TokenUsage::default(),
+                        repair_rounds: 0,
+                        mode_used: mode,
+                    });
+                }
+            }
+        } else {
+            rx.recv().await
+        };
+        let Some(event) = event else {
+            if let Some((value, raw_text, _)) = complete_candidate.take() {
+                cancel_token.cancel();
+                on_partial(&value);
+                return Ok(StructuredResult {
+                    object: value,
+                    raw_text: Some(raw_text),
+                    usage: TokenUsage::default(),
+                    repair_rounds: 0,
+                    mode_used: mode,
+                });
+            }
+            break;
+        };
         match event {
             StreamEvent::ToolUseInputDelta { delta, .. } if mode == StructuredMode::Tool => {
                 if final_response.is_some() {
@@ -334,6 +381,17 @@ pub async fn generate_streaming(
                     }
                     last_parse_len = json_buffer.len();
                 }
+                if complete_candidate.is_none() && (delta.contains('}') || delta.contains(']')) {
+                    complete_candidate = resolve_structured(
+                        std::slice::from_ref(&json_buffer),
+                        &req.schema,
+                        envelope,
+                    )
+                    .valid
+                    .map(|(value, raw_text)| {
+                        (value, raw_text, tokio::time::Instant::now() + DONE_GRACE)
+                    });
+                }
             }
             StreamEvent::TextDelta(delta) if mode != StructuredMode::Tool => {
                 if final_response.is_some() {
@@ -356,9 +414,21 @@ pub async fn generate_streaming(
                     }
                     last_parse_len = json_buffer.len();
                 }
+                if complete_candidate.is_none() && (delta.contains('}') || delta.contains(']')) {
+                    complete_candidate = resolve_structured(
+                        std::slice::from_ref(&json_buffer),
+                        &req.schema,
+                        envelope,
+                    )
+                    .valid
+                    .map(|(value, raw_text)| {
+                        (value, raw_text, tokio::time::Instant::now() + DONE_GRACE)
+                    });
+                }
             }
             StreamEvent::Done(resp) => {
                 final_response = Some(resp);
+                break;
             }
             _ => {}
         }

core/src/llm/structured.rs

@@ -155,6 +155,39 @@ impl LlmClient for MockStructuredClient {
     }
 }
 
+struct CompleteObjectWithoutDoneClient;
+
+#[async_trait]
+impl LlmClient for CompleteObjectWithoutDoneClient {
+    async fn complete(
+        &self,
+        _messages: &[Message],
+        _system: Option<&str>,
+        _tools: &[ToolDefinition],
+    ) -> anyhow::Result<LlmResponse> {
+        anyhow::bail!("blocking generation is not used in this test")
+    }
+
+    async fn complete_streaming(
+        &self,
+        _messages: &[Message],
+        _system: Option<&str>,
+        _tools: &[ToolDefinition],
+        cancel_token: CancellationToken,
+    ) -> anyhow::Result<mpsc::Receiver<StreamEvent>> {
+        let (tx, rx) = mpsc::channel(4);
+        tokio::spawn(async move {
+            tx.send(StreamEvent::TextDelta(
+                r#"{"color":"blue","count":3}"#.to_string(),
+            ))
+            .await
+            .ok();
+            cancel_token.cancelled().await;
+        });
+        Ok(rx)
+    }
+}
+
 // ========================================================================
 // extract_json_value tests
 // ========================================================================
@@ -1211,6 +1244,39 @@ async fn test_generate_streaming_text_mode() {
     let result = generate_streaming(&client, &req, callback).await.unwrap();
     assert_eq!(result.object["color"], "blue");
     assert_eq!(result.object["count"], 3);
+    assert_eq!(result.usage.total_tokens, 15);
+}
+
+#[tokio::test]
+async fn test_generate_streaming_accepts_a_complete_valid_object_without_done() {
+    let req = StructuredRequest {
+        prompt: "test".to_string(),
+        system: None,
+        schema: serde_json::json!({
+            "type": "object",
+            "additionalProperties": false,
+            "required": ["color", "count"],
+            "properties": {
+                "color": {"type": "string"},
+                "count": {"type": "integer"}
+            }
+        }),
+        schema_name: "result".to_string(),
+        schema_description: None,
+        mode: StructuredMode::Prompt,
+        max_repair_attempts: 0,
+    };
+
+    let result = tokio::time::timeout(
+        std::time::Duration::from_secs(1),
+        generate_streaming(&CompleteObjectWithoutDoneClient, &req, Box::new(|_| {})),
+    )
+    .await
+    .expect("a schema-valid streamed object must not wait forever for Done")
+    .expect("complete streamed object");
+
+    assert_eq!(result.object["color"], "blue");
+    assert_eq!(result.object["count"], 3);
 }
 
 #[tokio::test]