feat(04-01): finalize AgentApi migration

2026-04-17 16:31:48 +03:00 · 2026-04-17 16:31:48 +03:00 · 430c82dba1
commit 430c82dba1
parent cd59d89617
9 changed files with 225 additions and 350 deletions
--- a/sdk/agent_api_wrapper.py
+++ b/sdk/agent_api_wrapper.py
@ -0,0 +1,88 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import sys
+from pathlib import Path
+
+import aiohttp
+
+_api_root = Path(__file__).resolve().parents[1] / "external" / "platform-agent_api"
+if str(_api_root) not in sys.path:
+    sys.path.insert(0, str(_api_root))
+
+from lambda_agent_api.agent_api import AgentApi, AgentException
+from lambda_agent_api.server import (
+    MsgError,
+    MsgEventEnd,
+    MsgEventTextChunk,
+    MsgGracefulDisconnect,
+    ServerMessage,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class AgentApiWrapper(AgentApi):
+    """Capture tokens_used from MsgEventEnd without patching upstream code."""
+
+    def __init__(self, agent_id: str, url: str, **kwargs) -> None:
+        super().__init__(agent_id=agent_id, url=url, **kwargs)
+        self.last_tokens_used = 0
+
+    async def _listen(self):
+        try:
+            async for msg in self._ws:
+                if msg.type == aiohttp.WSMsgType.TEXT:
+                    try:
+                        outgoing_msg = ServerMessage.validate_json(msg.data)
+
+                        if isinstance(outgoing_msg, MsgEventTextChunk):
+                            if self._current_queue:
+                                await self._current_queue.put(outgoing_msg)
+                            elif self.callback:
+                                self.callback(outgoing_msg)
+                            else:
+                                logger.warning("[%s] AgentEvent without active request", self.id)
+
+                        elif isinstance(outgoing_msg, MsgEventEnd):
+                            self.last_tokens_used = outgoing_msg.tokens_used
+                            if self._current_queue:
+                                await self._current_queue.put(outgoing_msg)
+
+                        elif isinstance(outgoing_msg, MsgError):
+                            if self.callback:
+                                self.callback(outgoing_msg)
+                            error = AgentException(outgoing_msg.code, outgoing_msg.details)
+                            logger.error("[%s] Agent error: %s", self.id, error)
+                            if self._current_queue:
+                                await self._current_queue.put(error)
+
+                        elif isinstance(outgoing_msg, MsgGracefulDisconnect):
+                            if self.callback:
+                                self.callback(outgoing_msg)
+                            logger.info("[%s] Gracefully disconnecting", self.id)
+                            break
+
+                        else:
+                            logger.warning("[%s] Unknown message type: %s", self.id, outgoing_msg.type)
+                            if self.callback:
+                                self.callback(outgoing_msg)
+
+                    except Exception as exc:
+                        logger.error("[%s] Failed to deserialize message: %s", self.id, exc)
+                        if self._current_queue:
+                            await self._current_queue.put(
+                                AgentException("PARSE_ERROR", f"Validation failed: {exc}")
+                            )
+
+                elif msg.type in (aiohttp.WSMsgType.ERROR, aiohttp.WSMsgType.CLOSED):
+                    logger.error("[%s] WebSocket closed/error: %s", self.id, msg.type)
+                    break
+
+        except asyncio.CancelledError:
+            pass
+        except Exception as exc:
+            logger.error("[%s] Error in listen loop: %s", self.id, exc)
+        finally:
+            await self._cleanup()