231 lines
6.9 KiB
Python
231 lines
6.9 KiB
Python
import asyncio
|
|
|
|
import pytest
|
|
|
|
from core.protocol import SettingsAction
|
|
from sdk.interface import MessageChunk, MessageResponse, UserSettings
|
|
from sdk.prototype_state import PrototypeStateStore
|
|
from sdk.real import RealPlatformClient
|
|
|
|
|
|
class FakeChunk:
|
|
def __init__(self, text: str) -> None:
|
|
self.text = text
|
|
|
|
|
|
class FakeChatAgentApi:
|
|
def __init__(self, chat_id: str) -> None:
|
|
self.chat_id = chat_id
|
|
self.calls: list[str] = []
|
|
self.connect_calls = 0
|
|
self.close_calls = 0
|
|
self.last_tokens_used = 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):
|
|
self.calls.append(text)
|
|
midpoint = len(text) // 2
|
|
yield FakeChunk(text[:midpoint])
|
|
yield FakeChunk(text[midpoint:])
|
|
self.last_tokens_used = 3
|
|
|
|
|
|
class FakeAgentApiFactory:
|
|
def __init__(self) -> None:
|
|
self.created_chat_ids: list[str] = []
|
|
self.instances: dict[str, FakeChatAgentApi] = {}
|
|
|
|
def for_chat(self, chat_id: str) -> FakeChatAgentApi:
|
|
chat_api = FakeChatAgentApi(chat_id)
|
|
self.created_chat_ids.append(chat_id)
|
|
self.instances[chat_id] = chat_api
|
|
return chat_api
|
|
|
|
|
|
class LegacyAgentApi:
|
|
def __init__(self) -> None:
|
|
self.calls: list[str] = []
|
|
self.last_tokens_used = 0
|
|
|
|
async def send_message(self, text: str):
|
|
self.calls.append(text)
|
|
yield FakeChunk(text[:2])
|
|
yield FakeChunk(text[2:])
|
|
self.last_tokens_used = 7
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_platform_client_get_or_create_user_uses_local_state():
|
|
client = RealPlatformClient(
|
|
agent_api=FakeAgentApiFactory(),
|
|
prototype_state=PrototypeStateStore(),
|
|
)
|
|
|
|
first = await client.get_or_create_user("u1", "matrix", "Alice")
|
|
second = await client.get_or_create_user("u1", "matrix")
|
|
|
|
assert first.user_id == "usr-matrix-u1"
|
|
assert first.is_new is True
|
|
assert second.user_id == first.user_id
|
|
assert second.is_new is False
|
|
assert second.display_name == "Alice"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_platform_client_send_message_uses_chat_bound_client():
|
|
agent_api = FakeAgentApiFactory()
|
|
client = RealPlatformClient(
|
|
agent_api=agent_api,
|
|
prototype_state=PrototypeStateStore(),
|
|
platform="matrix",
|
|
)
|
|
|
|
result = await client.send_message("@alice:example.org", "chat-7", "hello")
|
|
|
|
assert result == MessageResponse(
|
|
message_id="@alice:example.org",
|
|
response="hello",
|
|
tokens_used=3,
|
|
finished=True,
|
|
)
|
|
assert agent_api.created_chat_ids == ["chat-7"]
|
|
assert agent_api.instances["chat-7"].chat_id == "chat-7"
|
|
assert agent_api.instances["chat-7"].calls == ["hello"]
|
|
assert agent_api.instances["chat-7"].connect_calls == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_platform_client_works_with_legacy_agent_api_without_for_chat():
|
|
legacy_api = LegacyAgentApi()
|
|
client = RealPlatformClient(
|
|
agent_api=legacy_api,
|
|
prototype_state=PrototypeStateStore(),
|
|
platform="matrix",
|
|
)
|
|
|
|
result = await client.send_message("@alice:example.org", "chat-legacy", "hello")
|
|
|
|
assert result == MessageResponse(
|
|
message_id="@alice:example.org",
|
|
response="hello",
|
|
tokens_used=7,
|
|
finished=True,
|
|
)
|
|
assert legacy_api.calls == ["hello"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_platform_client_reuses_cached_chat_client():
|
|
agent_api = FakeAgentApiFactory()
|
|
client = RealPlatformClient(
|
|
agent_api=agent_api,
|
|
prototype_state=PrototypeStateStore(),
|
|
platform="matrix",
|
|
)
|
|
|
|
await client.send_message("@alice:example.org", "chat-1", "hello")
|
|
await client.send_message("@alice:example.org", "chat-1", "again")
|
|
|
|
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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_platform_client_creates_chat_client_atomically_for_concurrent_requests():
|
|
agent_api = FakeAgentApiFactory()
|
|
client = RealPlatformClient(
|
|
agent_api=agent_api,
|
|
prototype_state=PrototypeStateStore(),
|
|
platform="matrix",
|
|
)
|
|
|
|
results = await asyncio.gather(
|
|
client.send_message("@alice:example.org", "chat-1", "hello"),
|
|
client.send_message("@alice:example.org", "chat-1", "again"),
|
|
)
|
|
|
|
assert [result.response for result in results] == ["hello", "again"]
|
|
assert agent_api.created_chat_ids == ["chat-1"]
|
|
assert agent_api.instances["chat-1"].connect_calls == 1
|
|
assert agent_api.instances["chat-1"].calls == ["hello", "again"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_platform_client_creates_distinct_clients_per_chat():
|
|
agent_api = FakeAgentApiFactory()
|
|
client = RealPlatformClient(
|
|
agent_api=agent_api,
|
|
prototype_state=PrototypeStateStore(),
|
|
platform="matrix",
|
|
)
|
|
|
|
await client.send_message("@alice:example.org", "chat-1", "hello")
|
|
await client.send_message("@alice:example.org", "chat-2", "world")
|
|
|
|
assert agent_api.created_chat_ids == ["chat-1", "chat-2"]
|
|
assert agent_api.instances["chat-1"] is not agent_api.instances["chat-2"]
|
|
assert agent_api.instances["chat-1"].calls == ["hello"]
|
|
assert agent_api.instances["chat-2"].calls == ["world"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_platform_client_stream_message_emits_final_tokens_chunk():
|
|
agent_api = FakeAgentApiFactory()
|
|
client = RealPlatformClient(
|
|
agent_api=agent_api,
|
|
prototype_state=PrototypeStateStore(),
|
|
platform="matrix",
|
|
)
|
|
|
|
chunks = []
|
|
async for chunk in client.stream_message("@alice:example.org", "chat-1", "hello"):
|
|
chunks.append(chunk)
|
|
|
|
assert chunks == [
|
|
MessageChunk(
|
|
message_id="@alice:example.org",
|
|
delta="he",
|
|
finished=False,
|
|
tokens_used=0,
|
|
),
|
|
MessageChunk(
|
|
message_id="@alice:example.org",
|
|
delta="llo",
|
|
finished=False,
|
|
tokens_used=0,
|
|
),
|
|
MessageChunk(
|
|
message_id="@alice:example.org",
|
|
delta="",
|
|
finished=True,
|
|
tokens_used=3,
|
|
),
|
|
]
|
|
assert agent_api.created_chat_ids == ["chat-1"]
|
|
assert agent_api.instances["chat-1"].calls == ["hello"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_real_platform_client_settings_are_local():
|
|
client = RealPlatformClient(
|
|
agent_api=FakeAgentApiFactory(),
|
|
prototype_state=PrototypeStateStore(),
|
|
platform="matrix",
|
|
)
|
|
|
|
await client.update_settings(
|
|
"usr-matrix-u1",
|
|
SettingsAction(action="toggle_skill", payload={"skill": "browser", "enabled": True}),
|
|
)
|
|
|
|
settings = await client.get_settings("usr-matrix-u1")
|
|
|
|
assert isinstance(settings, UserSettings)
|
|
assert settings.skills["browser"] is True
|
|
assert settings.skills["web-search"] is True
|