agent-lab/tests/test_websocket_api.py

import json
import asyncio
import logging
from collections.abc import AsyncIterator
from pathlib import Path

import httpx
import pytest
from fastapi.testclient import TestClient
from pydantic import ValidationError

from agent_lab.application.contracts import AgentParams, DebugRunRequest, EventAgentParams
from agent_lab.application.queues import RuntimeQueues
from agent_lab.application.runtime import DebugRuntime
from agent_lab.domain.events import ToolCallEvent
from agent_lab.domain.messages import ChatMessage, StreamItem
from agent_lab.infrastructure.chat_client import OpenAICompatibleChatClient
from agent_lab.presentation.web import create_app


class FakeRuntime:
    def __init__(self) -> None:
        self.requests: list[DebugRunRequest] = []

    async def run(self, request: DebugRunRequest) -> AsyncIterator[dict]:
        self.requests.append(request)
        yield {"type": "session_started"}
        yield {"type": "message_delta", "content": "hello"}
        yield {"type": "done"}


class QueueAwareRuntime:
    def __init__(self) -> None:
        self.requests: list[DebugRunRequest] = []
        self.queues: RuntimeQueues | None = None
        self.task: asyncio.Task | None = None

    def start(self, request: DebugRunRequest) -> RuntimeQueues:
        self.requests.append(request)
        self.queues = RuntimeQueues()
        self.task = asyncio.create_task(self._run())
        return self.queues

    async def _run(self) -> None:
        assert self.queues is not None
        await self.queues.output.put({"type": "session_started"})
        message = await self.queues.input.get()
        await self.queues.output.put(
            {"type": "message_delta", "content": message.content}
        )
        await self.queues.output.put({"type": "done"})

    async def aclose(self) -> None:
        if self.task is not None and not self.task.done():
            self.task.cancel()
            await asyncio.gather(self.task, return_exceptions=True)


def _request_payload() -> dict:
    return {
        "user_message": "debug this",
        "system_prompts": ["You are a debugger."],
        "pre_messages": [{"role": "user", "content": "previous turn"}],
        "chat_agent": {
            "model": "fake-model",
            "temperature": 0.1,
            "max_tokens": 200,
        },
        "event_agent": {
            "enabled_tools": ["handoff_note"],
            "max_event_loops": 2,
        },
    }


def test_agent_params_defaults_include_extra_body_and_event_agent_defaults_to_one_round():
    chat_params = AgentParams()
    event_params = EventAgentParams()

    assert chat_params.extra_body == {
        "thinking": {"type": "disabled"},
        "enable_search": False,
        "search_options": {"forced_search": False},
    }
    assert event_params.extra_body == chat_params.extra_body
    assert event_params.max_event_loops == 1
    assert event_params.system_prompt == ""


def test_health_returns_ok():
    app = create_app(runtime_factory=FakeRuntime)
    client = TestClient(app)

    response = client.get("/health")

    assert response.status_code == 200
    assert response.json() == {"status": "ok"}


def test_api_tools_returns_handoff_note_metadata():
    app = create_app(runtime_factory=FakeRuntime)
    client = TestClient(app)

    response = client.get("/api/tools")

    assert response.status_code == 200
    tools = response.json()
    assert [tool["name"] for tool in tools] == [
        "handoff_note",
        "mock_search",
        "mock_ticket",
    ]
    assert tools[0]["parameters"]["required"] == ["message"]


def test_websocket_debug_streams_runtime_messages():
    runtime = FakeRuntime()
    app = create_app(runtime_factory=lambda: runtime)
    client = TestClient(app)

    with client.websocket_connect("/ws/debug") as websocket:
        websocket.send_json(_request_payload())

        assert websocket.receive_json() == {"type": "session_started"}
        assert websocket.receive_json() == {
            "type": "message_delta",
            "content": "hello",
        }
        assert websocket.receive_json() == {"type": "done"}

    assert runtime.requests[0].user_message == "debug this"
    assert runtime.requests[0].pre_messages[0].content == "previous turn"


