"""Tests for chat prompt concurrency and cancellation contracts.""" from __future__ import annotations import asyncio import json from datetime import datetime, timedelta, timezone from types import SimpleNamespace import pytest from investibook.api import chat as chat_api from investibook.api import agent_sessions as sessions_api from investibook.core import agent_runtime as runtime_module from investibook.core.agent_runtime import AgentSession, AgentEventBridge from investibook.core import pubsub from investibook.db import queries as dbq async def acquire_lock_or_busy(lock: asyncio.Lock, session_id: str = "test"): """Mirror chat_stream's prompt-lock policy: reject instead of queueing.""" if lock.locked(): yield ( "data: " + json.dumps( { "type": "error", "code": "agent_busy", "error": "Agent is already responding in this workspace. Stop the current response before sending another message.", } ) + "\n\n" ) yield f'data: {json.dumps({"type": "complete", "session_id": session_id})}\n\n' return await lock.acquire() yield "__LOCK_ACQUIRED__" def _parse_events(chunks: list[str | bytes]) -> list[dict]: events = [] for chunk in chunks: text = chunk.decode() if isinstance(chunk, bytes) else chunk for line in text.split("\n"): line = line.strip() if line.startswith("data: "): events.append(json.loads(line[6:])) return events @pytest.mark.asyncio async def test_lock_free_acquires_immediately(): lock = asyncio.Lock() chunks = [] async for chunk in acquire_lock_or_busy(lock): chunks.append(chunk) assert "__LOCK_ACQUIRED__" in chunks assert lock.locked() lock.release() @pytest.mark.asyncio async def test_locked_prompt_returns_busy_without_waiting(): lock = asyncio.Lock() await lock.acquire() chunks = [] async for chunk in acquire_lock_or_busy(lock): chunks.append(chunk) assert "__LOCK_ACQUIRED__" not in chunks events = _parse_events(chunks) assert [event["type"] for event in events] == ["error", "complete"] assert events[0]["code"] == "agent_busy" assert lock.locked() lock.release() @pytest.mark.asyncio async def test_chat_stream_disconnect_only_detaches_from_durable_run(monkeypatch): execution_started = asyncio.Event() execution_finished = asyncio.Event() allow_execution_finish = asyncio.Event() execution_task: asyncio.Task | None = None async def execute_independently(): execution_started.set() await allow_execution_finish.wait() execution_finished.set() async def fake_submit_agent_turn(**_kwargs): nonlocal execution_task execution_task = asyncio.create_task(execute_independently()) return { "id": "run-durable", "assistant_message_id": "assistant-durable", }, True async def fake_events(_run_id, **_kwargs): yield 1, { "type": "session_created", "session_id": "runtime-session", "run_id": "run-durable", } await execution_finished.wait() yield 2, { "type": "complete", "run_id": "run-durable", "status": "completed", } monkeypatch.setattr( chat_api, "_db_get_user_by_token", lambda _token: {"id": "user-1", "email": "u@example.com"}, ) monkeypatch.setattr(sessions_api, "submit_agent_turn", fake_submit_agent_turn) monkeypatch.setattr(chat_api.agent_run_manager, "events", fake_events) request = SimpleNamespace( state=SimpleNamespace(trace_id="trace-disconnect"), base_url="http://testserver/", ) response = await chat_api.chat_stream( request=request, message="hello", workspace_id="ws-disconnect", user_message_id="user-message-1", token="token", authorization=None, ) iterator = response.body_iterator first_chunk = await anext(iterator) assert _parse_events([first_chunk])[0]["run_id"] == "run-durable" await asyncio.wait_for(execution_started.wait(), timeout=1) await iterator.aclose() assert not execution_finished.is_set() allow_execution_finish.set() await asyncio.wait_for(execution_finished.wait(), timeout=1) assert execution_task is not None await execution_task @pytest.mark.asyncio async def test_persist_assistant_snapshot_keeps_streaming_running_part(tmp_path, monkeypatch): workspace_dir = tmp_path / "ws-running" workspace_dir.mkdir() session = AgentSession("frontend-session", "opencode", workspace_dir=workspace_dir) client = AgentEventBridge(workspace_id="ws-running", acp_session=session) session.client = client client.start_assistant_message("assistant-running") client._msg_parts["tool-1"] = { "type": "tool", "id": "tool-1", "tool": "bash", "callID": "1", "startedAt": 123, "state": {"status": "running", "title": "Inspect logs"}, } client._msg_parts_order.append("tool-1") upserts = [] notifications = [] def fake_upsert(workspace_id, msg_id, role, content, parts=None, status="completed"): upserts.append( { "workspace_id": workspace_id, "msg_id": msg_id, "role": role, "content": content, "parts": parts, "status": status, } ) async def fake_notify(workspace_id, updated_at, count): notifications.append((workspace_id, updated_at, count)) monkeypatch.setattr(runtime_module, "_db_upsert_message_v2", fake_upsert) monkeypatch.setattr(dbq, "_db_get_messages_status_v2", lambda workspace_id: {"updated_at": "now", "count": 1}) monkeypatch.setattr(pubsub, "notify_message_subscribers", fake_notify) await session.persist_assistant_message_snapshot() assert upserts == [ { "workspace_id": "ws-running", "msg_id": "assistant-running", "role": "assistant", "content": "", "parts": [client._msg_parts["tool-1"]], "status": "streaming", } ] assert notifications == [("ws-running", "now", 1)] assert client._msg_id == "assistant-running" def test_messages_v2_compact_history_can_include_active_parts(tmp_path, monkeypatch): monkeypatch.setattr(dbq, "DB_PATH", tmp_path / "messages_active_parts.db") dbq._db_init() workspace_id = "ws-active-parts" dbq._db_append_messages_v2( workspace_id, [ { "id": "assistant-completed", "role": "assistant", "content": "done", "status": "completed", "timestamp": "2026-05-07T00:00:00+00:00", "parts": [ { "type": "tool", "id": "completed-tool", "state": {"status": "completed", "output": "large output"}, } ], }, { "id": "assistant-running", "role": "assistant", "content": "", "status": "streaming", "timestamp": "2026-05-07T00:00:01+00:00", "parts": [ { "type": "tool", "id": "running-tool", "state": {"status": "running"}, } ], }, ], ) compact = dbq._db_get_messages_v2(workspace_id, include_parts=False) active = dbq._db_get_messages_v2( workspace_id, include_parts=False, include_active_parts=True, ) assert "parts" not in compact[0] assert "parts" not in compact[1] assert "parts" not in active[0] assert active[1]["parts"][0]["id"] == "running-tool" def test_stale_streaming_messages_are_failed_after_timeout(tmp_path, monkeypatch): monkeypatch.setattr(dbq, "DB_PATH", tmp_path / "messages_stale_timeout.db") dbq._db_init() workspace_id = "ws-stale-timeout" base = datetime(2026, 5, 9, 4, 0, 0, tzinfo=timezone.utc) old_update = (base - timedelta(minutes=20)).isoformat() dbq._db_append_messages_v2( workspace_id, [ { "id": "assistant-stale", "role": "assistant", "content": "partial", "status": "streaming", "timestamp": (base - timedelta(minutes=30)).isoformat(), "parts": [ { "type": "tool", "id": "running-tool", "state": {"status": "running", "title": "Long check"}, } ], } ], ) with dbq._db_lock: conn = dbq._db_connect() conn.execute( "UPDATE messages_v2 SET updated_at = ? WHERE workspace_id = ? AND id = ?", (old_update, workspace_id, "assistant-stale"), ) conn.commit() conn.close() failed = dbq._db_fail_stale_streaming_messages_v2( workspace_id, timeout_seconds=15 * 60, now=base, ) messages = dbq._db_get_messages_v2(workspace_id) status = dbq._db_get_messages_status_v2(workspace_id) assert failed == 1 assert messages[0]["status"] == "failed" assert messages[0]["parts"][0]["state"]["status"] == "failed" assert "Timed out" in messages[0]["parts"][0]["state"]["output"]["error"] assert status["updated_at"] == base.isoformat() def test_fresh_streaming_messages_are_not_failed_before_timeout(tmp_path, monkeypatch): monkeypatch.setattr(dbq, "DB_PATH", tmp_path / "messages_fresh_timeout.db") dbq._db_init() workspace_id = "ws-fresh-timeout" base = datetime(2026, 5, 9, 4, 0, 0, tzinfo=timezone.utc) recent_update = (base - timedelta(minutes=5)).isoformat() dbq._db_append_messages_v2( workspace_id, [ { "id": "assistant-fresh", "role": "assistant", "content": "partial", "status": "streaming", "timestamp": (base - timedelta(minutes=30)).isoformat(), "parts": [ { "type": "tool", "id": "running-tool", "state": {"status": "running", "title": "Still active"}, } ], } ], ) with dbq._db_lock: conn = dbq._db_connect() conn.execute( "UPDATE messages_v2 SET updated_at = ? WHERE workspace_id = ? AND id = ?", (recent_update, workspace_id, "assistant-fresh"), ) conn.commit() conn.close() failed = dbq._db_fail_stale_streaming_messages_v2( workspace_id, timeout_seconds=15 * 60, now=base, ) messages = dbq._db_get_messages_v2(workspace_id) assert failed == 0 assert messages[0]["status"] == "streaming" assert messages[0]["parts"][0]["state"]["status"] == "running" @pytest.mark.asyncio async def test_cancel_prompt_waits_for_prompt_task(tmp_path): session = AgentSession("frontend-session", "opencode", workspace_dir=tmp_path) session.runtime_session_id = "acp-session" cancel_called = asyncio.Event() task_cleaned_up = asyncio.Event() class FakeConnection: async def cancel(self, session_id: str): assert session_id == "acp-session" cancel_called.set() async def prompt_task(): try: await asyncio.sleep(60) except asyncio.CancelledError: task_cleaned_up.set() raise task = asyncio.create_task(prompt_task()) await asyncio.sleep(0) session.conn = FakeConnection() session._current_prompt_task = task await session.cancel_prompt(timeout_s=1) assert cancel_called.is_set() assert task_cleaned_up.is_set() assert task.done() assert task.cancelled()