Surfaces_team/surfaces
4
0
Fork 0
Code Issues 5 Pull requests Projects Releases Packages Wiki Activity Actions

feat: finalize matrix platform audit and docs

Browse source
This commit is contained in:
Mikhail Putilovskij 2026-04-21 15:35:03 +03:00
parent 6422c7db58
commit 4524a6abc8
30 changed files with 3093 additions and 176 deletions
Download patch file Download diff file Expand all files Collapse all files

153
sdk/agent_api_wrapper.py
View file

@ -3,10 +3,12 @@ from __future__ import annotations
import asyncio
import inspect
import logging
import sys
import os
import re
from urllib.parse import urlsplit, urlunsplit
import sys
from collections.abc import AsyncIterator
from pathlib import Path
from urllib.parse import urlsplit, urlunsplit
import aiohttp
@ -14,16 +16,18 @@ _api_root = Path(__file__).resolve().parents[1] / "external" / "platform-agent_a
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,
)
from lambda_agent_api.agent_api import AgentApi, AgentBusyException, AgentException # noqa: E402
from lambda_agent_api.client import EClientMessage, MsgUserMessage # noqa: E402
from lambda_agent_api.server import AgentEventUnion, MsgEventEnd, ServerMessage # noqa: E402
logger = logging.getLogger(__name__)
_DEBUG_STREAM = os.environ.get("SURFACES_AGENT_DEBUG_STREAM", "").strip().lower() in {
"1",
"true",
"yes",
}
_POST_END_DRAIN_MS = int(os.environ.get("SURFACES_AGENT_POST_END_DRAIN_MS", "120"))
_STREAM_IDLE_TIMEOUT_MS = int(os.environ.get("SURFACES_AGENT_IDLE_TIMEOUT_MS", "60000"))
class AgentApiWrapper(AgentApi):
@ -78,7 +82,7 @@ class AgentApiWrapper(AgentApi):
def _build_ws_url(base_url: str, chat_id: int | str) -> str:
return base_url.rstrip("/") + f"/agent_ws/?thread_id={chat_id}"
def for_chat(self, chat_id: int | str) -> "AgentApiWrapper":
def for_chat(self, chat_id: int | str) -> AgentApiWrapper:
return type(self)(
agent_id=self.id,
base_url=self._base_url,
@ -133,7 +137,7 @@ class AgentApiWrapper(AgentApi):
if self.callback:
self.callback(event)
if self._current_queue and hasattr(event, "code") and hasattr(event, "details"):
await self._current_queue.put(AgentException(getattr(event, "code"), getattr(event, "details")))
await self._current_queue.put(AgentException(event.code, event.details))
async def _listen(self):
try:
@ -143,6 +147,13 @@ class AgentApiWrapper(AgentApi):
outgoing_msg = ServerMessage.validate_json(msg.data)
if self._is_text_event(outgoing_msg):
if _DEBUG_STREAM:
logger.warning(
"[%s] text chunk queue=%s text=%r",
self.id,
self._current_queue is not None,
getattr(outgoing_msg, "text", "")[:80],
)
if self._current_queue:
await self._current_queue.put(outgoing_msg)
elif self.callback:
@ -152,6 +163,13 @@ class AgentApiWrapper(AgentApi):
elif self._is_end_event(outgoing_msg):
self.last_tokens_used = outgoing_msg.tokens_used
if _DEBUG_STREAM:
logger.warning(
"[%s] end event queue=%s tokens=%s",
self.id,
self._current_queue is not None,
getattr(outgoing_msg, "tokens_used", None),
)
await self._publish_event(outgoing_msg)
elif self._is_kind(outgoing_msg, "ERROR"):
@ -184,3 +202,114 @@ class AgentApiWrapper(AgentApi):
logger.error("[%s] Error in listen loop: %s", self.id, exc)
finally:
await self._cleanup()
async def send_message(
self, text: str, attachments: list[str] | None = None
) -> AsyncIterator[AgentEventUnion]:
if not self._connected or not self._ws:
raise AgentException(
code="NOT_CONNECTED", details="Not connected. Call connect() first."
)
if self._request_lock.locked():
raise AgentBusyException("Agent is currently processing another request")
await self._request_lock.acquire()
try:
self._current_queue = asyncio.Queue()
message = MsgUserMessage(
type=EClientMessage.USER_MESSAGE,
text=text,
attachments=attachments or [],
)
await self._ws.send_str(message.model_dump_json())
logger.debug("[%s] Sent message: %s...", self.id, text[:50])
while True:
try:
chunk = await asyncio.wait_for(
self._current_queue.get(),
timeout=max(_STREAM_IDLE_TIMEOUT_MS, 0) / 1000,
)
except TimeoutError as exc:
raise AgentException(
"TIMEOUT",
(
"Timed out waiting for the next agent stream event "
f"after {max(_STREAM_IDLE_TIMEOUT_MS, 0)}ms"
),
) from exc
if isinstance(chunk, Exception):
raise chunk
if isinstance(chunk, MsgEventEnd):
self.last_tokens_used = chunk.tokens_used
async for late_chunk in self._drain_post_end_events():
yield late_chunk
break
yield chunk
finally:
if self._current_queue:
orphan_queue = self._current_queue
self._current_queue = None
while not orphan_queue.empty():
try:
orphan_msg = orphan_queue.get_nowait()
if isinstance(orphan_msg, Exception):
logger.debug(
"[%s] Dropped exception from queue during cleanup: %s",
self.id,
orphan_msg,
)
continue
if self.callback:
self.callback(orphan_msg)
else:
logger.debug("[%s] Dropped orphaned message during cleanup", self.id)
except asyncio.QueueEmpty:
break
if self._request_lock.locked():
self._request_lock.release()
async def _drain_post_end_events(self) -> AsyncIterator[AgentEventUnion]:
if self._current_queue is None:
return
timeout_s = max(_POST_END_DRAIN_MS, 0) / 1000
while True:
try:
chunk = await asyncio.wait_for(self._current_queue.get(), timeout=timeout_s)
except TimeoutError:
break
if isinstance(chunk, Exception):
logger.warning("[%s] dropping post-END exception: %s", self.id, chunk)
continue
if isinstance(chunk, MsgEventEnd):
self.last_tokens_used = chunk.tokens_used
if _DEBUG_STREAM:
logger.warning(
"[%s] dropped duplicate END tokens=%s",
self.id,
chunk.tokens_used,
)
continue
if _DEBUG_STREAM and self._is_text_event(chunk):
logger.warning(
"[%s] recovered post-END text chunk=%r",
self.id,
getattr(chunk, "text", "")[:80],
)
yield chunk