def test_websocket_debug_logs_session_lifecycle(caplog):
    runtime = FakeRuntime()
    app = create_app(runtime_factory=lambda: runtime)
    client = TestClient(app)

    with caplog.at_level(logging.INFO, logger="agent_lab.presentation.web"):
        with client.websocket_connect("/ws/debug") as websocket:
            websocket.send_json(_request_payload())
            assert websocket.receive_json() == {"type": "session_started"}
            assert websocket.receive_json() == {
                "type": "message_delta",
                "content": "hello",
            }
            assert websocket.receive_json() == {"type": "done"}

    assert "websocket session accepted" in caplog.text
    assert "websocket request accepted" in caplog.text
    assert "websocket session closed" in caplog.text


def test_websocket_debug_enqueues_user_messages_during_running_session():
    runtime = QueueAwareRuntime()
    app = create_app(runtime_factory=lambda: runtime)
    client = TestClient(app)

    with client.websocket_connect("/ws/debug") as websocket:
        websocket.send_json(_request_payload())

        assert websocket.receive_json() == {"type": "session_started"}
        websocket.send_json({"type": "user_message", "content": "follow-up"})
        assert websocket.receive_json() == {
            "type": "message_delta",
            "content": "follow-up",
        }
        assert websocket.receive_json() == {"type": "done"}

    assert runtime.requests[0].user_message == "debug this"


def test_websocket_debug_sends_error_for_invalid_request():
    app = create_app(runtime_factory=FakeRuntime)
    client = TestClient(app)

    with client.websocket_connect("/ws/debug") as websocket:
        websocket.send_json({"chat_agent": {"model": "fake-model"}})
        message = websocket.receive_json()

    assert message["type"] == "error"
    assert "user_message" in message["message"]


def test_debug_run_request_rejects_invalid_pre_message_role():
    payload = _request_payload()
    payload["pre_messages"] = [{"role": "developer", "content": "invalid"}]

    with pytest.raises(ValidationError) as exc_info:
        DebugRunRequest.model_validate(payload)

    assert "role" in str(exc_info.value)


def test_debug_run_request_rejects_tool_pre_messages():
    payload = _request_payload()
    payload["pre_messages"] = [
        {
            "role": "tool",
            "content": "orphan result",
            "tool_call_id": "call_1",
        }
    ]

    with pytest.raises(ValidationError, match="pre_messages cannot include tool messages"):
        DebugRunRequest.model_validate(payload)


def test_chat_message_allows_internal_tool_replies():
    message = ChatMessage(
        role="tool",
        content='{"message":"handled"}',
        name="handoff_note",
        tool_call_id="call_1",
    )

    assert message.role == "tool"
    assert message.tool_call_id == "call_1"


class HistoryCapturingChatClient:
    def __init__(self) -> None:
        self.calls = 0
        self.second_call_messages: list[ChatMessage] = []

    async def stream_chat(
        self,
        messages: list[ChatMessage],
        tools: list[dict],
        params: AgentParams,
    ) -> AsyncIterator[StreamItem]:
        self.calls += 1
        if self.calls == 1:
            yield StreamItem.message_delta("Need event help.")
            yield StreamItem.event(
                ToolCallEvent(
                    id="call_1",
                    name="handoff_note",
                    arguments={},
                    raw_arguments="{}",
                )
            )
            return

        self.second_call_messages = list(messages)
        yield StreamItem.message_delta("Final answer.")


@pytest.mark.asyncio
async def test_runtime_appends_assistant_message_before_tool_reply_history():
    request = DebugRunRequest(
        user_message="debug this",
        system_prompts=["You are a debugger."],
        pre_messages=[],
        chat_agent=AgentParams(model="fake-model", temperature=0.1, max_tokens=200),
        event_agent=EventAgentParams(enabled_tools=["handoff_note"], max_event_loops=2),
    )
    client = HistoryCapturingChatClient()
    runtime = DebugRuntime(client)

    outputs = [message async for message in runtime.run(request)]

    assert client.calls == 2
    assert [message.role for message in client.second_call_messages] == [
        "system",
        "user",
        "assistant",
        "tool",
    ]
    assert client.second_call_messages[2].content == "Need event help."
    assert outputs[-1] == {"type": "done"}


