ZNJJ-api-server/test/services/test_stream_and_session_safety.py

import asyncio

import pytest

from src.agent.collect import CollectReplyError, validated_reply
from src.agent.llm import FakeLLMService
from src.backends.chat import ChatInput, ChatResult, TextDelta
from src.backends.langgraph import LangGraphBackend
from src.services.app_service import AccidentAppService
from src.services.repository import InMemorySessionRepository


class ValidatingGraph:
    def __init__(self):
        self.committed = None

    async def ainvoke(self, data, *, config, context):
        raw, stage = await validated_reply(context.llm, 'system', 'user', '1002')
        self.committed = stage
        return {'response_text': raw, 'turn_patch': {}}


class ScriptedStream:
    def __init__(self, replies):
        self.replies = iter(replies)
        self.calls = 0

    async def complete(self, *args, **kwargs):
        return '{}'

    async def stream_tokens(self, *args, **kwargs):
        self.calls += 1
        for chunk in next(self.replies):
            yield chunk


@pytest.mark.asyncio
async def test_invalid_prefix_is_not_sent_and_repair_is_the_only_reply():
    graph = ValidatingGraph()
    llm = ScriptedStream([
        ['<state>9999</state>invalid'],
        ['<sta', 'te>1002</state>', 'repaired'],
    ])
    backend = LangGraphBackend(graph=graph, llm=llm)
    events = [e async for e in backend.stream(ChatInput('s', 'hello'))]
    assert ''.join(e.text for e in events if isinstance(e, TextDelta)) == '<state>1002</state>repaired'
    assert graph.committed == '1002'
    assert llm.calls == 2


@pytest.mark.asyncio
async def test_failure_after_output_aborts_without_appending_a_repair():
    graph = ValidatingGraph()
    llm = ScriptedStream([
        ['<state>1002</state>', 'partial', '<state>9999</state>'],
        ['<state>1002</state>must not appear'],
    ])
    backend = LangGraphBackend(graph=graph, llm=llm)
    events = []
    with pytest.raises(CollectReplyError):
        async for event in backend.stream(ChatInput('s', 'hello')):
            events.append(event)
    assert llm.calls == 1
    assert graph.committed is None
    assert 'must not appear' not in ''.join(e.text for e in events)


@pytest.mark.asyncio
async def test_repair_exhaustion_emits_no_invalid_text():
    llm = ScriptedStream([['<state>9999</state>x'], ['malformed']])
    backend = LangGraphBackend(graph=ValidatingGraph(), llm=llm)
    events = []
    with pytest.raises(CollectReplyError):
        async for event in backend.stream(ChatInput('s', 'hello')):
            events.append(event)
    assert events == []


class BlockingBackend:
    def __init__(self):
        self.entered = asyncio.Queue()
        self.release = asyncio.Event()
        self.active = 0
        self.peak = 0

    async def complete(self, request):
        self.active += 1
        self.peak = max(self.peak, self.active)
        await self.entered.put(request.text)
        try:
            await self.release.wait()
            return ChatResult('<state>1002</state>ok', form_update={'vehicle_count': 1})
        finally:
            self.active -= 1


@pytest.mark.asyncio
async def test_delete_waiter_and_new_chat_keep_one_lock():
    backend = BlockingBackend()
    service = AccidentAppService(backend, InMemorySessionRepository())
    dropping, drop_release = asyncio.Event(), asyncio.Event()

    async def drop(_):
        dropping.set()
        await drop_release.wait()

    deletion = asyncio.create_task(service.delete_session('s', drop))
    await dropping.wait()
    first = asyncio.create_task(service.chat(ChatInput('s', 'a')))
    await asyncio.sleep(0)
    drop_release.set()
    await deletion
    assert await asyncio.wait_for(backend.entered.get(), 1) == 'a'
    second = asyncio.create_task(service.chat(ChatInput('s', 'b')))
    await asyncio.sleep(0)
    assert backend.entered.empty()
    backend.release.set()
    await asyncio.wait_for(asyncio.gather(first, second), 2)
    assert backend.peak == 1
    assert not service._locks._entries


@pytest.mark.asyncio
async def test_set_waits_for_chat_and_cannot_be_overwritten_by_old_result():
    backend = BlockingBackend()
    service = AccidentAppService(backend, InMemorySessionRepository())
    chat = asyncio.create_task(service.chat(ChatInput('s', 'a')))
    await backend.entered.get()
    setting = asyncio.create_task(service.set_info('s', 'jdcsl', '2'))
    await asyncio.sleep(0)
    assert not setting.done()
    backend.release.set()
    await asyncio.gather(chat, setting)
    assert await service.get_info('s', 'jdcsl') == '"2"'


@pytest.mark.asyncio
async def test_cancelled_waiter_does_not_leak_or_replace_lock():
    backend = BlockingBackend()
    service = AccidentAppService(backend, InMemorySessionRepository())
    chat = asyncio.create_task(service.chat(ChatInput('s', 'a')))
    await backend.entered.get()
    waiting = asyncio.create_task(service.get_info('s', 'jdcsl'))
    await asyncio.sleep(0)
    waiting.cancel()
    with pytest.raises(asyncio.CancelledError):
        await waiting
    assert len(service._locks._entries) == 1
    backend.release.set()
    await chat
    assert not service._locks._entries