mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 13:16:24 +00:00
feat(subagents): preserve RAG evidence in durable batch exports (#6413)
* feat(subagents): preserve RAG evidence in durable batch exports * docs(subagents): clarify evidence omissions and migration head * docs(persistence): update the migration chain head * test(persistence): track the new batch evidence migration head * test(persistence): verify batch evidence migration on PostgreSQL * docs(subagents): retain evidence migration compatibility guidance * fix(subagents): propagate cited evidence forwarding omissions * fix(tests): adapt batch evidence PostgreSQL CI connection options * fix(tests): preserve TLS policy in batch evidence migration fixture
This commit is contained in:
Guo-Yixin
GitHub
parent
d073add71c
commit
3182ecb1e4
@ -1891,6 +1891,12 @@ Durable `batch_task` workers use one app-owned plugin snapshot for tool assembly
|
||||
|
||||
Ordinary `task` delegation and explicit durable `batch_task` execution share the startup-scoped `subagent_runtime` process capacity. Batch mode keeps large independent item sets in SQL with separate total, live, and running limits, restart recovery, bounded results, and a thread-scoped Web UI panel. The panel pages through bounded previews on demand; full stored result text is available only through the owner-scoped JSONL export, while internal execution and authorization context never enters owner-facing responses. If the batch worker is later stopped or disabled, threads with persisted batches retain read-only item inspection and JSONL export; execution controls remain disabled until the worker is running again. See `config.example.yaml` and [the implementation contract](docs/plans/2026-08-24-subagent-batch-capacity-implementation.md) for limits and recovery semantics.
|
||||
|
||||
Durable knowledge-research results also export nullable `result_artifact.knowledge_sources` snapshots. Each retained source pairs a citation ID in the stored report with the exact excerpt captured by `knowledge_search` or a delegated `task`, plus its RAGFlow dataset/document/chunk locator and pages. The existing authenticated `GET /api/threads/{thread_id}/subagent-batches/{batch_id}/results.jsonl` remains the entry point: find a report's `#knowledge-<id>` link in that row's `sources` by `id` to inspect the evidence after a Gateway restart or provider change. Compact item queries and the batch panel do not include these snapshots; this does not add clickable citations to batch previews.
|
||||
|
||||
Snapshots preserve whole excerpts, not freshly fetched provider content. They inherit the 100-source/1,000,000-text-character forwarding limits; the serialized snapshot has a separate character budget equal to `subagent_batches.max_result_chars`, in addition to the report's existing budget. `knowledge_sources.omitted_count` reports referenced records omitted by forwarding, snapshot projection checks or snapshot budgeting; malformed or unsupported producer artifacts filtered before candidate selection are not counted. Report truncation limits selection to complete citations in the stored text. Old/plain results have no snapshot (`null`); no historical evidence is reconstructed. Gateway startup adds a nullable JSON column through migration `0033_batch_result_artifact`. Back up the database before migrating; downgrading drops evidence snapshots while preserving report text. Cancelled, failed and stale attempts do not publish evidence, and explicit retries clear previous snapshots.
|
||||
|
||||
Explicitly enabled nested `task` delegation propagates omitted opaque source IDs in its captured artifact. Each outer report filters and deduplicates those IDs, and independently captured evidence can restore an omitted source. The durable snapshot publishes only the final count. Legacy task artifacts without omission IDs cannot reconstruct earlier forwarding losses; unknown prose citations are not treated as captured evidence.
|
||||
|
||||
If an item's final lease expires after exhausting the retry budget, the batch reaches a terminal state once all items are terminal: `failed` when none succeeded, or `completed` when at least one succeeded. The item remains `failed` in both cases so partial results stay visible.
|
||||
|
||||
Direct `create_deerflow_agent(...)` integrations can own the same boundary explicitly instead of relying on Gateway startup. Construct one `SubagentRuntime` and share it across every graph in that application; its `max_running`, ordinary per-run total, bound `task` tool, and optional durable-batch tools then use the same caller-owned snapshot and execution controller. A runtime with a batch repository owns a worker and must be started before graph construction and stopped during application shutdown:
|
||||
|
||||
@ -10,8 +10,16 @@ to the model and mark truncation; do not fetch a fresh chunk and present it as
|
||||
historical evidence. Direct `knowledge_search()` callers retain its string API.
|
||||
`sources.py` forwards only captured sources cited by ordinary subagent results,
|
||||
with count/text budgets. The source dialog uses stored thread messages and
|
||||
introduces no unauthenticated document proxy. Durable batch result storage and
|
||||
standalone Markdown do not include native source artifacts.
|
||||
introduces no unauthenticated document proxy. Durable batch exports retain
|
||||
schema-projected snapshots via `durable_source_artifact`; standalone Markdown
|
||||
does not include native source artifacts. Preserve entire excerpts and only
|
||||
select citations in the stored report. Bound the serialized JSON by the batch
|
||||
result cap, retain forwarding limits, and count omissions without refetching.
|
||||
Ordinary task artifacts carry omitted opaque source IDs, limited to complete
|
||||
citations in that task's report. Filter them again against each outer report,
|
||||
deduplicate, and let recaptured evidence supersede an omission. Never trust an
|
||||
incoming count or count unknown prose links. Durable exports expose only the
|
||||
final count; legacy artifacts without omission IDs cannot recover lost IDs.
|
||||
|
||||
Output budgeting retains complete evidence entries and their source records
|
||||
together, including delegated results and model-request history. Never shorten
|
||||
|
||||
@ -1,5 +1,6 @@
|
||||
"""Bounded source-artifact forwarding across ordinary subagent results."""
|
||||
"""Bounded captured-source forwarding for ordinary and durable subagent results."""
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
@ -71,9 +72,13 @@ def budget_source_artifact(content: str, artifact: object, max_chars: int, *, su
|
||||
return "\n\n".join(entries), updated or None
|
||||
|
||||
|
||||
def cited_source_artifact(messages: list[dict[str, Any]], content: str) -> dict[str, Any] | None:
|
||||
def cited_source_artifact(messages: list[dict[str, Any]], content: str, *, include_omissions: bool = False) -> dict[str, Any] | None:
|
||||
"""Carry only actual captured sources cited in the child's final result."""
|
||||
sources: dict[str, dict[str, Any]] = {}
|
||||
omitted_ids: set[str] = set()
|
||||
# Complete destinations bound omission metadata by this report's size;
|
||||
# never forward arbitrary input IDs or infer evidence from unknown links.
|
||||
referenced_ids = set(re.findall(r"\]\(#knowledge-([a-f0-9]{32}-[1-9][0-9]{0,2})\)", content)) if include_omissions else set()
|
||||
remaining = 1_000_000
|
||||
for message in messages:
|
||||
if message.get("type") != "tool" or message.get("name") not in {"knowledge_search", "task"}:
|
||||
@ -85,6 +90,9 @@ def cited_source_artifact(messages: list[dict[str, Any]], content: str) -> dict[
|
||||
raw_sources = payload.get("sources")
|
||||
if not isinstance(raw_sources, list):
|
||||
continue
|
||||
inherited_ids = payload.get("omitted_source_ids")
|
||||
if include_omissions and isinstance(inherited_ids, list):
|
||||
omitted_ids.update(source_id for source_id in inherited_ids if isinstance(source_id, str) and source_id in referenced_ids)
|
||||
for source in raw_sources[:100]:
|
||||
if not isinstance(source, dict):
|
||||
continue
|
||||
@ -95,7 +103,53 @@ def cited_source_artifact(messages: list[dict[str, Any]], content: str) -> dict[
|
||||
if source_id in sources:
|
||||
continue
|
||||
if len(sources) >= 100 or len(text) > remaining:
|
||||
if include_omissions and source_id in referenced_ids:
|
||||
omitted_ids.add(source_id)
|
||||
continue
|
||||
sources[source_id] = dict(source)
|
||||
remaining -= len(text)
|
||||
# Another captured result may carry a source previously omitted by a child.
|
||||
omitted_ids.difference_update(sources)
|
||||
if include_omissions and (sources or omitted_ids):
|
||||
return {"knowledge_sources": {"version": 1, "sources": list(sources.values()), "omitted_source_ids": sorted(omitted_ids), "omitted_count": len(omitted_ids)}}
|
||||
return {"knowledge_sources": {"version": 1, "sources": list(sources.values())}} if sources else None
|
||||
|
||||
|
||||
def durable_source_artifact(messages: list[dict[str, Any]], content: str, *, max_chars: int) -> dict[str, Any] | None:
|
||||
"""Project cited evidence into a bounded, self-contained batch export.
|
||||
|
||||
The report may already be truncated. Only complete references in that
|
||||
stored report select sources. Excerpts are never shortened under their IDs;
|
||||
records exceeding the separate JSON character budget are explicitly omitted.
|
||||
This does not persist arbitrary child artifacts or fetch provider data.
|
||||
"""
|
||||
artifact = cited_source_artifact(messages, content, include_omissions=True)
|
||||
if artifact is None:
|
||||
return None
|
||||
candidates = artifact["knowledge_sources"]["sources"]
|
||||
payload = {"version": 1, "sources": [], "omitted_count": len(candidates) + artifact["knowledge_sources"]["omitted_count"]}
|
||||
# Reserve the largest omission counter and the complete JSON envelope.
|
||||
remaining = max_chars - len(json.dumps({"knowledge_sources": payload}, ensure_ascii=False))
|
||||
if remaining < 0:
|
||||
return None
|
||||
fields = ("id", "provider", "dataset_id", "document_id", "chunk_id", "dataset_name", "document_name", "text")
|
||||
for source in candidates:
|
||||
if not all(isinstance(source.get(field), str) for field in fields):
|
||||
continue
|
||||
if source["provider"] != "ragflow" or not re.fullmatch(r"[a-f0-9]{32}-[1-9][0-9]{0,2}", source["id"]):
|
||||
continue
|
||||
if not all(0 < len(source[field]) <= 256 for field in ("dataset_id", "document_id", "chunk_id")):
|
||||
continue
|
||||
if any(len(source[field]) > 512 for field in ("dataset_name", "document_name")):
|
||||
continue
|
||||
record = {field: source[field] for field in fields}
|
||||
pages = source.get("pages")
|
||||
record["pages"] = sorted({page for page in pages[:100] if isinstance(page, int) and not isinstance(page, bool) and 1 <= page <= 1_000_000}) if isinstance(pages, list) else []
|
||||
record["truncated"] = source.get("truncated") is True
|
||||
cost = len(json.dumps(record, ensure_ascii=False)) + (2 if payload["sources"] else 0)
|
||||
if cost > remaining:
|
||||
continue
|
||||
payload["sources"].append(record)
|
||||
remaining -= cost
|
||||
payload["omitted_count"] -= 1
|
||||
return {"knowledge_sources": payload}
|
||||
|
||||
@ -28,11 +28,14 @@ The empty-DB path keeps using `create_all` because `Base.metadata` is the only a
|
||||
`0026_mcp_task_lease_tokens` → `0027_notification_deliveries` →
|
||||
`0028_parked_attempts` → `0029_scheduler_agent_tasks` →
|
||||
`0030_notification_claim_tokens` → `0031_scheduled_streak_boundary` →
|
||||
`0032_activity_and_task_events` (current head). The preference
|
||||
`0032_activity_and_task_events` → `0033_batch_result_artifact`
|
||||
(current head). The preference
|
||||
revision adds a separate owner/key table with a cascading users foreign key and
|
||||
does not alter users; the project-documents revision adds a new owner-scoped
|
||||
shelf table, and the MCP lease-token revision adds two nullable token columns to
|
||||
`mcp_tasks`, so the bootstrap forward-compat floor is unchanged.
|
||||
`mcp_tasks`. The batch-result revision adds a nullable JSON evidence snapshot
|
||||
to `subagent_batch_items`, without changing existing report text or backfilling
|
||||
historical evidence; the bootstrap forward-compat floor is unchanged.
|
||||
The incarnation revision deliberately retains the exact id audited by the
|
||||
rollback-floor binary; Alembic orders revisions by `down_revision`, not by the
|
||||
numeric prefix.
|
||||
|
||||
@ -0,0 +1,22 @@
|
||||
"""Preserve nullable evidence snapshots in durable batch results."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
|
||||
revision = "0033_batch_result_artifact"
|
||||
down_revision = "0032_activity_and_task_events"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
from deerflow.persistence.migrations._helpers import safe_add_column
|
||||
|
||||
safe_add_column("subagent_batch_items", sa.Column("result_artifact", sa.JSON(), nullable=True))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
from deerflow.persistence.migrations._helpers import safe_drop_column
|
||||
|
||||
safe_drop_column("subagent_batch_items", "result_artifact")
|
||||
@ -56,6 +56,7 @@ class SubagentBatchItemRow(Base):
|
||||
cancel_requested_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True)
|
||||
model_name: Mapped[str | None] = mapped_column(String(128), nullable=True)
|
||||
result: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
result_artifact: Mapped[dict[str, Any] | None] = mapped_column(JSON, nullable=True)
|
||||
result_preview: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
result_truncated: Mapped[bool] = mapped_column(default=False)
|
||||
error: Mapped[str | None] = mapped_column(Text, nullable=True)
|
||||
|
||||
@ -88,6 +88,7 @@ class SubagentBatchRepository:
|
||||
data["acceptance_verdict"] = validate_acceptance_verdict(row.acceptance_verdict)
|
||||
if include_result:
|
||||
data["result"] = row.result
|
||||
data["result_artifact"] = row.result_artifact
|
||||
for key in _ITEM_TIMESTAMP_FIELDS:
|
||||
if data.get(key) is not None:
|
||||
data[key] = coerce_iso(data[key])
|
||||
@ -412,6 +413,7 @@ class SubagentBatchRepository:
|
||||
model_name: str | None,
|
||||
completed_at: datetime,
|
||||
acceptance_verdict: AcceptanceVerdict | None = None,
|
||||
result_artifact: dict[str, Any] | None = None,
|
||||
) -> bool:
|
||||
async with self._sf() as session:
|
||||
item = (
|
||||
@ -436,6 +438,7 @@ class SubagentBatchRepository:
|
||||
item.token_usage = token_usage
|
||||
item.updated_at = completed_at
|
||||
item.acceptance_verdict = None
|
||||
item.result_artifact = None
|
||||
if cancelled:
|
||||
item.status = "cancelled"
|
||||
item.error = "Cancelled by user"
|
||||
@ -444,6 +447,7 @@ class SubagentBatchRepository:
|
||||
item.status = "succeeded"
|
||||
item.acceptance_verdict = validate_acceptance_verdict(acceptance_verdict)
|
||||
item.result = result
|
||||
item.result_artifact = result_artifact
|
||||
item.result_preview = result_preview
|
||||
item.result_truncated = result_truncated
|
||||
item.error = None
|
||||
@ -579,6 +583,7 @@ class SubagentBatchRepository:
|
||||
item.attempt = 0
|
||||
item.error = None
|
||||
item.result = None
|
||||
item.result_artifact = None
|
||||
item.result_preview = None
|
||||
item.result_truncated = False
|
||||
item.acceptance_verdict = None
|
||||
|
||||
@ -7,6 +7,15 @@ Registry overrides must not mutate `BUILTIN_SUBAGENTS`.
|
||||
Direct returns use the compiled tool registry, including middleware tools. Match current-turn IDs in call order; error ToolMessages fail the task with outputs preserved.
|
||||
Durable batch specs store overlays as JSON and restore them before execution.
|
||||
|
||||
Batch evidence uses nullable `result_artifact` (migration 0033), persisted with
|
||||
successful results under the same lease fence. Only full-result exports expose
|
||||
bounded, cited RAG snapshots; compact item projections omit them. Cancelled,
|
||||
failed and stale attempts cannot publish evidence; retry clears it. Cover the
|
||||
formatter/step capture -> worker -> reopened SQLite -> authorized export chain
|
||||
in `tests/test_batch_rag_evidence.py`; no provider fetch or preview UI is added.
|
||||
Explicitly enabled nested tasks propagate report-bound omission IDs; filter and
|
||||
deduplicate at each forwarding boundary before publishing the durable count.
|
||||
|
||||
**Remote empty files**: GNU stat's `regular file`/`regular empty file` are regular files. Empty files pass exists/file_written, fail non-empty; reject symlinks, directories, FIFOs.
|
||||
|
||||
**Context**: Capture after validation, before setup. Keep genuine replies, even hidden clarifications; exclude framework state and unpaired calls. Mark unserializable media as omitted.
|
||||
|
||||
@ -7,6 +7,7 @@ import uuid
|
||||
from datetime import UTC, datetime
|
||||
from typing import Any
|
||||
|
||||
from deerflow.community.ragflow.sources import durable_source_artifact
|
||||
from deerflow.config.app_config import AppConfig, get_app_config
|
||||
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
|
||||
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
||||
@ -321,6 +322,11 @@ class SubagentBatchService:
|
||||
truncated = len(raw_result) > self._config.max_result_chars
|
||||
stored_result = raw_result[: self._config.max_result_chars] if raw_result else None
|
||||
preview = raw_result[: self._config.result_preview_max_chars] if raw_result else None
|
||||
result_artifact = durable_source_artifact(
|
||||
getattr(result, "ai_messages", None) or [],
|
||||
stored_result or "",
|
||||
max_chars=self._config.max_result_chars,
|
||||
)
|
||||
acceptance_verdict = None
|
||||
if result.status is SubagentStatus.COMPLETED and item.get("acceptance_criteria"):
|
||||
try:
|
||||
@ -344,6 +350,7 @@ class SubagentBatchService:
|
||||
model_name=effective_model,
|
||||
completed_at=datetime.now(UTC),
|
||||
acceptance_verdict=acceptance_verdict,
|
||||
result_artifact=result_artifact,
|
||||
)
|
||||
except asyncio.CancelledError:
|
||||
if execution_id is not None:
|
||||
|
||||
@ -648,7 +648,7 @@ def _task_result_command(
|
||||
content=content,
|
||||
tool_call_id=tool_call_id,
|
||||
name="task",
|
||||
artifact=cited_source_artifact(source_messages or [], content),
|
||||
artifact=cited_source_artifact(source_messages or [], content, include_omissions=True),
|
||||
additional_kwargs=make_subagent_additional_kwargs(
|
||||
status,
|
||||
result=result,
|
||||
|
||||
@ -32,6 +32,22 @@ The local sandbox's UTF-8 subprocess guard inspects each text-mode call with
|
||||
`ast`, checking both `encoding` and `errors`; module-wide literal counts can
|
||||
hide unpinned calls behind unrelated settings.
|
||||
|
||||
## PostgreSQL batch fixtures
|
||||
|
||||
Batch fixtures use `support.postgres.asyncpg_test_url` to map libpq `sslmode`
|
||||
to asyncpg `ssl` and remove unsupported `channel_binding` before constructing
|
||||
database config. Preserve TLS modes, credentials and other query options;
|
||||
reject conflicting `ssl`/`sslmode` values. CI uses `?sslmode=disable`; validate
|
||||
that URI shape against a real test database, not only a parameter-free local
|
||||
URI. Reuse the normalized config for reopening and teardown, and drop only the
|
||||
fixture's UUID schema. This is test-only handling; production connection and
|
||||
TLS policy are unchanged.
|
||||
|
||||
The 0033 batch-evidence migration fixture uses the same adapter. Its connection
|
||||
contract probes execute the actual migration test setup through SQLAlchemy's
|
||||
dialect argument conversion, stopping before database acquisition; the real
|
||||
SQLite/PostgreSQL cases still exercise upgrade, downgrade and re-upgrade.
|
||||
|
||||
## Real Compose tests
|
||||
|
||||
`support/compose.py` probes `docker compose version --short` and requires Compose
|
||||
|
||||
15
backend/tests/support/postgres.py
Normal file
15
backend/tests/support/postgres.py
Normal file
@ -0,0 +1,15 @@
|
||||
"""Adapt libpq test URIs to SQLAlchemy's asyncpg connection options."""
|
||||
|
||||
from sqlalchemy.engine import make_url
|
||||
|
||||
|
||||
def asyncpg_test_url(uri: str) -> str:
|
||||
url = make_url(uri)
|
||||
query = dict(url.query)
|
||||
if "sslmode" in query:
|
||||
sslmode = query.pop("sslmode")
|
||||
if "ssl" in query and query["ssl"] != sslmode:
|
||||
raise ValueError("Conflicting ssl and sslmode in the test PostgreSQL URI")
|
||||
query["ssl"] = sslmode
|
||||
query.pop("channel_binding", None)
|
||||
return url.set(drivername="postgresql+asyncpg", query=query).render_as_string(hide_password=False)
|
||||
449
backend/tests/test_batch_rag_evidence.py
Normal file
449
backend/tests/test_batch_rag_evidence.py
Normal file
@ -0,0 +1,449 @@
|
||||
"""RAG producer -> durable worker -> reopened SQL -> owner JSONL evidence."""
|
||||
|
||||
import asyncio
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import uuid
|
||||
from copy import deepcopy
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from langchain_core.messages import ToolMessage
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
from deerflow.community.ragflow.formatting import format_retrieval_sources
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
from deerflow.config.subagent_batches_config import SubagentBatchesConfig
|
||||
from deerflow.config.subagent_runtime_config import SubagentRuntimeConfig
|
||||
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine_from_config
|
||||
from deerflow.persistence.subagent_batches import SubagentBatchRepository
|
||||
from deerflow.subagents import batch_service
|
||||
from deerflow.subagents.batch_runtime import BatchSubmitRequest
|
||||
from deerflow.subagents.step_events import capture_step_message
|
||||
|
||||
|
||||
class Status(Enum):
|
||||
RUNNING = "running"
|
||||
COMPLETED = "completed"
|
||||
FAILED = "failed"
|
||||
|
||||
@property
|
||||
def is_terminal(self):
|
||||
return self is not Status.RUNNING
|
||||
|
||||
|
||||
def retrieval(*, text="Original evidence", name="knowledge_search"):
|
||||
content, artifact = format_retrieval_sources(
|
||||
{"chunks": [{"id": "chunk", "dataset_id": "kb", "document_id": "doc", "document_keyword": "Manual.pdf", "content": text}]},
|
||||
dataset_names_by_id={"kb": "Engineering"},
|
||||
max_chars_per_chunk=20_000,
|
||||
max_total_chars=30_000,
|
||||
)
|
||||
captured = []
|
||||
capture_step_message(ToolMessage(content=content, artifact=artifact, name=name, tool_call_id="search", id="message"), captured, set())
|
||||
return content, captured
|
||||
|
||||
|
||||
@pytest_asyncio.fixture(params=["sqlite", "postgres"])
|
||||
async def env(monkeypatch, tmp_path, request):
|
||||
schema = None
|
||||
if request.param == "postgres":
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI (real Postgres batch evidence chain)")
|
||||
schema = f"batch_evidence_{uuid.uuid4().hex}"
|
||||
db = DatabaseConfig(backend="postgres", postgres_url=asyncpg_test_url(uri), postgres_schema=schema)
|
||||
else:
|
||||
db = DatabaseConfig(backend="sqlite", sqlite_dir=str(tmp_path / "db"))
|
||||
service = None
|
||||
try:
|
||||
await init_engine_from_config(db)
|
||||
if schema:
|
||||
from sqlalchemy import text
|
||||
|
||||
async with get_session_factory()() as session:
|
||||
assert (await session.execute(text("SELECT current_schema()"))).scalar_one() == schema
|
||||
repo = SubagentBatchRepository(get_session_factory())
|
||||
state = SimpleNamespace(repo=repo, db=db, result=SimpleNamespace(status=Status.COMPLETED, result="plain report", ai_messages=[], error=None, stop_reason=None, token_usage_records=[]))
|
||||
|
||||
class Executor:
|
||||
def __init__(self, **kwargs):
|
||||
state.owner = kwargs["user_id"]
|
||||
state.scope = kwargs["knowledge_scope"]
|
||||
|
||||
def execute_async(self, prompt, task_id=None):
|
||||
return task_id
|
||||
|
||||
monkeypatch.setattr(batch_service, "SubagentExecutor", Executor)
|
||||
monkeypatch.setattr(batch_service, "SubagentStatus", Status)
|
||||
monkeypatch.setattr(batch_service, "get_background_task_result", lambda _: state.result)
|
||||
monkeypatch.setattr(batch_service, "cleanup_background_task", lambda _: None)
|
||||
monkeypatch.setattr(batch_service, "resolve_subagent_model_name", lambda *args, **kwargs: "offline-model")
|
||||
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: [])
|
||||
state.service = batch_service.SubagentBatchService(repository=repo, config=SubagentBatchesConfig(), runtime_config=SubagentRuntimeConfig(), app_config=SimpleNamespace())
|
||||
service = state.service
|
||||
yield state
|
||||
finally:
|
||||
try:
|
||||
if service is not None:
|
||||
await service.stop()
|
||||
finally:
|
||||
try:
|
||||
await close_engine()
|
||||
finally:
|
||||
if schema:
|
||||
import sqlalchemy as sa
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
|
||||
engine = create_async_engine(db.app_sqlalchemy_url)
|
||||
try:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
async def execute(env):
|
||||
batch = await env.service.submit(
|
||||
BatchSubmitRequest(
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
run_id="run-1",
|
||||
tool_call_id="call-1",
|
||||
submission_key="submission",
|
||||
title="Research",
|
||||
subagent_type="general-purpose",
|
||||
items=[{"key": "one", "prompt": "Research the selected knowledge base"}],
|
||||
max_live_items=None,
|
||||
max_running_items=None,
|
||||
execution_spec={"subagent_config": {"name": "general-purpose", "description": "Researcher", "system_prompt": "Research"}, "knowledge_scope": {"version": 1, "mode": "selected", "dataset_ids": ["kb"]}},
|
||||
)
|
||||
)
|
||||
await env.service.run_once(now=datetime.now(UTC))
|
||||
await asyncio.gather(*list(env.service._executions.values()))
|
||||
return batch
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("tool_name", ["knowledge_search", "task"])
|
||||
async def test_cited_evidence_survives_worker_database_reopen_and_export(env, monkeypatch, tool_name):
|
||||
from app.gateway.routers import subagent_batches as router
|
||||
|
||||
env.result.result, env.result.ai_messages = retrieval(name=tool_name)
|
||||
source = env.result.ai_messages[0]["artifact"]["knowledge_sources"]["sources"][0]
|
||||
expected_source = deepcopy(source)
|
||||
batch = await execute(env)
|
||||
assert env.owner == "user-1" and env.scope["dataset_ids"] == ["kb"]
|
||||
# Actual engine teardown/reopen; neither the worker nor provider can serve evidence.
|
||||
await env.service.stop()
|
||||
await close_engine()
|
||||
await init_engine_from_config(env.db)
|
||||
repo = SubagentBatchRepository(get_session_factory())
|
||||
source["text"] = "Provider content changed after retrieval"
|
||||
env.result.ai_messages = []
|
||||
assert "result_artifact" not in (await repo.list_items(batch["id"], user_id="user-1"))[0]
|
||||
assert await repo.list_items(batch["id"], user_id="other", include_result=True) is None
|
||||
monkeypatch.setattr(router, "get_current_user", AsyncMock(return_value="user-1"))
|
||||
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(subagent_batch_repo=repo, subagent_batches_available=False)))
|
||||
response = await router.export_batch_results.__wrapped__(thread_id="thread-1", batch_id=batch["id"], request=request)
|
||||
rows = [json.loads(line) async for line in response.body_iterator]
|
||||
row = rows[0]
|
||||
assert row["result"] == env.result.result
|
||||
assert row["result_artifact"]["knowledge_sources"]["sources"] == [expected_source]
|
||||
assert "Original evidence" in row["result_artifact"]["knowledge_sources"]["sources"][0]["text"]
|
||||
assert "execution_spec" not in row and "prompt" not in row
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_real_subagent_graph_captures_artifact_for_durable_export(env, monkeypatch, tmp_path):
|
||||
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
|
||||
from langchain_core.messages import AIMessage
|
||||
from langchain_core.tools import StructuredTool
|
||||
|
||||
from deerflow.config.app_config import AppConfig
|
||||
from deerflow.subagents.config import SubagentConfig
|
||||
|
||||
# conftest stubs executor to break a circular import. Load the production
|
||||
# executor, keeping its actual graph, middleware, ToolNode and step capture.
|
||||
path = Path(__file__).parents[1] / "packages/harness/deerflow/subagents/executor.py"
|
||||
spec = importlib.util.spec_from_file_location("_batch_evidence_executor", path)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
monkeypatch.setitem(sys.modules, spec.name, module)
|
||||
spec.loader.exec_module(module)
|
||||
monkeypatch.setenv("DEER_FLOW_HOME", str(tmp_path / "graph"))
|
||||
content, messages = retrieval()
|
||||
artifact = messages[0]["artifact"]
|
||||
|
||||
class Model(GenericFakeChatModel):
|
||||
def bind_tools(self, tools, **kwargs):
|
||||
return self
|
||||
|
||||
model = Model(messages=iter([AIMessage(content="Searching", tool_calls=[{"name": "knowledge_search", "args": {}, "id": "search"}])]))
|
||||
monkeypatch.setattr(module, "create_chat_model", lambda **kwargs: model)
|
||||
tool = StructuredTool.from_function(lambda: (content, artifact), name="knowledge_search", description="Offline captured evidence", response_format="content_and_artifact", return_direct=True)
|
||||
config = AppConfig.model_validate({"models": [{"name": "offline", "use": "langchain_openai:ChatOpenAI", "model": "offline"}], "sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"}, "summarization": {"enabled": False}})
|
||||
executor = module.SubagentExecutor(config=SubagentConfig(name="researcher", description="Research", skills=[]), tools=[tool], parent_model="offline", app_config=config, thread_id="thread-1", user_id="user-1")
|
||||
try:
|
||||
env.result = await executor._aexecute("Return evidence")
|
||||
assert env.result.status is module.SubagentStatus.COMPLETED, env.result.error
|
||||
assert env.result.result == content
|
||||
captured = [m for m in env.result.ai_messages if m["type"] == "tool"]
|
||||
assert captured[0]["artifact"] == artifact
|
||||
monkeypatch.setattr(batch_service, "SubagentStatus", module.SubagentStatus)
|
||||
batch = await execute(env)
|
||||
row = (await env.repo.list_items(batch["id"], user_id="user-1", include_result=True))[0]
|
||||
assert row["status"] == "succeeded"
|
||||
assert row["result_artifact"]["knowledge_sources"]["sources"] == artifact["knowledge_sources"]["sources"]
|
||||
finally:
|
||||
module._shutdown_isolated_subagent_loop()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("change", ["uncited", "unsupported-tool", "future-version", "invalid-id", "invalid-provider", "invalid-text", "invalid-locator", "overlong-name"])
|
||||
def test_unrelated_or_invalid_evidence_is_not_persisted(change):
|
||||
from deerflow.community.ragflow.sources import durable_source_artifact
|
||||
|
||||
content, messages = retrieval()
|
||||
source = messages[0]["artifact"]["knowledge_sources"]["sources"][0]
|
||||
if change == "uncited":
|
||||
content = "A report with no citations"
|
||||
elif change == "unsupported-tool":
|
||||
messages[0]["name"] = "arbitrary_tool"
|
||||
elif change == "future-version":
|
||||
messages[0]["artifact"]["knowledge_sources"]["version"] = 2
|
||||
elif change == "invalid-id":
|
||||
content = content.replace(source["id"], "invalid")
|
||||
source["id"] = "invalid"
|
||||
elif change == "invalid-provider":
|
||||
source["provider"] = "other"
|
||||
elif change == "invalid-text":
|
||||
source["text"] = None
|
||||
elif change == "invalid-locator":
|
||||
source["document_id"] = ""
|
||||
else:
|
||||
source["document_name"] = "x" * 513
|
||||
result = durable_source_artifact(messages, content, max_chars=1000)
|
||||
assert result is None or result["knowledge_sources"]["sources"] == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("max_chars", [1, 1000, 1800, 10_000])
|
||||
def test_evidence_budget_keeps_whole_records_and_reports_omissions(max_chars):
|
||||
from deerflow.community.ragflow.sources import durable_source_artifact
|
||||
|
||||
content, messages = retrieval(text='Quoted "evidence" and 中文 ' * 200)
|
||||
original = deepcopy(messages)
|
||||
small_content, small_messages = retrieval(text="A small source fits")
|
||||
# Both citations occur in the report; a large early source cannot starve a small later one.
|
||||
result = durable_source_artifact(messages + small_messages, content + small_content, max_chars=max_chars)
|
||||
assert messages == original
|
||||
if result is None:
|
||||
assert max_chars == 1
|
||||
return
|
||||
assert len(json.dumps(result, ensure_ascii=False)) <= max_chars
|
||||
payload = result["knowledge_sources"]
|
||||
assert len(payload["sources"]) + payload["omitted_count"] == 2
|
||||
originals = [m["artifact"]["knowledge_sources"]["sources"][0] for m in messages + small_messages]
|
||||
assert all(source in originals for source in payload["sources"])
|
||||
if max_chars in (1000, 1800):
|
||||
assert payload["sources"] == [originals[1]]
|
||||
assert payload["omitted_count"] == 1
|
||||
|
||||
|
||||
def test_forwarding_limit_counts_duplicate_omissions_once():
|
||||
from deerflow.community.ragflow.sources import cited_source_artifact, durable_source_artifact
|
||||
|
||||
entries = [retrieval(text=f"Evidence {index}") for index in range(102)]
|
||||
content = "\n".join(entry[0] for entry in entries)
|
||||
messages = [message for entry in entries for message in entry[1]]
|
||||
result = durable_source_artifact(messages + messages, content, max_chars=100_000)
|
||||
assert len(result["knowledge_sources"]["sources"]) == 100
|
||||
assert result["knowledge_sources"]["omitted_count"] == 2
|
||||
# Ordinary task forwarding retains its existing artifact shape.
|
||||
assert "omitted_count" not in cited_source_artifact(messages, content)["knowledge_sources"]
|
||||
|
||||
|
||||
def forwarded_task(messages, content):
|
||||
from deerflow.tools.builtins.task_tool import _task_result_command
|
||||
|
||||
command = _task_result_command(tool_call_id="child", status="completed", result=content, source_messages=messages)
|
||||
captured = []
|
||||
capture_step_message(command.update["messages"][0], captured, set())
|
||||
return captured
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_nested_task_omissions_survive_worker_reopen_and_export(env, monkeypatch):
|
||||
from app.gateway.routers import subagent_batches as router
|
||||
|
||||
entries = [retrieval(text=f"Evidence {index}") for index in range(102)]
|
||||
report = "\n".join(entry[0] for entry in entries)
|
||||
messages = [message for entry in entries for message in entry[1]]
|
||||
forwarded = forwarded_task(messages, report)
|
||||
# Repeated child messages and a second delegation must not double-count.
|
||||
env.result.result = report
|
||||
env.result.ai_messages = forwarded_task(forwarded + forwarded, report)
|
||||
snapshot = deepcopy(env.result.ai_messages)
|
||||
batch = await execute(env)
|
||||
await env.service.stop()
|
||||
await close_engine()
|
||||
await init_engine_from_config(env.db)
|
||||
repo = SubagentBatchRepository(get_session_factory())
|
||||
monkeypatch.setattr(router, "get_current_user", AsyncMock(return_value="user-1"))
|
||||
request = SimpleNamespace(app=SimpleNamespace(state=SimpleNamespace(subagent_batch_repo=repo)))
|
||||
response = await router.export_batch_results.__wrapped__(thread_id="thread-1", batch_id=batch["id"], request=request)
|
||||
row = json.loads([line async for line in response.body_iterator][0])
|
||||
payload = row["result_artifact"]["knowledge_sources"]
|
||||
assert len(payload["sources"]) == 100
|
||||
assert payload["omitted_count"] == 2
|
||||
assert "omitted_source_ids" not in payload
|
||||
assert env.result.ai_messages == snapshot
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selection", ["retained", "omitted", "unknown", "partial"])
|
||||
def test_forwarded_omissions_follow_only_current_complete_references(selection):
|
||||
from deerflow.community.ragflow.sources import durable_source_artifact
|
||||
|
||||
entries = [retrieval(text=f"Evidence {index}") for index in range(101)]
|
||||
report = "\n".join(entry[0] for entry in entries)
|
||||
messages = [message for entry in entries for message in entry[1]]
|
||||
forwarded = forwarded_task(messages, report)
|
||||
selected = {"retained": entries[0][0], "omitted": entries[-1][0], "unknown": "[citation:1](#knowledge-" + "f" * 32 + "-1)", "partial": entries[-1][0].split(")")[0]}[selection]
|
||||
result = durable_source_artifact(forwarded, selected, max_chars=1000)
|
||||
if selection in {"unknown", "partial"}:
|
||||
assert result is None
|
||||
else:
|
||||
assert result["knowledge_sources"]["omitted_count"] == (selection == "omitted")
|
||||
assert len(result["knowledge_sources"]["sources"]) == (selection == "retained")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reverse", [False, True])
|
||||
def test_recaptured_source_supersedes_forwarded_omission(reverse):
|
||||
from deerflow.community.ragflow.sources import durable_source_artifact
|
||||
|
||||
content, messages = retrieval()
|
||||
source_id = messages[0]["artifact"]["knowledge_sources"]["sources"][0]["id"]
|
||||
omitted = {"type": "tool", "name": "task", "artifact": {"knowledge_sources": {"version": 1, "sources": [], "omitted_source_ids": [source_id, source_id, None, "invalid"], "omitted_count": 999}}}
|
||||
inputs = [omitted, *messages] if not reverse else [*messages, omitted]
|
||||
result = durable_source_artifact(inputs, content, max_chars=1000)
|
||||
assert result["knowledge_sources"]["sources"] == messages[0]["artifact"]["knowledge_sources"]["sources"]
|
||||
assert result["knowledge_sources"]["omitted_count"] == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_report_truncation_selects_only_complete_stored_citations(env):
|
||||
content, messages = retrieval()
|
||||
env.service._config = SubagentBatchesConfig(max_result_chars=1000, result_preview_max_chars=64)
|
||||
env.result.result = "x" * 1100 + content
|
||||
env.result.ai_messages = messages
|
||||
batch = await execute(env)
|
||||
row = (await env.repo.list_items(batch["id"], user_id="user-1", include_result=True))[0]
|
||||
assert row["result"] == "x" * 1000 and row["result_truncated"] is True
|
||||
assert row["result_artifact"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plain_results_keep_nullable_evidence(env):
|
||||
batch = await execute(env)
|
||||
row = (await env.repo.list_items(batch["id"], user_id="user-1", include_result=True))[0]
|
||||
assert row["result"] == "plain report" and row["result_artifact"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("action", ["cancel", "cancel-running", "stale-lease", "failure", "retry"])
|
||||
async def test_fenced_or_failed_attempts_cannot_publish_evidence(env, action):
|
||||
content, messages = retrieval()
|
||||
from deerflow.community.ragflow.sources import durable_source_artifact
|
||||
|
||||
artifact = durable_source_artifact(messages, content, max_chars=1000)
|
||||
batch = await env.repo.create_batch(
|
||||
batch_id="batch",
|
||||
user_id="user-1",
|
||||
thread_id="thread-1",
|
||||
run_id=None,
|
||||
tool_call_id=None,
|
||||
submission_key="key",
|
||||
title="R",
|
||||
subagent_type="general-purpose",
|
||||
items=[{"key": "item", "prompt": "Research"}],
|
||||
max_live_items=1,
|
||||
max_running_items=1,
|
||||
max_attempts=2 if action == "stale-lease" else 1,
|
||||
execution_spec={},
|
||||
)
|
||||
now = datetime.now(UTC)
|
||||
item = (await env.repo.claim_items(now=now, lease_owner="worker", lease_seconds=60, limit=1))[0]
|
||||
if action == "cancel-running":
|
||||
assert await env.repo.mark_item_running(item["id"], lease_owner="worker", now=now)
|
||||
if action in {"cancel", "cancel-running"}:
|
||||
await env.repo.cancel_batch(batch["id"], user_id="user-1")
|
||||
if action == "stale-lease":
|
||||
reclaimed = await env.repo.claim_items(now=now + timedelta(seconds=61), lease_owner="replacement", lease_seconds=60, limit=1)
|
||||
assert len(reclaimed) == 1 and reclaimed[0]["attempt"] == 2
|
||||
accepted = await env.repo.finalize_item(
|
||||
item["id"],
|
||||
lease_owner="worker",
|
||||
succeeded=action not in {"failure", "retry"},
|
||||
result=content,
|
||||
result_preview=content[:64],
|
||||
result_truncated=False,
|
||||
result_artifact=artifact,
|
||||
error="failed",
|
||||
stop_reason=None,
|
||||
token_usage=None,
|
||||
model_name=None,
|
||||
completed_at=now,
|
||||
)
|
||||
assert accepted is (action not in {"cancel", "cancel-running", "stale-lease"})
|
||||
if action == "retry":
|
||||
# A legacy/external failure row may have old evidence; explicit retry clears it.
|
||||
from deerflow.persistence.subagent_batches.model import SubagentBatchItemRow
|
||||
|
||||
async with get_session_factory()() as session:
|
||||
stored = await session.get(SubagentBatchItemRow, item["id"])
|
||||
stored.result_artifact = artifact
|
||||
await session.commit()
|
||||
assert await env.repo.retry_item(batch["id"], item["id"], user_id="other") is None
|
||||
await env.repo.retry_item(batch["id"], item["id"], user_id="user-1")
|
||||
row = (await env.repo.list_items(batch["id"], user_id="user-1", include_result=True))[0]
|
||||
assert row["result_artifact"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("user,permission,thread,status", [("user-1", True, "thread-1", 200), ("user-1", False, "thread-1", 403), ("other", True, "thread-1", 404), ("user-1", True, "other-thread", 404), (None, True, "thread-1", 401)])
|
||||
async def test_http_export_checks_permission_thread_and_batch_owner(env, monkeypatch, user, permission, thread, status):
|
||||
import httpx
|
||||
from fastapi import FastAPI
|
||||
|
||||
from app.gateway import deps
|
||||
from app.gateway.authz import AuthContext
|
||||
from app.gateway.routers import subagent_batches as router
|
||||
|
||||
env.result.result, env.result.ai_messages = retrieval()
|
||||
batch = await execute(env)
|
||||
app = FastAPI()
|
||||
app.include_router(router.router)
|
||||
app.state.subagent_batch_repo = env.repo
|
||||
|
||||
@app.middleware("http")
|
||||
async def auth_boundary(request, call_next):
|
||||
request.state.auth = AuthContext(user=SimpleNamespace(id=user) if user else None, permissions=["threads:read"] if permission else [])
|
||||
return await call_next(request)
|
||||
|
||||
# Authentication and the thread store are controlled; the actual permission decorator,
|
||||
# route, SQL batch ownership, streaming and serialization remain production code.
|
||||
monkeypatch.setattr(deps, "get_thread_store", lambda request: SimpleNamespace(check_access=AsyncMock(return_value=thread == "thread-1")))
|
||||
monkeypatch.setattr(router, "get_current_user", AsyncMock(return_value=user))
|
||||
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
|
||||
response = await client.get(f"/api/threads/{thread}/subagent-batches/{batch['id']}/results.jsonl")
|
||||
assert response.status_code == status
|
||||
if status == 200:
|
||||
row = json.loads(response.text)
|
||||
assert row["result_artifact"]["knowledge_sources"]["sources"][0]["text"] == "Original evidence"
|
||||
else:
|
||||
assert "Original evidence" not in response.text
|
||||
@ -23,7 +23,7 @@ OUTBOX = "0027_notification_deliveries"
|
||||
PARKED = "0028_parked_attempts"
|
||||
SCHEDULER = "0029_scheduler_agent_tasks"
|
||||
CLAIM_TOKENS = "0030_notification_claim_tokens"
|
||||
HEAD = "0032_activity_and_task_events"
|
||||
HEAD = "0033_batch_result_artifact"
|
||||
PREVIOUS = "0026_mcp_task_lease_tokens"
|
||||
TABLE = "notification_deliveries"
|
||||
|
||||
|
||||
@ -23,7 +23,7 @@ from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
|
||||
from deerflow.persistence.scheduled_tasks import ScheduledTaskRepository
|
||||
|
||||
REVISION = "0029_scheduler_agent_tasks"
|
||||
CURRENT_HEAD = "0032_activity_and_task_events"
|
||||
CURRENT_HEAD = "0033_batch_result_artifact"
|
||||
NEXT = "0030_notification_claim_tokens"
|
||||
PREVIOUS = "0028_parked_attempts"
|
||||
TASK_FIELDS = {"origin_thread_id", "goal_objective", "max_runs", "end_at", "standing_notes"}
|
||||
|
||||
@ -17,6 +17,7 @@ from deerflow.persistence import bootstrap
|
||||
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
|
||||
|
||||
REVISION = "0032_activity_and_task_events"
|
||||
CURRENT_HEAD = "0033_batch_result_artifact"
|
||||
PREVIOUS = "0031_scheduled_streak_boundary"
|
||||
TABLES = {"thread_read_markers", "thread_read_versions", "scheduled_task_events"}
|
||||
INDEXES = {
|
||||
@ -27,10 +28,11 @@ INDEXES = {
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
async def test_0032_is_the_single_head_after_0031():
|
||||
async def test_0032_remains_in_the_single_migration_chain():
|
||||
script = ScriptDirectory(str(bootstrap._MIGRATIONS_DIR))
|
||||
assert script.get_heads() == [REVISION]
|
||||
assert script.get_heads() == [CURRENT_HEAD]
|
||||
assert script.get_revision(REVISION).down_revision == PREVIOUS
|
||||
assert script.get_revision(CURRENT_HEAD).down_revision == REVISION
|
||||
# alembic_version.version_num is VARCHAR(32).
|
||||
assert len(REVISION) <= 32
|
||||
|
||||
|
||||
111
backend/tests/test_migration_0033_batch_result_artifact.py
Normal file
111
backend/tests/test_migration_0033_batch_result_artifact.py
Normal file
@ -0,0 +1,111 @@
|
||||
"""Nullable batch evidence migration preserves legacy report rows."""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
import sqlalchemy as sa
|
||||
from alembic import command
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
from deerflow.persistence import bootstrap
|
||||
from deerflow.persistence.postgres_schema import build_asyncpg_connect_args
|
||||
from deerflow.persistence.subagent_batches.model import SubagentBatchItemRow
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("backend", ["sqlite", "postgres"])
|
||||
async def test_upgrade_downgrade_and_reupgrade_preserve_report(tmp_path, backend):
|
||||
schema = None
|
||||
if backend == "postgres":
|
||||
uri = os.environ.get("TEST_POSTGRES_URI")
|
||||
if not uri:
|
||||
pytest.skip("requires TEST_POSTGRES_URI (real Postgres batch evidence migration)")
|
||||
schema = f"batch_evidence_{uuid.uuid4().hex}"
|
||||
engine = create_async_engine(asyncpg_test_url(uri), connect_args=build_asyncpg_connect_args(schema))
|
||||
else:
|
||||
engine = create_async_engine(f"sqlite+aiosqlite:///{tmp_path / 'batch.db'}")
|
||||
cfg = bootstrap._get_alembic_config(engine, postgres_schema=schema or "")
|
||||
try:
|
||||
if schema:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text(f'CREATE SCHEMA "{schema}"'))
|
||||
await asyncio.to_thread(bootstrap._upgrade, cfg, "0032_activity_and_task_events")
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
"INSERT INTO subagent_batches (id,user_id,thread_id,submission_key,title,subagent_type,status,total_items,max_live_items,max_running_items,max_attempts,execution_spec,created_at,updated_at) "
|
||||
"VALUES ('b','u','t','k','title','general-purpose','completed',1,1,1,2,'{}',CURRENT_TIMESTAMP,CURRENT_TIMESTAMP)"
|
||||
)
|
||||
)
|
||||
await conn.execute(
|
||||
sa.text(
|
||||
"INSERT INTO subagent_batch_items (id,batch_id,item_key,position,prompt,status,attempt,result,result_truncated,created_at,updated_at) "
|
||||
"VALUES ('i','b','k',0,'p','succeeded',1,'old report',FALSE,CURRENT_TIMESTAMP,CURRENT_TIMESTAMP)"
|
||||
)
|
||||
)
|
||||
await bootstrap.bootstrap_schema(engine, backend=backend, postgres_schema=schema or "")
|
||||
await bootstrap.bootstrap_schema(engine, backend=backend, postgres_schema=schema or "")
|
||||
async with engine.connect() as conn:
|
||||
row = (await conn.execute(sa.text("SELECT result,status,result_artifact FROM subagent_batch_items"))).one()
|
||||
assert tuple(row) == ("old report", "succeeded", None)
|
||||
snapshot = {"knowledge_sources": {"version": 1, "sources": [{"id": "a" * 32 + "-1", "provider": "ragflow", "dataset_id": "kb", "document_id": "doc", "chunk_id": "chunk", "text": "Captured evidence ☃"}], "omitted_count": 0}}
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.update(SubagentBatchItemRow).where(SubagentBatchItemRow.id == "i").values(result_artifact=snapshot))
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(sa.select(SubagentBatchItemRow.result_artifact).where(SubagentBatchItemRow.id == "i"))).scalar_one() == snapshot
|
||||
await asyncio.to_thread(command.downgrade, cfg, "0032_activity_and_task_events")
|
||||
async with engine.connect() as conn:
|
||||
columns = await conn.run_sync(lambda sync: {column["name"] for column in sa.inspect(sync).get_columns("subagent_batch_items")})
|
||||
assert "result_artifact" not in columns
|
||||
assert (await conn.execute(sa.text("SELECT result FROM subagent_batch_items"))).scalar_one() == "old report"
|
||||
await bootstrap.bootstrap_schema(engine, backend=backend, postgres_schema=schema or "")
|
||||
async with engine.connect() as conn:
|
||||
assert (await conn.execute(sa.text("SELECT result_artifact FROM subagent_batch_items"))).scalar_one() is None
|
||||
finally:
|
||||
if schema:
|
||||
async with engine.begin() as conn:
|
||||
await conn.execute(sa.text(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE'))
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
class _StopBeforeDatabase(Exception):
|
||||
"""End a connection-contract probe before acquiring database resources."""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("sslmode", [None, "disable", "allow", "prefer", "require", "verify-ca", "verify-full"])
|
||||
async def test_postgres_migration_setup_preserves_tls_policy(monkeypatch, tmp_path, sslmode):
|
||||
uri = "postgresql://user:password@localhost:5432/test?target_session_attrs=read-write&channel_binding=prefer"
|
||||
if sslmode is not None:
|
||||
uri += f"&sslmode={sslmode}"
|
||||
monkeypatch.setenv("TEST_POSTGRES_URI", uri)
|
||||
|
||||
def capture_engine(url, **kwargs):
|
||||
parsed = sa.engine.make_url(url)
|
||||
_args, connect_options = parsed.get_dialect()().create_connect_args(parsed)
|
||||
assert connect_options["target_session_attrs"] == "read-write"
|
||||
assert "sslmode" not in connect_options and "channel_binding" not in connect_options
|
||||
if sslmode is None:
|
||||
assert "ssl" not in connect_options
|
||||
else:
|
||||
assert connect_options.get("ssl") == sslmode
|
||||
raise _StopBeforeDatabase
|
||||
|
||||
monkeypatch.setitem(globals(), "create_async_engine", capture_engine)
|
||||
with pytest.raises(_StopBeforeDatabase):
|
||||
await test_upgrade_downgrade_and_reupgrade_preserve_report(tmp_path, "postgres")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_postgres_migration_setup_rejects_conflicting_tls_policy(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv("TEST_POSTGRES_URI", "postgresql://user:password@localhost:5432/test?ssl=require&sslmode=disable")
|
||||
|
||||
def reject_engine_creation(*args, **kwargs):
|
||||
pytest.fail("Conflicting TLS options must be rejected before engine creation")
|
||||
|
||||
monkeypatch.setitem(globals(), "create_async_engine", reject_engine_creation)
|
||||
with pytest.raises(ValueError, match="Conflicting ssl and sslmode"):
|
||||
await test_upgrade_downgrade_and_reupgrade_preserve_report(tmp_path, "postgres")
|
||||
35
backend/tests/test_support_postgres.py
Normal file
35
backend/tests/test_support_postgres.py
Normal file
@ -0,0 +1,35 @@
|
||||
"""Check the real asyncpg dialect's connection contract without a database."""
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.engine import make_url
|
||||
from support.postgres import asyncpg_test_url
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scheme", ["postgres", "postgresql", "postgresql+asyncpg"])
|
||||
@pytest.mark.parametrize("sslmode", [None, "disable", "allow", "prefer", "require", "verify-ca", "verify-full"])
|
||||
def test_asyncpg_test_url_preserves_tls_credentials_and_other_options(scheme, sslmode):
|
||||
uri = f"{scheme}://user:p%40ss%2Fword@localhost:5432/deerflow_test?target_session_attrs=read-write&channel_binding=prefer"
|
||||
if sslmode is not None:
|
||||
uri += f"&sslmode={sslmode}"
|
||||
url = make_url(asyncpg_test_url(uri))
|
||||
_args, kwargs = url.get_dialect()().create_connect_args(url)
|
||||
assert url.drivername == "postgresql+asyncpg"
|
||||
assert kwargs["password"] == "p@ss/word"
|
||||
assert kwargs["target_session_attrs"] == "read-write"
|
||||
assert "sslmode" not in kwargs and "channel_binding" not in kwargs
|
||||
if sslmode is None:
|
||||
assert "ssl" not in kwargs
|
||||
else:
|
||||
assert kwargs["ssl"] == sslmode
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query", ["ssl=require", "ssl=require&sslmode=require"])
|
||||
def test_asyncpg_test_url_preserves_explicit_ssl(query):
|
||||
url = make_url(asyncpg_test_url(f"postgresql://user:password@localhost/test?{query}"))
|
||||
_args, kwargs = url.get_dialect()().create_connect_args(url)
|
||||
assert kwargs["ssl"] == "require"
|
||||
|
||||
|
||||
def test_asyncpg_test_url_rejects_conflicting_tls_options():
|
||||
with pytest.raises(ValueError, match="Conflicting ssl and sslmode"):
|
||||
asyncpg_test_url("postgresql://user:password@localhost/test?ssl=require&sslmode=disable")
|
||||
Loading…
x
Reference in New Issue
Block a user