simbarag/tests/unit/test_chat_runs.py

import asyncio
import uuid
from datetime import timedelta

import pytest
import pytest_asyncio
from quart import Quart
from tortoise import Tortoise

from blueprints.conversation import runs
from blueprints.conversation.models import Conversation, ConversationMessage
from blueprints.conversation.run_models import ChatRun, ChatRunEvent
from blueprints.users.models import User


@pytest_asyncio.fixture
async def db(tmp_path):
    await Tortoise.init(
        db_url=f"sqlite://{tmp_path / 'runs.sqlite'}",
        modules={
            "models": [
                "blueprints.conversation.models",
                "blueprints.conversation.run_models",
                "blueprints.users.models",
            ]
        },
    )
    await Tortoise.generate_schemas()
    app = Quart(__name__)
    app.config["RUN_TEST_DB_URL"] = f"sqlite://{tmp_path / 'runs.sqlite'}"
    user = await User.create(username="owner", email="owner@test")
    other = await User.create(username="other", email="other@test")
    conv = await Conversation.create(name="Chat", user=user)
    yield app, user, other, conv
    await runs.drain_runs(app, 0)
    runs._draining_apps.discard(id(app))
    await Tortoise.close_connections()


async def reserved(user, conv, **kwargs):
    return await ChatRun.create(
        owner=user,
        conversation=conv,
        request_key=str(uuid.uuid4()),
        payload={"query": "hello"},
        active_key=str(conv.id),
        worker_token=uuid.uuid4(),
        lease_until=runs.now() + timedelta(seconds=30),
        **kwargs,
    )


async def wait_terminal(run):
    for _ in range(200):
        run = await ChatRun.get(id=run.id)
        if run.status != "running":
            return run
        await asyncio.sleep(0.01)
    pytest.fail("run did not terminate")


async def test_order_cursor_isolation_and_partial_persisted_once(db):
    app, user, other, conv = db
    run = await reserved(user, conv)
    await runs.append(run.id, run.worker_token, {"type": "content", "delta": "one"})
    await runs.append(run.id, run.worker_token, {"type": "content", "delta": "two"})
    await runs.finish(run.id, "completed")
    await runs.finish(run.id, "cancelled")
    events = await ChatRunEvent.filter(run=run).order_by("sequence")
    assert [e.sequence for e in events] == [1, 2, 3]
    assert (
        await ConversationMessage.filter(conversation=conv, speaker="simba").count()
        == 1
    )
    assert (await ConversationMessage.get(conversation=conv)).text == "onetwo"
    assert await runs.owned_run(other.id, conv.id, run.id) is None
    foreign = await Conversation.create(name="other", user=user)
    assert await runs.owned_run(user.id, foreign.id, run.id) is None
    async with app.app_context():
        response = await runs.replay_response(run, 1)
        text = (await response.get_data()).decode()
    assert "id: 2" in text and "id: 1" not in text
    assert text.index("id: 2") < text.index("[DONE]")


async def test_atomic_idempotency_and_concurrent_slot(db, monkeypatch):
    app, user, other, conv = db
    gate = asyncio.Event()

    async def fake(*args):
        await gate.wait()

    monkeypatch.setattr(runs, "_work", fake)
    data = {"conversation_id": str(conv.id), "query": "hi", "request_id": "same"}
    results = await asyncio.gather(
        *(runs.create_run(app, user.id, data) for _ in range(4))
    )
    assert len({str(run.id) for run, _ in results}) == 1
    assert sum(created for _, created in results) == 1
    assert await ChatRun.all().count() == 1
    with pytest.raises(runs.RunConflict, match="different input"):
        await runs.create_run(app, user.id, {**data, "query": "changed"})
    with pytest.raises(runs.RunConflict, match="already active"):
        await runs.create_run(app, user.id, {**data, "request_id": "another"})
    with pytest.raises(runs.RunConflict) as err:
        await runs.create_run(app, other.id, data)
    assert err.value.status == 404
    gate.set()
    await asyncio.sleep(0)


