from __future__ import annotations import asyncio from datetime import datetime, timedelta, timezone import pytest from investibook.db import queries as dbq from investibook.core import agent_turns as runs_module @pytest.fixture def run_db(tmp_path, monkeypatch): monkeypatch.setattr(dbq, "DB_PATH", tmp_path / "agent-runs.db") dbq._db_init() user = dbq._db_get_or_create_user("runner@example.com") now = dbq._utc_now() dbq._db_upsert_workspace( "ws-run", user["id"], "Run workspace", "", False, now, now, ) return user def _create_run(user: dict, *, run_id: str, user_message_id: str, workspace_id: str = "ws-run", code_context_id: str | None = None): return dbq._db_create_agent_run( run_id=run_id, workspace_id=workspace_id, user_id=user["id"], user_message_id=user_message_id, assistant_message_id=f"assistant-{run_id}", prompt="diagnose this", attachment_ids=[], message_attachments=[], agent_type="opencode", model_tier="standard", channel="web", api_base_url="https://investibook.example", code_context_id=code_context_id, ) def test_create_agent_run_is_idempotent_and_creates_canonical_messages(run_db): first, first_created = _create_run( run_db, run_id="run-1", user_message_id="user-message-1" ) second, second_created = _create_run( run_db, run_id="run-retry", user_message_id="user-message-1" ) assert first_created is True assert second_created is False assert second["id"] == first["id"] assert second["status"] == "queued" messages = dbq._db_get_messages_v2("ws-run") assert [(message["role"], message["status"]) for message in messages] == [ ("user", "completed"), ("assistant", "streaming"), ] def test_delete_workspace_removes_agent_run_transcripts(run_db): run, _ = _create_run( run_db, run_id="run-delete", user_message_id="user-message-delete" ) dbq._db_append_agent_transcript_record(run["id"], {"type": "text", "text": "hi"}) dbq._db_delete_workspace("ws-run") with dbq._db_lock: conn = dbq._db_connect() try: assert ( conn.execute( "SELECT COUNT(*) FROM agent_transcript_records" ).fetchone()[0] == 0 ) assert ( conn.execute( "SELECT COUNT(*) FROM agent_turns WHERE workspace_id = 'ws-run'" ).fetchone()[0] == 0 ) assert ( conn.execute( "SELECT COUNT(*) FROM workspaces WHERE id = 'ws-run'" ).fetchone()[0] == 0 ) finally: conn.close() def test_workspace_rejects_a_different_active_run(run_db): first, _ = _create_run( run_db, run_id="run-1", user_message_id="user-message-1" ) with pytest.raises(dbq.AgentTurnBusyError) as exc_info: _create_run(run_db, run_id="run-2", user_message_id="user-message-2") assert exc_info.value.active_turn["id"] == first["id"] assert len(dbq._db_get_messages_v2("ws-run")) == 2 def test_user_code_diagnosis_limit_allows_two_active_runs_across_workspaces(run_db): now = dbq._utc_now() dbq._db_upsert_workspace( "ws-run-two", run_db["id"], "Second workspace", "", False, now, now ) dbq._db_upsert_workspace( "ws-run-three", run_db["id"], "Third workspace", "", False, now, now ) _create_run( run_db, run_id="code-1", user_message_id="code-message-1", code_context_id="tidb:" + "a" * 40 + ":source", ) _create_run( run_db, run_id="code-2", user_message_id="code-message-2", workspace_id="ws-run-two", code_context_id="tidb:" + "b" * 40 + ":source", ) with pytest.raises(dbq.CodeDiagnosisConcurrencyLimitExceeded) as exc_info: _create_run( run_db, run_id="code-3", user_message_id="code-message-3", workspace_id="ws-run-three", code_context_id="tidb:" + "c" * 40 + ":source", ) assert exc_info.value.active_count == 2 assert exc_info.value.limit == 2 assert dbq._db_count_active_code_diagnosis_runs() == 2 def test_cancelling_a_queued_run_is_terminal_and_releases_workspace(run_db): run, _ = _create_run( run_db, run_id="run-queued", user_message_id="user-message-1" ) cancelled = dbq._db_request_agent_run_cancel(run["id"], run_db["id"]) assert cancelled is not None assert cancelled["status"] == "cancelled" assert cancelled["finished_at"] next_run, created = _create_run( run_db, run_id="run-next", user_message_id="user-message-2" ) assert created is True assert next_run["status"] == "queued" def test_claim_and_finish_require_the_same_worker_lease(run_db): _create_run(run_db, run_id="run-claim", user_message_id="user-message-1") claimed = dbq._db_claim_next_agent_run("worker-a", 30) assert claimed is not None assert claimed["status"] == "running" assert claimed["attempt"] == 1 assert dbq._db_finish_agent_run( claimed["id"], "worker-b", "completed" ) is False assert dbq._db_finish_agent_run( claimed["id"], "worker-a", "completed" ) is True assert dbq._db_get_agent_run(claimed["id"])["status"] == "completed" def test_running_cancel_transitions_through_cancelling(run_db): _create_run(run_db, run_id="run-cancel", user_message_id="user-message-1") claimed = dbq._db_claim_next_agent_run("worker-a", 30) cancelling = dbq._db_request_agent_run_cancel(claimed["id"], run_db["id"]) assert cancelling["status"] == "cancelling" assert cancelling["cancel_requested"] is True assert dbq._db_claim_next_agent_run("worker-b", 30) is None assert dbq._db_finish_agent_run( claimed["id"], "worker-a", "cancelled" ) is True def test_expired_dispatched_run_is_interrupted_instead_of_requeued(run_db): _create_run(run_db, run_id="run-expired", user_message_id="user-message-1") claimed = dbq._db_claim_next_agent_run("dead-worker", 30) expired = (datetime.now(timezone.utc) - timedelta(seconds=1)).isoformat() with dbq._db_lock: conn = dbq._db_connect() conn.execute( "UPDATE agent_turns SET lease_until = ? WHERE id = ?", (expired, claimed["id"]), ) conn.commit() conn.close() interrupted = dbq._db_interrupt_expired_agent_runs() assert [run["id"] for run in interrupted] == [claimed["id"]] stored = dbq._db_get_agent_run(claimed["id"]) assert stored["status"] == "interrupted" assert "not replayed" in stored["error"] def test_manager_supports_repeated_lifespans_on_different_event_loops(monkeypatch): manager = runs_module.AgentRunManager() monkeypatch.setattr(runs_module, "_db_interrupt_expired_agent_runs", lambda: []) monkeypatch.setattr(runs_module, "_db_claim_next_agent_run", lambda *_args: None) async def lifecycle(): await manager.start() await asyncio.sleep(0) await manager.stop() asyncio.run(lifecycle()) asyncio.run(lifecycle()) @pytest.mark.asyncio async def test_slow_observer_resync_does_not_advance_event_sequence(monkeypatch): manager = runs_module.AgentRunManager() monkeypatch.setattr( runs_module, "_db_append_agent_transcript_record", lambda *_args: 2, ) queue: asyncio.Queue = asyncio.Queue(maxsize=1) manager._subscribers["run-slow"].add(queue) queue.put_nowait((1, {"type": "text", "text": "old"})) await manager.publish("run-slow", {"type": "text", "text": "new"}) sequence, event = queue.get_nowait() assert sequence == 0 assert event == {"type": "resync", "run_id": "run-slow"} @pytest.mark.asyncio async def test_observer_disconnect_does_not_stop_worker_execution( run_db, tmp_path, monkeypatch ): manager = runs_module.AgentRunManager() _create_run(run_db, run_id="run-detach", user_message_id="user-message-1") claimed = dbq._db_claim_next_agent_run(manager.worker_id, 30) allow_prompt_finish = asyncio.Event() usage_rows: list[dict] = [] class FakeClient: def __init__(self): self.queue = asyncio.Queue() self._msg_id = None self.content = "" def start_assistant_message(self, message_id): self._msg_id = message_id def get_and_reset_usage(self): return { "input_tokens": 10, "output_tokens": 4, "reasoning_tokens": 0, "cache_read_tokens": 0, "cache_write_tokens": 0, "context_used": 14, "context_size": 100, "cost_amount": 0.01, "cost_currency": "USD", } class FakeSession: frontend_id = "frontend-session" runtime_session_id = "runtime-session" agent_type = "opencode" def __init__(self): self.client = FakeClient() self._prompt_lock = asyncio.Lock() self._current_prompt_task = None async def persist_assistant_message_snapshot(self): if self.client._msg_id: await asyncio.to_thread( dbq._db_upsert_message_v2, "ws-run", self.client._msg_id, "assistant", self.client.content, status="streaming", ) async def finalize_assistant_message(self, status): if self.client._msg_id: await asyncio.to_thread( dbq._db_upsert_message_v2, "ws-run", self.client._msg_id, "assistant", self.client.content, status=status, ) self.client._msg_id = None async def prompt(self, *_args, **_kwargs): self.client.content = "durable answer" await self.client.queue.put({"type": "text", "text": "durable answer"}) await allow_prompt_finish.wait() await self.finalize_assistant_message("completed") async def cancel_prompt(self, timeout_s=5.0): del timeout_s await self.finalize_assistant_message("cancelled") if self._current_prompt_task: self._current_prompt_task.cancel() session = FakeSession() async def fake_get_or_create_session(*_args, **_kwargs): return session async def no_cost_gate(*_args, **_kwargs): return None async def no_notify(*_args, **_kwargs): return None workspace_dir = tmp_path / "workspaces" / "ws-run" workspace_dir.mkdir(parents=True) monkeypatch.setattr(runs_module.cfg, "WORKSPACES_DIR", tmp_path / "workspaces") monkeypatch.setattr(runs_module, "get_or_create_session", fake_get_or_create_session) monkeypatch.setattr(runs_module, "effective_prompt_after_cost_gate", no_cost_gate) monkeypatch.setattr(runs_module, "_notify_workspace_messages", no_notify) monkeypatch.setattr(runs_module, "_populate_workspace_dotenv_from_server", lambda *_a, **_k: None) monkeypatch.setattr(runs_module, "_db_create_user_session", lambda *_a, **_k: "token") monkeypatch.setattr(runs_module, "_db_delete_user_session", lambda *_a, **_k: None) monkeypatch.setattr(runs_module, "schedule_workspace_evidence_postpass", lambda **_k: None) monkeypatch.setattr(runs_module, "record_ai_chat_completion", lambda **_k: None) monkeypatch.setattr(runs_module, "record_acp_session_interaction", lambda **_k: None) monkeypatch.setattr(runs_module, "record_acp_token_usage", lambda **_k: None) monkeypatch.setattr(runs_module, "set_acp_session_cumulative_cost", lambda **_k: None) monkeypatch.setattr( runs_module, "_db_insert_token_usage", lambda **kwargs: usage_rows.append(kwargs) ) observer = manager.events(claimed["id"]) first_event_task = asyncio.create_task(anext(observer)) execution_task = asyncio.create_task(manager._execute(claimed)) _, first_event = await asyncio.wait_for(first_event_task, timeout=1) assert first_event["type"] == "session_created" await observer.aclose() assert not execution_task.done() allow_prompt_finish.set() await asyncio.wait_for(execution_task, timeout=1) stored = dbq._db_get_agent_run(claimed["id"]) assert stored["status"] == "completed" messages = dbq._db_get_messages_v2("ws-run") assert messages[-1]["content"] == "durable answer" assert messages[-1]["status"] == "completed" assert usage_rows[0]["input_tokens"] == 10 assert usage_rows[0]["output_tokens"] == 4 @pytest.mark.asyncio async def test_worker_failure_terminalizes_run_and_assistant_message( run_db, tmp_path, monkeypatch ): manager = runs_module.AgentRunManager() _create_run(run_db, run_id="run-failed", user_message_id="user-message-1") claimed = dbq._db_claim_next_agent_run(manager.worker_id, 30) async def fail_session_start(*_args, **_kwargs): raise RuntimeError("runtime unavailable") async def no_notify(*_args, **_kwargs): return None workspace_dir = tmp_path / "workspaces" / "ws-run" workspace_dir.mkdir(parents=True) monkeypatch.setattr(runs_module.cfg, "WORKSPACES_DIR", tmp_path / "workspaces") monkeypatch.setattr(runs_module, "get_or_create_session", fail_session_start) monkeypatch.setattr(runs_module, "_notify_workspace_messages", no_notify) monkeypatch.setattr(runs_module, "_populate_workspace_dotenv_from_server", lambda *_a, **_k: None) monkeypatch.setattr(runs_module, "_db_create_user_session", lambda *_a, **_k: "token") monkeypatch.setattr(runs_module, "_db_delete_user_session", lambda *_a, **_k: None) monkeypatch.setattr(runs_module, "record_ai_chat_completion", lambda **_k: None) observer = manager.events(claimed["id"]) next_event_task = asyncio.create_task(anext(observer)) execution_task = asyncio.create_task(manager._execute(claimed)) events = [] while True: _, event = await asyncio.wait_for(next_event_task, timeout=1) events.append(event) if event["type"] == "complete": break next_event_task = asyncio.create_task(anext(observer)) await execution_task assert [event["type"] for event in events] == ["error", "complete"] assert events[-1]["status"] == "failed" stored = dbq._db_get_agent_run(claimed["id"]) assert stored["status"] == "failed" assert stored["error"] == "runtime unavailable" messages = dbq._db_get_messages_v2("ws-run") assert messages[-1]["status"] == "failed"