349 lines
14 KiB
Python
349 lines
14 KiB
Python
"""Unit tests for chat.py WebSocket route — message handling flows.
|
|
|
|
Tests the WebSocket endpoint at /api/v1/chat/ws/{session_id}, focusing on
|
|
the message handling protocol: connected event, ping/pong, cancel,
|
|
API key auth, session validation, happy-path message → final_answer,
|
|
confirmation request/reply flow, invalid JSON handling, session isolation,
|
|
and connection cleanup.
|
|
|
|
The LLM gateway and agent execution are mocked — these tests focus on the
|
|
WebSocket protocol layer, not real LLM calls.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
from fastapi import FastAPI, WebSocket
|
|
from fastapi.testclient import TestClient
|
|
|
|
from agentkit.core.protocol import CancellationToken
|
|
from agentkit.session.manager import SessionManager
|
|
from agentkit.session.store import InMemorySessionStore
|
|
|
|
|
|
# ── Stub handlers (match _handle_chat_message signature) ────────────────
|
|
#
|
|
# 10 params — all required by the call site in chat_websocket which passes
|
|
# them positionally. Full signature avoids *args (no Any usage).
|
|
|
|
|
|
async def _stub_final_answer(
|
|
websocket: WebSocket,
|
|
session_id: str,
|
|
content: str,
|
|
sm: SessionManager,
|
|
cancellation_token: CancellationToken,
|
|
pending_replies: dict[str, asyncio.Future],
|
|
pending_confirmations: dict[str, asyncio.Future] | None = None,
|
|
pending_spec_reviews: dict[str, asyncio.Future] | None = None,
|
|
pending_autonomy_resumes: dict[str, asyncio.Future] | None = None,
|
|
model_override: str | None = None,
|
|
) -> None:
|
|
"""Stub that immediately sends a final_answer event."""
|
|
await websocket.send_json(
|
|
{"type": "final_answer", "content": f"echo: {content}", "is_final": True}
|
|
)
|
|
|
|
|
|
async def _stub_with_confirmation(
|
|
websocket: WebSocket,
|
|
session_id: str,
|
|
content: str,
|
|
sm: SessionManager,
|
|
cancellation_token: CancellationToken,
|
|
pending_replies: dict[str, asyncio.Future],
|
|
pending_confirmations: dict[str, asyncio.Future] | None = None,
|
|
pending_spec_reviews: dict[str, asyncio.Future] | None = None,
|
|
pending_autonomy_resumes: dict[str, asyncio.Future] | None = None,
|
|
model_override: str | None = None,
|
|
) -> None:
|
|
"""Stub: send confirmation_request, wait for reply, send final_answer."""
|
|
conf_id = "test-conf-1"
|
|
# Register future BEFORE sending request so the WS loop can resolve it
|
|
# even if the reply arrives before we await (race-safe).
|
|
loop = asyncio.get_running_loop()
|
|
future: asyncio.Future[bool] = loop.create_future()
|
|
if pending_confirmations is not None:
|
|
pending_confirmations[conf_id] = future
|
|
await websocket.send_json(
|
|
{
|
|
"type": "confirmation_request",
|
|
"data": {
|
|
"confirmation_id": conf_id,
|
|
"command": "rm -rf /tmp/test",
|
|
"reason": "dangerous command",
|
|
},
|
|
}
|
|
)
|
|
try:
|
|
approved = await asyncio.wait_for(future, timeout=10.0)
|
|
except asyncio.TimeoutError:
|
|
approved = False
|
|
await websocket.send_json(
|
|
{"type": "final_answer", "content": f"approved={approved}", "is_final": True}
|
|
)
|
|
|
|
|
|
# ── Fixtures ────────────────────────────────────────────────────────────
|
|
|
|
|
|
@pytest.fixture
|
|
def app_with_chat_ws() -> FastAPI:
|
|
"""Create a FastAPI app with Chat routes and mocked dependencies."""
|
|
from agentkit.server.routes.chat import router
|
|
|
|
app = FastAPI()
|
|
app.include_router(router, prefix="/api/v1")
|
|
app.state.session_manager = SessionManager(store=InMemorySessionStore())
|
|
app.state.llm_gateway = MagicMock()
|
|
app.state.agent_pool = MagicMock()
|
|
app.state.server_config = MagicMock()
|
|
app.state.server_config.api_key = None
|
|
return app
|
|
|
|
|
|
@pytest.fixture
|
|
def client(app_with_chat_ws: FastAPI) -> TestClient:
|
|
return TestClient(app_with_chat_ws)
|
|
|
|
|
|
def _create_session(client: TestClient, agent_name: str = "test-agent") -> str:
|
|
"""Create a chat session via REST and return its id."""
|
|
resp = client.post("/api/v1/chat/sessions", json={"agent_name": agent_name})
|
|
assert resp.status_code == 200
|
|
return resp.json()["session_id"]
|
|
|
|
|
|
# ── Connection & session validation ─────────────────────────────────────
|
|
|
|
|
|
class TestWSConnection:
|
|
"""Connection establishment and session validation."""
|
|
|
|
def test_connected_event_with_session_id(self, client: TestClient) -> None:
|
|
session_id = _create_session(client)
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
|
msg = ws.receive_json()
|
|
assert msg["type"] == "connected"
|
|
assert msg["session_id"] == session_id
|
|
|
|
def test_session_not_found_rejected(self, client: TestClient) -> None:
|
|
with client.websocket_connect("/api/v1/chat/ws/nonexistent-session") as ws:
|
|
msg = ws.receive_json()
|
|
assert msg["type"] == "error"
|
|
assert "not found" in msg["data"]["message"].lower()
|
|
|
|
def test_closed_session_rejected(self, client: TestClient) -> None:
|
|
session_id = _create_session(client)
|
|
client.delete(f"/api/v1/chat/sessions/{session_id}")
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
|
msg = ws.receive_json()
|
|
assert msg["type"] == "error"
|
|
assert "closed" in msg["data"]["message"].lower()
|
|
|
|
|
|
# ── Authentication ──────────────────────────────────────────────────────
|
|
|
|
|
|
class TestWSAuth:
|
|
"""API key authentication via ?api_key= query param."""
|
|
|
|
def test_missing_api_key_rejected(self, app_with_chat_ws: FastAPI) -> None:
|
|
app_with_chat_ws.state.server_config.api_key = "secret123"
|
|
client = TestClient(app_with_chat_ws)
|
|
session_id = _create_session(client)
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
|
msg = ws.receive_json()
|
|
assert msg["type"] == "error"
|
|
assert "api_key" in msg["data"]["message"].lower()
|
|
|
|
def test_wrong_api_key_rejected(self, app_with_chat_ws: FastAPI) -> None:
|
|
app_with_chat_ws.state.server_config.api_key = "secret123"
|
|
client = TestClient(app_with_chat_ws)
|
|
session_id = _create_session(client)
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}?api_key=wrong") as ws:
|
|
msg = ws.receive_json()
|
|
assert msg["type"] == "error"
|
|
assert "api_key" in msg["data"]["message"].lower()
|
|
|
|
def test_valid_api_key_accepted(self, app_with_chat_ws: FastAPI) -> None:
|
|
app_with_chat_ws.state.server_config.api_key = "secret123"
|
|
client = TestClient(app_with_chat_ws)
|
|
session_id = _create_session(client)
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}?api_key=secret123") as ws:
|
|
msg = ws.receive_json()
|
|
assert msg["type"] == "connected"
|
|
assert msg["session_id"] == session_id
|
|
|
|
|
|
# ── Message handling ────────────────────────────────────────────────────
|
|
|
|
|
|
class TestWSMessages:
|
|
"""Message type handling: ping, cancel, message, confirmation, invalid JSON."""
|
|
|
|
def test_ping_returns_pong(self, client: TestClient) -> None:
|
|
session_id = _create_session(client)
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
|
ws.receive_json() # connected
|
|
ws.send_json({"type": "ping"})
|
|
msg = ws.receive_json()
|
|
assert msg["type"] == "pong"
|
|
|
|
def test_cancel_sends_result_and_stays_open(self, client: TestClient) -> None:
|
|
session_id = _create_session(client)
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
|
ws.receive_json() # connected
|
|
ws.send_json({"type": "cancel"})
|
|
msg = ws.receive_json()
|
|
assert msg["type"] == "result"
|
|
assert msg["data"]["status"] == "cancelled"
|
|
# Connection stays open — ping still works after cancel
|
|
ws.send_json({"type": "ping"})
|
|
assert ws.receive_json()["type"] == "pong"
|
|
|
|
def test_message_happy_path(
|
|
self, app_with_chat_ws: FastAPI, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
from agentkit.server.routes import chat as chat_module
|
|
|
|
monkeypatch.setattr(chat_module, "_handle_chat_message", _stub_final_answer)
|
|
client = TestClient(app_with_chat_ws)
|
|
session_id = _create_session(client)
|
|
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
|
ws.receive_json() # connected
|
|
ws.send_json({"type": "message", "content": "hello"})
|
|
final = ws.receive_json()
|
|
assert final["type"] == "final_answer"
|
|
assert final["content"] == "echo: hello"
|
|
assert final["is_final"] is True
|
|
|
|
def test_confirmation_flow(
|
|
self, app_with_chat_ws: FastAPI, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""Server sends confirmation_request → client replies → server continues."""
|
|
from agentkit.server.routes import chat as chat_module
|
|
|
|
monkeypatch.setattr(chat_module, "_handle_chat_message", _stub_with_confirmation)
|
|
client = TestClient(app_with_chat_ws)
|
|
session_id = _create_session(client)
|
|
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
|
ws.receive_json() # connected
|
|
ws.send_json({"type": "message", "content": "do something risky"})
|
|
|
|
conf_req = ws.receive_json()
|
|
assert conf_req["type"] == "confirmation_request"
|
|
assert conf_req["data"]["confirmation_id"] == "test-conf-1"
|
|
|
|
ws.send_json(
|
|
{
|
|
"type": "confirmation_reply",
|
|
"confirmation_id": "test-conf-1",
|
|
"approved": True,
|
|
}
|
|
)
|
|
|
|
final = ws.receive_json()
|
|
assert final["type"] == "final_answer"
|
|
assert final["content"] == "approved=True"
|
|
|
|
def test_confirmation_flow_rejected(
|
|
self, app_with_chat_ws: FastAPI, monkeypatch: pytest.MonkeyPatch
|
|
) -> None:
|
|
"""Client rejects confirmation → final_answer reflects approved=False."""
|
|
from agentkit.server.routes import chat as chat_module
|
|
|
|
monkeypatch.setattr(chat_module, "_handle_chat_message", _stub_with_confirmation)
|
|
client = TestClient(app_with_chat_ws)
|
|
session_id = _create_session(client)
|
|
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
|
ws.receive_json() # connected
|
|
ws.send_json({"type": "message", "content": "risky"})
|
|
|
|
conf_req = ws.receive_json()
|
|
assert conf_req["type"] == "confirmation_request"
|
|
|
|
ws.send_json(
|
|
{
|
|
"type": "confirmation_reply",
|
|
"confirmation_id": "test-conf-1",
|
|
"approved": False,
|
|
}
|
|
)
|
|
|
|
final = ws.receive_json()
|
|
assert final["type"] == "final_answer"
|
|
assert final["content"] == "approved=False"
|
|
|
|
def test_invalid_json_ignored(self, client: TestClient) -> None:
|
|
"""Invalid JSON text is silently ignored — connection stays open.
|
|
|
|
Note: the implementation catches JSONDecodeError with `continue`
|
|
rather than sending an error event. This is the actual behavior.
|
|
"""
|
|
session_id = _create_session(client)
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
|
ws.receive_json() # connected
|
|
ws.send_text("not valid json {{{")
|
|
# Connection stays open — ping still works
|
|
ws.send_json({"type": "ping"})
|
|
msg = ws.receive_json()
|
|
assert msg["type"] == "pong"
|
|
|
|
def test_unknown_message_type_ignored(self, client: TestClient) -> None:
|
|
"""Unknown message type is silently ignored — no error sent."""
|
|
session_id = _create_session(client)
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
|
ws.receive_json() # connected
|
|
ws.send_json({"type": "unknown_type", "data": "whatever"})
|
|
# Connection stays open — ping still works
|
|
ws.send_json({"type": "ping"})
|
|
msg = ws.receive_json()
|
|
assert msg["type"] == "pong"
|
|
|
|
|
|
# ── Isolation & cleanup ────────────────────────────────────────────────
|
|
|
|
|
|
class TestWSIsolationAndCleanup:
|
|
"""Session isolation and connection cleanup."""
|
|
|
|
def test_session_isolation(self, client: TestClient) -> None:
|
|
"""Different session_ids get separate connected events."""
|
|
sid1 = _create_session(client)
|
|
sid2 = _create_session(client)
|
|
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{sid1}") as ws1:
|
|
msg1 = ws1.receive_json()
|
|
assert msg1["session_id"] == sid1
|
|
ws1.send_json({"type": "ping"})
|
|
assert ws1.receive_json()["type"] == "pong"
|
|
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{sid2}") as ws2:
|
|
msg2 = ws2.receive_json()
|
|
assert msg2["session_id"] == sid2
|
|
ws2.send_json({"type": "ping"})
|
|
assert ws2.receive_json()["type"] == "pong"
|
|
|
|
def test_connection_cleanup_on_disconnect(self, app_with_chat_ws: FastAPI) -> None:
|
|
"""After disconnect, chat_manager no longer tracks the connection."""
|
|
from agentkit.server.routes.chat import chat_manager
|
|
|
|
client = TestClient(app_with_chat_ws)
|
|
session_id = _create_session(client)
|
|
|
|
assert len(chat_manager.get_connections(session_id)) == 0
|
|
|
|
with client.websocket_connect(f"/api/v1/chat/ws/{session_id}") as ws:
|
|
ws.receive_json() # connected
|
|
assert len(chat_manager.get_connections(session_id)) == 1
|
|
|
|
# After disconnect, the finally block removes the connection
|
|
assert len(chat_manager.get_connections(session_id)) == 0
|