@pytest.mark.asyncio
async def test_openai_chat_client_streams_sse_chunks_through_parser():
    requests: list[httpx.Request] = []

    def handler(request: httpx.Request) -> httpx.Response:
        requests.append(request)
        payload = json.loads(request.content)
        assert payload["stream"] is True
        assert payload["model"] == "model-x"
        assert payload["messages"] == [{"role": "user", "content": "hi"}]
        assert payload["tools"][0]["function"]["name"] == "handoff_note"
        assert payload["temperature"] == 0.3
        assert payload["max_tokens"] == 50
        assert request.headers["authorization"] == "Bearer test-key"
        return httpx.Response(
            200,
            content=(
                b'data: {"choices":[{"delta":{"content":"hi"},"finish_reason":null}]}\n\n'
                b"data: [DONE]\n\n"
            ),
        )

    async with httpx.AsyncClient(
        transport=httpx.MockTransport(handler),
        base_url="https://llm.test/v1",
    ) as http_client:
        client = OpenAICompatibleChatClient(
            api_key="test-key",
            base_url="https://llm.test/v1",
            default_model="default-model",
            request_timeout_seconds=5,
            http_client=http_client,
        )
        items = [
            item
            async for item in client.stream_chat(
                messages=[ChatMessage(role="user", content="hi")],
                tools=[
                    {
                        "type": "function",
                        "function": {"name": "handoff_note", "parameters": {}},
                    }
                ],
                params=AgentParams(model="model-x", temperature=0.3, max_tokens=50),
            )
        ]

    assert requests[0].url.path == "/v1/chat/completions"
    assert [item.content for item in items] == ["hi"]


@pytest.mark.asyncio
async def test_openai_chat_client_uses_default_model_when_request_model_is_blank():
    captured_payloads: list[dict] = []

    def handler(request: httpx.Request) -> httpx.Response:
        captured_payloads.append(json.loads(request.content))
        return httpx.Response(
            200,
            content=b'data: {"choices":[{"delta":{"content":"hi"},"finish_reason":null}]}\n\n',
        )

    async with httpx.AsyncClient(
        transport=httpx.MockTransport(handler),
        base_url="https://llm.test/v1",
    ) as http_client:
        client = OpenAICompatibleChatClient(
            api_key="",
            base_url="https://llm.test/v1",
            default_model="provider-default",
            request_timeout_seconds=5,
            http_client=http_client,
        )
        [
            item
            async for item in client.stream_chat(
                messages=[ChatMessage(role="user", content="hi")],
                tools=[],
                params=AgentParams(model="   ", temperature=0.3, max_tokens=50),
            )
        ]

    assert captured_payloads[0]["model"] == "provider-default"


@pytest.mark.asyncio
async def test_openai_chat_client_serializes_tool_reply_without_name():
    captured_payloads: list[dict] = []
    assistant_tool_call = {
        "id": "call_1",
        "type": "function",
        "function": {
            "name": "handoff_note",
            "arguments": '{"message":"inspect"}',
        },
    }

    def handler(request: httpx.Request) -> httpx.Response:
        captured_payloads.append(json.loads(request.content))
        return httpx.Response(
            200,
            content=b'data: {"choices":[{"delta":{"content":"done"},"finish_reason":null}]}\n\n',
        )

    async with httpx.AsyncClient(
        transport=httpx.MockTransport(handler),
        base_url="https://llm.test/v1",
    ) as http_client:
        client = OpenAICompatibleChatClient(
            api_key="",
            base_url="https://llm.test/v1",
            default_model="provider-default",
            request_timeout_seconds=5,
            http_client=http_client,
        )
        [
            item
            async for item in client.stream_chat(
                messages=[
                    ChatMessage(
                        role="assistant",
                        content="",
                        tool_calls=[assistant_tool_call],
                    ),
                    ChatMessage(
                        role="tool",
                        content="noted",
                        name="handoff_note",
                        tool_call_id="call_1",
                    ),
                ],
                tools=[],
                params=AgentParams(model="provider-default", temperature=0.3, max_tokens=50),
            )
        ]

    assert captured_payloads[0]["messages"] == [
        {
            "role": "assistant",
            "content": "",
            "tool_calls": [assistant_tool_call],
        },
        {
            "role": "tool",
            "content": "noted",
            "tool_call_id": "call_1",
        },
    ]


