@@ -59,7 +59,7 @@ def __init__(self, runtime_id: str, storage: UiPathRuntimeStorageProtocol | None
5959 self .runtime_id = runtime_id
6060 self .storage = storage
6161 self .current_message : AIMessageChunk | AIMessage
62- self .tool_confirmation_schemas : dict [str , Any ] = {}
62+ self .tools_requiring_confirmation : dict [str , Any ] = {}
6363 self .seen_message_ids : set [str ] = set ()
6464 self ._storage_lock = asyncio .Lock ()
6565 self ._citation_stream_processor = CitationStreamProcessor ()
@@ -320,6 +320,7 @@ async def map_ai_message_chunk_to_events(
320320
321321 events : list [UiPathConversationMessageEvent ] = []
322322
323+ # For every new message_id, start a new message
323324 if message .id not in self .seen_message_ids :
324325 self .current_message = message
325326 self .seen_message_ids .add (message .id )
@@ -338,6 +339,7 @@ async def map_ai_message_chunk_to_events(
338339 self ._chunk_to_message_event (message .id , chunk )
339340 )
340341 case "tool_call_chunk" :
342+ # Accumulate the message chunk. Note that we assume no interweaving of AIMessage and AIMessageChunks for a given message.
341343 # Skip the first chunk — it's already assigned as current_message above,
342344 # so accumulating it with itself would duplicate fields via string concat
343345 # (e.g. tool name "search_web" becomes "search_websearch_web").
@@ -431,9 +433,9 @@ async def map_current_message_to_start_tool_call_events(self):
431433
432434 tool_name = tool_call ["name" ]
433435 require_confirmation = (
434- tool_name in self .tool_confirmation_schemas
436+ tool_name in self .tools_requiring_confirmation
435437 )
436- input_schema = self .tool_confirmation_schemas .get (tool_name )
438+ input_schema = self .tools_requiring_confirmation .get (tool_name )
437439 events .append (
438440 self .map_tool_call_to_tool_call_start_event (
439441 self .current_message .id ,
0 commit comments