301 lines
12 KiB
Python
301 lines
12 KiB
Python
"""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})
|