lambda_agent_api/agent_api.py

@@ -226,20 +226,19 @@ class AgentApi:
                         outgoing_msg = ServerMessage.validate_json(
                             msg.data)
 
-                        if isinstance(outgoing_msg, MsgEventTextChunk):
+                        if isinstance(outgoing_msg, (MsgEventTextChunk,
+                                                     MsgEventToolCallChunk,
+                                                     MsgEventToolResult,
+                                                     MsgEventCustomUpdate,
+                                                     MsgEventEnd)):
                             if self._current_queue:
                                 await self._current_queue.put(outgoing_msg)
-                            # Если очереди нет (клиент отменил запрос), но токены идут — шлем их в коллбек
                             elif self.callback:
                                 self.callback(outgoing_msg)
                             else:
                                 logger.warning(
                                     f"[{self.id}] AgentEvent without active request")
 
-                        elif isinstance(outgoing_msg, MsgEventEnd):
-                            if self._current_queue:
-                                await self._current_queue.put(outgoing_msg)
-
                         elif isinstance(outgoing_msg, MsgError):
                             if self.callback:
                                 self.callback(outgoing_msg)