async def test_expired_restart_never_regenerates_and_releases_slot(db):
    app, user, _, conv = db
    run = await reserved(user, conv)
    await runs.append(run.id, run.worker_token, {"type": "content", "delta": "partial"})
    await ChatRun.filter(id=run.id).update(
        lease_until=runs.now() - timedelta(seconds=1)
    )
    # Close/reopen the real file database, discarding process-local ORM state.
    await Tortoise.close_connections()
    await Tortoise.init(
        db_url=app.config["RUN_TEST_DB_URL"],
        modules={
            "models": [
                "blueprints.conversation.models",
                "blueprints.conversation.run_models",
                "blueprints.users.models",
            ]
        },
    )
    await runs.expire_runs(str(conv.id))
    await runs.expire_runs(str(conv.id))
    latest = await ChatRun.get(id=run.id)
    assert latest.status == "interrupted" and latest.active_key is None
    assert await ConversationMessage.filter(conversation=conv).count() == 1
    assert (await ChatRunEvent.filter(run=run).order_by("sequence"))[1].payload[
        "type"
    ] == "error"
    with pytest.raises(asyncio.CancelledError):
        await runs.append(
            run.id, run.worker_token, {"type": "content", "delta": "zombie"}
        )


async def test_background_context_detach_and_cross_process_cancel(db, monkeypatch):
    app, user, _, conv = db
    gate = asyncio.Event()
    started = asyncio.Event()
    from quart import current_app
    import blueprints.conversation as conversation

    async def fake(data, user, conv, save_user, save_answer):
        assert current_app._get_current_object() is app
        await save_user("hi", None)
        yield 'data: {"type":"content","delta":"partial"}\n\n'
        started.set()
        await gate.wait()
        await save_answer("complete")
        yield 'data: {"type":"response","message":"complete"}\n\n'
        yield "data: [DONE]\n\n"

    monkeypatch.setattr(conversation, "_chat_run_stream", fake)
    run, _ = await runs.create_run(
        app,
        user.id,
        {"conversation_id": str(conv.id), "query": "hi", "request_id": "turn"},
    )
    await started.wait()
    async with app.app_context():
        response = await runs.replay_response(run, 0)
        subscriber = response.response.__aiter__()
        frame = ""
        while "partial" not in frame:  # skip keepalives until the delta is written
            frame = await asyncio.wait_for(subscriber.__anext__(), 2)
        await subscriber.aclose()  # browser disconnect
    assert (await ChatRun.get(id=run.id)).status == "running"
    # Same write made by another process, no local task cancellation API involved.
    await ChatRun.filter(id=run.id).update(cancel_requested=True)
    latest = await wait_terminal(run)
    assert latest.status == "cancelled"
    assert (
        await ConversationMessage.filter(conversation=conv, speaker="user").count() == 1
    )
    assert (
        await ConversationMessage.filter(conversation=conv, speaker="simba").count()
        == 1
    )
    assert (
        await ConversationMessage.get(conversation=conv, speaker="simba")
    ).text == "partial"


async def test_completed_answer_once_and_drain(db, monkeypatch):
    app, user, _, conv = db
    import blueprints.conversation as conversation

    async def fake(data, user, conv, save_user, save_answer):
        await save_user("hi", None)
        yield 'data: {"type":"content","delta":"answer"}\n\n'
        await save_answer("answer")
        yield 'data: {"type":"response","message":"answer"}\n\n'

    monkeypatch.setattr(conversation, "_chat_run_stream", fake)
    run, _ = await runs.create_run(
        app,
        user.id,
        {"conversation_id": str(conv.id), "query": "hi", "request_id": "turn"},
    )
    await runs.drain_runs(app, 2)
    latest = await ChatRun.get(id=run.id)
    assert latest.status == "completed"
    assert (
        await ConversationMessage.filter(conversation=conv, speaker="simba").count()
        == 1
    )
    with pytest.raises(runs.RunConflict) as err:
        await runs.create_run(
            app, user.id, {"conversation_id": str(conv.id), "query": "new"}
        )
    assert err.value.status == 503


