mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
75245e4b73
* fix(subagents): leave full results out of the polled batch item list The item list endpoint returns only previews, but list_items selected whole rows, so every page loaded and decoded each item's full report, evidence JSON and prompt before discarding them. The workspace polls this list every 2 seconds while a batch runs; with the default max_result_chars a 100-item page held about 34 MB and blocked the event loop for about 70 ms (170 MB and 300 ms at limit=500). Defer those columns unless the caller asks for them, with raiseload so an accidental read fails loudly instead of lazy-loading. A 100-item page now peaks under 1 MB; responses are unchanged. * test(subagents): pin the item list route to compact columns Drive list_batch_items against a real SQLite repository so a route that starts requesting full results fails, not only a repository change. * docs(changelog): reference #6781 * test(subagents): scope the route test's seed update to its batch
516 lines
18 KiB
Python
516 lines
18 KiB
Python
import re
|
|
from datetime import UTC, datetime, timedelta
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
from sqlalchemy import event
|
|
|
|
from deerflow.config.database_config import DatabaseConfig
|
|
from deerflow.mcp_scope import THREAD_INCARNATION_CONTEXT_KEY
|
|
from deerflow.persistence.engine import close_engine, get_engine, get_session_factory, init_engine_from_config
|
|
from deerflow.persistence.subagent_batches import SubagentBatchRepository
|
|
|
|
_MISSING = object()
|
|
|
|
|
|
@pytest_asyncio.fixture(autouse=True)
|
|
async def _close_engine() -> None:
|
|
yield
|
|
await close_engine()
|
|
|
|
|
|
async def _repo(tmp_path) -> SubagentBatchRepository:
|
|
await init_engine_from_config(DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path)))
|
|
sf = get_session_factory()
|
|
assert sf is not None
|
|
return SubagentBatchRepository(sf)
|
|
|
|
|
|
async def _create(
|
|
repo: SubagentBatchRepository,
|
|
*,
|
|
count: int = 4,
|
|
max_live: int = 2,
|
|
max_running: int = 1,
|
|
max_attempts: int = 2,
|
|
) -> dict:
|
|
return await repo.create_batch(
|
|
batch_id="batch-1",
|
|
user_id="user-1",
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
tool_call_id="call-1",
|
|
submission_key="run-1:call-1",
|
|
title="Research records",
|
|
subagent_type="general-purpose",
|
|
items=[{"key": f"item-{i}", "prompt": f"Process {i}"} for i in range(count)],
|
|
max_live_items=max_live,
|
|
max_running_items=max_running,
|
|
max_attempts=max_attempts,
|
|
execution_spec={
|
|
"subagent_config": {
|
|
"name": "general-purpose",
|
|
"description": "test",
|
|
"system_prompt": "private instructions",
|
|
},
|
|
"authz_attributes": {"tenant": "private-tenant"},
|
|
},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_claim_separates_total_live_leased_and_running(tmp_path) -> None:
|
|
repo = await _repo(tmp_path)
|
|
created = await _create(repo)
|
|
assert created["counts"]["pending"] == 4
|
|
|
|
now = datetime.now(UTC)
|
|
claimed = await repo.claim_items(now=now, lease_owner="worker-1", lease_seconds=60, limit=10)
|
|
assert len(claimed) == 1
|
|
assert claimed[0]["status"] == "leased"
|
|
|
|
batch = await repo.get_batch("batch-1", user_id="user-1")
|
|
assert batch is not None
|
|
assert batch["counts"] == {
|
|
"pending": 2,
|
|
"queued": 1,
|
|
"leased": 1,
|
|
"running": 0,
|
|
"succeeded": 0,
|
|
"failed": 0,
|
|
"cancelled": 0,
|
|
}
|
|
assert await repo.mark_item_running(claimed[0]["id"], lease_owner="worker-1", now=now)
|
|
|
|
while_full = await repo.claim_items(
|
|
now=now + timedelta(seconds=1),
|
|
lease_owner="worker-2",
|
|
lease_seconds=60,
|
|
limit=10,
|
|
)
|
|
assert while_full == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("thread_incarnation", "expected_present"),
|
|
[
|
|
("incarnation-1", True),
|
|
(None, True),
|
|
(_MISSING, False),
|
|
],
|
|
)
|
|
async def test_claim_preserves_thread_incarnation_presence(
|
|
tmp_path,
|
|
thread_incarnation,
|
|
expected_present,
|
|
) -> None:
|
|
repo = await _repo(tmp_path)
|
|
execution_spec = {
|
|
"subagent_config": {
|
|
"name": "general-purpose",
|
|
"description": "test",
|
|
},
|
|
}
|
|
if thread_incarnation is not _MISSING:
|
|
execution_spec[THREAD_INCARNATION_CONTEXT_KEY] = thread_incarnation
|
|
await repo.create_batch(
|
|
batch_id="batch-1",
|
|
user_id="user-1",
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
tool_call_id="call-1",
|
|
submission_key="run-1:call-1",
|
|
title="Incarnation persistence",
|
|
subagent_type="general-purpose",
|
|
items=[{"key": "item-1", "prompt": "Process item"}],
|
|
max_live_items=1,
|
|
max_running_items=1,
|
|
max_attempts=1,
|
|
execution_spec=execution_spec,
|
|
)
|
|
|
|
claimed = await repo.claim_items(
|
|
now=datetime.now(UTC),
|
|
lease_owner="worker-1",
|
|
lease_seconds=60,
|
|
limit=1,
|
|
)
|
|
|
|
assert len(claimed) == 1
|
|
restored_spec = claimed[0]["batch"]["execution_spec"]
|
|
assert (THREAD_INCARNATION_CONTEXT_KEY in restored_spec) is expected_present
|
|
if expected_present:
|
|
assert restored_spec[THREAD_INCARNATION_CONTEXT_KEY] == thread_incarnation
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_expired_lease_is_recovered_with_stable_item_identity(tmp_path) -> None:
|
|
repo = await _repo(tmp_path)
|
|
await _create(repo, count=1, max_live=1, max_running=1)
|
|
now = datetime.now(UTC)
|
|
first = await repo.claim_items(now=now, lease_owner="worker-1", lease_seconds=30, limit=1)
|
|
|
|
reclaimed = await repo.claim_items(
|
|
now=now + timedelta(seconds=31),
|
|
lease_owner="worker-2",
|
|
lease_seconds=30,
|
|
limit=1,
|
|
)
|
|
assert len(reclaimed) == 1
|
|
assert reclaimed[0]["id"] == first[0]["id"]
|
|
assert reclaimed[0]["item_key"] == "item-0"
|
|
assert reclaimed[0]["attempt"] == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("succeed_sibling", "expected_status"),
|
|
[(False, "failed"), (True, "completed")],
|
|
)
|
|
async def test_exhausted_expired_lease_terminalizes_batch(tmp_path, succeed_sibling: bool, expected_status: str) -> None:
|
|
repo = await _repo(tmp_path)
|
|
count = 2 if succeed_sibling else 1
|
|
await _create(repo, count=count, max_live=count, max_running=count, max_attempts=1)
|
|
now = datetime.now(UTC)
|
|
claimed = await repo.claim_items(now=now, lease_owner="worker-1", lease_seconds=30, limit=count)
|
|
assert len(claimed) == count
|
|
|
|
if succeed_sibling:
|
|
assert await repo.finalize_item(
|
|
claimed[0]["id"],
|
|
lease_owner="worker-1",
|
|
succeeded=True,
|
|
result="done",
|
|
result_preview="done",
|
|
result_truncated=False,
|
|
error=None,
|
|
stop_reason=None,
|
|
token_usage=None,
|
|
model_name=None,
|
|
completed_at=now + timedelta(seconds=1),
|
|
)
|
|
|
|
assert (
|
|
await repo.claim_items(
|
|
now=now + timedelta(seconds=31),
|
|
lease_owner="worker-2",
|
|
lease_seconds=30,
|
|
limit=count,
|
|
)
|
|
== []
|
|
)
|
|
batch = await repo.get_batch("batch-1", user_id="user-1")
|
|
assert batch is not None
|
|
assert batch["counts"]["failed"] == 1
|
|
assert batch["status"] == expected_status
|
|
assert batch["completed_at"] is not None
|
|
|
|
retried = await repo.retry_item("batch-1", claimed[-1]["id"], user_id="user-1")
|
|
assert retried is not None and retried["status"] == "pending"
|
|
reopened = await repo.get_batch("batch-1", user_id="user-1")
|
|
assert reopened is not None
|
|
assert reopened["status"] == "queued"
|
|
assert reopened["completed_at"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_finalize_retries_then_terminalizes_and_completes_batch(tmp_path) -> None:
|
|
repo = await _repo(tmp_path)
|
|
await _create(repo, count=1, max_live=1, max_running=1)
|
|
now = datetime.now(UTC)
|
|
first = (await repo.claim_items(now=now, lease_owner="worker-1", lease_seconds=60, limit=1))[0]
|
|
await repo.finalize_item(
|
|
first["id"],
|
|
lease_owner="worker-1",
|
|
succeeded=False,
|
|
result=None,
|
|
result_preview=None,
|
|
result_truncated=False,
|
|
error="temporary",
|
|
stop_reason=None,
|
|
token_usage=None,
|
|
model_name="model-a",
|
|
completed_at=now,
|
|
)
|
|
item = (await repo.list_items("batch-1", user_id="user-1"))[0]
|
|
assert item["status"] == "queued"
|
|
|
|
second = (await repo.claim_items(now=now + timedelta(seconds=1), lease_owner="worker-2", lease_seconds=60, limit=1))[0]
|
|
await repo.finalize_item(
|
|
second["id"],
|
|
lease_owner="worker-2",
|
|
succeeded=True,
|
|
result="done",
|
|
result_preview="done",
|
|
result_truncated=False,
|
|
error=None,
|
|
stop_reason=None,
|
|
token_usage={"total_tokens": 12},
|
|
model_name="model-a",
|
|
completed_at=now + timedelta(seconds=2),
|
|
)
|
|
batch = await repo.get_batch("batch-1", user_id="user-1")
|
|
assert batch is not None
|
|
assert batch["status"] == "completed"
|
|
assert batch["counts"]["succeeded"] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_pause_resume_cancel_and_owner_scope(tmp_path) -> None:
|
|
repo = await _repo(tmp_path)
|
|
await _create(repo, count=2, max_live=2, max_running=1)
|
|
paused = await repo.pause_batch("batch-1", user_id="user-1")
|
|
assert paused is not None and paused["status"] == "paused"
|
|
assert await repo.claim_items(now=datetime.now(UTC), lease_owner="worker", lease_seconds=60, limit=1) == []
|
|
resumed = await repo.resume_batch("batch-1", user_id="user-1")
|
|
assert resumed is not None and resumed["status"] == "queued"
|
|
cancelled = await repo.cancel_batch("batch-1", user_id="user-1")
|
|
assert cancelled is not None and cancelled["status"] == "cancelled"
|
|
assert cancelled["counts"]["cancelled"] == 2
|
|
assert await repo.get_batch("batch-1", user_id="other") is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retry_item_keeps_paused_batch_paused_until_resume(tmp_path) -> None:
|
|
repo = await _repo(tmp_path)
|
|
await _create(repo, count=3, max_live=3, max_running=3, max_attempts=1)
|
|
now = datetime.now(UTC)
|
|
failed = (await repo.claim_items(now=now, lease_owner="worker-1", lease_seconds=60, limit=1))[0]
|
|
await repo.finalize_item(
|
|
failed["id"],
|
|
lease_owner="worker-1",
|
|
succeeded=False,
|
|
result=None,
|
|
result_preview=None,
|
|
result_truncated=False,
|
|
error="permanent",
|
|
stop_reason=None,
|
|
token_usage=None,
|
|
model_name="model-a",
|
|
completed_at=now,
|
|
)
|
|
paused = await repo.pause_batch("batch-1", user_id="user-1")
|
|
assert paused is not None and paused["status"] == "paused"
|
|
assert await repo.claim_items(now=now + timedelta(seconds=1), lease_owner="worker-2", lease_seconds=60, limit=10) == []
|
|
|
|
retried = await repo.retry_item("batch-1", failed["id"], user_id="user-1")
|
|
|
|
assert retried is not None
|
|
assert retried["status"] == "pending"
|
|
assert retried["attempt"] == 0
|
|
assert retried["error"] is None
|
|
batch = await repo.get_batch("batch-1", user_id="user-1")
|
|
assert batch is not None and batch["status"] == "paused"
|
|
assert batch["counts"]["pending"] == 1
|
|
assert batch["counts"]["queued"] == 2
|
|
assert await repo.claim_items(now=now + timedelta(seconds=2), lease_owner="worker-2", lease_seconds=60, limit=10) == []
|
|
|
|
await close_engine()
|
|
repo = await _repo(tmp_path)
|
|
batch = await repo.get_batch("batch-1", user_id="user-1")
|
|
assert batch is not None and batch["status"] == "paused"
|
|
|
|
resumed = await repo.resume_batch("batch-1", user_id="user-1")
|
|
assert resumed is not None and resumed["status"] == "queued"
|
|
claimed = await repo.claim_items(now=now + timedelta(seconds=3), lease_owner="worker-3", lease_seconds=60, limit=10)
|
|
assert len(claimed) == 3
|
|
assert failed["id"] in {item["id"] for item in claimed}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retry_item_in_running_batch_keeps_batch_schedulable(tmp_path) -> None:
|
|
repo = await _repo(tmp_path)
|
|
await _create(repo, count=3, max_live=3, max_running=3, max_attempts=1)
|
|
now = datetime.now(UTC)
|
|
failed = (await repo.claim_items(now=now, lease_owner="worker-1", lease_seconds=60, limit=1))[0]
|
|
await repo.finalize_item(
|
|
failed["id"],
|
|
lease_owner="worker-1",
|
|
succeeded=False,
|
|
result=None,
|
|
result_preview=None,
|
|
result_truncated=False,
|
|
error="permanent",
|
|
stop_reason=None,
|
|
token_usage=None,
|
|
model_name="model-a",
|
|
completed_at=now,
|
|
)
|
|
batch = await repo.get_batch("batch-1", user_id="user-1")
|
|
assert batch is not None and batch["status"] == "running"
|
|
|
|
retried = await repo.retry_item("batch-1", failed["id"], user_id="user-1")
|
|
|
|
assert retried is not None and retried["status"] == "pending"
|
|
batch = await repo.get_batch("batch-1", user_id="user-1")
|
|
assert batch is not None and batch["status"] == "queued"
|
|
claimed = await repo.claim_items(now=now + timedelta(seconds=1), lease_owner="worker-2", lease_seconds=60, limit=10)
|
|
assert len(claimed) == 3
|
|
assert failed["id"] in {item["id"] for item in claimed}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_terminalizes_in_flight_items_and_fences_stale_completion(tmp_path) -> None:
|
|
repo = await _repo(tmp_path)
|
|
await _create(repo, count=1, max_live=1, max_running=1)
|
|
now = datetime.now(UTC)
|
|
claimed = (await repo.claim_items(now=now, lease_owner="worker-1", lease_seconds=60, limit=1))[0]
|
|
assert await repo.mark_item_running(claimed["id"], lease_owner="worker-1", now=now)
|
|
|
|
cancelled = await repo.cancel_batch("batch-1", user_id="user-1")
|
|
|
|
assert cancelled is not None
|
|
assert cancelled["counts"]["cancelled"] == 1
|
|
item = (await repo.list_items("batch-1", user_id="user-1"))[0]
|
|
assert item["status"] == "cancelled"
|
|
assert not await repo.finalize_item(
|
|
claimed["id"],
|
|
lease_owner="worker-1",
|
|
succeeded=True,
|
|
result="late result",
|
|
result_preview="late result",
|
|
result_truncated=False,
|
|
error=None,
|
|
stop_reason=None,
|
|
token_usage=None,
|
|
model_name="model-a",
|
|
completed_at=now + timedelta(seconds=1),
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_executor_admission_failure_requeues_without_consuming_attempt(tmp_path) -> None:
|
|
repo = await _repo(tmp_path)
|
|
await _create(repo, count=1, max_live=1, max_running=1)
|
|
now = datetime.now(UTC)
|
|
first = (await repo.claim_items(now=now, lease_owner="worker-1", lease_seconds=60, limit=1))[0]
|
|
assert first["attempt"] == 1
|
|
|
|
assert await repo.requeue_item_after_admission_failure(
|
|
first["id"],
|
|
lease_owner="worker-1",
|
|
error="Process-wide subagent capacity is full",
|
|
now=now,
|
|
)
|
|
queued = (await repo.list_items("batch-1", user_id="user-1"))[0]
|
|
assert queued["status"] == "queued"
|
|
assert queued["attempt"] == 0
|
|
|
|
second = (await repo.claim_items(now=now + timedelta(seconds=1), lease_owner="worker-2", lease_seconds=60, limit=1))[0]
|
|
assert second["attempt"] == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_all_failed_items_mark_batch_failed(tmp_path) -> None:
|
|
repo = await _repo(tmp_path)
|
|
await _create(repo, count=1, max_live=1, max_running=1, max_attempts=1)
|
|
now = datetime.now(UTC)
|
|
item = (await repo.claim_items(now=now, lease_owner="worker-1", lease_seconds=60, limit=1))[0]
|
|
|
|
assert await repo.finalize_item(
|
|
item["id"],
|
|
lease_owner="worker-1",
|
|
succeeded=False,
|
|
result=None,
|
|
result_preview=None,
|
|
result_truncated=False,
|
|
error="permanent",
|
|
stop_reason=None,
|
|
token_usage=None,
|
|
model_name="model-a",
|
|
completed_at=now,
|
|
)
|
|
batch = await repo.get_batch("batch-1", user_id="user-1")
|
|
assert batch is not None
|
|
assert batch["status"] == "failed"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_public_projections_omit_execution_context_and_full_results(tmp_path) -> None:
|
|
repo = await _repo(tmp_path)
|
|
created = await _create(repo, count=1, max_live=1, max_running=1)
|
|
assert "execution_spec" not in created
|
|
assert "user_id" not in created
|
|
assert "submission_key" not in created
|
|
assert "run_id" not in created
|
|
assert "tool_call_id" not in created
|
|
|
|
now = datetime.now(UTC)
|
|
item = (await repo.claim_items(now=now, lease_owner="worker-1", lease_seconds=60, limit=1))[0]
|
|
assert await repo.finalize_item(
|
|
item["id"],
|
|
lease_owner="worker-1",
|
|
succeeded=True,
|
|
result="full private result",
|
|
result_preview="preview",
|
|
result_truncated=False,
|
|
error=None,
|
|
stop_reason=None,
|
|
token_usage=None,
|
|
model_name="model-a",
|
|
completed_at=now,
|
|
)
|
|
|
|
public_batch = await repo.get_batch("batch-1", user_id="user-1")
|
|
assert public_batch is not None
|
|
assert "execution_spec" not in public_batch
|
|
public_item = (await repo.list_items("batch-1", user_id="user-1"))[0]
|
|
assert public_item["result_preview"] == "preview"
|
|
assert "result" not in public_item
|
|
assert "lease_owner" not in public_item
|
|
export_item = (await repo.list_items("batch-1", user_id="user-1", include_result=True))[0]
|
|
assert export_item["result"] == "full private result"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_item_list_selects_only_the_heavy_columns_it_returns(tmp_path) -> None:
|
|
# The workspace polls this list every 2 s. Selecting whole rows decoded
|
|
# every full report and evidence artifact only to return their previews.
|
|
repo = await _repo(tmp_path)
|
|
await _create(repo, count=1)
|
|
statements = []
|
|
sync_engine = get_engine().sync_engine
|
|
listener = lambda *args, **kwargs: statements.append(args[2]) # noqa: E731
|
|
event.listen(sync_engine, "before_cursor_execute", listener)
|
|
try:
|
|
public_items = await repo.list_items("batch-1", user_id="user-1")
|
|
full_items = await repo.list_items("batch-1", user_id="user-1", include_result=True, include_prompt=True)
|
|
finally:
|
|
event.remove(sync_engine, "before_cursor_execute", listener)
|
|
|
|
item_columns = [set(re.findall(r"subagent_batch_items\.(\w+)", statement)) for statement in statements if "FROM subagent_batch_items" in statement]
|
|
heavy_columns = {"prompt", "result", "result_artifact"}
|
|
assert len(item_columns) == 2
|
|
assert not heavy_columns & item_columns[0]
|
|
assert heavy_columns <= item_columns[1]
|
|
assert "result_preview" in public_items[0]
|
|
assert full_items[0]["prompt"] == "Process 0"
|
|
assert full_items[0]["result"] is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_duplicate_submission_key_returns_original_batch(tmp_path) -> None:
|
|
repo = await _repo(tmp_path)
|
|
original = await _create(repo, count=1)
|
|
duplicate = await repo.create_batch(
|
|
batch_id="batch-2",
|
|
user_id="user-1",
|
|
thread_id="thread-1",
|
|
run_id="run-1",
|
|
tool_call_id="call-1",
|
|
submission_key="run-1:call-1",
|
|
title="Duplicate retry",
|
|
subagent_type="general-purpose",
|
|
items=[{"key": "different", "prompt": "Must not be inserted"}],
|
|
max_live_items=1,
|
|
max_running_items=1,
|
|
max_attempts=2,
|
|
execution_spec={"subagent_config": {"name": "general-purpose", "description": "test"}},
|
|
)
|
|
|
|
assert duplicate["id"] == original["id"] == "batch-1"
|
|
items = await repo.list_items("batch-1", user_id="user-1", include_prompt=True)
|
|
assert items is not None
|
|
assert [item["item_key"] for item in items] == ["item-0"]
|