@pytest.mark.asyncio
async def test_openai_chat_client_rejects_tool_message_without_tool_call_id_before_network():
    network_called = False

    def handler(request: httpx.Request) -> httpx.Response:
        nonlocal network_called
        network_called = True
        return httpx.Response(
            200,
            content=b'data: {"choices":[{"delta":{"content":"done"},"finish_reason":null}]}\n\n',
        )

    async with httpx.AsyncClient(
        transport=httpx.MockTransport(handler),
        base_url="https://llm.test/v1",
    ) as http_client:
        client = OpenAICompatibleChatClient(
            api_key="",
            base_url="https://llm.test/v1",
            default_model="provider-default",
            request_timeout_seconds=5,
            http_client=http_client,
        )

        with pytest.raises(ValueError, match="tool messages require tool_call_id"):
            [
                item
                async for item in client.stream_chat(
                    messages=[ChatMessage(role="tool", content="orphan tool result")],
                    tools=[],
                    params=AgentParams(model="provider-default", temperature=0.3, max_tokens=50),
                )
            ]

    assert network_called is False


def test_static_pre_message_role_selector_does_not_offer_tool():
    html = Path("src/agent_lab/presentation/static/index.html").read_text()

    assert '<option value="tool">tool</option>' not in html


def test_static_tools_ui_is_dynamic_and_not_hardcoded_to_handoff_note_checkbox():
    html = Path("src/agent_lab/presentation/static/index.html").read_text()
    js = Path("src/agent_lab/presentation/static/app.js").read_text()

    assert 'id="tool-handoff-note"' not in html
    assert "#tool-handoff-note" not in js
    assert 'id="tool-list"' in html
    assert '"/api/tools"' in js
    assert "fetch" in js


def test_static_prompt_workspace_controls_are_available():
    html = Path("src/agent_lab/presentation/static/index.html").read_text()

    assert 'id="prompt-set-name"' in html
    assert 'id="saved-prompt-sets"' in html
    assert 'id="save-prompt-set"' in html
    assert 'id="load-prompt-set"' in html
    assert 'id="delete-prompt-set"' in html


def test_static_prompt_workspace_is_opened_from_a_modal():
    html = Path("src/agent_lab/presentation/static/index.html").read_text()
    js = Path("src/agent_lab/presentation/static/app.js").read_text()

    assert 'id="open-prompt-config"' in html
    assert '<dialog id="prompt-dialog"' in html
    assert 'id="close-prompt-config"' in html
    assert "promptDialog.showModal()" in js
    assert "promptDialog.close()" in js


def test_static_agent_config_controls_are_aligned_for_chat_and_event_agents():
    html = Path("src/agent_lab/presentation/static/index.html").read_text()
    js = Path("src/agent_lab/presentation/static/app.js").read_text()

    for control_id in [
        "model",
        "temperature",
        "max-tokens",
        "chat-system-prompt",
        "chat-thinking-disabled",
        "chat-enable-search",
        "chat-forced-search",
        "event-model",
        "event-temperature",
        "event-max-tokens",
        "event-system-prompt",
        "event-thinking-disabled",
        "event-enable-search",
        "event-forced-search",
    ]:
        assert f'id="{control_id}"' in html

    assert 'extra_body: buildExtraBody("chat")' in js
    assert 'extra_body: buildExtraBody("event")' in js
    assert 'system_prompt: document.querySelector("#chat-system-prompt")' in js
    assert 'system_prompt: document.querySelector("#event-system-prompt")' in js


def test_static_tools_ui_shows_tool_metadata_and_selection_actions():
    html = Path("src/agent_lab/presentation/static/index.html").read_text()
    js = Path("src/agent_lab/presentation/static/app.js").read_text()

    assert 'id="tool-count"' in html
    assert 'id="select-all-tools"' in html
    assert 'id="clear-tools"' in html
    assert "tool-card" in js
    assert "tool-description" in js
    assert "tool-required" in js
    assert "function updateToolCount()" in js


def test_static_app_handles_audit_messages():
    js = Path("src/agent_lab/presentation/static/app.js").read_text()
    css = Path("src/agent_lab/presentation/static/styles.css").read_text()

    assert 'message.type === "audit"' in js
    assert 'appendLog("audit"' in js
    assert ".audit" in css


def test_static_app_shows_wait_state_and_immediate_user_echo():
    html = Path("src/agent_lab/presentation/static/index.html").read_text()
    js = Path("src/agent_lab/presentation/static/app.js").read_text()
    css = Path("src/agent_lab/presentation/static/styles.css").read_text()

    assert 'id="run-button"' in html
    assert 'const runButton = document.querySelector("#run-button")' in js
    assert 'appendLog("user", submittedMessage)' in js
    assert 'statusEl.textContent = "Waiting for first token"' in js
    assert "function startElapsedTimer()" in js
    assert "function stopElapsedTimer()" in js
    assert "function setRunState(" in js
    assert "if (firstTokenAt) {\n    hasBackendRoundStats = true;" in js
    assert "hasBackendRoundStats && firstTokenAt" in js
    assert ".user" in css


def test_static_prompt_workspace_uses_stable_storage_hooks():
    js = Path("src/agent_lab/presentation/static/app.js").read_text()

    assert "agent-lab.prompt-sets.v1" in js
    assert "function savePromptSet()" in js
    assert "function loadPromptSet()" in js
    assert "function deletePromptSet()" in js


def test_static_app_handles_backend_round_stats_messages():
    js = Path("src/agent_lab/presentation/static/app.js").read_text()

    assert 'message.type === "round_stats"' in js
    assert "function updateRoundStats(" in js
    assert "#stat-ttft" in js
    assert "#stat-elapsed" in js


@pytest.mark.asyncio
async def test_openai_chat_client_omits_stream_options_when_usage_disabled():
    captured_payloads: list[dict] = []

    def handler(request: httpx.Request) -> httpx.Response:
        captured_payloads.append(json.loads(request.content))
        return httpx.Response(
            200,
            content=b'data: {"choices":[{"delta":{"content":"hi"},"finish_reason":null}]}\n\n',
        )

    async with httpx.AsyncClient(
        transport=httpx.MockTransport(handler),
        base_url="https://llm.test/v1",
    ) as http_client:
        client = OpenAICompatibleChatClient(
            api_key="",
            base_url="https://llm.test/v1",
            default_model="provider-default",
            request_timeout_seconds=5,
            include_usage=False,
            http_client=http_client,
        )
        [
            item
            async for item in client.stream_chat(
                messages=[ChatMessage(role="user", content="hi")],
                tools=[],
                params=AgentParams(model="provider-default", temperature=0.3, max_tokens=50),
            )
        ]

    assert "stream_options" not in captured_payloads[0]


@pytest.mark.asyncio
async def test_openai_chat_client_merges_agent_extra_body_into_payload():
    captured_payloads: list[dict] = []

    def handler(request: httpx.Request) -> httpx.Response:
        captured_payloads.append(json.loads(request.content))
        return httpx.Response(
            200,
            content=b'data: {"choices":[{"delta":{"content":"hi"},"finish_reason":null}]}\n\n',
        )

    async with httpx.AsyncClient(
        transport=httpx.MockTransport(handler),
        base_url="https://llm.test/v1",
    ) as http_client:
        client = OpenAICompatibleChatClient(
            api_key="",
            base_url="https://llm.test/v1",
            default_model="provider-default",
            request_timeout_seconds=5,
            http_client=http_client,
        )
        [
            item
            async for item in client.stream_chat(
                messages=[ChatMessage(role="user", content="hi")],
                tools=[],
                params=AgentParams(
                    model="provider-default",
                    temperature=0.3,
                    max_tokens=50,
                    extra_body={
                        "thinking": {"type": "disabled"},
                        "enable_search": False,
                        "search_options": {"forced_search": False},
                    },
                ),
            )
        ]

    assert captured_payloads[0]["thinking"] == {"type": "disabled"}
    assert captured_payloads[0]["enable_search"] is False
    assert captured_payloads[0]["search_options"] == {"forced_search": False}