async def test_authenticated_routes_validate_owner_conversation_cursor_and_cancel(db):
    from quart_jwt_extended import JWTManager, create_access_token

    app, owner, foreign, conv = db
    app.config.update(
        TESTING=True, JWT_SECRET_KEY="chat-runs-test-secret-at-least-32-characters"
    )
    JWTManager(app)
    app.register_blueprint(runs.runs_blueprint)
    async with app.app_context():
        headers = {
            "Authorization": "Bearer " + create_access_token(identity=str(owner.id))
        }
        other_headers = {
            "Authorization": "Bearer " + create_access_token(identity=str(foreign.id))
        }
    client = app.test_client()
    run = await reserved(owner, conv)
    base = f"/api/conversation/{conv.id}/runs/{run.id}"
    assert (await client.get(base + "/events")).status_code == 401
    assert (
        await client.get(base + "/events", headers=other_headers)
    ).status_code == 404
    assert (
        await client.post(base + "/cancel", headers=other_headers)
    ).status_code == 404
    assert not (await ChatRun.get(id=run.id)).cancel_requested
    for cursor in ("-1", "bad", "999"):
        assert (
            await client.get(base + f"/events?cursor={cursor}", headers=headers)
        ).status_code == 400
    assert (await client.post(base + "/cancel", headers=headers)).status_code == 200
    assert (await ChatRun.get(id=run.id)).cancel_requested
    assert (
        await client.get(f"/api/conversation/{conv.id}/runs", headers=other_headers)
    ).status_code == 404
    discovery = await (
        await client.get(f"/api/conversation/{conv.id}/runs", headers=headers)
    ).get_json()
    assert discovery["run"]["id"] == str(run.id)
    assert (
        await client.post(
            "/api/conversation/runs",
            headers=other_headers,
            json={"conversation_id": str(conv.id), "query": "hi"},
        )
    ).status_code == 404
    assert (
        await client.post("/api/conversation/runs", headers=headers, json=[])
    ).status_code == 400


async def test_real_shared_runtime_legacy_endpoint_persists_and_replays(
    db, monkeypatch
):
    from quart_jwt_extended import JWTManager, create_access_token
    from langchain_core.messages import AIMessageChunk, AIMessage
    import blueprints.conversation as conversation

    app, owner, _, conv = db

    class Agent:
        async def astream_events(self, payload, version, config):
            assert config["configurable"]["user_id"] == str(owner.id)
            yield {"event": "on_chat_model_start", "run_id": "llm", "name": "model"}
            yield {
                "event": "on_chat_model_stream",
                "run_id": "llm",
                "data": {"chunk": AIMessageChunk(content="answer")},
            }
            yield {
                "event": "on_chat_model_end",
                "run_id": "llm",
                "data": {"output": AIMessage(content="answer")},
            }
            yield {
                "event": "on_chain_end",
                "data": {"output": {"messages": [AIMessage(content="answer")]}},
            }

    app.main_agent = Agent()
    from unittest.mock import AsyncMock

    monkeypatch.setattr(
        conversation.runtime,
        "build_system_prompt_with_memories",
        AsyncMock(return_value="system"),
    )
    app.config.update(
        TESTING=True, JWT_SECRET_KEY="chat-runs-test-secret-at-least-32-characters"
    )
    JWTManager(app)
    app.register_blueprint(conversation.conversation_blueprint)
    app.register_blueprint(runs.runs_blueprint)
    async with app.app_context():
        headers = {
            "Authorization": "Bearer " + create_access_token(identity=str(owner.id))
        }
    response = await app.test_client().post(
        "/api/conversation/stream-query",
        headers=headers,
        json={"conversation_id": str(conv.id), "query": "hi", "request_id": "legacy"},
    )
    body = (await response.get_data()).decode()
    assert (
        response.status_code == 200
        and '"type": "response"' in body
        and "[DONE]" in body
    )
    run = await ChatRun.get(request_key="legacy")
    assert run.status == "completed"
    assert (
        await ConversationMessage.filter(conversation=conv, speaker="simba").count()
        == 1
    )
    replay = await app.test_client().get(
        f"/api/conversation/{conv.id}/runs/{run.id}/events?cursor=0", headers=headers
    )
    assert (await replay.get_data()).decode() == body.replace(": keepalive\n\n", "")


async def test_nonstream_query_validates_input_and_respects_active_run(db):
    from quart_jwt_extended import JWTManager, create_access_token

    from blueprints.conversation import conversation_blueprint

    app, owner, _, conv = db
    app.config.update(
        TESTING=True, JWT_SECRET_KEY="chat-runs-test-secret-at-least-32-characters"
    )
    JWTManager(app)
    app.register_blueprint(conversation_blueprint)
    async with app.app_context():
        headers = {
            "Authorization": "Bearer " + create_access_token(identity=str(owner.id))
        }
    client = app.test_client()
    for body in (
        ["not", "an", "object"],
        {"conversation_id": str(conv.id), "query": 123},
        {"conversation_id": str(conv.id), "query": "   "},
    ):
        resp = await client.post("/api/conversation/query", headers=headers, json=body)
        assert resp.status_code == 400, body
    resp = await client.post(
        "/api/conversation/query",
        headers=headers,
        json={"conversation_id": ["x"], "query": "hi"},
    )
    assert resp.status_code == 404

    # A durable run owns the checkpoint thread: reject before saving anything.
    await reserved(owner, conv)
    resp = await client.post(
        "/api/conversation/query",
        headers=headers,
        json={"conversation_id": str(conv.id), "query": "hi"},
    )
    assert resp.status_code == 409
    assert await ConversationMessage.filter(conversation=conv).count() == 0


