@@ -137,6 +137,7 @@ def __init__(self, runtime: dict[str, Any], tools) -> None:
137137 # Cache the serialized tool list — tools rarely change within a session
138138 self ._cached_tools_json : list [dict [str , Any ]] | None = None
139139 self ._tools_cache_key : int = 0 # hash of tool list for invalidation
140+ self ._thinking_blocks : list [dict [str , Any ]] = [] # Preserve thinking blocks for round-trip
140141
141142 def _get_serialized_tools (self ) -> list [dict [str , Any ]]:
142143 """Get serialized tool list with caching."""
@@ -156,6 +157,18 @@ def _get_serialized_tools(self) -> list[dict[str, Any]]:
156157
157158 def next (self , messages : list [dict [str , Any ]], on_stream_chunk : Callable [[str ], None ] | None = None , store : Store [AppState ] | None = None ) -> AgentStep :
158159 system_message , converted_messages = _to_anthropic_messages (messages )
160+
161+ # Inject stored thinking blocks into the last assistant message
162+ if self ._thinking_blocks :
163+ for i in range (len (converted_messages ) - 1 , - 1 , - 1 ):
164+ if converted_messages [i ].get ("role" ) == "assistant" :
165+ existing = converted_messages [i ].get ("content" , [])
166+ if isinstance (existing , list ):
167+ converted_messages [i ] = dict (converted_messages [i ])
168+ converted_messages [i ]["content" ] = list (self ._thinking_blocks ) + existing
169+ break
170+ self ._thinking_blocks = []
171+
159172 request_body = {
160173 "model" : self .runtime ["model" ],
161174 "system" : system_message ,
@@ -251,6 +264,8 @@ def next(self, messages: list[dict[str, Any]], on_stream_chunk: Callable[[str],
251264 text_parts .append (block ["text" ])
252265 elif block_type == "tool_use" and isinstance (block .get ("id" ), str ) and isinstance (block .get ("name" ), str ):
253266 tool_calls .append ({"id" : block ["id" ], "toolName" : block ["name" ], "input" : block .get ("input" )})
267+ elif block_type == "thinking" :
268+ self ._thinking_blocks .append (block ) # Preserve for round-trip
254269 else :
255270 ignored_block_types .append (str (block_type ))
256271
@@ -277,6 +292,7 @@ def next(self, messages: list[dict[str, Any]], on_stream_chunk: Callable[[str],
277292 block_types = []
278293 ignored_block_types = []
279294 active_tool_call = None
295+ active_thinking_block = None
280296 stop_reason = None
281297
282298 # Streaming cost tracking
@@ -315,6 +331,8 @@ def next(self, messages: list[dict[str, Any]], on_stream_chunk: Callable[[str],
315331 "name" : cb .get ("name" ),
316332 "input_json" : ""
317333 }
334+ elif c_type == "thinking" :
335+ active_thinking_block = {"type" : "thinking" , "thinking" : "" }
318336 elif etype == "content_block_delta" :
319337 delta = event .get ("delta" , {})
320338 d_type = delta .get ("type" )
@@ -325,6 +343,12 @@ def next(self, messages: list[dict[str, Any]], on_stream_chunk: Callable[[str],
325343 elif d_type == "input_json_delta" :
326344 if active_tool_call :
327345 active_tool_call ["input_json" ] += delta .get ("partial_json" , "" )
346+ elif d_type == "thinking_delta" :
347+ if active_thinking_block :
348+ active_thinking_block ["thinking" ] += delta .get ("thinking" , "" )
349+ elif d_type == "signature_delta" :
350+ if active_thinking_block :
351+ active_thinking_block ["signature" ] = active_thinking_block .get ("signature" , "" ) + delta .get ("signature" , "" )
328352 elif etype == "content_block_stop" :
329353 if active_tool_call :
330354 try :
@@ -337,6 +361,9 @@ def next(self, messages: list[dict[str, Any]], on_stream_chunk: Callable[[str],
337361 "input" : parsed_input
338362 })
339363 active_tool_call = None
364+ if active_thinking_block :
365+ self ._thinking_blocks .append (active_thinking_block )
366+ active_thinking_block = None
340367 elif etype == "message_delta" :
341368 delta = event .get ("delta" , {})
342369 if "stop_reason" in delta :
0 commit comments