kamidso-back/tests/kamidso/test_concurrent_runs.py

"""Different conversations run in parallel; the same conversation stays serial."""

from __future__ import annotations

import asyncio
from types import SimpleNamespace
from typing import Any, AsyncIterator, Dict, List
from unittest.mock import AsyncMock

import pytest

from kamidso.application.orchestrator import KamidsoOrchestrator
from kamidso.application.run_registry import ActiveRunRegistry
from kamidso.contracts.chat import KamidsoChatRequest
from kamidso.contracts.errors import KamidsoError, KamidsoErrorCode
from kamidso.contracts.run import KamidsoRun
from kamidso.contracts.user import CurrentUser
from kamidso.data_graph.explorer.agent import DataMapperAgent, DataMapperMode
from kamidso.data_graph.llm import bind_llm_conv_uid, reset_llm_conv_uid
from kamidso.engine.app.openapi.api_v1 import agentic_data_api
from kamidso.integrations.engine.engine import KamidsoEngine


USER = CurrentUser(
    id="u1",
    email="a@example.com",
    name="A",
    role="owner",
    workspace_id="ws1",
)


def _request(conversation_id: str, message: str = "hello") -> KamidsoChatRequest:
    return KamidsoChatRequest.from_body(
        {
            "conversation_id": conversation_id,
            "message": message,
            "resources": [],
            "mode": "auto",
            "temperature": 0.6,
            "max_new_tokens": 4000,
            "reasoning_mode": "fast",
        }
    )


@pytest.mark.asyncio
async def test_chat_rejects_second_message_while_run_is_live():
    active = ActiveRunRegistry()
    live = KamidsoRun(
        id="run_live",
        conversation_id="conv-a",
        workspace_id="ws1",
        user_id="u1",
        status="running",
        input="first",
    )
    active.register(live)

    runs = SimpleNamespace(
        active_run_id=active.active_run_id,
        cancel_active_for_conversation=AsyncMock(),
        create_run=AsyncMock(),
        active_count=lambda: 1,
    )
    conversations = SimpleNamespace(
        get_or_create=AsyncMock(return_value=SimpleNamespace(id="conv-a")),
        add_message=AsyncMock(),
        replace_resources=AsyncMock(),
    )
    orchestrator = KamidsoOrchestrator(
        engines=SimpleNamespace(default=lambda: None),
        resources=SimpleNamespace(authorize_many=AsyncMock()),
        conversations=conversations,
        runs=runs,
    )

    with pytest.raises(KamidsoError) as exc_info:
        async for _ in orchestrator.chat(USER, _request("conv-a", "second")):
            pass

    assert exc_info.value.code == KamidsoErrorCode.VALIDATION_ERROR
    assert exc_info.value.run_id == "run_live"
    runs.cancel_active_for_conversation.assert_not_called()
    runs.create_run.assert_not_called()
    conversations.add_message.assert_not_called()


