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.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):
    usage = (
        {
            "input_tokens": prompt_tokens,
            "output_tokens": 5,
            "total_tokens": prompt_tokens + 5,
        }
        if prompt_tokens
        else None
    )
    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):
    user, conversation = owner
    monkeypatch.setattr(conv, "MODEL_STALL_TIMEOUT", LIMIT)
    monkeypatch.setattr(conv, "get_agent_for", lambda u: _TimedAgent(script))

    async def _prompt(user_id):
        return "system"

    monkeypatch.setattr(conv, "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))
    resp = await app.test_client().post(
        "/api/conversation/stream-query",
        json={"query": "hi", "conversation_id": str(conversation.id)},
        headers={"Authorization": f"Bearer {token}"},
    )
    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)),
        (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["prefill_ms"] is not None and call["ms"] is not None