feat: finalize matrix platform audit and docs
This commit is contained in:
parent
6422c7db58
commit
4524a6abc8
30 changed files with 3093 additions and 176 deletions
|
|
@ -1,11 +1,12 @@
|
|||
import asyncio
|
||||
|
||||
import pytest
|
||||
from lambda_agent_api.server import MsgEventEnd, MsgEventTextChunk
|
||||
|
||||
from core.protocol import SettingsAction
|
||||
import sdk.agent_api_wrapper as agent_api_wrapper_module
|
||||
from core.protocol import SettingsAction
|
||||
from sdk.agent_api_wrapper import AgentApiWrapper
|
||||
from sdk.interface import Attachment, MessageChunk, MessageResponse, UserSettings
|
||||
from sdk.interface import Attachment, MessageChunk, MessageResponse, PlatformError, UserSettings
|
||||
from sdk.prototype_state import PrototypeStateStore
|
||||
from sdk.real import RealPlatformClient
|
||||
|
||||
|
|
@ -110,6 +111,23 @@ class AttachmentTrackingChatAgentApi:
|
|||
self.last_tokens_used = 5
|
||||
|
||||
|
||||
class FlakyChatAgentApi:
|
||||
def __init__(self, chat_id: str) -> None:
|
||||
self.chat_id = chat_id
|
||||
self.connect_calls = 0
|
||||
self.close_calls = 0
|
||||
|
||||
async def connect(self) -> None:
|
||||
self.connect_calls += 1
|
||||
|
||||
async def close(self) -> None:
|
||||
self.close_calls += 1
|
||||
|
||||
async def send_message(self, text: str, attachments: list[str] | None = None):
|
||||
raise ConnectionError("Connection closed")
|
||||
yield
|
||||
|
||||
|
||||
class SendFileEvent:
|
||||
def __init__(self, *, workspace_path: str, mime_type: str, filename: str, size: int) -> None:
|
||||
self.type = "AGENT_EVENT_SEND_FILE"
|
||||
|
|
@ -180,6 +198,26 @@ class FakeWebSocket:
|
|||
return self._messages.pop(0)
|
||||
|
||||
|
||||
class QueueFeedingWebSocket:
|
||||
def __init__(self, owner, queued_events: list[object]) -> None:
|
||||
self.owner = owner
|
||||
self.queued_events = list(queued_events)
|
||||
self.sent_payloads: list[str] = []
|
||||
|
||||
async def send_str(self, payload: str) -> None:
|
||||
self.sent_payloads.append(payload)
|
||||
for event in self.queued_events:
|
||||
await self.owner._current_queue.put(event)
|
||||
|
||||
|
||||
class SilentWebSocket:
|
||||
def __init__(self) -> None:
|
||||
self.sent_payloads: list[str] = []
|
||||
|
||||
async def send_str(self, payload: str) -> None:
|
||||
self.sent_payloads.append(payload)
|
||||
|
||||
|
||||
class MessageResponseWithAttachments(MessageResponse):
|
||||
attachments: list[Attachment] = []
|
||||
|
||||
|
|
@ -271,6 +309,68 @@ def test_agent_api_wrapper_falls_back_to_legacy_url_constructor(monkeypatch):
|
|||
assert wrapper.last_tokens_used == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_api_wrapper_recovers_late_text_after_first_end(monkeypatch):
|
||||
def fake_init(self, agent_id, base_url=None, chat_id=0, **kwargs):
|
||||
self.id = agent_id
|
||||
self.url = base_url
|
||||
self.callback = kwargs.get("callback")
|
||||
self.on_disconnect = kwargs.get("on_disconnect")
|
||||
|
||||
monkeypatch.setattr(agent_api_wrapper_module.AgentApi, "__init__", fake_init)
|
||||
|
||||
wrapper = AgentApiWrapper(
|
||||
agent_id="agent-1",
|
||||
base_url="https://agent.example.com/v1/agent_ws",
|
||||
chat_id="chat-1",
|
||||
)
|
||||
wrapper._connected = True
|
||||
wrapper._request_lock = asyncio.Lock()
|
||||
wrapper._current_queue = None
|
||||
wrapper._ws = QueueFeedingWebSocket(
|
||||
wrapper,
|
||||
[
|
||||
MsgEventTextChunk(text="Иллюстра"),
|
||||
MsgEventEnd(tokens_used=5),
|
||||
MsgEventTextChunk(text="ция"),
|
||||
MsgEventEnd(tokens_used=5),
|
||||
],
|
||||
)
|
||||
|
||||
chunks = []
|
||||
async for chunk in wrapper.send_message("hello"):
|
||||
chunks.append(chunk)
|
||||
|
||||
assert [chunk.text for chunk in chunks] == ["Иллюстра", "ция"]
|
||||
assert wrapper.last_tokens_used == 5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_agent_api_wrapper_times_out_on_idle_stream(monkeypatch):
|
||||
def fake_init(self, agent_id, base_url=None, chat_id=0, **kwargs):
|
||||
self.id = agent_id
|
||||
self.url = base_url
|
||||
self.callback = kwargs.get("callback")
|
||||
self.on_disconnect = kwargs.get("on_disconnect")
|
||||
|
||||
monkeypatch.setattr(agent_api_wrapper_module.AgentApi, "__init__", fake_init)
|
||||
monkeypatch.setattr(agent_api_wrapper_module, "_STREAM_IDLE_TIMEOUT_MS", 10)
|
||||
|
||||
wrapper = AgentApiWrapper(
|
||||
agent_id="agent-1",
|
||||
base_url="https://agent.example.com/v1/agent_ws",
|
||||
chat_id="chat-1",
|
||||
)
|
||||
wrapper._connected = True
|
||||
wrapper._request_lock = asyncio.Lock()
|
||||
wrapper._current_queue = None
|
||||
wrapper._ws = SilentWebSocket()
|
||||
|
||||
with pytest.raises(agent_api_wrapper_module.AgentException, match="Timed out waiting"):
|
||||
async for _ in wrapper.send_message("hello"):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_platform_client_get_or_create_user_uses_local_state():
|
||||
client = RealPlatformClient(
|
||||
|
|
@ -418,6 +518,58 @@ async def test_real_platform_client_reuses_cached_chat_client():
|
|||
assert agent_api.created_chat_ids == ["chat-1"]
|
||||
assert agent_api.instances["chat-1"].calls == ["hello", "again"]
|
||||
assert agent_api.instances["chat-1"].connect_calls == 1
|
||||
assert agent_api.instances["chat-1"].close_calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_platform_client_wraps_connection_closed_as_platform_error():
|
||||
agent_api = FakeAgentApiFactory()
|
||||
agent_api.instances["chat-1"] = FlakyChatAgentApi("chat-1")
|
||||
agent_api.for_chat = lambda chat_id: agent_api.instances.setdefault(
|
||||
chat_id, FlakyChatAgentApi(chat_id)
|
||||
)
|
||||
client = RealPlatformClient(
|
||||
agent_api=agent_api,
|
||||
prototype_state=PrototypeStateStore(),
|
||||
platform="matrix",
|
||||
)
|
||||
|
||||
with pytest.raises(PlatformError, match="Connection closed") as exc_info:
|
||||
await client.send_message("@alice:example.org", "chat-1", "hello")
|
||||
|
||||
assert exc_info.value.code == "PLATFORM_CONNECTION_ERROR"
|
||||
assert "chat-1" not in client._chat_apis
|
||||
assert agent_api.instances["chat-1"].close_calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_platform_client_reconnects_after_closed_chat_api():
|
||||
agent_api = FakeAgentApiFactory()
|
||||
flaky = FlakyChatAgentApi("chat-1")
|
||||
healthy = AttachmentTrackingChatAgentApi("chat-1")
|
||||
provided = iter([flaky, healthy])
|
||||
|
||||
def for_chat(chat_id: str):
|
||||
chat_api = next(provided)
|
||||
agent_api.created_chat_ids.append(chat_id)
|
||||
agent_api.instances[chat_id] = chat_api
|
||||
return chat_api
|
||||
|
||||
agent_api.for_chat = for_chat
|
||||
client = RealPlatformClient(
|
||||
agent_api=agent_api,
|
||||
prototype_state=PrototypeStateStore(),
|
||||
platform="matrix",
|
||||
)
|
||||
|
||||
with pytest.raises(PlatformError, match="Connection closed"):
|
||||
await client.send_message("@alice:example.org", "chat-1", "hello")
|
||||
|
||||
result = await client.send_message("@alice:example.org", "chat-1", "again")
|
||||
|
||||
assert result.response == "again"
|
||||
assert agent_api.created_chat_ids == ["chat-1", "chat-1"]
|
||||
assert healthy.calls == [("again", None)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -462,7 +614,9 @@ async def test_real_platform_client_creates_distinct_clients_per_chat():
|
|||
async def test_real_platform_client_serializes_same_chat_streams_across_send_paths():
|
||||
agent_api = FakeAgentApiFactory()
|
||||
agent_api.instances["chat-1"] = BlockingChatAgentApi("chat-1")
|
||||
agent_api.for_chat = lambda chat_id: agent_api.instances.setdefault(chat_id, BlockingChatAgentApi(chat_id))
|
||||
agent_api.for_chat = lambda chat_id: agent_api.instances.setdefault(
|
||||
chat_id, BlockingChatAgentApi(chat_id)
|
||||
)
|
||||
client = RealPlatformClient(
|
||||
agent_api=agent_api,
|
||||
prototype_state=PrototypeStateStore(),
|
||||
|
|
@ -587,10 +741,12 @@ async def test_agent_api_wrapper_transparently_surfaces_modern_events(monkeypatc
|
|||
monkeypatch.setattr(
|
||||
agent_api_wrapper_module.AgentApi,
|
||||
"__init__",
|
||||
lambda self, agent_id, base_url=None, chat_id=0, **kwargs: setattr(self, "id", agent_id)
|
||||
or setattr(self, "callback", kwargs.get("callback"))
|
||||
or setattr(self, "on_disconnect", kwargs.get("on_disconnect"))
|
||||
or setattr(self, "_current_queue", None),
|
||||
lambda self, agent_id, base_url=None, chat_id=0, **kwargs: (
|
||||
setattr(self, "id", agent_id)
|
||||
or setattr(self, "callback", kwargs.get("callback"))
|
||||
or setattr(self, "on_disconnect", kwargs.get("on_disconnect"))
|
||||
or setattr(self, "_current_queue", None)
|
||||
),
|
||||
)
|
||||
|
||||
wrapper = AgentApiWrapper(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue