simbarag/tests/unit/test_stream_stall.py

"""Tests for the stream-query stall timeout and per-call model metrics."""

import asyncio
import json
import logging

import pytest
from langchain_core.messages import AIMessage, AIMessageChunk
from quart import Quart
from quart_jwt_extended import JWTManager, create_access_token
from tortoise import Tortoise

import blueprints.conversation as conv
from blueprints.conversation.models import Conversation
from blueprints.users.models import User

# Shortened stall limit for the tests (seconds).
LIMIT = 0.3


@pytest.fixture
async def owner():
    await Tortoise.init(
        db_url="sqlite://:memory:",
        modules={
            "models": [
                "blueprints.users.models",
                "blueprints.conversation.models",
                "blueprints.conversation.run_models",
                "blueprints.artifacts.models",
            ]
        },
    )
    await Tortoise.generate_schemas()
    user = await User.create(username="owner", email="owner@example.com")
    conversation = await Conversation.create(name="c", user=user)
    yield user, conversation
    await Tortoise.close_connections()


class _TimedAgent:
    """Yields (delay_seconds, event) pairs, sleeping before each event."""

    def __init__(self, script):
        self._script = script

    async def astream_events(self, payload, version, config):
        for delay, event in self._script:
            await asyncio.sleep(delay)
            yield event


def _start(run_id):
    return {"event": "on_chat_model_start", "run_id": run_id, "name": "llm", "data": {}}


def _chunk(run_id, text):
    return {
        "event": "on_chat_model_stream",
        "run_id": run_id,
        "data": {"chunk": AIMessageChunk(content=text)},
    }


def _end(run_id, prompt_tokens=None, cached_tokens=None):
    usage = (
        {
            "input_tokens": prompt_tokens,
            "output_tokens": 5,
            "total_tokens": prompt_tokens + 5,
        }
        if prompt_tokens
        else None
    )
    if usage and cached_tokens is not None:
        usage["input_token_details"] = {"cache_read": cached_tokens}
    return {
        "event": "on_chat_model_end",
        "run_id": run_id,
        "data": {"output": AIMessage(content="", usage_metadata=usage)},
    }


def _final(text):
    return {
        "event": "on_chain_end",
        "name": "agent",
        "data": {"output": {"messages": [AIMessage(content=text)]}},
    }


async def _run(owner, monkeypatch, script, *, response_timeout=60, direct=False):
    user, conversation = owner
    monkeypatch.setattr(conv, "MODEL_STALL_TIMEOUT", LIMIT)
    monkeypatch.setattr(conv.runtime, "get_agent_for", lambda u: _TimedAgent(script))

    async def _prompt(user_id, query=""):
        return "system"

    monkeypatch.setattr(conv.runtime, "build_system_prompt_with_memories", _prompt)

    app = Quart(__name__)
    app.config["RESPONSE_TIMEOUT"] = response_timeout
    app.config["JWT_SECRET_KEY"] = "test-secret"
    JWTManager(app)
    app.register_blueprint(conv.conversation_blueprint)
    async with app.app_context():
        token = create_access_token(identity=str(user.id))
    kwargs = {
        "json": {"query": "hi", "conversation_id": str(conversation.id)},
        "headers": {"Authorization": f"Bearer {token}"},
    }
    if direct:
        async with app.test_request_context(
            "/api/conversation/stream-query", method="POST", **kwargs
        ):
            resp = await app.full_dispatch_request()
            body = (await resp.get_data()).decode()
    else:
        resp = await app.test_client().post("/api/conversation/stream-query", **kwargs)
        body = (await resp.get_data()).decode()
    return [
        json.loads(line[len("data: ") :])
        for line in body.splitlines()
        if line.startswith("data: {")
    ]


def test_stall_timeout_only_while_a_model_call_is_in_flight(monkeypatch):
    monkeypatch.setattr(conv, "MODEL_STALL_TIMEOUT", 30.0)
    assert conv.stall_timeout(None, 100.0) is None
    assert conv.stall_timeout(100.0, 110.0) == 20.0
    assert conv.stall_timeout(100.0, 200.0) == 1.0  # floor


async def test_active_stream_outlasts_response_timeout(owner, monkeypatch):
    script = [
        (0, _start("m1")),
        (0, _chunk("m1", "<think>")),
        (0.15, _chunk("m1", "working")),
        (0.15, _chunk("m1", "</think>Answer")),
        (0, _end("m1")),
        (0, _final("Answer")),
    ]

    frames = await _run(owner, monkeypatch, script, response_timeout=0.1)

    assert "error" not in [f["type"] for f in frames]
    assert frames[-1] == {"type": "response", "message": "Answer"}


async def test_long_thinking_is_not_a_stall(owner, monkeypatch):
    # The whole call takes ~4x the limit, but output never pauses for longer
    # than half of it; the old time-to-first-content deadline killed this.
    script = [(0, _start("m1"))]
    script += [(LIMIT / 2, _chunk("m1", "<think>hmm ")) for _ in range(8)]
    script += [
        (0, _chunk("m1", "</think>Answer")),
        (0, _end("m1")),
        (0, _final("Answer")),
    ]

    frames = await _run(owner, monkeypatch, script)

    types = [f["type"] for f in frames]
    assert "error" not in types
    assert frames[-1] == {"type": "response", "message": "Answer"}


async def test_tool_time_between_model_calls_does_not_count(owner, monkeypatch):
    tool = {"event": "on_tool_start", "name": "web_search", "run_id": "t1", "data": {}}
    tool_end = {
        "event": "on_tool_end",
        "name": "web_search",
        "run_id": "t1",
        "data": {},
    }
    script = [
        (0, _start("m1")),
        (0, _chunk("m1", "")),
        (0, _end("m1")),
        (0, tool),
        (LIMIT * 3, tool_end),  # slow tool, no model call in flight
        (0, _start("m2")),
        (LIMIT / 2, _chunk("m2", "Done")),
        (0, _end("m2")),
        (0, _final("Done")),
    ]

    frames = await _run(owner, monkeypatch, script)

    assert "error" not in [f["type"] for f in frames]
    assert frames[-1]["message"] == "Done"


async def test_silent_model_call_is_a_stall(owner, monkeypatch):
    script = [
        (0, _start("m1")),
        (0, _chunk("m1", "<think>start")),
        (LIMIT * 4, _chunk("m1", "never arrives")),
    ]

    frames = await _run(owner, monkeypatch, script)

    assert frames[-1]["type"] == "error"
    assert frames[-1]["stalled"] is True


async def test_logs_per_call_prompt_tokens(owner, monkeypatch, caplog):
    script = [
        (0, _start("m1")),
        (0, _chunk("m1", "Hi")),
        (0, _end("m1", prompt_tokens=27262, cached_tokens=25000)),
        (0, _final("Hi")),
    ]

    with caplog.at_level(logging.INFO):
        await _run(owner, monkeypatch, script)

    summaries = [
        json.loads(r.getMessage())
        for r in caplog.records
        if r.getMessage().startswith('{"event": "stream_request"')
    ]
    assert summaries, "stream_request summary not logged"
    summary = summaries[-1]
    assert summary["max_prompt_tokens"] == 27262
    (call,) = summary["model_calls"]
    assert call["prompt_tokens"] == 27262
    assert call["cached_prompt_tokens"] == 25000
    assert call["prefill_ms"] is not None and call["ms"] is not None


async def test_token_counts_use_usage_not_chunk_count(owner, monkeypatch, caplog):
    script = [
        (0, _start("m1")),
        (0, _chunk("m1", "Hello there")),
        (0, _chunk("m1", " world again")),
        (0, _end("m1", prompt_tokens=100)),
        (0, _final("Hello there world again")),
    ]
    with caplog.at_level(logging.INFO):
        await _run(owner, monkeypatch, script)
    summary = [
        json.loads(r.getMessage())
        for r in caplog.records
        if r.getMessage().startswith('{"event": "stream_request"')
    ][-1]
    assert summary["tokens"] == 5
    assert summary["content_chunks"] == 2
    assert summary["usage_complete"] is True
    assert "thinking_ms" not in summary
    assert summary["model_to_first_content_ms"] is not None


async def test_stalled_call_reports_partial_usage_and_outcome(
    owner, monkeypatch, caplog
):
    results = []
    durations = []
    monkeypatch.setattr(
        conv, "count_result", lambda **kwargs: results.append(kwargs["result"])
    )
    monkeypatch.setattr(
        conv, "observe_response", lambda **kwargs: durations.append(kwargs["seconds"])
    )
    script = [
        (0, _start("m1")),
        (0, _chunk("m1", "Hi")),
        (0, _end("m1", prompt_tokens=100)),
        (0, _start("m2")),
        (LIMIT * 4, _chunk("m2", "late")),
    ]
    with caplog.at_level(logging.INFO):
        await _run(owner, monkeypatch, script)
    summary = [
        json.loads(r.getMessage())
        for r in caplog.records
        if r.getMessage().startswith('{"event": "stream_request"')
    ][-1]
    assert summary["tokens"] is None
    assert summary["completed_output_tokens"] == 5
    assert summary["usage_complete"] is False
    assert results == ["stalled"]
    assert len(durations) == 1


async def test_no_response_is_counted_as_error(owner, monkeypatch):
    results = []
    monkeypatch.setattr(
        conv, "count_result", lambda **kwargs: results.append(kwargs["result"])
    )
    frames = await _run(
        owner, monkeypatch, [(0, _start("m1")), (0, _end("m1", prompt_tokens=100))]
    )
    assert frames[-1]["type"] == "error"
    assert results == ["error"]


@pytest.mark.parametrize("cancelled", [False, True])
async def test_final_persistence_exit_has_exactly_one_outcome(
    owner, monkeypatch, caplog, cancelled
):
    # Answers are persisted by the detached run worker, so the request only
    # observes the terminal event; cancellation must not escape the worker.
    from blueprints.conversation import runs

    results, durations, writes = [], [], []
    original = runs.add_message_to_conversation

    async def persist(conversation, message, speaker, *args, **kwargs):
        if speaker == "simba":
            writes.append(message)
            if cancelled:
                raise asyncio.CancelledError()
            raise RuntimeError("persistence failed")
        return await original(conversation, message, speaker, *args, **kwargs)

    monkeypatch.setattr(runs, "add_message_to_conversation", persist)
    monkeypatch.setattr(
        conv, "count_result", lambda **kwargs: results.append(kwargs["result"])
    )
    monkeypatch.setattr(
        conv, "observe_response", lambda **kwargs: durations.append(kwargs["seconds"])
    )
    script = [
        (0, _start("m1")),
        (0, _chunk("m1", "Answer")),
        (0, _end("m1", prompt_tokens=100)),
        (0, _final("Answer")),
    ]
    with caplog.at_level(logging.INFO):
        frames = await _run(owner, monkeypatch, script, direct=True)
    assert frames[-1]["type"] == "error"
    assert len(writes) == 1
    assert results == ["cancelled" if cancelled else "error"]
    assert len(durations) == 1
    summaries = [
        json.loads(r.getMessage())
        for r in caplog.records
        if r.getMessage().startswith('{"event": "stream_request"')
    ]
    assert len(summaries) == 1
    assert summaries[0]["cancelled"] is cancelled