"""Cross-user isolation, fair admission and durable event retention.""" import asyncio import threading from contextvars import copy_context from types import SimpleNamespace import pytest from deerflow.config.runtime_contention_config import RuntimeContentionConfig from deerflow.runtime.events.store.jsonl import JsonlRunEventStore from deerflow.runtime.events.store.spool import SpoolingRunEventStore from deerflow.runtime.fair_queue import FairAsyncGate, FairBackgroundScheduler from deerflow.utils.file_io import FileIOPool async def until(predicate, timeout=2): async with asyncio.timeout(timeout): while not predicate(): await asyncio.sleep(0.005) @pytest.mark.asyncio async def test_cancelled_disk_caller_does_not_release_a_still_running_job(): pool = FileIOPool(1, 1, "file-test") entered, release, second = threading.Event(), threading.Event(), threading.Event() def blocked(): entered.set() assert release.wait(3) task = asyncio.create_task(pool(blocked)) await until(entered.is_set) task.cancel() with pytest.raises(asyncio.CancelledError): await task queued = asyncio.create_task(pool(second.set)) await asyncio.sleep(0.02) assert not second.is_set() release.set() await asyncio.wait_for(queued, 2) assert second.is_set() pool._pool.shutdown() @pytest.mark.asyncio @pytest.mark.parametrize("kind", ["spool", "jsonl"]) async def test_event_store_concurrent_cursors_restart_pagination_and_explicit_deletion(tmp_path, kind): def create(): return SpoolingRunEventStore(tmp_path / "events.db") if kind == "spool" else JsonlRunEventStore(tmp_path) store = create() content = "完整消息" * 6000 rows = await asyncio.gather(*(store.put(thread_id="thread", run_id=f"r{i % 3}", event_type="answer", category="message", content=content, metadata={"id": i}) for i in range(60))) assert sorted(row["seq"] for row in rows) == list(range(1, 61)) reopened = create() latest = await reopened.list_messages("thread", limit=5) assert [row["seq"] for row in latest] == [56, 57, 58, 59, 60] assert latest[-1]["content"] == content assert [row["seq"] for row in await reopened.list_messages("thread", after_seq=3, limit=4)] == [4, 5, 6, 7] assert [row["seq"] for row in await reopened.list_messages("thread", before_seq=8, limit=3)] == [5, 6, 7] assert len(await reopened.list_messages_by_run("thread", "r1", limit=50)) == 20 assert await reopened.count_messages("thread") == 60 assert await reopened.list_events("thread", "r1", event_types=["missing"]) == [] assert await reopened.delete_by_run("thread", "r1") == 20 assert await reopened.count_messages("thread") == 40 last = await reopened.put(thread_id="thread", run_id="new", event_type="answer", category="message", content="new") assert last["seq"] == 61 assert await reopened.delete_by_thread("thread") == 41 @pytest.mark.asyncio async def test_spool_failed_batch_rolls_back_events_and_cursor(tmp_path): store = SpoolingRunEventStore(tmp_path / "events.db") with pytest.raises(KeyError): await store.put_batch([ dict(thread_id="t", run_id="r", event_type="answer", category="message", content="must rollback"), dict(thread_id="t", run_id="r", event_type="invalid"), ]) assert await store.count_messages("t") == 0 saved = await store.put(thread_id="t", run_id="r", event_type="answer", category="message", content="retained") assert saved["seq"] == 1 @pytest.mark.asyncio async def test_event_backend_memory_factory_spools_without_losing_messages(tmp_path, monkeypatch): from deerflow.config import runtime_paths from deerflow.runtime.events.store import make_run_event_store monkeypatch.setattr(runtime_paths, "runtime_home", lambda: tmp_path) first = make_run_event_store(SimpleNamespace(backend="memory")) await first.put(thread_id="t", run_id="r", event_type="answer", category="message", content="saved") second = make_run_event_store(SimpleNamespace(backend="memory")) assert (await second.list_messages("t"))[0]["content"] == "saved" @pytest.mark.asyncio async def test_async_gate_reserves_capacity_for_other_user_and_removes_cancelled_waiters(): gate = FairAsyncGate(3, 2, 20) entered = [] release = asyncio.Event() async def work(user, index): async with gate.slot(user): entered.append((user, index)) await release.wait() tasks = [asyncio.create_task(work("A", i)) for i in range(4)] await until(lambda: len(entered) == 2) tasks.append(asyncio.create_task(work("B", 0))) await until(lambda: len(entered) == 3) assert ("B", 0) in entered tasks[3].cancel() await asyncio.gather(tasks[3], return_exceptions=True) release.set() await asyncio.gather(*tasks, return_exceptions=True) assert gate._active == {} and gate._queues == {} assert ("A", 3) not in entered @pytest.mark.asyncio async def test_scheduler_fairness_queued_timeout_and_cancellation(): scheduler = FairBackgroundScheduler(3, 2, 30) release, cancel = threading.Event(), threading.Event() entered, expired = [], [] lock = threading.Lock() def run(name): def operation(remaining): assert remaining > 0 with lock: entered.append(name) assert release.wait(3) return operation try: for i in range(4): scheduler.submit("A", run(f"A{i}"), timeout=2, cancelled=lambda: False, expired=lambda cancelled: expired.append(cancelled)) await until(lambda: len(entered) == 2) scheduler.submit("B", run("B"), timeout=2, cancelled=lambda: False, expired=lambda cancelled: expired.append(cancelled)) await until(lambda: len(entered) == 3) assert "B" in entered scheduler.submit("C", run("expired"), timeout=0.02, cancelled=lambda: False, expired=lambda cancelled: expired.append("timeout")) scheduler.submit("D", run("cancelled"), timeout=2, cancelled=cancel.is_set, expired=lambda cancelled: expired.append("cancel" if cancelled else "wrong")) cancel.set() await until(lambda: len(expired) == 2) assert set(expired) == {"timeout", "cancel"} assert "expired" not in entered and "cancelled" not in entered finally: release.set() await asyncio.to_thread(scheduler.shutdown) @pytest.mark.asyncio async def test_subagent_user_context_and_cancellation_ack_keep_real_slot(monkeypatch): # Initialize the normal agent/tool import order before loading the executor. import deerflow.agents # noqa: F401 from deerflow.runtime.user_context import get_effective_user_id, reset_current_user, set_current_user from deerflow.subagents import executor scheduler = FairBackgroundScheduler(1, 1, 20) monkeypatch.setattr(executor, "_scheduler_pool", scheduler) release = threading.Event() started, cancelled = [], threading.Event() async def fake(self, task, holder): started.append((task, get_effective_user_id())) try: await asyncio.sleep(10) except asyncio.CancelledError: cancelled.set() while not release.is_set(): await asyncio.sleep(0.005) holder.status = executor.SubagentStatus.COMPLETED return holder monkeypatch.setattr(executor.SubagentExecutor, "_aexecute", fake) first = object.__new__(executor.SubagentExecutor) first.config = SimpleNamespace(name="fixture", timeout_seconds=0.05) first.trace_id = "test" token = set_current_user(SimpleNamespace(id="A")) try: first_id = first.execute_async("first") finally: reset_current_user(token) second = object.__new__(executor.SubagentExecutor) second.config, second.trace_id = SimpleNamespace(name="fixture", timeout_seconds=2), "test" try: await until(cancelled.is_set) assert executor.get_background_task_result(first_id).status == executor.SubagentStatus.TIMED_OUT executor.cleanup_background_task(first_id) assert executor.get_background_task_result(first_id) is not None token = set_current_user(SimpleNamespace(id="B")) try: second_id = second.execute_async("second") finally: reset_current_user(token) await asyncio.sleep(0.02) assert started == [("first", "A")] release.set() await until(lambda: len(started) == 2) assert started[-1] == ("second", "B") executor.request_cancel_background_task(second_id) # The fake intentionally waits for its deadline to exercise the bridge. await until(lambda: executor.get_background_task_result(second_id).execution_finished.is_set(), timeout=3) executor.cleanup_background_task(second_id) finally: release.set() await asyncio.to_thread(scheduler.shutdown) @pytest.mark.asyncio async def test_isolated_cancellation_before_first_coroutine_step_acknowledges(): import deerflow.agents # noqa: F401 from deerflow.subagents import executor stopped = threading.Event() future = executor._submit_to_isolated_loop_in_context(copy_context(), lambda: asyncio.sleep(5), on_finished=stopped.set) future.cancel() await until(stopped.is_set) @pytest.mark.asyncio async def test_background_roundtable_cancel_waits_for_actual_run_before_next_user(monkeypatch): from app.gateway import roundtable_inprocess_gateway as gateway from app.gateway import services gate = FairAsyncGate(1, 1) started, cancelled = [], [] release = asyncio.Event() async def run_child(thread_id): started.append(thread_id) try: await asyncio.Event().wait() except asyncio.CancelledError: cancelled.append(thread_id) await release.wait() async def start(body, thread_id, request, **kwargs): return SimpleNamespace(task=asyncio.create_task(run_child(thread_id))) async def watch(*args, **kwargs): await asyncio.Event().wait() monkeypatch.setattr(gateway, "_background_gate", lambda: gate) monkeypatch.setattr(gateway, "_await_run_with_idle_guard", watch) monkeypatch.setattr(services, "start_run", start) app = SimpleNamespace(state=SimpleNamespace()) def launch(user): return asyncio.create_task(gateway._run_agent_turn(app, user_id=user, thread_id=user, agent_name="seat", message="question", model="test", excluded_tools=[], skill_stop_names=[], thinking_enabled=False, reasoning_effort="low")) first = launch("A") second = None try: await until(lambda: started == ["A"]) first.cancel() await until(lambda: cancelled == ["A"]) second = launch("B") await asyncio.sleep(0.02) assert started == ["A"] and not first.done() release.set() await until(lambda: started == ["A", "B"]) with pytest.raises(asyncio.CancelledError): await first finally: release.set() for task in [first, second]: if task and not task.done(): task.cancel() await asyncio.gather(*(task for task in [first, second] if task), return_exceptions=True) assert gate._active == {} @pytest.mark.asyncio async def test_failed_event_disk_write_preserves_journal_buffer_for_retry(tmp_path, monkeypatch): from deerflow.runtime.journal import RunJournal store = SpoolingRunEventStore(tmp_path / "events.db") await store.put(thread_id="t", run_id="old", event_type="answer", category="message", content="old answer") journal = RunJournal("new", "t", store, flush_threshold=100) journal._put(event_type="answer", category="message", content="new answer" * 6000) original = store._connect with monkeypatch.context() as patch: patch.setattr(store, "_connect", lambda: (_ for _ in ()).throw(OSError("disk unavailable"))) with pytest.raises(OSError): await journal.flush() assert len(journal._buffer) == 1 assert store._connect == original await journal.flush() rows = await store.list_messages("t") assert [row["seq"] for row in rows] == [1, 2] assert [row["content"] for row in rows] == ["old answer", "new answer" * 6000] @pytest.mark.parametrize("field", ["subagent_concurrency", "subagent_per_user", "background_concurrency", "queue_limit"]) def test_contention_budget_validation(field): with pytest.raises(ValueError): RuntimeContentionConfig(**{field: 0})