deer-flow/backend/scripts/migrate_agents_to_db.py
Aari 0d4d0cb17d
feat(agents): database-backed storage for custom agent definitions (#4359)
* feat(agents): database-backed storage for custom agent definitions

Add an agent_storage.backend switch (default file, behaviour-unchanged) with a
db backend that stores each custom agent as a row in the shared SQL persistence
layer, so a multi-instance deployment sees the same agents on every node
(#4331, #4357). Introduces an AgentStore interface routing all read/write
surfaces, an agents table + migration 0006, startup validation, and a file->db
importer. Follows the thread_meta store / run_events backend-switch /
0003_scheduled_tasks migration patterns; no new dependency.

* fix(agents): make db storage path production-ready (review round 1)

Addresses review feedback on the db/sync agent-storage path:

- sql.py: mirror the async engine's per-connection SQLite PRAGMAs on the sync
  engine (busy_timeout=30000, synchronous=NORMAL, foreign_keys=ON, WAL) so both
  engines behave identically against the shared DB; guard the engine cache with
  a lock (double-checked) so concurrent first-touch cannot build duplicate
  engines or register the connect listener twice.
- routers/agents.py + routers/assistants_compat.py: offload the sync-store reads
  that ran on the event loop (list/get/check, update's pre-read + legacy guard +
  refresh, and assistants_compat's four list routes) via asyncio.to_thread — on
  db+postgres each was a network round trip stalling the loop. Writes were
  already offloaded.
- file.py: translate the create() mkdir(exist_ok=False) race FileExistsError
  into AgentExistsError (router 409, matching SqlAgentStore's IntegrityError
  path); correct the _write docstring — per-file atomic replace, two commits
  sequential not transactional.

Tests: sync-engine PRAGMA + engine-cache reuse assertions; file create-race ->
AgentExistsError; strict Blockbuster anchor over the read endpoints so a
regression back onto the loop fails CI.

* fix(agents): address round-2 review on the db store path

- update_agent tool: align the docstring/inline comment with FileAgentStore._write.
  Cross-field write atomicity is db-only; the file backend commits config then
  soul via two sequential os.replace (a crash between them can leave a fresh
  config.yaml beside a stale SOUL.md). The dropped partial-write *reporting* is
  an intentional tradeoff — the stage-then-replace safety is preserved
  (test_update_agent_soul_failure_does_not_replace_config still holds).
- SqlAgentStore.update(): true upsert. Catch IntegrityError on the
  insert-on-missing branch, re-fetch and apply, so two concurrent first-time
  writes (e.g. two setup_agent handshakes) converge instead of surfacing a raw
  UNIQUE(user_id, name) violation as a 500. Symmetric with create().
- get_agent_store(): document the graph-subprocess config-resolution invariant
  (the except->file fallback is a genuine no-config path, not a mask for a
  misconfigured graph process) and pin it with two tests driving the real
  get_app_config() file resolution: db resolves from an on-disk config.yaml,
  file fallback when config is unresolvable.

* test(agents): cover SqlAgentStore.update() write-race upsert recovery

Mandatory-TDD test for the round-2 fix in 0680340a: two concurrent first-time
update()s where the loser's insert hits UNIQUE(user_id, name). Deterministically
forces the IntegrityError recovery path by making the first _row probe miss the
committed winner, and asserts last-writer-wins instead of a surfaced 500.
2026-07-23 08:03:21 +08:00

94 lines
3.6 KiB
Python

#!/usr/bin/env python
"""One-shot importer: copy file-backed custom agents into the ``db`` agent store.
For operators switching ``agent_storage.backend`` from ``file`` to ``db``. Reads
every agent from the on-disk layout (both the per-user
``{base_dir}/users/{user_id}/agents/`` and the legacy shared
``{base_dir}/agents/``, with the same shadowing rule the file store uses) and
writes each as a row in the shared ``agents`` table.
Design (mirrors ``scripts/migrate_user_isolation.py``):
- Explicit, operator-run. Nothing auto-imports on boot.
- Idempotent: an agent already present in the db is skipped, so re-running is safe.
- Non-destructive: the on-disk files are left untouched, so unsetting
``agent_storage.backend`` (back to ``file``) is a clean rollback.
Usage::
python scripts/migrate_agents_to_db.py [--dry-run]
Requires ``database.backend`` to be ``sqlite`` or ``postgres`` in config.yaml.
"""
from __future__ import annotations
import argparse
import asyncio
import logging
import sys
from deerflow.config.app_config import get_app_config
from deerflow.persistence.agents.base import AgentExistsError
from deerflow.persistence.agents.file import FileAgentStore
from deerflow.persistence.agents.sql import SqlAgentStore
logger = logging.getLogger("migrate_agents_to_db")
def main() -> int:
parser = argparse.ArgumentParser(description="Import file-backed custom agents into the db agent store.")
parser.add_argument("--dry-run", action="store_true", help="List what would be imported without writing to the database.")
args = parser.parse_args()
logging.basicConfig(level=logging.INFO, format="%(message)s")
config = get_app_config()
db_backend = getattr(config.database, "backend", None)
if db_backend not in ("sqlite", "postgres"):
logger.error(
"database.backend is %r; this importer needs 'sqlite' or 'postgres'. Set it in config.yaml (the same database the gateway uses).",
db_backend,
)
return 1
source = FileAgentStore()
agents = source.list_all()
if not agents:
logger.info("No file-backed agents found — nothing to import.")
return 0
if args.dry_run:
for user_id, cfg in agents:
logger.info("[dry-run] would import %s/%s", user_id, cfg.name)
logger.info("[dry-run] %d agent(s) would be imported. Source files are left in place.", len(agents))
return 0
# Ensure the schema exists (creates the ``agents`` table via the same
# Alembic bootstrap the gateway runs) before the sync store writes rows.
from deerflow.persistence.engine import init_engine_from_config
asyncio.run(init_engine_from_config(config.database))
dest = SqlAgentStore(config.database.app_sync_sqlalchemy_url)
imported = 0
skipped = 0
for user_id, cfg in agents:
soul = source.get_soul(cfg.name, user_id=user_id) or ""
# exclude_unset keeps the stored document as sparse as the source file
# (only the keys the operator actually wrote), matching the file layout.
document = cfg.model_dump(exclude_unset=True)
try:
dest.create(cfg.name, document, soul, user_id=user_id)
imported += 1
logger.info("imported %s/%s", user_id, cfg.name)
except AgentExistsError:
skipped += 1
logger.info("skip %s/%s: already present in db", user_id, cfg.name)
logger.info("Done: %d imported, %d already present. Source files left in place (rollback: revert agent_storage.backend to 'file').", imported, skipped)
return 0
if __name__ == "__main__":
sys.exit(main())