README.md

@@ -31,14 +31,6 @@ async def main():
         async for chunk in response:
             if isinstance(chunk, MsgEventTextChunk):
                 print(chunk.text, end="", flush=True)
-            elif isinstance(chunk, MsgEventToolCallChunk):
-                print(f"Tool call started: {chunk.tool_name}")
-            elif isinstance(chunk, MsgEventToolResult):
-                print(f"Tool result: {chunk.result}")
-            elif isinstance(chunk, MsgEventCustomUpdate):
-                print(f"Progress update: {chunk.payload}")
-            elif isinstance(chunk, MsgEventEnd):
-                print(f"Generation ended, tokens used: {chunk.tokens_used}")
 
     finally:
         await api.close()
@@ -47,8 +39,6 @@ async def main():
 asyncio.run(main())
 ```
 
-> `AgentApi.send_message()` возвращает стриминг-итерируемый объект, который может выдавать не только текстовые чанки, но и события инструментов (`MsgEventToolCallChunk`, `MsgEventToolResult`, `MsgEventCustomUpdate`) и финальный `MsgEventEnd`.
-
 ## Предполагаемое использование
 
 ```python
@@ -115,97 +105,37 @@ async def on_telegram_message(from_user: int, text: str):
 }
 ```
 
-#### AGENT_EVENT_TEXT_CHUNK
+#### AGENT_EVENT
+
+Базовый класс для ивентов, которые стримит агент во время генерации ответа. Конкретный класс для ивента определяется по `subtype`.
+
+##### TEXT_CHUNK
 
 Чанк текста ответа агента.
 
 ```json
 {
-  "type": "AGENT_EVENT_TEXT_CHUNK",
+  "type": "AGENT_EVENT",
-  "text": "Фрагмент текста",
+  "subtype": "TEXT_CHUNK",
-  "source": "main"
+  "text": "Фрагмент текста"
 }
 ```
 
-| Поле   | Тип    | Описание                                      |
+##### END
-|--------|--------|-----------------------------------------------|
-| type   | string | Всегда `AGENT_EVENT_TEXT_CHUNK`              |
-| text   | string | Фрагмент текста ответа агента                |
-| source | string | Источник события (по умолчанию "main")       |
-
-#### AGENT_EVENT_TOOL_CALL_CHUNK
-
-Агент решил использовать инструмент и генерирует аргументы.
-
-```json
-{
-  "type": "AGENT_EVENT_TOOL_CALL_CHUNK",
-  "tool_name": "имя_инструмента",
-  "args_chunk": "{\"key\": \"value\"}",
-  "source": "main"
-}
-```
-
-| Поле        | Тип     | Описание                                      |
-|-------------|---------|-----------------------------------------------|
-| type        | string  | Всегда `AGENT_EVENT_TOOL_CALL_CHUNK`         |
-| tool_name   | string  | Имя инструмента (может быть null в первом чанке) |
-| args_chunk  | string  | Кусок JSON-аргументов (может быть null)      |
-| source      | string  | Источник события (по умолчанию "main")       |
-
-#### AGENT_EVENT_TOOL_RESULT
-
-Инструмент отработал и вернул результат.
-
-```json
-{
-  "type": "AGENT_EVENT_TOOL_RESULT",
-  "tool_name": "имя_инструмента",
-  "result": "результат выполнения",
-  "source": "main"
-}
-```
-
-| Поле       | Тип    | Описание                                      |
-|------------|--------|-----------------------------------------------|
-| type       | string | Всегда `AGENT_EVENT_TOOL_RESULT`             |
-| tool_name  | string | Имя инструмента                              |
-| result     | any    | Результат выполнения (строка, объект или массив) |
-| source     | string | Источник события (по умолчанию "main")       |
-
-#### AGENT_EVENT_CUSTOM_UPDATE
-
-Кастомный прогресс (например, скачивание файла) изнутри инструмента.
-
-```json
-{
-  "type": "AGENT_EVENT_CUSTOM_UPDATE",
-  "payload": {"status": "in_progress", "progress": 50},
-  "source": "main"
-}
-```
-
-| Поле     | Тип             | Описание                                      |
-|----------|-----------------|-----------------------------------------------|
-| type     | string          | Всегда `AGENT_EVENT_CUSTOM_UPDATE`           |
-| payload  | object          | Любые данные о прогрессе                      |
-| source   | string          | Источник события (по умолчанию "main")       |
-
-#### AGENT_EVENT_END
 
 Агент закончил генерацию ответа.
 
 ```json
 {
-  "type": "AGENT_EVENT_END",
+  "type": "AGENT_EVENT",
+  "subtype": "END",
   "tokens_used": 42
 }
 ```
 
-| Поле        | Тип    | Описание                                      |
+| Поле        | Тип    | Описание              |
-|-------------|--------|-----------------------------------------------|
+|-------------|--------|-----------------------|
-| type        | string | Всегда `AGENT_EVENT_END`                     |
+| tokens_used | int    | Количество использованных токенов |
-| tokens_used | int    | Количество использованных токенов            |
 
 #### ERROR
 
@@ -234,29 +164,4 @@ async def on_telegram_message(from_user: int, text: str):
 }
 ```
 
-Неопределенная ошибка в работе агента.
-
-```json
-{
-  "type": "ERROR",
-  "code": "error_code",
-  "details": "Описание ошибки"
-}
-```
-
-| Поле    | Тип   | Описание       |
-|---------|-------|----------------|
-| code    | string | Код ошибки    |
-| details | string | Подробности   |
-
-#### GRACEFUL_DISCONNECT
-
-Отправляется перед завершением работы контейнера с агентом. Например, при долгом бездействии. Нужно, чтобы отделять обрыв соединения из-за ошибки с необходимостью повторного подключения. Приход этого сообщения означает, что агент осознанно завершает работу с клиентом по какой-то причине. Для дальнейшего взаимодействия нужно снова обратиться к мастеру.
-
-```json
-{
-  "type": "GRACEFUL_DISCONNECT"
-}
-```
-
 ![Схема взаимодействия](docs/schema.png)

lambda_agent_api/agent_api.py

@@ -226,19 +226,20 @@ 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)

tests/test_models.py

@@ -33,9 +33,6 @@ def test_client_message_invalid(data):
     [
         ({"type": "STATUS"}, MsgStatus),
         ({"type": "AGENT_EVENT_TEXT_CHUNK", "text": "hi"}, MsgEventTextChunk),
-        ({"type": "AGENT_EVENT_TOOL_CALL_CHUNK", "tool_name": "search", "args_chunk": "{\"q\": \"hello\"}"}, MsgEventToolCallChunk),
-        ({"type": "AGENT_EVENT_TOOL_RESULT", "tool_name": "search", "result": {"items": [1, 2, 3]}}, MsgEventToolResult),
-        ({"type": "AGENT_EVENT_CUSTOM_UPDATE", "payload": {"status": "in_progress", "progress": 50}}, MsgEventCustomUpdate),
         ({"type": "AGENT_EVENT_END", "tokens_used": 10}, MsgEventEnd),
         ({"type": "ERROR", "code": "E1", "details": "fail"}, MsgError),
         ({"type": "GRACEFUL_DISCONNECT"}, MsgGracefulDisconnect),
@@ -50,8 +47,6 @@ def test_server_message_valid(data, expected_type):
     "data",
     [
         {"type": "AGENT_EVENT_TEXT_CHUNK"},  # нет text
-        {"type": "AGENT_EVENT_TOOL_RESULT", "tool_name": "search"},  # нет result
-        {"type": "AGENT_EVENT_CUSTOM_UPDATE"},  # нет payload
         {"type": "AGENT_EVENT_END"},  # нет tokens_used
         {"type": "ERROR", "code": "E1"},  # нет details
         {"type": "UNKNOWN"},