From b840fff80c4e2f93d367fae6a6b2e25b61cecbee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9B=B7=E7=94=B5=E8=8A=BD=E8=A1=A3?= Date: Tue, 29 Sep 2026 13:32:03 -0400 Subject: [PATCH] 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. --- .../kigi-sampler/src/stream/messages.rs | 24 +++--- .../kigi-sampler/src/stream/messages_tests.rs | 85 +++++++++++++++++++ 2 files changed, 99 insertions(+), 10 deletions(-) diff --git a/crates/codegen/kigi-sampler/src/stream/messages.rs b/crates/codegen/kigi-sampler/src/stream/messages.rs index 729e053..11d8ab1 100644 --- a/crates/codegen/kigi-sampler/src/stream/messages.rs +++ b/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 = None; let mut final_stop_message: Option = 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 = Vec::new(); - let mut assistant_reasoning: Option = None; + let mut assistant_reasoning: Vec = 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 = 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(); diff --git a/crates/codegen/kigi-sampler/src/stream/messages_tests.rs b/crates/codegen/kigi-sampler/src/stream/messages_tests.rs index b6c8f3f..26aba42 100644 --- a/crates/codegen/kigi-sampler/src/stream/messages_tests.rs +++ b/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> { + 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<(String, Option)> { + 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::(); + (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 {