@pytest.mark.asyncio
async def test_chat_allows_new_run_on_different_conversation():
    active = ActiveRunRegistry()
    live = KamidsoRun(
        id="run_a",
        conversation_id="conv-a",
        workspace_id="ws1",
        user_id="u1",
        status="running",
        input="first",
    )
    active.register(live)

    created: List[str] = []

    async def create_run(user, conversation_id, request):
        run = KamidsoRun(
            id=f"run_{conversation_id}",
            conversation_id=conversation_id,
            workspace_id=user.workspace_id,
            user_id=user.id,
            status="created",
            input=request.message,
        )
        created.append(conversation_id)
        active.register(run)
        return run

    async def persist_event(event, _workspace_id):
        return SimpleNamespace(sequence=1)

    async def update_from_event(run, event):
        if event.get("type") == "run.completed":
            run.status = "completed"
            run.answer = str(event.get("answer") or "")
        return run

    runs = SimpleNamespace(
        active_run_id=active.active_run_id,
        cancel_active_for_conversation=AsyncMock(),
        create_run=create_run,
        active_count=lambda: active.active_count(),
        materialize_resources=AsyncMock(return_value=[]),
        persist_event=persist_event,
        update_from_event=update_from_event,
        get_run=AsyncMock(return_value=None),
    )
    conversations = SimpleNamespace(
        get_or_create=AsyncMock(side_effect=lambda _user, cid: SimpleNamespace(id=cid or "new")),
        add_message=AsyncMock(return_value=SimpleNamespace(id="msg1")),
        replace_resources=AsyncMock(),
        list_messages=AsyncMock(return_value=[]),
    )

    class _FakeEngine(KamidsoEngine):
        def __init__(self) -> None:
            pass

        async def execute(self, **kwargs) -> AsyncIterator[Dict[str, Any]]:
            yield {
                "type": "run.completed",
                "run_id": kwargs["run"].id,
                "answer": "ok",
            }

    orchestrator = KamidsoOrchestrator(
        engines=SimpleNamespace(default=lambda: _FakeEngine()),
        resources=SimpleNamespace(authorize_many=AsyncMock()),
        conversations=conversations,
        runs=runs,
    )

    events: List[str] = []
    async for chunk in orchestrator.chat(USER, _request("conv-b", "parallel")):
        events.append(chunk)

    assert created == ["conv-b"]
    runs.cancel_active_for_conversation.assert_not_called()
    assert any("run.created" in chunk for chunk in events)
    assert active.active_run_id("conv-a") == "run_a"


@pytest.mark.asyncio
async def test_second_stream_does_not_cancel_other_conversation(monkeypatch) -> None:
    cancelled: set[str] = set()
    started = asyncio.Event()
    started_ids: set[str] = set()
    gate = asyncio.Event()

    async def _fake_stream_impl(
        dialogue, tool_mode="full", attachment_ctx=None, agent_task_holder=None
    ):
        del tool_mode, attachment_ctx

        async def work() -> None:
            try:
                await gate.wait()
            except asyncio.CancelledError:
                cancelled.add(dialogue.conv_uid)
                raise

        task = asyncio.create_task(work())
        if agent_task_holder is not None:
            agent_task_holder.append(task)
        agentic_data_api.REACT_CONVERSATION_AGENT_TASKS[str(dialogue.conv_uid)] = task
        started_ids.add(dialogue.conv_uid)
        if len(started_ids) >= 2:
            started.set()
        yield f"data: start-{dialogue.conv_uid}\n\n"
        try:
            await task
        except asyncio.CancelledError:
            return
        yield f"data: end-{dialogue.conv_uid}\n\n"

    monkeypatch.setattr(agentic_data_api, "_react_agent_stream_impl", _fake_stream_impl)

    async def consume(conv_id: str):
        events = []
        async for event in agentic_data_api._react_agent_stream_inner(
            SimpleNamespace(conv_uid=conv_id, ext_info={})
        ):
            events.append(event)
        return events

    first = asyncio.create_task(consume("conv-a"))
    second = asyncio.create_task(consume("conv-b"))
    await asyncio.wait_for(started.wait(), timeout=2)
    assert cancelled == set()
    gate.set()
    await asyncio.gather(first, second)
    assert cancelled == set()


def test_data_mapper_agents_use_distinct_conv_ids():
    one = DataMapperAgent(DataMapperMode.STRUCTURE, tools=[])
    two = DataMapperAgent(DataMapperMode.STRUCTURE, tools=[])
    assert one.conv_id != two.conv_id
    assert one.conv_id.startswith("data-mapper-")
    fixed = DataMapperAgent(DataMapperMode.STRUCTURE, tools=[], conv_id="data-mapper-abc")
    assert fixed.conv_id == "data-mapper-abc"


def test_bind_llm_conv_uid_scopes_context():
    from kamidso.data_graph import llm as llm_mod

    token = bind_llm_conv_uid("data-mapper-xyz")
    try:
        assert llm_mod._llm_conv_uid.get() == "data-mapper-xyz"
    finally:
        reset_llm_conv_uid(token)
    assert llm_mod._llm_conv_uid.get() is None