fix(sampler): Messages streams keep every thinking block and join split signatures

A turn with interleaved thinking lost all but its last signed block on
replay, and a signature split across deltas kept only its last part.

crates/codegen/kigi-sampler/src/stream/messages.rs

@@ -46,6 +46,8 @@ struct BlockState {
     args_acc: String,
     thinking_acc: String,
     signature: String,
+    /// The first `SignatureDelta` replaces a start-seeded signature; later ones append.
+    signature_delta_seen: bool,
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -109,13 +111,10 @@ pub fn stream_messages<'a>(
         let mut final_stop_reason: Option<StopReason> = None;
         let mut final_stop_message: Option<String> = None;
 
-        // Assistant-response accumulators (built up as ContentBlockStop
-        // events fire). Reasoning is collected into a synthesized
-        // `rs::ReasoningItem` and emitted as a sibling
-        // `ConversationItem::Reasoning` before the trailing Assistant.
+        // Assistant-response accumulators, filled as ContentBlockStop events fire.
         let mut assistant_text = String::new();
         let mut assistant_tool_calls: Vec<ToolCall> = Vec::new();
-        let mut assistant_reasoning: Option<rs::ReasoningItem> = None;
+        let mut assistant_reasoning: Vec<rs::ReasoningItem> = Vec::new();
 
         // Index counters
         let mut chunk_index: u64 = 0;
@@ -185,6 +184,7 @@ pub fn stream_messages<'a>(
                                 args_acc: String::new(),
                                 thinking_acc: thinking.clone(),
                                 signature: signature.clone(),
+                                signature_delta_seen: false,
                             },
                         );
                         if !first_token_emitted {
@@ -205,6 +205,7 @@ pub fn stream_messages<'a>(
                                 args_acc: String::new(),
                                 thinking_acc: String::new(),
                                 signature: String::new(),
+                                signature_delta_seen: false,
                             },
                         );
                         if !first_token_emitted {
@@ -237,6 +238,7 @@ pub fn stream_messages<'a>(
                                 args_acc: String::new(),
                                 thinking_acc: String::new(),
                                 signature: String::new(),
+                                signature_delta_seen: false,
                             },
                         );
 
@@ -276,7 +278,11 @@ pub fn stream_messages<'a>(
                                 }
                             }
                             StreamDelta::SignatureDelta { signature } => {
-                                state.signature = signature;
+                                if !state.signature_delta_seen {
+                                    state.signature_delta_seen = true;
+                                    state.signature.clear();
+                                }
+                                state.signature.push_str(&signature);
                             }
                             StreamDelta::TextDelta { text } => {
                                 if !text.is_empty() {
@@ -347,7 +353,7 @@ pub fn stream_messages<'a>(
                                     } else {
                                         Some(state.signature)
                                     };
-                                    assistant_reasoning = Some(rs::ReasoningItem {
+                                    assistant_reasoning.push(rs::ReasoningItem {
                                         id: String::new(),
                                         summary,
                                         content: None,
@@ -513,9 +519,7 @@ pub fn stream_messages<'a>(
         });
 
         let mut items: Vec<ConversationItem> = Vec::new();
-        if let Some(r) = assistant_reasoning {
-            items.push(ConversationItem::Reasoning(r));
-        }
+        items.extend(assistant_reasoning.into_iter().map(ConversationItem::Reasoning));
         items.push(assistant_item);
 
         let stream_end = Instant::now();

crates/codegen/kigi-sampler/src/stream/messages_tests.rs

@@ -208,6 +208,91 @@ async fn thinking_block_emits_reasoning_channel_and_preserved_in_response() {
     }
 }
 
+fn thinking_block(
+    index: u32,
+    start_signature: &str,
+    thinking: &str,
+    signature_deltas: &[&str],
+) -> Vec<Result<MessageStreamEvent, SamplingError>> {
+    let mut events = vec![
+        Ok(MessageStreamEvent::ContentBlockStart {
+            index,
+            content_block: ContentBlock::Thinking {
+                thinking: String::new(),
+                signature: start_signature.into(),
+            },
+        }),
+        Ok(MessageStreamEvent::ContentBlockDelta {
+            index,
+            delta: StreamDelta::ThinkingDelta {
+                thinking: thinking.into(),
+            },
+        }),
+    ];
+    for signature in signature_deltas {
+        events.push(Ok(MessageStreamEvent::ContentBlockDelta {
+            index,
+            delta: StreamDelta::SignatureDelta {
+                signature: (*signature).into(),
+            },
+        }));
+    }
+    events.push(Ok(block_stop(index)));
+    events
+}
+
+async fn reasoning_signatures(
+    blocks: Vec<Vec<Result<MessageStreamEvent, SamplingError>>>,
+) -> Vec<(String, Option<String>)> {
+    let mut events = vec![Ok(message_start())];
+    events.extend(blocks.into_iter().flatten());
+    events.push(Ok(MessageStreamEvent::MessageStop));
+    let raw = stream::iter(events).boxed();
+    let evs = collect(stream_messages(raw, None, rid(), Duration::from_secs(60))).await;
+    match evs.last().unwrap() {
+        SamplingEvent::Completed { response, .. } => response
+            .reasoning_items()
+            .map(|r| {
+                let text = r
+                    .summary
+                    .iter()
+                    .map(|rs::SummaryPart::SummaryText(t)| t.text.as_str())
+                    .collect::<String>();
+                (text, r.encrypted_content.clone())
+            })
+            .collect(),
+        other => panic!("expected Completed, got {other:?}"),
+    }
+}
+
+#[tokio::test]
+async fn every_thinking_block_is_kept_in_order() {
+    let items = reasoning_signatures(vec![
+        thinking_block(0, "", "first", &["sig-a"]),
+        thinking_block(1, "", "second", &["sig-b"]),
+    ])
+    .await;
+    assert_eq!(
+        items,
+        vec![
+            ("first".to_string(), Some("sig-a".to_string())),
+            ("second".to_string(), Some("sig-b".to_string())),
+        ]
+    );
+}
+
+#[tokio::test]
+async fn split_signature_deltas_join_into_one_signature() {
+    let items = reasoning_signatures(vec![thinking_block(0, "", "t", &["abc", "def"])]).await;
+    assert_eq!(items, vec![("t".to_string(), Some("abcdef".to_string()))]);
+}
+
+#[tokio::test]
+async fn first_signature_delta_replaces_a_start_seeded_signature() {
+    let items = reasoning_signatures(vec![thinking_block(0, "sig", "t", &["sig"])]).await;
+    assert_eq!(items, vec![("t".to_string(), Some("sig".to_string()))]);
+}
+
 #[tokio::test]
 async fn tool_use_block_assembles_into_tool_call() {
     let tool_start = MessageStreamEvent::ContentBlockStart {