async def test_deltas_coalesce_and_local_subscribers_wake_immediately(db, monkeypatch):
    app, user, _, conv = db
    import blueprints.conversation as conversation

    release = asyncio.Event()

    async def fake(data, user, conv, save_user, save_answer):
        await save_user("hi", None)
        for word in ("one ", "two ", "three"):
            yield f'data: {{"type":"content","delta":"{word}"}}\n\n'
        yield 'data: {"type":"tool_start","tool":"t"}\n\n'
        await release.wait()
        await save_answer("one two three")
        yield 'data: {"type":"response","message":"one two three"}\n\n'

    monkeypatch.setattr(conversation, "_chat_run_stream", fake)
    monkeypatch.setattr(runs, "POLL_SECONDS", 30)  # only a wakeup can be fast
    run, _ = await runs.create_run(
        app, user.id, {"conversation_id": str(conv.id), "query": "hi"}
    )
    async with app.app_context():
        response = await runs.replay_response(run, 0)
        frames = response.response.__aiter__()
        seen = []
        while not any('"tool_start"' in f for f in seen):
            seen.append(await asyncio.wait_for(frames.__anext__(), 2))
        release.set()
        while "[DONE]" not in seen[-1]:
            seen.append(await asyncio.wait_for(frames.__anext__(), 2))
    deltas = [
        e.payload["delta"]
        for e in await ChatRunEvent.filter(run=run).order_by("sequence")
        if e.payload["type"] == "content"
    ]
    # First delta is written immediately; the burst behind it shares one write.
    assert "".join(deltas) == "one two three"
    assert len(deltas) == 2
    await wait_terminal(run)


async def _content_deltas(run):
    return [
        e.payload["delta"]
        for e in await ChatRunEvent.filter(run=run).order_by("sequence")
        if e.payload["type"] == "content"
    ]


async def test_buffered_delta_flushes_while_model_is_silent(db, monkeypatch):
    app, user, _, conv = db
    import blueprints.conversation as conversation

    gate = asyncio.Event()

    async def fake(data, user, conv, save_user, save_answer):
        await save_user("hi", None)
        yield 'data: {"type":"content","delta":"a"}\n\n'
        yield 'data: {"type":"content","delta":"b"}\n\n'
        await gate.wait()  # e.g. the model is writing tool-call arguments
        yield 'data: {"type":"response","message":"ab"}\n\n'

    monkeypatch.setattr(conversation, "_chat_run_stream", fake)
    monkeypatch.setattr(runs, "FLUSH_SECONDS", 0.05)
    run, _ = await runs.create_run(
        app, user.id, {"conversation_id": str(conv.id), "query": "hi"}
    )
    for _ in range(100):
        if "".join(await _content_deltas(run)) == "ab":
            break
        await asyncio.sleep(0.01)
    assert "".join(await _content_deltas(run)) == "ab"
    gate.set()
    await wait_terminal(run)


async def test_stop_keeps_buffered_text_in_saved_answer(db, monkeypatch):
    app, user, _, conv = db
    import blueprints.conversation as conversation

    buffered = asyncio.Event()

    async def fake(data, user, conv, save_user, save_answer):
        await save_user("hi", None)
        yield 'data: {"type":"content","delta":"first "}\n\n'
        yield 'data: {"type":"content","delta":"tail"}\n\n'
        buffered.set()
        await asyncio.Event().wait()

    monkeypatch.setattr(conversation, "_chat_run_stream", fake)
    monkeypatch.setattr(runs, "FLUSH_SECONDS", 60)  # "tail" stays buffered
    run, _ = await runs.create_run(
        app, user.id, {"conversation_id": str(conv.id), "query": "hi"}
    )
    await buffered.wait()
    await ChatRun.filter(id=run.id).update(cancel_requested=True)
    assert (await wait_terminal(run)).status == "cancelled"
    answer = await ConversationMessage.get(conversation=conv, speaker="simba")
    assert answer.text == "first tail"