@@ -4638,13 +4638,21 @@ async fn stream_with_provider(
46384638 let mut stream = client. stream_message ( message_request) . await ?;
46394639 let mut events = Vec :: new ( ) ;
46404640 let mut pending_tools: BTreeMap < u32 , ( String , String , String ) > = BTreeMap :: new ( ) ;
4641+ let mut pending_thinking: BTreeMap < u32 , ( String , Option < String > ) > = BTreeMap :: new ( ) ;
46414642 let mut saw_stop = false ;
46424643
46434644 while let Some ( event) = stream. next_event ( ) . await ? {
46444645 match event {
46454646 ApiStreamEvent :: MessageStart ( start) => {
46464647 for block in start. message . content {
4647- push_output_block ( block, 0 , & mut events, & mut pending_tools, true ) ;
4648+ push_output_block (
4649+ block,
4650+ 0 ,
4651+ & mut events,
4652+ & mut pending_tools,
4653+ & mut pending_thinking,
4654+ true ,
4655+ ) ;
46484656 }
46494657 }
46504658 ApiStreamEvent :: ContentBlockStart ( start) => {
@@ -4653,6 +4661,7 @@ async fn stream_with_provider(
46534661 start. index ,
46544662 & mut events,
46554663 & mut pending_tools,
4664+ & mut pending_thinking,
46564665 true ,
46574666 ) ;
46584667 }
@@ -4667,10 +4676,23 @@ async fn stream_with_provider(
46674676 input. push_str ( & partial_json) ;
46684677 }
46694678 }
4670- ContentBlockDelta :: ThinkingDelta { .. }
4671- | ContentBlockDelta :: SignatureDelta { .. } => { }
4679+ ContentBlockDelta :: ThinkingDelta { thinking } => {
4680+ if let Some ( ( pending, _) ) = pending_thinking. get_mut ( & delta. index ) {
4681+ pending. push_str ( & thinking) ;
4682+ }
4683+ }
4684+ ContentBlockDelta :: SignatureDelta { signature } => {
4685+ if let Some ( ( _, pending_signature) ) = pending_thinking. get_mut ( & delta. index ) {
4686+ pending_signature
4687+ . get_or_insert_with ( String :: new)
4688+ . push_str ( & signature) ;
4689+ }
4690+ }
46724691 } ,
46734692 ApiStreamEvent :: ContentBlockStop ( stop) => {
4693+ if let Some ( ( thinking, signature) ) = pending_thinking. remove ( & stop. index ) {
4694+ events. push ( AssistantEvent :: Thinking { thinking, signature } ) ;
4695+ }
46744696 if let Some ( ( id, name, input) ) = pending_tools. remove ( & stop. index ) {
46754697 events. push ( AssistantEvent :: ToolUse { id, name, input } ) ;
46764698 }
@@ -4767,8 +4789,12 @@ fn convert_messages(messages: &[ConversationMessage]) -> Vec<InputMessage> {
47674789 . iter ( )
47684790 . map ( |block| match block {
47694791 ContentBlock :: Text { text } => InputContentBlock :: Text { text : text. clone ( ) } ,
4770- ContentBlock :: Thinking { .. } => InputContentBlock :: Text {
4771- text : String :: new ( ) ,
4792+ ContentBlock :: Thinking {
4793+ thinking,
4794+ signature,
4795+ } => InputContentBlock :: Thinking {
4796+ thinking : thinking. clone ( ) ,
4797+ signature : signature. clone ( ) ,
47724798 } ,
47734799 ContentBlock :: ToolUse { id, name, input } => InputContentBlock :: ToolUse {
47744800 id : id. clone ( ) ,
@@ -4806,6 +4832,7 @@ fn push_output_block(
48064832 block_index : u32 ,
48074833 events : & mut Vec < AssistantEvent > ,
48084834 pending_tools : & mut BTreeMap < u32 , ( String , String , String ) > ,
4835+ pending_thinking : & mut BTreeMap < u32 , ( String , Option < String > ) > ,
48094836 streaming_tool_input : bool ,
48104837) {
48114838 match block {
@@ -4825,17 +4852,35 @@ fn push_output_block(
48254852 } ;
48264853 pending_tools. insert ( block_index, ( id, name, initial_input) ) ;
48274854 }
4828- OutputContentBlock :: Thinking { .. } | OutputContentBlock :: RedactedThinking { .. } => { }
4855+ OutputContentBlock :: Thinking {
4856+ thinking,
4857+ signature,
4858+ } => {
4859+ if streaming_tool_input {
4860+ pending_thinking. insert ( block_index, ( thinking, signature) ) ;
4861+ } else {
4862+ events. push ( AssistantEvent :: Thinking { thinking, signature } ) ;
4863+ }
4864+ }
4865+ OutputContentBlock :: RedactedThinking { .. } => { }
48294866 }
48304867}
48314868
48324869fn response_to_events ( response : MessageResponse ) -> Vec < AssistantEvent > {
48334870 let mut events = Vec :: new ( ) ;
48344871 let mut pending_tools = BTreeMap :: new ( ) ;
4872+ let mut pending_thinking = BTreeMap :: new ( ) ;
48354873
48364874 for ( index, block) in response. content . into_iter ( ) . enumerate ( ) {
48374875 let index = u32:: try_from ( index) . expect ( "response block index overflow" ) ;
4838- push_output_block ( block, index, & mut events, & mut pending_tools, false ) ;
4876+ push_output_block (
4877+ block,
4878+ index,
4879+ & mut events,
4880+ & mut pending_tools,
4881+ & mut pending_thinking,
4882+ false ,
4883+ ) ;
48394884 if let Some ( ( id, name, input) ) = pending_tools. remove ( & index) {
48404885 events. push ( AssistantEvent :: ToolUse { id, name, input } ) ;
48414886 }
@@ -7259,6 +7304,7 @@ mod tests {
72597304 fn pending_tools_preserve_multiple_streaming_tool_calls_by_index ( ) {
72607305 let mut events = Vec :: new ( ) ;
72617306 let mut pending_tools = BTreeMap :: new ( ) ;
7307+ let mut pending_thinking = BTreeMap :: new ( ) ;
72627308
72637309 push_output_block (
72647310 OutputContentBlock :: ToolUse {
@@ -7269,6 +7315,7 @@ mod tests {
72697315 1 ,
72707316 & mut events,
72717317 & mut pending_tools,
7318+ & mut pending_thinking,
72727319 true ,
72737320 ) ;
72747321 push_output_block (
@@ -7280,6 +7327,7 @@ mod tests {
72807327 2 ,
72817328 & mut events,
72827329 & mut pending_tools,
7330+ & mut pending_thinking,
72837331 true ,
72847332 ) ;
72857333
0 commit comments