fix prototype backend review issues
This commit is contained in:
parent
94bdb44b93
commit
37643a9695
9 changed files with 182 additions and 46 deletions
|
|
@ -1,9 +1,58 @@
|
|||
import sys
|
||||
from pathlib import Path
|
||||
from types import ModuleType
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
from sdk.interface import MessageChunk, MessageResponse
|
||||
from sdk.agent_session import AgentSessionClient, AgentSessionConfig, build_thread_key
|
||||
|
||||
AGENT_ROOT = Path(__file__).resolve().parents[2] / "external" / "platform-agent"
|
||||
AGENT_API_ROOT = Path(__file__).resolve().parents[2] / "external" / "platform-agent_api"
|
||||
for path in (AGENT_ROOT, AGENT_API_ROOT):
|
||||
if str(path) not in sys.path:
|
||||
sys.path.insert(0, str(path))
|
||||
|
||||
if "fastapi" not in sys.modules:
|
||||
fastapi = ModuleType("fastapi")
|
||||
|
||||
class _Router:
|
||||
def websocket(self, _path: str):
|
||||
def decorator(fn):
|
||||
return fn
|
||||
|
||||
return decorator
|
||||
|
||||
class _WebSocketDisconnect(Exception):
|
||||
pass
|
||||
|
||||
def _depends(value):
|
||||
return value
|
||||
|
||||
fastapi.APIRouter = _Router
|
||||
fastapi.WebSocket = object
|
||||
fastapi.WebSocketDisconnect = _WebSocketDisconnect
|
||||
fastapi.Depends = _depends
|
||||
sys.modules["fastapi"] = fastapi
|
||||
|
||||
if "src.agent" not in sys.modules:
|
||||
agent_module = ModuleType("src.agent")
|
||||
|
||||
class _AgentService:
|
||||
async def astream(self, text: str, thread_id: str):
|
||||
yield text
|
||||
|
||||
def _get_agent_service():
|
||||
return _AgentService()
|
||||
|
||||
agent_module.AgentService = _AgentService
|
||||
agent_module.get_agent_service = _get_agent_service
|
||||
sys.modules["src.agent"] = agent_module
|
||||
|
||||
from lambda_agent_api.client import MsgUserMessage # noqa: E402
|
||||
from src.api.external import process_message # noqa: E402
|
||||
|
||||
|
||||
def test_build_thread_key_uses_platform_user_and_chat_id():
|
||||
assert build_thread_key("matrix", "@alice:example.org", "C1") == "6:matrix18:@alice:example.org2:C1"
|
||||
|
|
@ -18,11 +67,13 @@ def test_build_thread_key_does_not_collide_when_user_id_contains_colons():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_message_yields_text_chunks_and_end(aiohttp_server):
|
||||
thread_key = build_thread_key("matrix", "@alice:example.org", "C1")
|
||||
|
||||
async def handler(request):
|
||||
ws = web.WebSocketResponse()
|
||||
await ws.prepare(request)
|
||||
|
||||
assert request.query["thread_id"] == "matrix:@alice:example.org:C1"
|
||||
assert request.query["thread_id"] == thread_key
|
||||
|
||||
await ws.send_json({"type": "STATUS"})
|
||||
|
||||
|
|
@ -43,25 +94,27 @@ async def test_stream_message_yields_text_chunks_and_end(aiohttp_server):
|
|||
|
||||
chunks = []
|
||||
async for chunk in client.stream_message(
|
||||
thread_key="matrix:@alice:example.org:C1",
|
||||
thread_key=thread_key,
|
||||
text="hello",
|
||||
):
|
||||
chunks.append(chunk)
|
||||
|
||||
assert chunks == [
|
||||
MessageChunk(message_id="matrix:@alice:example.org:C1", delta="hel", finished=False, tokens_used=0),
|
||||
MessageChunk(message_id="matrix:@alice:example.org:C1", delta="lo", finished=False, tokens_used=0),
|
||||
MessageChunk(message_id="matrix:@alice:example.org:C1", delta="", finished=True, tokens_used=7),
|
||||
MessageChunk(message_id=thread_key, delta="hel", finished=False, tokens_used=0),
|
||||
MessageChunk(message_id=thread_key, delta="lo", finished=False, tokens_used=0),
|
||||
MessageChunk(message_id=thread_key, delta="", finished=True, tokens_used=7),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_message_collects_streamed_chunks_and_tokens(aiohttp_server):
|
||||
thread_key = build_thread_key("matrix", "@alice:example.org", "C1")
|
||||
|
||||
async def handler(request):
|
||||
ws = web.WebSocketResponse()
|
||||
await ws.prepare(request)
|
||||
|
||||
assert request.query["thread_id"] == "matrix:@alice:example.org:C1"
|
||||
assert request.query["thread_id"] == thread_key
|
||||
|
||||
await ws.send_json({"type": "STATUS"})
|
||||
|
||||
|
|
@ -81,13 +134,60 @@ async def test_send_message_collects_streamed_chunks_and_tokens(aiohttp_server):
|
|||
client = AgentSessionClient(AgentSessionConfig(base_ws_url=str(server.make_url("/agent_ws/"))))
|
||||
|
||||
result = await client.send_message(
|
||||
thread_key="matrix:@alice:example.org:C1",
|
||||
thread_key=thread_key,
|
||||
text="hello world",
|
||||
)
|
||||
|
||||
assert result == MessageResponse(
|
||||
message_id="matrix:@alice:example.org:C1",
|
||||
message_id=thread_key,
|
||||
response="hello world",
|
||||
tokens_used=11,
|
||||
finished=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_message_requires_thread_id_query_param():
|
||||
class FakeWebSocket:
|
||||
query_params = {}
|
||||
|
||||
async def send_text(self, text: str) -> None:
|
||||
raise AssertionError(f"send_text should not be called: {text}")
|
||||
|
||||
class FakeAgentService:
|
||||
async def astream(self, text: str, thread_id: str):
|
||||
yield text
|
||||
|
||||
with pytest.raises(ValueError, match="thread_id query parameter is required"):
|
||||
await process_message(
|
||||
FakeWebSocket(),
|
||||
MsgUserMessage(text="hello"),
|
||||
FakeAgentService(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_message_passes_thread_id_to_agent_service():
|
||||
class FakeWebSocket:
|
||||
def __init__(self) -> None:
|
||||
self.query_params = {"thread_id": "6:matrix18:@alice:example.org2:C1"}
|
||||
self.sent_messages: list[str] = []
|
||||
|
||||
async def send_text(self, text: str) -> None:
|
||||
self.sent_messages.append(text)
|
||||
|
||||
class FakeAgentService:
|
||||
def __init__(self) -> None:
|
||||
self.calls: list[tuple[str, str]] = []
|
||||
|
||||
async def astream(self, text: str, thread_id: str):
|
||||
self.calls.append((text, thread_id))
|
||||
yield "hello"
|
||||
|
||||
ws = FakeWebSocket()
|
||||
agent_service = FakeAgentService()
|
||||
await process_message(ws, MsgUserMessage(text="hello"), agent_service)
|
||||
|
||||
assert agent_service.calls == [("hello", "6:matrix18:@alice:example.org2:C1")]
|
||||
assert any("AGENT_EVENT_TEXT_CHUNK" in message for message in ws.sent_messages)
|
||||
assert any("AGENT_EVENT_END" in message for message in ws.sent_messages)
|
||||
|
|
|
|||
|
|
@ -9,18 +9,12 @@ from sdk.prototype_state import PrototypeStateStore
|
|||
async def test_get_or_create_user_is_stable_per_surface_identity():
|
||||
store = PrototypeStateStore()
|
||||
|
||||
first = await store.get_or_create_user(
|
||||
external_id="@alice:example.org",
|
||||
platform="matrix",
|
||||
display_name="Alice",
|
||||
)
|
||||
second = await store.get_or_create_user(
|
||||
external_id="@alice:example.org",
|
||||
platform="matrix",
|
||||
)
|
||||
first = await store.get_or_create_user("@alice:example.org", "matrix", "Alice")
|
||||
second = await store.get_or_create_user("@alice:example.org", "matrix")
|
||||
|
||||
assert first.user_id == "usr-matrix-@alice:example.org"
|
||||
assert first.is_new is True
|
||||
assert store._users["matrix:@alice:example.org"].is_new is False
|
||||
|
||||
first.display_name = "Mallory"
|
||||
first.is_new = False
|
||||
|
|
@ -56,6 +50,22 @@ async def test_settings_defaults_match_existing_mock_shape():
|
|||
assert settings.plan == {"name": "Beta", "tokens_used": 0, "tokens_limit": 1000}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_settings_returns_connectors_copy():
|
||||
store = PrototypeStateStore()
|
||||
store._settings["usr-matrix-@alice:example.org"] = {
|
||||
"connectors": {"github": {"enabled": True}},
|
||||
}
|
||||
|
||||
settings = await store.get_settings("usr-matrix-@alice:example.org")
|
||||
settings.connectors["github"]["enabled"] = False
|
||||
settings.connectors["slack"] = {"enabled": True}
|
||||
|
||||
assert store._settings["usr-matrix-@alice:example.org"]["connectors"] == {
|
||||
"github": {"enabled": True},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_settings_supports_toggle_skill_and_setters():
|
||||
store = PrototypeStateStore()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import pytest
|
||||
|
||||
from core.protocol import SettingsAction
|
||||
from sdk.agent_session import build_thread_key
|
||||
from sdk.interface import MessageChunk, MessageResponse, UserSettings
|
||||
from sdk.prototype_state import PrototypeStateStore
|
||||
from sdk.real import RealPlatformClient
|
||||
|
|
@ -31,13 +32,12 @@ async def test_real_platform_client_get_or_create_user_uses_local_state():
|
|||
client = RealPlatformClient(
|
||||
agent_sessions=FakeAgentSessionClient(),
|
||||
prototype_state=PrototypeStateStore(),
|
||||
platform="telegram",
|
||||
)
|
||||
|
||||
first = await client.get_or_create_user("u1", "telegram", "Alice")
|
||||
second = await client.get_or_create_user("u1", "telegram")
|
||||
first = await client.get_or_create_user("u1", "matrix", "Alice")
|
||||
second = await client.get_or_create_user("u1", "matrix")
|
||||
|
||||
assert first.user_id == "usr-telegram-u1"
|
||||
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
|
||||
|
|
@ -45,57 +45,55 @@ async def test_real_platform_client_get_or_create_user_uses_local_state():
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_platform_client_send_message_uses_configured_platform():
|
||||
async def test_real_platform_client_send_message_uses_surface_user_thread_identity():
|
||||
agent_sessions = FakeAgentSessionClient()
|
||||
client = RealPlatformClient(
|
||||
agent_sessions=agent_sessions,
|
||||
prototype_state=PrototypeStateStore(),
|
||||
platform="telegram",
|
||||
platform="matrix",
|
||||
)
|
||||
|
||||
result = await client.send_message("usr-telegram-u1", "C1", "hello")
|
||||
thread_key = build_thread_key("matrix", "@alice:example.org", "C1")
|
||||
result = await client.send_message("@alice:example.org", "C1", "hello")
|
||||
|
||||
assert result == MessageResponse(
|
||||
message_id="8:telegram15:usr-telegram-u12:C1",
|
||||
message_id=thread_key,
|
||||
response="echo:hello",
|
||||
tokens_used=3,
|
||||
finished=True,
|
||||
)
|
||||
assert agent_sessions.send_calls == [
|
||||
("8:telegram15:usr-telegram-u12:C1", "hello")
|
||||
]
|
||||
assert agent_sessions.send_calls == [(thread_key, "hello")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_platform_client_stream_message_uses_configured_platform():
|
||||
async def test_real_platform_client_stream_message_uses_surface_user_thread_identity():
|
||||
agent_sessions = FakeAgentSessionClient()
|
||||
client = RealPlatformClient(
|
||||
agent_sessions=agent_sessions,
|
||||
prototype_state=PrototypeStateStore(),
|
||||
platform="telegram",
|
||||
platform="matrix",
|
||||
)
|
||||
|
||||
thread_key = build_thread_key("matrix", "@alice:example.org", "C1")
|
||||
chunks = []
|
||||
async for chunk in client.stream_message("usr-telegram-u1", "C1", "hello"):
|
||||
async for chunk in client.stream_message("@alice:example.org", "C1", "hello"):
|
||||
chunks.append(chunk)
|
||||
|
||||
assert chunks == [
|
||||
MessageChunk(
|
||||
message_id="8:telegram15:usr-telegram-u12:C1",
|
||||
message_id=thread_key,
|
||||
delta="he",
|
||||
finished=False,
|
||||
tokens_used=0,
|
||||
),
|
||||
MessageChunk(
|
||||
message_id="8:telegram15:usr-telegram-u12:C1",
|
||||
message_id=thread_key,
|
||||
delta="llo",
|
||||
finished=True,
|
||||
tokens_used=3,
|
||||
),
|
||||
]
|
||||
assert agent_sessions.stream_calls == [
|
||||
("8:telegram15:usr-telegram-u12:C1", "hello")
|
||||
]
|
||||
assert agent_sessions.stream_calls == [(thread_key, "hello")]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue