Merge branch 'main' into rayhpeng/hexagonal-feedback-slice

Rebase the feedback migration onto main's new chain tip again: main added
0009_webhook_dedupe (also chained after 0008_thread_operation_kind), so
0009_feedback_tags becomes 0010_feedback_tags with
down_revision=0009_webhook_dedupe, keeping the chain linear.

Conflicts resolved:
- Five head-pin tests take main's version with the pin bumped to
  0010_feedback_tags.
- test_thread_messages_page.py keeps this branch's feedback_service
  wiring over main's feedback_repo wiring, but stubs both service
  methods so the thread-grouped path main added coverage for stays
  stubbed (latest_per_run_in_thread alongside latest_for_runs).
- test_thread_messages_feedback.py keeps both sides' imports; each is
  used (Feedback for the fixture, EditReplayVisibility for the run
  manager stub).

Verified: 79 tests across the conflicted files plus the migration and
bootstrap suites, and 54 feedback tests, all pass.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
rayhpeng 2026-07-28 07:13:25 +08:00
commit fd74553f91
129 changed files with 13742 additions and 724 deletions

View File

@ -83,4 +83,10 @@ jobs:
- name: Run unit tests of backend
working-directory: backend
env:
# Expose the job's Postgres service to the cross-pod dedupe integration
# tests (issue #4120), which read DEDUPE_TEST_POSTGRES_URL. Without this
# mapping those tests silently skip and the headline capability is never
# exercised in CI.
DEDUPE_TEST_POSTGRES_URL: ${{ env.TEST_POSTGRES_URI }}
run: make test

View File

@ -255,6 +255,10 @@ single-label cluster hosts, and Docker/Podman internal hostnames do not inherit
honor environment proxy settings.
Backend processes automatically pick up `config.yaml` changes on the next config access, so model metadata updates do not require a manual restart during development.
The checkpoint storage settings `database.checkpoint_channel_mode` and
`database.checkpoint_delta_snapshot_frequency` are exceptions: both are frozen
when the process first builds an agent (including through `DeerFlowClient`) and
require a process restart to change safely.
> [!TIP]
> On Linux, if Docker-based commands fail with `permission denied while trying to connect to the Docker daemon socket at unix:///var/run/docker.sock`, add your user to the `docker` group and re-login before retrying. See [CONTRIBUTING.md](CONTRIBUTING.md#linux-docker-daemon-permission-denied) for the full fix.
@ -284,7 +288,7 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser
>
> With lease heartbeat enabled, a transient RunStore renewal error is retried only until the last confirmed lease expires; the stale worker then cancels local execution and suppresses checkpoint, completion-hook, delivery-receipt, and thread-status finalization. A remote tool side effect already in flight may still be outside local cancellation.
>
> Reconciliation uses an atomic takeover claim that re-checks the lease after candidate selection, so a successful owner renewal wins over orphan recovery and only one reconciler can report a run as recovered.
> Reconciliation uses an atomic takeover claim that re-checks the lease after candidate selection, so a successful owner renewal wins over orphan recovery and only one reconciler can report a run as recovered. When multiple Gateway workers share the Docker/AIO or E2B sandbox backend, also configure `sandbox.ownership.type: redis`; E2B uses the leases during background startup and periodic reconciliation so duplicate/orphan cleanup cannot terminate a live peer's sandbox.
See [CONTRIBUTING.md](CONTRIBUTING.md) for detailed Docker development guide.
@ -341,6 +345,8 @@ DeerFlow runs the agent runtime inside the Gateway API. Development mode enables
Gateway owns `/api/langgraph/*` and translates those public LangGraph-compatible paths to its native `/api/*` routers behind nginx.
Gateway runs automatically enforce native delivery for artifacts created or modified under `/mnt/user-data/outputs`: `present_files` must present at least one output produced by the current run, and the terminal `run.delivery` receipt must be durably recorded. Runs that do not produce output artifacts keep ordinary conversational behavior.
DeerFlow's built-in custom events are available through both LangGraph streaming interfaces: native clients can continue subscribing to `stream_mode="custom"`, while callback-based integrations can consume the same payloads as `on_custom_event` records from `astream_events(version="v2")`. The callback event name matches the payload's `type` field.
#### Docker Production Deployment
@ -434,6 +440,8 @@ channels:
telegram:
enabled: true
bot_token: $TELEGRAM_BOT_TOKEN
# Optional: render final Markdown replies as Telegram Rich Messages.
rich_messages: false
allowed_users: [] # empty = allow all
wechat:
@ -759,7 +767,7 @@ uv run python -m deerflow.skills.review.cli ../skills/public/data-analysis --for
Tools follow the same philosophy. DeerFlow comes with a core toolset — web search, web fetch, rendered web capture, file operations, bash execution — and supports custom tools via MCP servers and Python functions. Swap anything. Add anything.
Advanced deployments can enable pluggable tool authorization with `authorization.enabled` in `config.yaml`. A configured `AuthorizationProvider` filters denied tools before they reach the model or deferred-tool catalog, then the same provider is checked again before every business-tool execution through the existing guardrail middleware. A generated `tool_search` may bypass that second check only when it fronts the current build's already-filtered deferred catalog. The built-in RBAC provider supports per-role tool allow/deny policies and validates that `default_role` names a configured role; authorization is disabled by default. See `config.example.yaml` and the [authorization RFC](docs/plans/2026-07-10-pluggable-authorization-rfc.md).
Advanced deployments can enable pluggable authorization with `authorization.enabled` in `config.yaml`. A configured `AuthorizationProvider` filters denied tools before they reach the model or deferred-tool catalog, then the same provider is checked again before every business-tool execution through the existing guardrail middleware. Gateway `threads:*` and `runs:*` route permissions are derived from the same provider, while existing owner checks and admin-only management gates remain in force. A generated `tool_search` may bypass the second tool check only when it fronts the current build's already-filtered deferred catalog. The built-in RBAC provider supports per-role `tools` and `routes` allow/deny policies and validates that `default_role` names a configured role; authorization is disabled by default. See `config.example.yaml` and the [authorization RFC](docs/plans/2026-07-10-pluggable-authorization-rfc.md).
Advanced deployments can also extend the agent runtime itself by declaring zero-argument `AgentMiddleware` classes under `extensions.middlewares` in `config.yaml` or `extensions_config.json`. DeerFlow loads the same configured class list into the lead-agent and subagent pipelines after their built-in runtime middlewares and loop/token guards, but before the terminal-response/safety/clarification tail, so enterprise forks can add domain guardrails, tool-call governance, or observability hooks without patching the built-in middleware builders. Missing packages, invalid classes, and broken modules fail loudly at agent creation. Treat `config.yaml` and `extensions_config.json` as trusted operator-controlled files: middleware paths are code execution, just like custom tool, model, sandbox, guardrail, MCP server, and MCP interceptor declarations. Gateway skill/MCP toggle endpoints preserve this field but do not expose an API write path for `extensions.middlewares`. Per-context parameterization and separate lead-only/subagent-only middleware lists are not supported yet.
@ -778,10 +786,12 @@ Interrupted first-turn runs still persist a fallback conversation title, so stop
Streaming Markdown responses animate only newly arrived words; text that is already visible is not faded out and replayed when the next chunk extends the same block.
In the Web UI, completed assistant turns can be branched into a new main conversation. The new thread starts from that turn's checkpoint and keeps the preceding replay checkpoint, so the branched response can be regenerated immediately. Legacy or imported histories without checkpoint parent links use a bounded chronological fallback; if no earlier replay checkpoint exists, branching still succeeds with the legacy single-checkpoint shape, while regeneration remains unavailable for that inherited response. Existing single-checkpoint branches are left unchanged rather than attempting an unsafe checkpoint copy. Because workspace files are not checkpointed, the branch only receives a best-effort copy of the current workspace when you branch from the latest turn; branching from an older turn keeps just the restored message history so the branch never inherits files that were created in a later part of the conversation.
In the Web UI, completed assistant turns can be branched into a new main conversation. The new thread starts from that turn's checkpoint and keeps the preceding replay checkpoint, so the branched response can be regenerated immediately. Regenerating the latest response preserves the thread's current title, including a title you renamed manually after the original response. Legacy or imported histories without checkpoint parent links use a bounded chronological fallback; if no earlier replay checkpoint exists, branching still succeeds with the legacy single-checkpoint shape, while regeneration remains unavailable for that inherited response. Existing single-checkpoint branches are left unchanged rather than attempting an unsafe checkpoint copy. Because workspace files are not checkpointed, the branch only receives a best-effort copy of the current workspace when you branch from the latest turn; branching from an older turn keeps just the restored message history so the branch never inherits files that were created in a later part of the conversation.
The Web UI reports completed task time once per run. This is total wall-clock time—including model reasoning, tool calls, and waiting—not a per-step or model-only thinking duration. Reasoning content remains available through its own separate disclosure.
In the Web UI, the latest completed user turn can also be edited and rerun from the message toolbar. DeerFlow restores the conversation checkpoint before that user message, submits the edited text as a new user message, and hides the superseded turn once the replay is in progress or succeeds. This is a conversation-state replay only: files, memory updates, and external tool side effects are not undone.
Web UI chat links percent-encode custom thread identifiers before placing them in route segments, so reserved URL characters such as `#` and `?` do not change which conversation is opened.
```
@ -931,7 +941,7 @@ Across sessions, DeerFlow builds a persistent memory of your profile, preference
Memory updates now skip duplicate fact entries at apply time, so repeated preferences and context do not accumulate endlessly across sessions.
File-backed memory now separates global user context from agent facts. Each user has one `memory.json` containing only the project-independent `user` and `history` summaries; every fact is a canonical Markdown file below `agents/{agent_name}/facts/`. Existing lead-agent middleware, API, Settings, import/export, and embedded-client calls that omit `agent_name` resolve inside DeerMem to the reserved `__default__` bucket. That bucket is outside the valid custom-agent name grammar, so a real custom agent named `lead-agent` has a separate fact repository and deleting a custom agent cannot delete a memory-only directory without `config.yaml`. Public agent identifiers are case-insensitive and canonicalized to lowercase. Runtime/API readers still receive a compatibility `facts` array for the selected/default agent, so the frontend does not read agent facts from `memory.json`; structured Markdown `source` metadata is projected to the historical string field at the MemoryManager boundary. An unscoped Clear All first migrates facts from unread legacy per-agent JSON without adopting its soon-to-be-cleared summaries, then removes shared summaries and facts from every agent bucket while preserving agent configuration files, so a later read cannot resurrect skipped legacy facts; an explicitly agent-scoped clear removes only that agent's facts. On first normal read, old facts embedded in the user JSON are migrated automatically to `__default__`; facts written to the earlier implicit `lead-agent` bucket are also moved when that directory is not a real custom agent. Migration and normal writes notify a configured retrieval adapter only after durable storage locks are released. Journaled writes, a shared user lock, and optimistic user-memory revisions prevent silent lost updates. Retrieval engines remain optional behind the memory retrieval-adapter contract, with an explicit local substring fallback when no adapter is configured.
File-backed memory now separates global user context from agent facts. Each user has one `memory.json` containing only the project-independent `user` and `history` summaries; every fact is a canonical Markdown file below `agents/{agent_name}/facts/`. Existing lead-agent middleware, API, Settings, import/export, and embedded-client calls that omit `agent_name` resolve inside DeerMem to the reserved `__default__` bucket. That bucket is outside the valid custom-agent name grammar, so a real custom agent named `lead-agent` has a separate fact repository and deleting a custom agent cannot delete a memory-only directory without `config.yaml`. Public agent identifiers are case-insensitive and canonicalized to lowercase. Runtime/API readers still receive a compatibility `facts` array for the selected/default agent, so the frontend does not read agent facts from `memory.json`; structured Markdown `source` metadata is projected to the historical string field at the MemoryManager boundary. An unscoped Clear All first migrates facts from unread legacy per-agent JSON without adopting its soon-to-be-cleared summaries, then removes shared summaries and facts from every agent bucket while preserving agent configuration files, so a later read cannot resurrect skipped legacy facts; an explicitly agent-scoped clear removes only that agent's facts. On first normal read, old facts embedded in the user JSON are migrated automatically to `__default__`; facts written to the earlier implicit `lead-agent` bucket are also moved when that directory is not a real custom agent. Migration and normal writes notify the configured retrieval adapter only after durable storage locks are released. DeerMem uses a scope-aware SQLite FTS5/BM25 adapter by default, stores only rebuildable derived index data under `.retrieval/`, and rebuilds it in the background during Gateway startup or lazily on the first scoped search. A corrupt derived index is recreated automatically. Set `memory.backend_config.retrieval_adapter` to an empty string to disable it and use the local substring fallback. Chinese tokenization is optional; install the backend `memory-zh` extra (`uv sync --extra memory-zh`) for jieba-assisted sub-phrase search. Journaled writes, a shared user lock, and optimistic user-memory revisions prevent silent lost updates.
Single-fact repository operations are genuinely incremental: an upsert/delete reads, journals, writes, and re-indexes only the addressed fact files, and returns an explicit incomplete delta rather than a cache-dependent fake full document. Summary change sets merge the supplied `user`/`history` child keys over the persisted sections so a partial update cannot erase omitted siblings; full imports normalize both sections to the complete compatibility schema before applying replacement values. Manager/API compatibility methods materialize a fresh full document only when their public response contract requires one. Fact-level point operations use separate expected user-memory and fact revisions and may explicitly rebase when every addressed fact precondition still holds. Snapshot-derived operations such as scoped clear, capped create, consolidation, and trimming never replay stale delete/trim sets: a manifest conflict reloads the complete document and recomputes the operation, with a bounded retry. Fact paths use the first two hexadecimal characters of `SHA-256(fact_id)` so generated `fact_*` IDs distribute across shards. The cache token combines the shared JSON's nanosecond mtime, size, and persisted revision; this prevents coarse-mtime same-size writes from returning stale data without scanning fact files. Direct out-of-band Markdown edits require an explicit reload. Storage-specific conflicts and corruption are translated at the MemoryManager boundary; the Gateway returns conflict as HTTP 409 and a stable, non-sensitive corruption error as HTTP 500. Full-document `save()` remains a compatibility API and computes a diff before writing; malformed or missing `facts` can no longer silently erase an agent's Markdown files. Legacy migration preserves non-empty `user`/`history` before deleting an agent `memory.json`; conflicting summaries keep the legacy file and fail loudly instead of choosing a winner.

View File

@ -340,6 +340,8 @@ Lead-agent middlewares are assembled in strict order across three functions: the
Authorization identity plumbing is independent of whether authorization enforcement is enabled. Gateway removes client-supplied `is_internal` / `authz_attributes` / `channel_user_id`, derives `is_internal` only from the server-owned `request.state.auth_source`, and accepts `channel_user_id` only from an internally authenticated IM caller's top-level `body.context`; free-form `body.config` can never supply it. `build_principal_from_context` is the shared Principal builder for assembly-time authorization and `GuardrailAuthorizationAdapter`; it applies `default_role`, strict-boolean internal provenance, and copy-on-read `authz_attributes`. The built-in RBAC provider validates `authorization.default_role` during provider resolution so an unknown fallback role fails agent construction instead of degrading into an empty tool set. Task delegation carries `is_internal` plus copied attributes through `SubagentExecutor`, while `GuardrailMiddleware` maps the same runtime fields into `GuardrailRequest`. Phase 1B applies Layer 1 before deferred-tool assembly on the lead, native-subagent, and embedded-client paths, then passes the same provider instance into Layer 2. Framework-provided `describe_skill` and memory tools are included in Layer 1 but restored to their legacy post-`tool_search` ordering afterward. `DeerFlowClient.stream()` treats its in-process caller as trusted and accepts the same identity fields as keyword overrides; it includes the complete Principal in its agent cache key and deep-copies nested attributes so caller mutation cannot make a stale tool set look current.
Gateway route authorization uses `authz.py::resolve_route_permissions()` as the single provider integration point for both `AuthMiddleware` and decorator-only authentication. When enabled, it evaluates the six registered `threads:*` / `runs:*` permissions as `resource="route"` requests whose targets are the full `resource:action` strings. Decisions use the async provider API and are cached for the request in `AuthContext`; decorators do not call the provider again. Provider resolution or decision errors follow `authorization.fail_closed`, scoped per permission for decision errors. When authorization is disabled, the legacy complete permission set is returned without resolving a provider. Existing `owner_check` enforcement and `require_admin_user()` management gates remain independent and unchanged. Tests: `tests/test_authorization_route_permissions.py`, `tests/test_auth.py`, and `tests/test_auth_middleware.py`.
Before changing a later authorization phase, read the [authorization RFC](../docs/plans/2026-07-10-pluggable-authorization-rfc.md) and its [implementation notes](../docs/plans/2026-07-10-pluggable-authorization-implementation-notes.md). The notes are the cumulative handoff record for merged PR behavior, reviewer feedback, trust-boundary decisions, deferred scope, and required regression coverage.
**Lead-only middlewares** (`build_middlewares`, appended after the base):
@ -365,7 +367,7 @@ Before changing a later authorization phase, read the [authorization RFC](../doc
32. **TerminalResponseMiddleware** - When a provider returns an empty terminal `AIMessage` after tool execution, injects a hidden recovery prompt and retries the model once; a second empty response is replaced in checkpoint state by a visible error fallback marked for the run worker, so the run finishes as an error instead of a silent success
33. **ModelLengthFinishReasonMiddleware** - Records `stop_reason=model_length_capped` when provider-specific length detectors match a terminal `AIMessage` without tool-call intent (`finish_reason=length` / `MAX_TOKENS`, or `stop_reason=max_tokens`), preserving the original assistant content and never reparsing textual tool-call-like envelopes
34. **SafetyFinishReasonMiddleware** - *(optional, if `safety_finish_reason.enabled`)* Suppresses tool execution when the provider safety-terminated the response (e.g. `finish_reason=content_filter`); registered after terminal-response/custom/configured middlewares so LangChain's reverse-order `after_model` dispatch runs it first
35. **ClarificationMiddleware** - Intercepts `ask_clarification` tool calls, writes a readable `ToolMessage.content` fallback plus structured `ToolMessage.artifact.human_input` request payload, and interrupts via `Command(goto=END)` (must be last). Because this middleware can short-circuit tool execution before LangChain emits `on_tool_end`, `RunJournal` performs a root-run final reconciliation for allowlisted clarification `ToolMessage`s whose `tool_call_id` was produced by the current run, so human-input request cards remain recoverable from `run_events` after checkpoint compaction. Human Input Card replies are submitted as `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden response sources (currently `ask_clarification`) as `llm.human.input`, which preserves answered-card state after compaction without exposing generic internal hidden context.
35. **ClarificationMiddleware** - Intercepts `ask_clarification` tool calls, writes a readable `ToolMessage.content` fallback plus structured `ToolMessage.artifact.human_input` request payload, and interrupts via `Command(goto=END)` (must be last). Payloads are versioned: legacy modes (`free_text` / `choice_with_other`) keep `version: 1` unchanged, while the v2 `form` mode (from `fields`) carries `version: 2` so older frontends reject the payload and degrade to the plain-text fallback. Field normalization is deterministic and lives in the middleware, not the tool schema — the middleware short-circuits before tool execution, so tool-arg typing alone provides no runtime validation. Validation is atomic: any structurally broken entry (non-dict, bad/duplicate name, a name colliding with a JS `Object.prototype` member like `__proto__`/`constructor`, exceeding the caps of 16 fields / 24 options per field / 200 chars per text, or the whole normalized definition exceeding `MAX_FORM_SERIALIZED_BYTES` = 16KB UTF-8 — the per-item caps alone admit forms whose IM text fallback would blow channel delivery limits and truncate away trailing fields) degrades the whole form to the legacy option/free-text modes, so a card can never render "complete" while silently missing a business field; benign issues keep local degradation (unknown types — including unhashable JSON like `type: []`, which must never raise from the membership probe — and option-less selects become `text`), and options are trimmed/deduped with blanks dropped (both form-level and top-level) because the frontend parser rejects blank option labels. Checkbox fields are booleans that default to an explicit "no"; `required` on a checkbox means must-agree/consent semantics. The response protocol is deliberately unchanged (v1 `text`/`option` only): form cards submit a readable text summary as `response_kind: "text"`, so journal persistence and answered-card recovery need no new allowlist entries. Because this middleware can short-circuit tool execution before LangChain emits `on_tool_end`, `RunJournal` performs a root-run final reconciliation for allowlisted clarification `ToolMessage`s whose `tool_call_id` was produced by the current run, so human-input request cards remain recoverable from `run_events` after checkpoint compaction. Human Input Card replies are submitted as `hide_from_ui` `HumanMessage`s with `additional_kwargs.human_input_response`; `RunJournal` persists only allowlisted hidden response sources (currently `ask_clarification`) as `llm.human.input`, which preserves answered-card state after compaction without exposing generic internal hidden context.
### Configuration System
@ -439,7 +441,7 @@ Localhost persistence deliberately reads the direct request `Host` and ignores `
| **Artifacts** (`/api/threads/{id}/artifacts`) | `GET /{path}` - serve artifacts; active content types (`text/html`, `application/xhtml+xml`, `image/svg+xml`) are always forced as download attachments to reduce XSS risk; `?download=true` still forces download for other file types |
| **Suggestions** (`/api/suggestions`) | `GET /config` - returns global suggestions config boolean; `POST /threads/{id}/suggestions` - generate follow-up questions; rich list/block model content is normalized and inline reasoning (`<think>...</think>`, including unclosed/truncated blocks from reasoning models like MiniMax-M3) is stripped before JSON parsing |
| **Input Polish** (`/api/input-polish`) | `POST /` - rewrite a composer draft before it is sent. This is a short authenticated `runs:create` LLM request using `input_polish` config; it does not create a LangGraph run, persist a message, or modify thread state. Shares the non-graph one-shot LLM path (`deerflow.utils.oneshot_llm.run_oneshot_llm`) with the suggestions route so model build + Langfuse metadata + invoke stay in one place; validates the same stripped view of the draft it sends to the model, and preserves literal `<think>` substrings in the rewrite (`strip_think_blocks(truncate_unclosed=False)`) |
| **Thread Runs** (`/api/threads/{id}/runs`) | `POST /` - create background run; `POST /stream` - create + SSE stream; `POST /wait` - create + block; `POST /regenerate/prepare` - prepare clean input + checkpoint metadata for regenerating the latest assistant answer; `GET /` - list runs; `GET /{rid}` - run details; `POST /{rid}/cancel` - cancel; `GET /{rid}/join` - join SSE; `GET /{rid}/messages` - paginated per-run messages `{data, has_more}`; `GET /{rid}/events` - full event stream; `GET /{rid}/workspace-changes` - workspace/output file change summary and optional diffs; `GET /../messages` - legacy thread message array; `GET /../messages/page` - backward thread-global `seq` history page with middleware/subagent-AI/successful-regenerate filtering and page-run-scoped feedback enrichment; subagent AI callbacks remain available through run events while parent `task` ToolMessages stay visible for card restoration; `GET /../token-usage` - aggregate tokens |
| **Thread Runs** (`/api/threads/{id}/runs`) | `POST /` - create background run; `POST /stream` - create + SSE stream; `POST /wait` - create + block; `POST /regenerate/prepare` - prepare clean input + checkpoint metadata for regenerating the latest assistant answer; carrying the latest non-empty thread title in graph input so resuming an older checkpoint cannot roll back a later manual rename (#4457); `POST /edit-regenerate/prepare` - prepare a checkpoint replay from the latest editable human turn with a replacement user message and edit replay metadata; `GET /` - list runs; `GET /{rid}` - run details; `POST /{rid}/cancel` - cancel; `GET /{rid}/join` - join SSE; `GET /{rid}/messages` - paginated per-run messages `{data, has_more}`; `GET /{rid}/events` - full event stream; `GET /{rid}/workspace-changes` - workspace/output file change summary and optional diffs; `GET /../messages` - legacy thread message array; `GET /../messages/page` - backward thread-global `seq` history page with middleware/subagent-AI/successful-regenerate/edit-replay filtering and page-run-scoped feedback enrichment; subagent AI callbacks remain available through run events while parent `task` ToolMessages stay visible for card restoration; `GET /../token-usage` - aggregate tokens |
| **Feedback** (`/api/threads/{id}/runs/{rid}/feedback`) | `PUT /` - upsert feedback; `DELETE /` - delete user feedback; `POST /` - create feedback; `GET /` - list feedback; `GET /stats` - aggregate stats; `DELETE /{fid}` - delete specific |
| **Runs** (`/api/runs`) | `POST /stream` - stateless run + SSE; `POST /wait` - stateless run + block; `GET /{rid}/messages` - paginated messages by run_id `{data, has_more}` (cursor: `after_seq`/`before_seq`); `GET /{rid}/feedback` - list feedback by run_id |
| **GitHub Webhooks** (`/api/webhooks/github`) | `POST /` - receive GitHub App / repo webhook deliveries. Verifies `X-Hub-Signature-256` against `GITHUB_WEBHOOK_SECRET`; exempt from auth + CSRF because authenticity is enforced by HMAC. The route is fail-closed: mounted only when `GITHUB_WEBHOOK_SECRET` is set, or when explicit dev opt-in `DEER_FLOW_ALLOW_UNVERIFIED_GITHUB_WEBHOOKS=1` is set. Recognized events include `ping`, `issues`, `issue_comment`, `pull_request`, `pull_request_review`, and `pull_request_review_comment`; unknown events return 200 with `handled=false`. Fan-out runtime failures return 503, keeping the delivery recorded as failed for manual/API/scripted redelivery (GitHub does not automatically retry any failed delivery, 5xx included); permanent/non-retryable conditions such as `channels.github.enabled: false`, unknown events, malformed payloads, or unavailable channel service return 200 with a skipped/handled response. |
@ -468,10 +470,18 @@ compatibility failures and cancellation while waiting for prior finalization
still emit a zero-delivery receipt. The worker flushes ordinary journal events,
idempotently persists the run-scoped receipt, and only then persists the staged
terminal run status. A receipt failure is retried on a short bounded schedule
while the owning worker still knows the real outcome and holds the lease. If
all attempts fail, the worker persists that real terminal status instead of
leaving a successful run inflight for orphan recovery to rewrite as an error;
the receipt remains best-effort in that outage case. Orphan recovery first
while the owning worker still knows the real outcome and holds the lease. The
worker derives delivery requirements from the run's workspace snapshots rather
than a client request option: every regular file created or modified under
`/mnt/user-data/outputs` is a candidate produced artifact. At least one candidate
must be covered by a path attributed by the journal to `present_files`;
presenting only an unrelated pre-existing path does not satisfy delivery.
Receipts for such runs add `produced_paths`, `presented_paths`, `matched_paths`,
`verification`, `stage`, and `satisfied` to the Slice 1 fact fields. Missing a
matching presentation becomes a run error; a successful presentation is also
downgraded to error if its receipt cannot be durably verified. Runs without
changed outputs preserve ordinary chat behavior and the original receipt shape.
Orphan recovery first
atomically claims an expired lease, then uses the same singleton write to
backfill a zero-delivery receipt. This ordering prevents a stale recovery scan
from overwriting a live run's later detailed receipt; an event-store outage
@ -494,7 +504,8 @@ JSONL event stores when `GATEWAY_WORKERS > 1`.
**RunManager / RunStore contract**:
- LangGraph-compatible run requests validate their supported subset before creating a run. `runtime/stream_modes.py` is the shared backend contract for public stream modes and the worker's `graph.astream` mapping; the public `messages-tuple` mode maps to LangGraph's internal `messages` mode, while public `messages`, `events`, and other unsupported modes are rejected instead of being dropped or replaced with `values`. `app/gateway/run_models.py::RunCreateRequest` is shared by HTTP and internal scheduled launch paths, retains only truthful compatibility defaults for unimplemented options (`if_not_exists="create"` plus `None` placeholders), returns 422 for unsupported values including `on_completion="complete"`, `on_completion="continue"`, and `multitask_strategy="enqueue"`, and forbids undeclared SDK options so fields such as `checkpoint_during` and `durability` cannot be silently discarded. A placeholder must still accept the stock SDK's own default: `langgraph_sdk` drops only `None` from its run payload, so `stream_resumable=False` reaches every request and means "non-resumable", which is what DeerFlow serves — rejecting it 422'd every IM channel run (#4466). `tests/test_run_request_validation.py::test_gateway_accepts_langgraph_sdk_default_payload` pins the real SDK payload against this boundary; channel tests mock the SDK client and cannot catch this class of drift.
- `RunManager.get()` is async; direct callers must `await` it.
- The history batch helpers `list_successful_regenerate_sources()` and `get_many_by_thread()` default to `user_id=AUTO`: they resolve the request user and fail closed when no user context exists. Migration/admin callers that intentionally need an unscoped read must pass `user_id=None` explicitly.
- The history batch helpers `list_successful_regenerate_sources()`, `list_edit_regenerate_runs()`, and `get_many_by_thread()` default to `user_id=AUTO`: they resolve the request user and fail closed when no user context exists. Migration/admin callers that intentionally need an unscoped read must pass `user_id=None` explicitly.
- Edit-and-rerun visibility is derived from edit replay runs (`metadata.replay_kind="edit"` plus `regenerate_from_run_id`) by `RunManager.list_edit_replay_visibility()`: the newest attempt for each source run is authoritative. Pending/running/success attempts hide the original source run; failed, timed-out, or interrupted attempts hide only the failed attempt so the original conversation reappears.
- When a persistent `RunStore` is configured, `get()` and `list_by_thread()` hydrate historical runs from the store. In-memory records win for the same `run_id` so task, abort, and stream-control state stays attached to active local runs.
- Thread metadata status switches to `running` only after `RunManager.try_start()` succeeds. Pending-cancelled runs therefore skip the old `running` projection, while clients may observe the prior thread status during the short worker-startup window.
- `cancel()` returns a :class:`~deerflow.runtime.CancelOutcome` enum: `cancelled` (local cancel), `taken_over` (non-owning worker claimed the run because the owner's lease expired — marks it as `error`), `lease_valid_elsewhere` (owner's lease is still alive — caller should return 409 + `Retry-After`), `not_active_locally` (heartbeat disabled, preserving the old 409 path), `not_cancellable` (terminal state), or `unknown` (not found in memory or store). `create_or_reject(..., multitask_strategy="interrupt"|"rollback")` persists interrupted status through `RunStore.update_status()`, matching normal `set_status()` transitions.
@ -547,7 +558,15 @@ copying a raw checkpoint because delta state is not self-contained in one tuple.
failed kill through a new client. An unconfirmed remote ID stays tracked.
`reset()` uses full shutdown semantics. It destroys tracked active and warm
VMs. It wakes capacity waiters. Callers cannot reuse the old provider
instance.
instance. A background startup pass and periodic reconciliation list
provider-tagged remote sandboxes within page/item/time budgets, probe every
candidate until a healthy canonical sandbox is found, adopt only through the
shared ownership store, and reap duplicates/orphans only after their
configured grace/TTL and an atomic `del:` claim. Lease renewal is independent
of reconciliation. Failed canonical adoption clears both its capacity
reservation and acquire-intent marker even if a peer takes ownership between
the initial claim and bootstrap cleanup. Shutdown kills only IDs whose leases
are owned by this provider instance; peer-owned clients are merely closed.
- **Cross-instance ownership store** (`aio_sandbox/ownership/`, #4206): gateway instances sharing a container backend coordinate container ownership through a pluggable lease store, selected by `sandbox.ownership.type` (`memory` | `redis`) and resolved like `stream_bridge` (`factory.py`, lazy per-branch import, `redis` optional extra, `DEER_FLOW_SANDBOX_OWNERSHIP_REDIS_URL` env escape hatch; a set `DEER_FLOW_STREAM_BRIDGE_REDIS_URL` implies a multi-instance deployment and infers `redis`). `memory` is single-instance only and declares `supports_cross_process = False`.
- **A lease answers "who reaps this container", not "who may use it".** That splits the interface in two: `take()` transfers ownership on the **acquire** path (a container is deterministic per user/thread, so consecutive turns legitimately land on different instances — a conditional claim there would strand the thread until the previous lease expired), while `claim()` succeeds only if the container is unowned or already ours and gates every **adopt/reap** path. `release()` never clears a peer's lease.
- **A lease carries a state, and that is what makes the destroy window safe.** `own:` = responsible for this container; `del:` = tearing it down (`claim(..., for_destroy=True)`). `take()` is refused against a `del:` lease, so a container cannot be re-acquired between a destroy path's claim and its container stop. Without the two states an unconditional `take()` would silently overwrite the destroyer's claim and the peer's stop would land on a container the new owner had already handed to an agent — i.e. #4206 again. That pairing is what replaced the previous same-host `flock` guard, which is gone; Redis makes the scope genuinely multi-instance instead of same-host. A destroyer that dies mid-stop leaves a `del:` marker that lapses with the TTL. On the acquire path a refused take raises `SandboxBeingDestroyedError`: the reuse/reclaim paths drop the container and cold-start, and the discover path propagates (falling through to create would collide with the not-yet-removed container name).
@ -573,9 +592,10 @@ copying a raw checkpoint because delta state is not self-contained in one tuple.
- `BoxliteProvider` (`packages/harness/deerflow/community/boxlite/`) - BoxLite micro-VM isolation. The `boxlite` runtime is optional (`deerflow-harness[boxlite]`) and lazy-imported only when this provider is selected. The provider owns one private asyncio event loop on a daemon thread because BoxLite handles are loop-affine; sync `Sandbox` calls marshal onto that loop with `run_coroutine_threadsafe`.
Boxes are named deterministically from `user_id:thread_id`, released into an in-process warm pool after each agent turn, and reclaimed only by the same user/thread. Warm-pool health checks use a short explicit timeout and forward that timeout through both BoxLite `exec(timeout=...)` and the private-loop `.result(timeout)` bridge so a hung VM cannot pin the per-thread acquire lock indefinitely.
`sandbox.replicas` caps active + warm VMs per gateway process; if capacity is exhausted, only warm-pool VMs are evicted. `sandbox.idle_timeout` stops idle warm VMs after the configured seconds. `reset()` is intentionally a lightweight registry clear for `reset_sandbox_provider()` and does not close boxes, stop the idle reaper, or close the private loop; full teardown remains `shutdown()`.
- `TenkiSandboxProvider` (`packages/harness/deerflow/community/tenki/`) - Tenki cloud microVM isolation. The `tenki-sandbox` SDK is optional (`deerflow-harness[tenki]`) and lazy-imported (`_import_client`) only when this provider is selected. Unlike Boxlite, the SDK is synchronous, so the adapter calls it directly with no event-loop bridge. File transport uses Tenki's native `sandbox.fs` API (`read_text`/`read_stream`/`write_stream`/`mkdir`/`stat`) — binary-safe and streaming, no base64/shell hop; only directory/content *search* (`list_dir`/`glob`/`grep`) shells out to busybox-portable `find`/`grep`, parsed with the shared `deerflow.sandbox.search` helpers like `community/e2b_sandbox`. Sandboxes run as the unprivileged `tenki` user, so DeerFlow's `/mnt/user-data` prefix is remapped under a writable HOME (`_resolve_path`) and best-effort `sudo`-symlinked at bootstrap. Boxes are named deterministically from `sha256(user_id:thread_id)[:16]` (64-bit, matching E2B; the warm pool is keyed by this id alone with no full-seed fallback), released into an in-process warm pool, and reclaimed only by the same user/thread after a liveness check. A terminal session error (named SDK errors plus builtin `ConnectionError`/`BrokenPipeError`/`EOFError`) routes through `_invalidate_sandbox` to evict the dead microVM. Cross-process orphan reconciliation is a follow-up (single-process warm pool today).
**Shared warm-pool lifecycle:** community sandbox providers that keep released sandboxes alive for fast reuse share `deerflow.community.warm_pool_lifecycle.WarmPoolLifecycleMixin`. The mixin owns the common `DEFAULT_IDLE_TIMEOUT=600`, `IDLE_CHECK_INTERVAL=60`, `DEFAULT_REPLICAS=3`, idle-checker loop, warm-pool expiry, oldest-warm eviction, replica counting, and soft-cap logging. Providers remain responsible for their own active registries, creation/discovery, health checks, and destroy hook (`_destroy_warm_entry`): AIO destroys `SandboxInfo` through its backend; Boxlite closes loop-affine `BoxliteBox` handles. AIO keeps active-idle cleanup outside the mixin and delegates only warm-pool expiry to the shared helper.
**Shared warm-pool lifecycle:** community sandbox providers that keep released sandboxes alive for fast reuse share `deerflow.community.warm_pool_lifecycle.WarmPoolLifecycleMixin`. The mixin owns the common `DEFAULT_IDLE_TIMEOUT=600`, `IDLE_CHECK_INTERVAL=60`, `DEFAULT_REPLICAS=3`, idle-checker loop, warm-pool expiry, oldest-warm eviction, replica counting, and soft-cap logging. Providers remain responsible for their own active registries, creation/discovery, health checks, and destroy hook (`_destroy_warm_entry`): AIO destroys `SandboxInfo` through its backend; Boxlite closes loop-affine `BoxliteBox` handles; Tenki closes the microVM session (`TenkiSandbox.close`, which terminates the remote sandbox). AIO keeps active-idle cleanup outside the mixin and delegates only warm-pool expiry to the shared helper.
**Virtual Path System**:
- Agent sees: `/mnt/user-data/{workspace,uploads,outputs}`, `/mnt/skills`
@ -613,7 +633,7 @@ copying a raw checkpoint because delta state is not self-contained in one tuple.
2. **MCP tools** - From enabled MCP servers (lazy initialized, cached with resolved-path + content-signature invalidation)
3. **Built-in tools**:
- `present_files` - Make output files visible to user (only `/mnt/user-data/outputs`)
- `ask_clarification` - Request clarification (intercepted by ClarificationMiddleware, which preserves text fallback and adds `artifact.human_input` for Web UI Human Input Cards)
- `ask_clarification` - Request clarification (intercepted by ClarificationMiddleware, which preserves text fallback and adds `artifact.human_input` for Web UI Human Input Cards). Beyond free text and single choice, the request-side v2 protocol supports `fields` (structured form card collecting several values at once; field types: text/textarea/number/select/multi_select/checkbox/date, validated and normalized server-side in the middleware — invalid entries are dropped, unknown types degrade to `text`; a standalone multi-select question is a one-field form). Replies stay on the v1 response protocol (`text`/`option`): the form card submits a readable text summary
- `view_image` - Read image as base64 (added only if model supports vision)
- `setup_agent` - Bootstrap-only: persist a brand-new custom agent's `SOUL.md` and `config.yaml`. Bound only when `is_bootstrap=True`.
- `update_agent` - Custom-agent-only: persist self-updates to the current agent's `SOUL.md` / `config.yaml` from inside a normal chat (partial update + atomic write). Bound when `agent_name` is set and `is_bootstrap=False`.
@ -632,7 +652,7 @@ Scheduled-task runtime note:
- `browser_automation/` - Agentic browser control (stateful `navigate → observe → click/type` loop) via Playwright, distinct from the read-only `web_fetch`/`web_capture` tools. Tools: `browser_navigate`, `browser_snapshot`, `browser_click`, `browser_type`, `browser_get_text`, `browser_back`, `browser_screenshot`, `browser_close` (config `group: browser`). A process-local `BrowserSessionManager` owns one private, loop-affine Playwright event-loop thread (same pattern as the BoxLite provider) so a per-thread browser session survives across turns regardless of the caller's loop (Gateway / TUI / test). Each action returns a fresh page snapshot whose interactive elements are addressed by a stable numeric `[ref]` index (stamped as `data-df-ref` during snapshot), so the model acts on what it just observed instead of holding stale handles or guessing selectors. URLs are SSRF-screened via the shared `validate_public_http_url` (opt-out `allow_private_addresses` only for intentional internal targets). CDP attachment cannot install the request guard on an existing Chrome context, so `cdp_url` fails closed unless the operator explicitly sets `allow_unguarded_cdp: true` for a trusted local browser. Browser REST/Live access also requires an exact non-NULL thread owner, rather than the general legacy shared-thread policy, because retained pages may contain authenticated state. Session admission is a hard `max_sessions` cap: pinned Live/operation sessions are never evicted, and a new thread is rejected when no unpinned session can be closed; one Live viewer owns a session at a time. Optional dependency: `cd backend && uv sync --extra browser && uv run playwright install chromium`; `scripts/detect_uv_extras.py` preserves the extra when `config.yaml` enables `browser_navigate`, and Gateway startup fails fast if configured browser control cannot import Playwright. Tests: `tests/test_browser_automation.py` (mocked tools + a real-Chromium integration test guarded by `importorskip`); `tests/manual_browser_live_check.py` is a manual DeepSeek-driven end-to-end check (not collected by pytest).
Live UI input dispatch is kept independent from JPEG capture: non-move actions start a rate-limited background refresh loop, so pointer, wheel, or keyboard input stays responsive while continuous gestures still produce frames throughout the interaction.
Additional providers also live here (`boxlite`, `brave`, `browserless`, `crawl4ai`, `ddg_search`, `e2b_sandbox`, `exa`, `fastcrw`, `groundroute`, `infoquest`, `searxng`, `serper`); see each subpackage for specifics. E2B bootstrap is required. If it fails, the provider kills and closes the unusable remote sandbox. New sandbox creation raises an error. Warm-pool reclaim and remote discovery discard the sandbox and continue acquisition. E2B mounts remain optional.
Additional providers also live here (`boxlite`, `brave`, `browserless`, `crawl4ai`, `ddg_search`, `e2b_sandbox`, `exa`, `fastcrw`, `groundroute`, `infoquest`, `searxng`, `serper`, `tenki`); see each subpackage for specifics. E2B bootstrap is required. If it fails, the provider kills and closes the unusable remote sandbox. New sandbox creation raises an error. Warm-pool reclaim and remote discovery discard the sandbox and continue acquisition. E2B mounts remain optional.
E2B output sync records remote file versions and actual host file metadata in a thread-local manifest. The manifest binds to the remote sandbox ID. A complete output listing removes entries for deleted files. This avoids repeat downloads when the host filesystem rounds modification times. A single release-time sync pass is bounded by aggregate ceilings (`_MAX_SYNC_TOTAL_BYTES`, `_MAX_SYNC_FILES`, `_SYNC_DEADLINE_SECONDS`) on top of the per-file `_MAX_DOWNLOAD_SIZE` cap, so a pathological outputs tree cannot make release download unboundedly; a truncated pass logs what it dropped and leaves the manifest un-pruned (only entries observed in that pass are reconciled), so files it never reached are retried on the next release rather than being forgotten.
@ -719,7 +739,7 @@ Bridges external messaging platforms (Feishu, Slack, Telegram, Discord, DingTalk
A swallowed streaming failure publishes its final outbound before releasing the inbound dedupe key, so a provider redelivery can retry without overtaking the terminal reply.
- `base.py` - Abstract `Channel` base class (start/stop/send lifecycle)
- `service.py` - Manages lifecycle of all configured channels from `config.yaml`
- `slack.py` / `feishu.py` / `telegram.py` / `discord.py` / `dingtalk.py` - Platform-specific implementations (`feishu.py` tracks the running card `message_id` in memory and patches the same card in place; `telegram.py` accepts inbound text/photos/documents, preserves media captions, hands token-free attachment bytes to the shared upload pipeline, and edits the "Working on it..." stream target in place via `editMessageText`; `dingtalk.py` optionally uses AI Card streaming for in-place updates when `card_template_id` is configured)
- `slack.py` / `feishu.py` / `telegram.py` / `discord.py` / `dingtalk.py` - Platform-specific implementations (`feishu.py` tracks the running card `message_id` in memory and patches the same card in place; `telegram.py` accepts inbound text/photos/documents, preserves media captions, hands token-free attachment bytes to the shared upload pipeline, edits the "Working on it..." stream target in place via `editMessageText`, and can optionally send final Markdown replies as Rich Messages through `channels.telegram.rich_messages`; `dingtalk.py` optionally uses AI Card streaming for in-place updates when `card_template_id` is configured, and overrides `receive_file` to download inbound images (`picture`/`richText`) and documents (`file`) by `downloadCode` into the thread uploads bucket, mirroring `feishu.py`)
- `github.py` - Webhook-driven GitHub channel. Inbound messages come from `POST /api/webhooks/github`; outbound is log-only because GitHub agents post explicitly with `gh` from their sandbox when they choose to comment or create a PR
- `app/gateway/routers/channel_connections.py` - Browser-facing user connection and disconnect APIs
- `deerflow.persistence.channel_connections` - SQL-backed user-owned connection, optional credential, connect state, and conversation store
@ -751,7 +771,7 @@ The cached value is reused for both the blocking (`runs.wait`) and streaming (`_
- `langgraph_url` - LangGraph-compatible Gateway API base URL (default: `http://localhost:8001/api`)
- `gateway_url` - Gateway API URL for auxiliary commands (default: `http://localhost:8001`)
- In Docker Compose, IM channels run inside the `gateway` container, so `localhost` points back to that container. Use `http://gateway:8001/api` for `langgraph_url` and `http://gateway:8001` for `gateway_url`, or set `DEER_FLOW_CHANNELS_LANGGRAPH_URL` / `DEER_FLOW_CHANNELS_GATEWAY_URL`.
- Per-channel configs: `feishu` (app_id, app_secret), `slack` (bot_token, app_token), `telegram` (bot_token), `dingtalk` (client_id, client_secret, optional `card_template_id` for AI Card streaming), `github` (operator kill-switch `enabled`, plus `default_mention_login` for mention-required GitHub triggers)
- Per-channel configs: `feishu` (app_id, app_secret), `slack` (bot_token, app_token), `telegram` (bot_token, optional `rich_messages` for final Markdown Rich Messages), `dingtalk` (client_id, client_secret, optional `card_template_id` for AI Card streaming), `github` (operator kill-switch `enabled`, plus `default_mention_login` for mention-required GitHub triggers)
**User-owned channel connections** (`config.yaml` -> `channel_connections`):
- Disabled by default. It is a user-binding layer on top of the existing `channels.*` runtime config, not a replacement for provider bot credentials.
@ -780,7 +800,8 @@ The cached value is reused for both the blocking (`runs.wait`) and streaming (`_
- `updater.py` - LLM-based memory updates with fact extraction, whitespace-normalized fact deduplication, optimistic revision checks, and repository change sets
- `queue.py` - Debounced update queue (per-thread deduplication, configurable wait time); captures `user_id` at enqueue time so it survives the `threading.Timer` boundary
- `prompt.py` - Prompt templates for memory updates
- `storage.py` - File repository with one user-global summary JSON, agent-owned single-fact Markdown, target-only journaled changes, strict fact validation, shared-user plus per-fact optimistic revisions, lock-protected migration, deep-copy caching, and an optional retrieval adapter
- `storage.py` - File repository with one user-global summary JSON, agent-owned single-fact Markdown, target-only journaled changes, strict fact validation, shared-user plus per-fact optimistic revisions, lock-protected migration, deep-copy caching, and a RetrievalPort adapter boundary
- `retrieval.py` - Built-in scope-aware SQLite FTS5/BM25 adapter; it stores only rebuildable derived data and can be disabled with an empty `retrieval_adapter`. Chinese jieba tokenization is optional via the backend `memory-zh` extra; without it the adapter uses SQLite unicode tokenization and the substring fallback. A corrupt persistent derived database is deleted and recreated once before falling back to substring retrieval. The Gateway closes the derived SQLite connection after its shutdown flush; reads and writes remain serialized by the adapter lock, with connection pooling deferred as a performance follow-up.
- `tools.py` - Tool-driven memory mode (`memory_search`, `memory_add`, `memory_update`, `memory_delete`) using the same storage/update primitives
**Per-User Isolation**:
@ -809,7 +830,7 @@ The cached value is reused for both the blocking (`runs.wait`) and streaming (`_
- Both modes share `FileMemoryStorage`, per-user/per-agent isolation, prompt injection, manual CRUD primitives, and the updater backend.
- Middleware mode queue debounces (30s default), batches updates, and commits global summaries plus the selected/default agent's fact delta through a user-level lock, optimistic user-memory revisions, per-fact revisions, and a recoverable target-file journal. Only explicitly marked point operations may rebase a stale shared revision, and only while every addressed fact still satisfies its original absent/revision precondition. Snapshot-derived clear/trim/consolidation operations instead reload the complete document and recompute their intent on a manifest conflict, with a bounded retry. Typed manifest/fact conflict subclasses keep that decision independent of exception text, and same-ID creates and stale same-fact writes fail. Scope-lock objects are weakly cached so inactive users do not grow a process-lifetime map. Cache validation does not scale with the fact-file count: its token combines the shared JSON's `(mtime_ns, size, revision)`, so the persisted revision invalidates stale caches even when a coarse-mtime filesystem reports identical metadata for same-size writes; direct out-of-band Markdown edits require `reload()`. Atomic replacement also syncs the parent directory on POSIX so the rename is durable. DeerMem translates private storage conflict/corruption exceptions to the backend-neutral MemoryManager contract; the Gateway maps them to HTTP 409 and a stable HTTP 500 response respectively. A normal default-manager read automatically migrates legacy facts from the global JSON into `__default__`; it also adopts the earlier implicit `lead-agent` fact bucket only when that directory has no custom-agent `config.yaml`, and rejects unexpected files instead of deleting them. The v1-to-v2 migration is one-way for the running application: operators must stop DeerFlow and snapshot the configured storage root before upgrade. Before any destructive v2 write, every migrated JSON source is durably retained as `{manifest_filename}.v1.bak`; a missing-write or mismatched existing backup aborts without modifying v1 data. Legacy per-agent JSON is deleted only after its non-empty summaries are safely adopted or confirmed identical; summary conflicts keep the source file and fail loudly.
- **Proactive Markdown migration CLI**: from `backend/`, run `PYTHONPATH=. python scripts/migrate_memory_markdown.py --all-users --dry-run` to audit and omit `--dry-run` to migrate before serving traffic. Use repeated `--user-id` values when selecting exact original identities, especially standalone raw IDs containing `@` or other characters that are normalized in directory names; `--storage-path` selects a non-default DeerMem root. The CLI reuses `FileMemoryStorage.migrate`, is idempotent, continues across per-user failures, and exits non-zero if any user fails. It is optional because the first normal read still performs the same migration automatically.
- A configured `retrieval_adapter` owns indexing and semantic retrieval. File storage sends upsert/remove notifications for normal writes and both explicit and lazy migrations after releasing durable storage locks, then delegates search; without an adapter it declares and uses `substring_fallback`.
- `retrieval_adapter` owns indexing and retrieval. `fts5` is the DeerMem default and uses a persistent derived SQLite index under `.retrieval/`; an empty value disables the adapter and selects `substring_fallback`. File storage sends upsert/remove notifications for normal writes and both explicit and lazy migrations after releasing durable storage locks, then delegates search. Gateway startup schedules `DeerMem.warm_retrieval()` as a background full rebuild so readiness is not delayed, while a first search lazily rebuilds its exact scope until warm-up completes. Individual malformed facts are logged and skipped without triggering repeated full scans; only a fatal adapter rebuild failure keeps lazy retry enabled. During shutdown, the Gateway waits at most one second for this derived rebuild and leaves the full configured timeout to the canonical memory flush; if the rebuild is still active, its adapter remains open until process exit. Adapter failures mark the scope dirty and fall back to canonical substring search until rebuilding succeeds. `FileMemoryStorage` owns and closes the adapter so higher layers do not reach into private storage state.
- Staleness pass (same LLM invocation as the regular updater, no extra API call): when `staleness_review_enabled` is `true` and at least `staleness_min_candidates` aged facts exist, `_select_stale_candidates` selects facts older than their individual review window (`expected_valid_days`, or the global `staleness_age_days` fallback) that are not in `staleness_protected_categories` (default: `correction`), surfaces them in the prompt with a `valid:Nd` annotation, and the LLM judges each as KEEP, REMOVE, or EXTEND. REMOVE entries go in `staleFactsToRemove`; EXTEND entries go in `staleFactsToExtend` with an `extend_by_days` value, which sets the fact's `expected_valid_days` to `min(days_since_created + extend_by_days, staleness_max_extension_days)`. The LLM assigns `expected_valid_days` when creating a fact; it is clamped at write time to `staleness_age_days × staleness_max_lifetime_multiplier` (creation cap). `_apply_updates` enforces the guardrail unconditionally at apply time: it intersects both the removal and extension sets with `_select_stale_candidates` output before applying the per-cycle cap (`staleness_max_removals_per_cycle`), so protected and non-aged facts can never be targeted regardless of model behavior or the feature flag setting. Facts the LLM proposed for removal are excluded from extension even if the per-cycle cap prevented their actual deletion that cycle. Extensions use an absolute ceiling (`staleness_max_extension_days`) rather than the creation multiplier so a deliberate review decision can advance the window beyond the initial cap while preventing `timedelta` overflow from a malformed `extend_by_days`.
- Consolidation pass (same LLM invocation as the regular updater, no extra API call): when `consolidation_enabled` is `true` and at least one category holds `consolidation_min_facts` or more facts, `_select_consolidation_candidates` identifies fragmented categories and surfaces at most `consolidation_max_groups_per_cycle` of them (largest first) in the prompt. The LLM decides which groups to merge and proposes a synthesised fact per group. `_apply_updates` enforces guardrails: source IDs must exist and must not overlap across groups, group size is capped at `consolidation_max_sources`, the merged fact's confidence cannot exceed the source maximum, and facts below `fact_confidence_threshold` are not written. The merged fact carries the newest source's `createdAt` (so the staleness clock reflects the underlying information, not synthesis time) and inherits `expected_valid_days` set so the merged fact is re-reviewed at the earliest source review deadline (`min(createdAt + effective_lifetime)` across sources, where a source's effective lifetime is its `expected_valid_days` or the global `staleness_age_days` fallback for legacy facts without one - so a legacy source's default window is not swallowed by a long-lived sibling), relative to the merged `createdAt`, clamped to a minimal positive window if a source is already past its deadline, then capped at the creation-time `staleness_max_lifetime_multiplier`; this keeps a volatile or legacy sub-detail from inheriting a stable source's long window and escaping staleness review for years, while a merge of uniformly stable sources does not re-enter review prematurely.
- Next interaction injects selected facts + context into `<memory>` tags in the system prompt when `injection_enabled` is true.
@ -835,9 +856,9 @@ Focused regression coverage for the updater lives in `backend/tests/test_memory_
- `strict_user_scope` - Require `user_id` for all storage access (default `false` for no-auth/legacy compatibility)
- `manifest_filename` - User-global summary JSON filename (kept for configuration compatibility)
- `file_lock_timeout_seconds` - Scope-lock wait; Markdown facts and the recovery journal are required storage invariants rather than configurable modes
- `retrieval_adapter` - Optional dotted factory receiving `DeerMemConfig` and returning a retrieval-port implementation
- `retrieval_adapter` - `fts5` by default, empty to disable, or a dotted factory receiving `DeerMemConfig` and returning a retrieval-port implementation
- `debounce_seconds` - Wait time before processing (default: 30)
- `shutdown_flush_timeout_seconds` - Host-shared hard budget (seconds) to drain the memory backend's pending-update buffer on Gateway graceful shutdown (default: 30; 1300). Each pending item does one LLM call, so large IM batches may need more. The Gateway lifespan calls `MemoryManager.shutdown_flush(timeout)` after channels/scheduler stop; the backend short-circuits on an idle buffer, so the host calls it unconditionally (no pending/processing gate). Must fit inside the pod's K8s `terminationGracePeriodSeconds` (gateway Helm chart sets this; default 45s) or K8s SIGKILLs the drain mid-flight.
- `shutdown_flush_timeout_seconds` - Hard budget (seconds) reserved for draining the memory backend's pending-update buffer on Gateway graceful shutdown (default: 30; 1300). Each pending item does one LLM call, so large IM batches may need more. The Gateway lifespan calls `MemoryManager.shutdown_flush(timeout)` after channels/scheduler stop and after waiting at most one additional second for the derived retrieval warm-up; the backend short-circuits on an idle buffer, so the host calls it unconditionally (no pending/processing gate). The retrieval wait does not reduce this canonical flush budget. The combined shutdown hooks, brief retrieval wait, flush budget, and scheduling margin must fit inside the pod's K8s `terminationGracePeriodSeconds` (gateway Helm chart default: 45s) or K8s SIGKILLs the drain mid-flight.
- `model_name` - LLM for updates (null = default model)
- `max_facts` / `fact_confidence_threshold` - Fact storage limits (100 / 0.7)
- `max_injection_tokens` - Token limit for prompt injection (2000)
@ -902,7 +923,7 @@ This invokes `alembic revision --autogenerate` against the live ORM models. Revi
Checkpointer storage runs in one of two channel modes, selected by `checkpoint_channel_mode` in `config.yaml` (default `full`). `delta` mode adopts LangGraph 1.2's `DeltaChannel` for `messages`: checkpoints store a sentinel + per-step writes instead of the full message list, so storage/serde grows O(N) instead of O(N²) in turns. All checkpointer backends (memory/sqlite/postgres) serve both modes unchanged — the semantics live in the compiled graph's channel table, not in the saver.
**Mode is process-frozen and restart-required.** `make_lead_agent` freezes the resolved mode (`runtime/checkpoint_mode.py::freeze_checkpoint_channel_mode`) before compiling the graph with the mode-matched schema (`agents/thread_state.py::get_thread_state_schema`, plus `adapt_state_schema_for_mode` / `normalize_middleware_state_schemas` for middleware state). A second, different mode in the same process raises `CheckpointModeReconfigurationError`. To switch: edit config, restart.
**Mode is process-frozen and restart-required.** `make_lead_agent` and the embedded `DeerFlowClient` freeze the resolved mode (`runtime/checkpoint_mode.py::freeze_checkpoint_channel_mode`) and the delta snapshot frequency (`agents/thread_state.py::freeze_delta_snapshot_frequency`, from the `checkpoint_delta_snapshot_frequency` config knob, default 1000 — restart-required like the mode) before compiling the graph with the mode-matched schema (`agents/thread_state.py::get_thread_state_schema`, plus `adapt_state_schema_for_mode` / `normalize_middleware_state_schemas` for middleware state). Adapted middleware schemas are cached by schema, mode, and resolved snapshot frequency so a pre-freeze ephemeral graph cannot leave a stale default-frequency schema behind. A second, different mode or frequency in the same process raises `CheckpointModeReconfigurationError`. To switch: edit config, restart.
**Compatibility is asymmetric and fail-closed.** Every checkpoint written in delta mode carries metadata marker `deerflow_checkpoint_channel_mode: "delta"` (injected via `inject_checkpoint_mode`; absence of marker = full, so pre-feature checkpoints need no migration). Before any state read/write, `ensure_checkpoint_mode_compatible` rejects a full-mode process opening a delta thread with `CheckpointModeMismatchError` (surfaced as HTTP 409 with the cause and thread id by the threads router; `CheckpointModeReconfigurationError` maps to 503) — a full-mode raw read of a delta blob would silently return empty/partial `messages`. The reverse direction is allowed: delta-mode processes read full checkpoints transparently (old full checkpoints seed the delta channel), so full → delta is the smooth migration path; delta → full requires materializing/converting the data first. Detection also honors upstream's `counters_since_delta_snapshot.messages` metadata, and an explicit config marker takes precedence over any ambient context value.
@ -914,17 +935,17 @@ Checkpointer storage runs in one of two channel modes, selected by `checkpoint_c
**Wholesale state replacement uses a state-only mutation graph + `Overwrite`.** `update_state` values pass through channel reducers (`add_messages` merge in full, append in delta), so replacing reducer values requires `Overwrite` rather than an ordinary update. Full-mode rollback and context compaction replace `messages`; delta resume and delta rollback replace every materialized channel and reset current-head-only channels to their schema default (or `None`). These writes go through `build_state_mutation_graph(as_node, mode, state_schema)`, and `state_schema` MUST be the thread's effective schema (`graph_state_schema(assistant_graph)`), because the base-ThreadState fallback silently discards written channels contributed by custom `AgentMiddleware.state_schema`. Channels absent from a full-mode fork write inherit the parent's channel blobs, so middleware channels survive rollback/compaction (locked by `test_rollback_preserves_middleware_contributed_channels` and `test_compact_thread_context_preserves_middleware_contributed_channels`). The compiled mutation graph has one no-op node (entry = finish) whose checkpoint machinery (channels/versions/metadata) is identical to the agent graph's but schedules no pending tasks, so the restored/compacted head stays idle instead of re-triggering the agent. Never hand-write checkpoints via `checkpointer.aput` for this; raw writers elsewhere must preserve checkpoint parentage — severed ancestry breaks delta replay (see `runtime/runs/worker.py` writer parenting and `checkpoint_patches.py`).
**Run rollback flow** (`runtime/runs/worker.py`): `_capture_rollback_point` materializes the complete pre-run state via the accessor and captures raw `pending_writes` via `aget_tuple` into an immutable `RollbackPoint` before the run starts — capture failure disables rollback (fail-closed), never restores partial state. In `full` mode, cancel-with-rollback forks from the pre-run checkpoint via the mutation graph and inherits non-message channels from that parent. In `delta` mode, forking is unsafe once the cancelled path has attached sibling writes to the pre-run checkpoint, so rollback replaces every captured channel on the current head, using `Overwrite` for reducers and schema defaults for current-head-only channels. Both modes reattach only the captured pre-run pending writes to the restored checkpoint.
**Run rollback flow** (`runtime/runs/worker.py`): `_capture_rollback_point` materializes the complete pre-run state via the accessor and captures raw `pending_writes` via `aget_tuple` into an immutable `RollbackPoint` before the run starts — capture failure disables rollback (fail-closed), never restores partial state. In `full` mode, cancel-with-rollback forks from the pre-run checkpoint via the mutation graph and inherits non-message channels from that parent. In `delta` mode, forking is unsafe once the cancelled path has attached sibling writes to the pre-run checkpoint, so rollback replaces every captured channel on the current head, using `Overwrite` for reducers and schema defaults for current-head-only channels. Both modes reattach only the captured pre-run pending writes to the restored checkpoint. Edit replay runs (`metadata.replay_kind="edit"`) also restore the pre-run checkpoint on failed, timed-out, or interrupted completion and publish the restored `values` snapshot to the stream before `end`, so clients do not remain on a transient edited branch when the replay did not produce a successful replacement.
**Where things live**:
- `runtime/checkpoint_mode.py` — mode freeze, marker injection, delta detection, compatibility gate, both error types
- `runtime/checkpoint_state.py``CheckpointStateAccessor`, `build_state_mutation_graph`, `RollbackPoint`
- `checkpoint_patches.py` (package root) — checkpoint-machinery patches: delta-history folding for `InMemorySaver` (delegating to the base walk), stable message IDs across materialization, upstream first-write drop fix, and `BinaryOperatorAggregate` unwrapping an `Overwrite` first write into an empty (MISSING) channel — Union-typed reducer channels (`sandbox`/`goal`/`todos`/`promoted`) have no constructible default, so a replace-style write into a fresh branch thread or a never-written channel stored the wrapper literally and crashed the next consumer (#4380; probe-guarded, stands down if upstream fixes it)
- `agents/thread_state.py``ThreadState`/`DeltaThreadState`, `DELTA_MESSAGES_FIELD` (`DeltaChannel` with `snapshot_frequency=1000`), schema adaptation helpers
- `agents/thread_state.py``ThreadState`/`DeltaThreadState`, `DELTA_MESSAGES_FIELD` (`DeltaChannel` whose `snapshot_frequency` resolves from the `checkpoint_delta_snapshot_frequency` config knob via `freeze_delta_snapshot_frequency`/`resolved_delta_snapshot_frequency`; 1000 is only the default), schema adaptation helpers
- `runtime/context_compaction.py` — compaction via accessor + mutation graph (reference consumer)
- Tests: `tests/test_checkpoint_mode.py` (freeze/detect/gate), `tests/test_checkpoint_state.py` (accessor/mutation graph), `tests/test_delta_channel_checkpointers.py` (saver parity), `tests/test_threads_checkpoint_mode.py`, `tests/test_gateway_checkpoint_mode.py` (dual-mode e2e parity), `tests/test_context_compaction.py` (mutation-graph write, no scheduling), `tests/test_run_worker_rollback.py`
**Checkpoint channel benchmark**: `scripts/benchmark/bench_checkpoint_channels.py`
**Checkpoint channel benchmark**: `scripts/benchmark/checkpoint/bench_channels.py`
runs paired `full`/`delta` message-only StateGraphs in a fresh child process per
case, using sync `InMemorySaver` or `SqliteSaver` so reducer, serialization, and
saver costs stay separate from Gateway/async scheduling. It reports deterministic
@ -937,7 +958,7 @@ estimated cumulative full-payload cap skips both modes of an oversized pair when
full-payload cap, so size those runs explicitly. Use `--allow-large-cases` only
on a provisioned machine. Duplicate CSV matrix values are ignored with a warning;
use `--repetitions` for repeated samples. Summarize paired successful repetitions
with `scripts/benchmark/summarize_checkpoint_channels.py` (all ratios are
with `scripts/benchmark/checkpoint/summarize_channels.py` (all ratios are
`delta/full`). `--profile-dir /tmp/checkpoint-profiles` writes one cProfile
artifact per case for attribution. Profiled rows carry `profiled: true`, and the
summarizer automatically excludes them from baseline summaries with a warning.
@ -948,19 +969,61 @@ Example:
```bash
cd backend
PYTHONPATH=. uv run python scripts/benchmark/bench_checkpoint_channels.py \
PYTHONPATH=. uv run python scripts/benchmark/checkpoint/bench_channels.py \
--backends sqlite --updates 100,500,999,1000,1001 --payload-bytes 128 \
--repetitions 7 --output /tmp/checkpoint-bench.jsonl
PYTHONPATH=. uv run python scripts/benchmark/summarize_checkpoint_channels.py \
PYTHONPATH=. uv run python scripts/benchmark/checkpoint/summarize_channels.py \
/tmp/checkpoint-bench.jsonl
```
The sync storage benchmark is not an end-to-end Gateway benchmark. Complete
`ThreadState`/`DeltaThreadState`, async saver scheduling, history, mutation,
rollback, migration, and branch-heavy cases belong to the production-shaped
follow-up layer. Harness tests live in `tests/test_bench_checkpoint_channels.py`
and `tests/test_summarize_checkpoint_channels.py`; timing thresholds are not CI
gates.
The production-shaped layer lives in
`scripts/benchmark/checkpoint/bench_production.py`: per-case child processes
run graph-level `ainvoke` turns through the real lead-agent graph (scripted
deterministic model, real `AsyncSqliteSaver`), then measure
`GET /threads/{id}/state` and `POST /threads/{id}/history` through the real
Gateway route stack in the same event loop (httpx ASGITransport), split into
cold/warm accessor-graph-cache samples. It sweeps `snapshot_frequency`
(config: `checkpoint_delta_snapshot_frequency`, process-frozen like the
mode), pairs every delta frequency against the same full row, and fails both
rows of a pair when materialized or wire digests diverge. Each case must have
more than the two discarded warm-up turns, and SQLite DB/WAL/SHM sizes are
captured while the saver is still open so they represent the online storage
footprint. Summarize with
`scripts/benchmark/checkpoint/summarize_production.py` (ratios are
`delta/full`; it also emits `snapshot_write_spike` and `cache_effect_ms`,
the decision inputs for the production snapshot-frequency and accessor-cache
defaults). Harness tests live in `tests/test_bench_checkpoint_production.py`
and `tests/test_summarize_checkpoint_production.py`; timing thresholds are
not CI gates. The matrix test pins that every `(repetition, turns)` group
contains both modes and that their execution order flips between consecutive
groups, including across repetition boundaries.
Operational limits learned from the first runs (the default matrix is too
large to run blindly):
- The default `--timeout-seconds 900` is insufficient for delta mode at
`snapshot_frequency=1000` once turns reach 500 (measured: delta-500 takes
~1100-1200s; delta-2000 takes ~45min). Pass an explicit
`--timeout-seconds` for any large matrix, and treat the turns=2000 corner
as practical only at small snapshot frequencies.
- Full-mode 2000-turn runs produce a ~33GB sqlite DB. Point `TMPDIR` at real
disk, not tmpfs (the benchmark uses `tempfile.TemporaryDirectory`, which
honors `TMPDIR`), or the run dies mid-case.
- The history route clamps `limit` to 100 (`le=100` on
`ThreadHistoryRequest.limit`), so `--history-limits` values above 100 are
measured and reported by their effective (clamped) limit.
Example:
```bash
cd backend
PYTHONPATH=. uv run python scripts/benchmark/checkpoint/bench_production.py \
--turns 10,100,500,1000,2000 --payload-bytes 128 \
--snapshot-frequencies 10,50,100,500,1000 \
--repetitions 7 --output /tmp/production-bench.jsonl
PYTHONPATH=. uv run python scripts/benchmark/checkpoint/summarize_production.py \
/tmp/production-bench.jsonl
```
### Terminal Workbench / TUI (`packages/harness/deerflow/tui/`)

View File

@ -0,0 +1,276 @@
"""Inbound webhook dedupe store.
The manager-level inbound dedupe (``ChannelManager._inbound_dedupe_key``) guards an
agent run / final answer against provider redeliveries. The default store is an
in-process ``OrderedDict`` (backward compatible, single-pod). A shared store (e.g.
Postgres) can be injected for multi-pod deployments so a redelivery landing on a
different pod is still dropped as a duplicate. See issue #4120.
"""
from __future__ import annotations
import logging
import os
import time
from collections import OrderedDict
from typing import Any, Protocol
from sqlalchemy import text
logger = logging.getLogger(__name__)
INBOUND_DEDUPE_TTL_SECONDS = 10 * 60
INBOUND_DEDUPE_MAX_ENTRIES = 4096
# Key tuple matches ChannelManager._inbound_dedupe_key:
# (channel_name, workspace_id, chat_id, message_id).
InboundDedupeKey = tuple[str, str, str, str]
class InboundDedupeStore(Protocol):
"""Async contract for recording / releasing inbound dedupe keys.
``try_record`` returns ``True`` if the key already existed (duplicate -> drop)
and ``False`` if it was newly recorded or its prior entry had expired (proceed).
Shared-state implementations must be atomic; the Postgres variant uses a single
conditional upsert (``INSERT ... ON CONFLICT DO UPDATE ... WHERE first_seen < TTL``).
"""
async def try_record(self, key: InboundDedupeKey) -> bool: ...
async def release(self, key: InboundDedupeKey) -> None: ...
class MemoryInboundDedupeStore:
"""Process-local ``OrderedDict`` store — preserves the pre-#4120 behavior exactly."""
def __init__(
self,
ttl_seconds: int = INBOUND_DEDUPE_TTL_SECONDS,
max_entries: int = INBOUND_DEDUPE_MAX_ENTRIES,
) -> None:
self._ttl = ttl_seconds
self._max = max_entries
# Insertion order == chronological (keys are never re-inserted), so an
# OrderedDict lets us evict expired/overflow entries from the front in
# O(k) instead of scanning all entries on every inbound message.
self._store: OrderedDict[InboundDedupeKey, float] = OrderedDict()
async def try_record(self, key: InboundDedupeKey) -> bool:
now = time.monotonic()
# Entries are in chronological insertion order, so expired ones cluster at
# the front: pop from the front until we hit a still-live entry.
while self._store:
_, oldest_at = next(iter(self._store.items()))
if now - oldest_at > self._ttl:
self._store.popitem(last=False)
else:
break
while len(self._store) > self._max:
self._store.popitem(last=False)
if key in self._store:
return True
self._store[key] = now
return False
async def release(self, key: InboundDedupeKey) -> None:
self._store.pop(key, None)
class PostgresInboundDedupeStore:
"""Shared Postgres-backed dedupe store (issue #4120).
One row per dispatched inbound webhook (keyed by the 4-tuple). A redelivery
routed to a different gateway pod hits the same table. The acquire is a single
atomic conditional upsert:
INSERT ... ON CONFLICT (4-tuple) DO UPDATE SET first_seen = now()
WHERE first_seen < now() - TTL RETURNING channel
- No conflict -> row inserted -> proceed.
- Conflict + expired row -> DO UPDATE refreshes ``first_seen``, row RETURNED ->
proceed (honors the 10-minute ceiling and re-admits a never-released/expired
redelivery, e.g. a manual provider "Redeliver").
- Conflict + live row -> WHERE fails, no row RETURNED -> drop as duplicate.
Because the upsert is a single row-locked statement, two pods racing on the same
expired key cannot both proceed (one wins the update; the other sees the fresh row
and is dropped). Lazy cleanup (``DELETE`` of rows older than the TTL) runs in the
same transaction, amortized into the proceed path with no background task.
Fail-open: any DB error is logged and treated as "allow" so a storage
outage never drops a webhook or returns 5xx to the provider.
"""
def __init__(self, session_factory: Any | None = None) -> None:
# Injected in tests; otherwise resolved lazily from the app engine.
self._session_factory = session_factory
def _resolve_session_factory(self) -> Any:
if self._session_factory is not None:
return self._session_factory
from deerflow.persistence.engine import get_session_factory
sf = get_session_factory()
if sf is None:
raise RuntimeError("PostgresInboundDedupeStore requires a Postgres session factory")
return sf
async def try_record(self, key: InboundDedupeKey) -> bool:
channel, workspace_id, chat_id, message_id = key
try:
sf = self._resolve_session_factory()
async with sf() as session:
async with session.begin():
# Atomic acquire + TTL reclamation in ONE row-locked statement.
#
# - No conflict -> new row inserted, RETURNING a row -> proceed.
# - Conflict + the existing row is EXPIRED (first_seen older than
# the TTL) -> DO UPDATE refreshes first_seen to now() and the row
# is RETURNED, so the redelivery is re-admitted (proceed). This is
# the cross-pod-safe equivalent of the memory store's "evict
# expired entries before the membership check", and it honors the
# 10-minute ceiling of issue #4120 even for a row that was never
# released (e.g. a crashed run): a manual provider "Redeliver" of
# an old message is re-admitted instead of being dropped forever
# in a quiet deployment.
# - Conflict + the existing row is still LIVE -> the WHERE fails, no
# UPDATE, no row RETURNED -> duplicate -> drop.
#
# A single conditional upsert has no TOCTOU window (unlike a separate
# DELETE-then-INSERT), so two pods racing on the same expired key
# cannot both proceed: one wins the update and proceeds, the other
# sees the now-fresh row and is dropped.
result = await session.execute(
text(
"INSERT INTO webhook_deliveries "
"(channel, workspace_id, chat_id, message_id, first_seen) "
"VALUES (:c, :w, :ch, :m, now()) "
"ON CONFLICT (channel, workspace_id, chat_id, message_id) "
"DO UPDATE SET first_seen = now() "
"WHERE webhook_deliveries.first_seen < now() - make_interval(secs => :ttl) "
"RETURNING channel"
),
{
"c": channel,
"w": workspace_id,
"ch": chat_id,
"m": message_id,
"ttl": INBOUND_DEDUPE_TTL_SECONDS,
},
)
# A returned row means the key was admitted (new delivery, or an
# expired row that was re-admitted). No row means a live duplicate
# that must be dropped. RETURNING (not rowcount) is used because
# rowcount reliability for ON CONFLICT DO NOTHING/DO UPDATE varies
# across DB drivers.
inserted = result.fetchone() is not None
# Lazy cleanup in the same transaction: drop rows older than the
# TTL. Only when a row was admitted (proceed path) so the periodic
# sweep is amortized into normal inbound traffic and keys never
# re-accessed still get reclaimed.
if inserted:
await session.execute(
text("DELETE FROM webhook_deliveries WHERE first_seen < now() - make_interval(secs => :ttl)"),
{"ttl": INBOUND_DEDUPE_TTL_SECONDS},
)
# inserted=True -> admitted (proceed, not a duplicate).
return not inserted
except Exception:
# Fail-open: if the store is unavailable we must NOT drop the
# message. Return False so the caller treats it as a new delivery
# and proceeds (at worst a possible duplicate, never silent loss).
logger.exception("PostgresInboundDedupeStore.try_record failed; proceeding without dedupe (fail-open)")
return False
async def release(self, key: InboundDedupeKey) -> None:
channel, workspace_id, chat_id, message_id = key
try:
sf = self._resolve_session_factory()
async with sf() as session:
async with session.begin():
await session.execute(
text("DELETE FROM webhook_deliveries WHERE channel = :c AND workspace_id = :w AND chat_id = :ch AND message_id = :m"),
{"c": channel, "w": workspace_id, "ch": chat_id, "m": message_id},
)
except Exception:
logger.exception("PostgresInboundDedupeStore.release failed; key left for TTL expiry (fail-open)")
def _gateway_workers() -> int:
"""Mirror deps._enforce_postgres_for_multi_worker's worker detection."""
try:
return int(os.environ.get("GATEWAY_WORKERS", "1") or 1)
except (TypeError, ValueError):
return 1
def _build_postgres_store() -> InboundDedupeStore:
"""Build the shared Postgres dedupe store."""
return PostgresInboundDedupeStore()
def make_inbound_dedupe_store(app_config: Any | None = None) -> InboundDedupeStore:
"""Resolve the inbound dedupe store from app config.
- ``memory`` -> in-process store (per-pod; a redelivery routed to a different
replica is NOT deduped).
- ``postgres`` -> shared Postgres store. Requires ``database.backend='postgres'``;
if the application DB is not Postgres the store falls back to the in-process
memory store and logs a WARNING (otherwise cross-pod dedupe would be silently
disabled).
- ``auto`` (default) -> shared Postgres store whenever the application DB is
Postgres. This is the recommended setting for any deployment that may run more
than one replica, including Kubernetes with ``GATEWAY_WORKERS=1`` per pod where
multiple pods still share the single Postgres DB. For non-Postgres DBs ``auto``
falls back to the cheaper in-process memory store (correct for a single-DB,
single-pod deployment).
Emits a WARNING when a multi-worker/multi-replica deployment cannot use the shared
Postgres store (the cross-pod dedupe gap becomes an explicit misconfiguration
rather than silent default behavior).
"""
backend = "auto"
db_is_postgres = False
db_backend = None
if app_config is not None:
dedupe_cfg = getattr(app_config, "dedupe_storage", None)
if dedupe_cfg is not None:
backend = str(dedupe_cfg.backend.value if hasattr(dedupe_cfg.backend, "value") else dedupe_cfg.backend)
db = getattr(app_config, "database", None)
db_backend = getattr(db, "backend", None)
db_is_postgres = db_backend == "postgres"
multi_worker = _gateway_workers() > 1
if backend == "postgres":
if not db_is_postgres:
logger.warning(
"dedupe_storage=postgres requires database.backend='postgres' (got '%s'). Falling back to the in-process memory store; inbound webhook dedupe is per-pod and cross-pod redeliveries will NOT be deduped. See issue #4120.",
db_backend,
)
return MemoryInboundDedupeStore()
return _build_postgres_store()
if backend == "memory":
if multi_worker:
logger.warning(
"dedupe_storage=memory with GATEWAY_WORKERS>1: inbound webhook dedupe "
"is per-pod and will NOT drop redeliveries routed to a different replica. "
"Use dedupe_storage=postgres (or remove the setting to let 'auto' pick it) "
"for multi-worker deployments. See issue #4120."
)
return MemoryInboundDedupeStore()
# auto
if db_is_postgres:
logger.info("dedupe_storage=auto resolved to the shared Postgres store (database.backend=postgres); inbound webhook dedupe is shared across pods.")
return _build_postgres_store()
if multi_worker:
logger.warning(
"Multi-worker deployment detected but dedupe_storage=auto resolved to the "
"in-process memory store (application database is not Postgres). Inbound "
"webhook dedupe is per-pod and will NOT drop redeliveries routed to a different "
"replica. Set database.backend=postgres (required for multi-worker) so dedupe "
"shares state across pods. See issue #4120."
)
return MemoryInboundDedupeStore()

View File

@ -17,6 +17,10 @@ from app.channels.base import Channel
from app.channels.commands import is_known_channel_command
from app.channels.connection_identity import attach_connection_identity
from app.channels.message_bus import InboundMessage, InboundMessageType, MessageBus, OutboundMessage, ResolvedAttachment
from deerflow.config.paths import VIRTUAL_PATH_PREFIX, get_paths
from deerflow.runtime.user_context import get_effective_user_id
from deerflow.sandbox.sandbox_provider import get_sandbox_provider
from deerflow.uploads.manager import UnsafeUploadPathError, claim_unique_filename, normalize_filename, write_upload_file_no_symlink
logger = logging.getLogger(__name__)
@ -29,6 +33,11 @@ _CONVERSATION_TYPE_GROUP = "2"
_MAX_UPLOAD_SIZE_BYTES = 20 * 1024 * 1024
# Inbound attachments are buffered in memory before being persisted and synced
# into the sandbox, so bound them. DingTalk chats accept far larger files than
# an agent can usefully read; oversized ones surface as a failed-load marker.
_MAX_INBOUND_FILE_SIZE_BYTES = 50 * 1024 * 1024
def _normalize_conversation_type(raw: Any) -> str:
"""Normalize ``conversationType`` to ``"1"`` (P2P) or ``"2"`` (group).
@ -63,6 +72,16 @@ def _is_dingtalk_command(text: str) -> bool:
return is_known_channel_command(text)
def _display_filename(filename: str) -> str:
"""Collapse whitespace and cap length for embedding a filename in message text.
``fileName`` is webhook-supplied data: embedded raw in a failed-load marker,
a newline could forge a standalone ``/mnt/user-data/uploads/...`` line and an
over-long name would bloat the message text.
"""
return re.sub(r"\s+", " ", filename).strip()[:80]
def _extract_text_from_rich_text(rich_text_list: list) -> str:
parts: list[str] = []
for item in rich_text_list:
@ -137,6 +156,9 @@ class DingTalkChannel(Channel):
self._incoming_messages: dict[str, Any] = {}
self._incoming_messages_lock = threading.Lock()
self._card_repliers: dict[str, Any] = {}
# Serialize inbound-file writes into the uploads directory to avoid
# racing writers clobbering one another (mirrors FeishuChannel).
self._file_write_lock = threading.Lock()
@property
def supports_streaming(self) -> bool:
@ -366,8 +388,9 @@ class DingTalkChannel(Channel):
sender_nick = message.sender_nick or ""
text = self._extract_text(message)
if not text:
logger.info("[DingTalk] empty text, ignoring message")
files = self._extract_files(message)
if not text and not files:
logger.info("[DingTalk] empty message with no files, ignoring")
return
connect_code = self._pending_connect_code(text)
@ -396,12 +419,13 @@ class DingTalkChannel(Channel):
# INFO logs, and only after the allowed_users gate so blocked senders are
# not logged at all.
logger.info(
"[DingTalk] parsed message: conv_type=%s, msg_id=%s, sender=%s(%s), text_len=%d",
"[DingTalk] parsed message: conv_type=%s, msg_id=%s, sender=%s(%s), text_len=%d, files=%d",
conversation_type,
msg_id,
sender_staff_id,
sender_nick,
len(text or ""),
len(files),
)
if _is_dingtalk_command(text):
@ -435,6 +459,7 @@ class DingTalkChannel(Channel):
text=text,
msg_type=msg_type,
thread_ts=msg_id,
files=files,
metadata={
"conversation_type": conversation_type,
"conversation_id": conversation_id,
@ -471,6 +496,237 @@ class DingTalkChannel(Channel):
return _extract_text_from_rich_text(message.rich_text_content.rich_text_list).strip()
return ""
# -- inbound: file attachments -----------------------------------------
@staticmethod
def _extract_files(message: Any) -> list[dict[str, Any]]:
"""Extract inbound file/image descriptors from a DingTalk message.
Returns a list of dicts shaped for ``InboundMessage.files``; each carries
``type`` (``"image"`` or ``"file"``), ``download_code``, and ``filename``.
Images arrive as ``picture`` messages (a single ``downloadCode``) or as
inline images inside a ``richText`` message. Documents arrive as ``file``
messages, which ``dingtalk_stream.ChatbotMessage.from_dict`` does not
parse the descriptor (``downloadCode`` / ``fileName``) is read from the
raw callback payload stashed on the message as ``_df_raw_data`` by
``_DingTalkMessageHandler.process``.
"""
msg_type = getattr(message, "message_type", None)
files: list[dict[str, Any]] = []
if msg_type == "picture":
image_content = getattr(message, "image_content", None)
code = getattr(image_content, "download_code", None)
if code:
files.append({"type": "image", "download_code": code, "filename": "image.png"})
elif msg_type == "richText":
try:
codes = message.get_image_list() or []
except Exception:
# Log rather than swallow: without this, a richText message whose
# SDK image parse fails looks identical to one that simply has no
# inline images, which is hard to diagnose if the SDK shape changes.
logger.warning("[DingTalk] failed to read inline images from richText message", exc_info=True)
codes = []
for idx, code in enumerate(codes):
if code:
files.append({"type": "image", "download_code": code, "filename": f"image_{idx}.png"})
elif msg_type == "file":
raw = getattr(message, "_df_raw_data", None)
content = raw.get("content") if isinstance(raw, dict) else None
if isinstance(content, dict):
code = content.get("downloadCode")
if code:
filename = content.get("fileName") or "file.bin"
files.append({"type": "file", "download_code": code, "filename": filename})
return files
async def receive_file(self, msg: InboundMessage, thread_id: str, *, user_id: str | None = None) -> InboundMessage:
"""Download inbound DingTalk files into the thread uploads directory.
Mirrors :meth:`FeishuChannel.receive_file`: each descriptor in
``msg.files`` is downloaded by its ``downloadCode``, persisted under the
thread's uploads bucket, synced into a non-local sandbox, and its sandbox
virtual path is prepended to ``msg.text`` so downstream models can read
the file by path. Descriptors are then cleared so the generic
URL-based ``_ingest_inbound_files`` pass does not try to re-fetch them
(DingTalk files are downloaded by code, not by URL).
An attachment that fails to download contributes a short
``[failed to load ...]`` marker instead of a path, so the agent can tell
the user something was sent but could not be read rather than silently
ignoring it.
"""
if not msg.files:
return msg
virtual_paths: list[str] = []
failures: list[str] = []
for f in msg.files:
if not isinstance(f, dict):
continue
download_code = f.get("download_code")
if not download_code:
continue
file_type = "image" if f.get("type") == "image" else "file"
filename = f.get("filename") if isinstance(f.get("filename"), str) else ""
# Per-attachment isolation: the manager awaits receive_file without a
# try, so anything escaping here would kill the chat turn with no
# reply — the silent-drop failure mode this feature exists to fix.
try:
virtual_path = await self._receive_single_file(download_code, file_type, filename, thread_id, user_id=user_id)
except Exception:
logger.exception("[DingTalk] unexpected error receiving inbound %s", file_type)
virtual_path = ""
if virtual_path:
virtual_paths.append(virtual_path)
else:
display = _display_filename(filename)
failures.append(f"[failed to load {file_type}: {display}]" if display else f"[failed to load {file_type}]")
prefix_lines = virtual_paths + failures
if prefix_lines:
block = "\n".join(prefix_lines)
msg.text = f"{block}\n\n{msg.text}".strip() if msg.text else block
msg.files = []
return msg
async def _receive_single_file(
self,
download_code: str,
file_type: str,
filename: str,
thread_id: str,
*,
user_id: str | None = None,
) -> str:
"""Download one file by ``downloadCode`` and persist it into uploads.
Returns the sandbox virtual path on success, or ``""`` on failure.
"""
content = await self._download_by_code(download_code)
if not content:
return ""
paths = get_paths()
effective_user_id = user_id or get_effective_user_id()
default_ext = "png" if file_type == "image" else "bin"
# ``download_code`` is attacker-controllable webhook data, so restrict it to
# a safe character set before embedding it in the fallback name. This keeps
# the fallback safe by construction instead of relying on a later write
# failure to reject a name that escaped the uploads directory.
code_token = re.sub(r"[^A-Za-z0-9_-]", "", download_code)[-12:] or "attachment"
fallback_name = f"dingtalk_{code_token}.{default_ext}"
# normalize_filename strips directory components and rejects traversal
# patterns ("..", backslash paths, over-long names); fall back to the
# sanitized generated name if the platform-supplied filename is rejected.
try:
safe_filename = normalize_filename(filename or fallback_name)
except ValueError:
safe_filename = fallback_name
def _persist() -> Path:
# Directory prep, the uniqueness claim, and the write are blocking
# filesystem IO — the whole sequence stays off the event loop. The
# claim and the write share one lock because generated names repeat
# across messages ("image.png" for every picture message): without a
# claim a later attachment silently overwrites an earlier one whose
# path was already handed to the agent, and letting the claim and
# write interleave would resolve two attachments to the same free name.
paths.ensure_thread_dirs(thread_id, user_id=effective_user_id)
uploads_dir = paths.sandbox_uploads_dir(thread_id, user_id=effective_user_id).resolve()
with self._file_write_lock:
seen = {entry.name for entry in uploads_dir.iterdir() if entry.is_file()}
unique_name = claim_unique_filename(safe_filename, seen)
# write_upload_file_no_symlink refuses a symlinked destination:
# uploads dirs can be mounted into local sandboxes, so a sandbox
# process could otherwise redirect this privileged write outside
# the bucket.
return write_upload_file_no_symlink(uploads_dir, unique_name, content)
try:
resolved_target = await asyncio.to_thread(_persist)
except (OSError, UnsafeUploadPathError):
logger.exception("[DingTalk] failed to persist downloaded file: %s", safe_filename)
return ""
virtual_path = f"{VIRTUAL_PATH_PREFIX}/uploads/{resolved_target.name}"
try:
sandbox_provider = get_sandbox_provider()
# acquire_async keeps provider lifecycle work (Docker discovery,
# readiness polls) off the event loop; update_file is blocking
# transport IO on remote sandboxes, so it is offloaded too.
sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=effective_user_id)
if sandbox_id != "local":
sandbox = sandbox_provider.get(sandbox_id)
if sandbox is None:
# Mirror Feishu: the agent's non-local sandbox cannot see this
# file, so returning the virtual path would hand the model a
# path that reads as nothing — surface a failed-load marker.
logger.warning("[DingTalk] sandbox %s not found after acquire, dropping attachment: %s", sandbox_id, virtual_path)
return ""
await asyncio.to_thread(sandbox.update_file, virtual_path, content)
except Exception:
# Same failure mode as the sandbox-is-None branch: the bytes never
# reached the agent's sandbox, so the virtual path would read as
# nothing. Mirror Feishu and surface a failed-load marker.
logger.exception("[DingTalk] failed to sync downloaded file into non-local sandbox: %s", virtual_path)
return ""
return virtual_path
async def _download_by_code(self, download_code: str) -> bytes | None:
"""Exchange a DingTalk ``downloadCode`` for the raw file bytes.
Two steps per the DingTalk robot OpenAPI:
1. ``POST /v1.0/robot/messageFiles/download`` -> ``{downloadUrl}``
2. ``GET downloadUrl`` -> binary content
"""
try:
# Token acquisition stays inside the try: the manager awaits
# receive_file without a try, so a token failure escaping here would
# abort the whole chat turn instead of degrading to a failed-load marker.
token = await self._get_access_token()
async with httpx.AsyncClient(timeout=httpx.Timeout(30.0)) as client:
response = await client.post(
f"{DINGTALK_API_BASE}/v1.0/robot/messageFiles/download",
headers={
"x-acs-dingtalk-access-token": token,
"Content-Type": "application/json",
},
json={"downloadCode": download_code, "robotCode": self._client_id},
)
if response.status_code != 200:
logger.warning("[DingTalk] messageFiles/download failed: status=%d, body=%s", response.status_code, response.text[:300])
return None
download_url = response.json().get("downloadUrl")
if not download_url:
logger.warning("[DingTalk] messageFiles/download returned no downloadUrl")
return None
# Stream with a size cap: the bytes are buffered in memory before
# being persisted, so an oversized attachment must be refused
# before it is fully read, not after.
chunks: list[bytes] = []
total = 0
async with client.stream("GET", download_url, follow_redirects=True) as file_response:
file_response.raise_for_status()
async for chunk in file_response.aiter_bytes():
total += len(chunk)
if total > _MAX_INBOUND_FILE_SIZE_BYTES:
logger.warning("[DingTalk] inbound file exceeds %d bytes, dropping", _MAX_INBOUND_FILE_SIZE_BYTES)
return None
chunks.append(chunk)
return b"".join(chunks)
except (httpx.HTTPError, ValueError):
logger.exception("[DingTalk] failed to download file by code")
return None
async def _prepare_inbound(self, chat_id: str, inbound: InboundMessage) -> None:
inbound = await self._attach_connection_identity(inbound)
# Running reply must finish before publish_inbound so AI card tracks are
@ -831,5 +1087,9 @@ class _DingTalkMessageHandler:
import dingtalk_stream
incoming_message = dingtalk_stream.ChatbotMessage.from_dict(callback.data)
# Stash the raw callback payload: dingtalk_stream does not parse ``file``
# (document) messages, so DingTalkChannel._extract_files reads the file
# descriptor (downloadCode / fileName) from it.
incoming_message._df_raw_data = callback.data
self._channel._on_chatbot_message(incoming_message)
return dingtalk_stream.AckMessage.STATUS_OK, "OK"

View File

@ -20,6 +20,7 @@ from langgraph_sdk.errors import ConflictError
from app.channels import feishu_run_policy as _feishu_run_policy # noqa: F401
from app.channels.commands import KNOWN_CHANNEL_COMMANDS
from app.channels.dedupe_store import InboundDedupeStore, MemoryInboundDedupeStore
from app.channels.message_bus import (
INBOUND_FILE_CONTENT_KEY,
PENDING_CLARIFICATION_METADATA_KEY,
@ -75,15 +76,14 @@ MESSAGE_STREAM_EVENTS = ("messages-tuple", "messages")
THREAD_BUSY_MESSAGE = "This conversation is already processing another request. Please wait for it to finish and try again."
BOUND_IDENTITY_REQUIRED_MESSAGE = "Connect this channel from DeerFlow Settings, complete the in-channel connect step, then send your message again."
BOUND_IDENTITY_UNAVAILABLE_MESSAGE = "Channel connection verification is temporarily unavailable. Please try again later or contact the DeerFlow operator."
# Inbound-redelivery dedup window. ``_recent_inbound_events`` (below) is a
# process-local, in-memory-only OrderedDict — it is never persisted to
# ``ChannelStore`` — so a recorded key survives only for this TTL (or until
# evicted by the entry cap below) and is gone entirely across a Gateway
# restart. 10 minutes is a deliberately bounded window: long enough to
# absorb a near-term redelivery of the same event — whether a provider's
# own automatic retry (where the provider implements one) or an operator
# explicitly triggering a resend — without keeping a growing in-memory
# ledger.
# Inbound-redelivery dedup window. The dedupe state lives in
# ``self._inbound_dedupe_store``: the default in-process Memory store is
# local to this Gateway process (a recorded key survives only for the store's
# TTL / entry cap and is gone across a restart), while a Postgres-backed store
# is shared across pods. 10 minutes is a deliberately bounded window: long
# enough to absorb a near-term redelivery of the same event — whether a
# provider's own automatic retry or an operator resend — without keeping a
# growing ledger.
#
# For GitHub specifically: GitHub does NOT automatically retry or redeliver
# a failed delivery (non-2xx response, timeout, or connection error) — it
@ -114,8 +114,6 @@ BOUND_IDENTITY_UNAVAILABLE_MESSAGE = "Channel connection verification is tempora
FOLLOWUP_BUFFER_MAX_PER_THREAD = 20
FOLLOWUP_DRAIN_BATCH_SIZE = 10
FOLLOWUP_BLOCK_TAG = "followups-while-busy"
INBOUND_DEDUPE_TTL_SECONDS = 10 * 60
INBOUND_DEDUPE_MAX_ENTRIES = 4096
# Only server-stable provider message ids: client-generated ids (client_msg_id,
# client_id) are not guaranteed identical across a provider's own redelivery, so
# keying dedupe on them would miss exactly the retries we want to absorb.
@ -917,6 +915,7 @@ class ChannelManager:
channel_sessions: dict[str, Any] | None = None,
connection_repo: Any | None = None,
require_bound_identity: bool = False,
inbound_dedupe_store: InboundDedupeStore | None = None,
get_stream_bridge: Callable[[], StreamBridge | None] | None = None,
) -> None:
self.bus = bus
@ -956,14 +955,14 @@ class ChannelManager:
# stop() has run, for the follow-up drain guard below.
self._stopped = False
self._task: asyncio.Task | None = None
# Insertion order == chronological (keys are never re-inserted), so an
# OrderedDict lets us evict expired/overflow entries from the front in
# O(k) instead of scanning all entries on every inbound message.
self._recent_inbound_events: OrderedDict[tuple[str, str, str, str], float] = OrderedDict()
# Inbound webhook dedupe store. Defaults to the in-process Memory store
# (pre-#4120 behavior). Multi-pod deployments inject a shared store so
# duplicate deliveries landing on different pods are collapsed.
self._inbound_dedupe_store = inbound_dedupe_store if inbound_dedupe_store is not None else MemoryInboundDedupeStore()
# Per-thread follow-up buffers for busy fire_and_forget channels that
# opted into ChannelRunPolicy.buffer_followups_on_busy (issue #4121
# Slice 2). Keyed by thread_id -> OrderedDict[dedupe_key -> entry],
# oldest-first, mirroring _recent_inbound_events's shape but scoped
# oldest-first, mirroring the dedupe store's shape but scoped
# per-thread with a hard cap instead of a global TTL (see
# _buffer_followup / _enforce_followup_cap).
self._followup_buffers: dict[str, OrderedDict[str, _FollowupEntry]] = {}
@ -1549,7 +1548,7 @@ class ChannelManager:
# answer. Provider adapters may emit ack side-effects (a "Working on
# it…" reply, an "eyes" reaction) before publish_inbound, so those are
# intentionally not deduped here.
if self._is_duplicate_inbound(msg):
if await self._is_duplicate_inbound(msg):
continue
logger.info(
"[Manager] received inbound: channel=%s, chat_id=%s, type=%s, text_len=%d, files=%d",
@ -1598,36 +1597,25 @@ class ChannelManager:
return None
return (msg.channel_name, str(workspace_id), msg.chat_id, message_id)
def _is_duplicate_inbound(self, msg: InboundMessage) -> bool:
async def _is_duplicate_inbound(self, msg: InboundMessage) -> bool:
key = self._inbound_dedupe_key(msg)
if key is None:
return False
now = time.monotonic()
# Entries are in chronological insertion order, so expired ones cluster at
# the front: pop from the front until we hit a still-live entry.
while self._recent_inbound_events:
_, oldest_at = next(iter(self._recent_inbound_events.items()))
if now - oldest_at > INBOUND_DEDUPE_TTL_SECONDS:
self._recent_inbound_events.popitem(last=False)
else:
break
while len(self._recent_inbound_events) > INBOUND_DEDUPE_MAX_ENTRIES:
self._recent_inbound_events.popitem(last=False)
if key in self._recent_inbound_events:
# Delegated to the shared/per-pod dedupe store. The store owns TTL eviction
# and capacity bounds; try_record returns True when the key was already
# present (i.e. this is a duplicate delivery to drop).
is_duplicate = await self._inbound_dedupe_store.try_record(key)
if is_duplicate:
logger.info(
"[Manager] duplicate inbound ignored: channel=%s, chat_id=%s, message_id=%s",
msg.channel_name,
msg.chat_id,
key[-1],
)
return True
return is_duplicate
self._recent_inbound_events[key] = now
return False
def _release_inbound_dedupe_key(self, msg: InboundMessage) -> None:
async def _release_inbound_dedupe_key(self, msg: InboundMessage) -> None:
"""Drop a recorded dedupe key so a provider redelivery can be reprocessed.
Called only on transient/unexpected handling failures: the key was
@ -1637,7 +1625,7 @@ class ChannelManager:
"""
key = self._inbound_dedupe_key(msg)
if key is not None:
self._recent_inbound_events.pop(key, None)
await self._inbound_dedupe_store.release(key)
@staticmethod
def _log_task_error(task: asyncio.Task) -> None:
@ -1692,7 +1680,7 @@ class ChannelManager:
# Transient/unexpected failure: release the dedupe key so a provider
# redelivery of the same message can recover instead of being dropped
# for the dedupe TTL.
self._release_inbound_dedupe_key(msg)
await self._release_inbound_dedupe_key(msg)
await self._send_error(msg, "An internal error occurred. Please try again.")
# -- chat handling -----------------------------------------------------
@ -2064,7 +2052,7 @@ class ChannelManager:
# Swallowed like the generic handler would not be: release the
# key so the provider's redelivery can retry once the thread
# frees, instead of being dropped for the dedupe TTL.
self._release_inbound_dedupe_key(msg)
await self._release_inbound_dedupe_key(msg)
await self._send_error(msg, THREAD_BUSY_MESSAGE)
return
raise
@ -2084,7 +2072,7 @@ class ChannelManager:
logger.warning("[Manager] thread busy (concurrent run rejected): thread_id=%s", thread_id)
# Same reason as the fire-and-forget branch above: this error is
# handled here rather than re-raised, so release explicitly.
self._release_inbound_dedupe_key(msg)
await self._release_inbound_dedupe_key(msg)
await self._send_error(msg, THREAD_BUSY_MESSAGE)
return
else:
@ -2261,7 +2249,7 @@ class ChannelManager:
# handler never runs and never releases the key. Release only
# after publishing the final outbound so a provider redelivery
# cannot overtake this attempt's terminal reply.
self._release_inbound_dedupe_key(msg)
await self._release_inbound_dedupe_key(msg)
# -- command handling --------------------------------------------------

View File

@ -99,6 +99,7 @@ class ChannelService:
*,
connection_repo: Any | None = None,
require_bound_identity: bool = False,
app_config: AppConfig | None = None,
get_stream_bridge: Callable[[], StreamBridge | None] | None = None,
) -> None:
self.bus = MessageBus()
@ -110,6 +111,8 @@ class ChannelService:
gateway_url = _resolve_service_url(config, "gateway_url", _CHANNELS_GATEWAY_URL_ENV, DEFAULT_GATEWAY_URL)
default_session = config.pop("session", None)
channel_sessions = {name: channel_config.get("session") for name, channel_config in config.items() if isinstance(channel_config, dict)}
from app.channels.dedupe_store import make_inbound_dedupe_store
self.manager = ChannelManager(
bus=self.bus,
store=self.store,
@ -119,6 +122,7 @@ class ChannelService:
channel_sessions=channel_sessions,
connection_repo=connection_repo,
require_bound_identity=require_bound_identity,
inbound_dedupe_store=make_inbound_dedupe_store(app_config),
get_stream_bridge=get_stream_bridge,
)
self._channels: dict[str, Any] = {} # name -> Channel instance
@ -157,6 +161,7 @@ class ChannelService:
channels_config=channels_config,
connection_repo=_make_connection_repo(connection_config),
require_bound_identity=require_bound_identity,
app_config=app_config,
get_stream_bridge=get_stream_bridge,
)

View File

@ -24,6 +24,7 @@ from deerflow.uploads.manager import is_upload_staging_file, normalize_filename
logger = logging.getLogger(__name__)
TELEGRAM_MAX_MESSAGE_LENGTH = 4096
TELEGRAM_MAX_RICH_MESSAGE_LENGTH = 32768
# Telegram's hosted Bot API documents this as 20 MB (decimal bytes).
TELEGRAM_MAX_INBOUND_FILE_BYTES = 20_000_000
# Keep Telegram cleanup inside the Gateway's five-second shutdown-hook budget.
@ -188,10 +189,23 @@ class TelegramChannel(Channel):
state = self._stream_messages.pop(key, None)
if state is not None:
if self._can_send_rich(msg.text) and await self._edit_rich_message(chat_id, state["message_id"], msg.text):
self._last_bot_message[msg.chat_id] = state["message_id"]
return
await self._finalize_stream_message(chat_id, msg.chat_id, state, msg.text)
return
await self._send_new_message(chat_id, msg.chat_id, msg.text, _max_retries=_max_retries)
if self._can_send_rich(msg.text):
try:
message_id = await self._send_new_rich_message(chat_id, msg.chat_id, msg.text, _max_retries=_max_retries)
except Exception as exc:
logger.warning("[Telegram] Rich Message send failed in chat=%s; falling back to plain text: %s", chat_id, exc)
message_id = None
if message_id is not None:
return
for chunk in self._split_message(msg.text):
await self._send_new_message(chat_id, msg.chat_id, chunk, _max_retries=_max_retries)
async def _send_stream_update(self, chat_id: int, key: str, text: str, reply_to: int | None = None) -> None:
"""Edit the in-flight streamed message with accumulated text.
@ -285,6 +299,51 @@ class TelegramChannel(Channel):
return False
return False
def _can_send_rich(self, text: str) -> bool:
return bool(self.config.get("rich_messages")) and 0 < len(text) <= TELEGRAM_MAX_RICH_MESSAGE_LENGTH
async def _edit_rich_message(self, chat_id: int, message_id: int, text: str) -> bool:
"""Replace a streamed preview with a persistent Telegram Rich Message."""
from telegram.error import BadRequest, EndPointNotFound
bot = self._application.bot
data = {"chat_id": chat_id, "message_id": message_id, "rich_message": {"markdown": text}}
for attempt in range(2):
try:
await bot.do_api_request("editMessageText", api_kwargs=data)
return True
except Exception as exc:
if self._is_retry_after(exc) and attempt == 0:
await asyncio.sleep(self._retry_after_seconds(exc))
continue
if isinstance(exc, (BadRequest, EndPointNotFound)):
logger.warning("[Telegram] Rich Message rejected in chat=%s; falling back to plain text: %s", chat_id, exc)
return False
logger.warning("[Telegram] final rich edit failed in chat=%s: %s", chat_id, exc)
return False
return False
async def _send_new_rich_message(self, chat_id: int, chat_key: str, text: str, *, _max_retries: int) -> int | None:
"""Send raw agent Markdown through Bot API 10.1 Rich Messages."""
from telegram.error import BadRequest, EndPointNotFound
bot = self._application.bot
async def send_message() -> int | None:
try:
result = await bot.do_api_request(
"sendRichMessage",
api_kwargs={"chat_id": chat_id, "rich_message": {"markdown": text}},
)
except (BadRequest, EndPointNotFound) as exc:
logger.warning("[Telegram] Rich Message rejected in chat=%s; falling back to plain text: %s", chat_id, exc)
return None
message_id = int(result["message_id"])
self._last_bot_message[chat_key] = message_id
return message_id
return await self._send_with_retry(send_message, max_retries=_max_retries, log_prefix="[Telegram]")
async def _send_new_message(self, chat_id: int, chat_key: str, text: str, *, _max_retries: int = 3) -> int | None:
"""Send a fresh message with retry/backoff. Returns the sent message_id."""
kwargs: dict[str, Any] = {"chat_id": chat_id, "text": text}

View File

@ -60,6 +60,10 @@ logger = logging.getLogger(__name__)
# firing signals into a worker that is stuck waiting for shutdown cleanup.
_SHUTDOWN_HOOK_TIMEOUT_SECONDS = 5.0
# The retrieval index is derived state, so shutdown only waits briefly for its
# startup rebuild. The canonical memory flush keeps its full configured budget.
_RETRIEVAL_WARM_SHUTDOWN_TIMEOUT_SECONDS = 1.0
async def _ensure_admin_user(app: FastAPI) -> None:
"""Startup hook: handle first boot and migrate orphan threads otherwise.
@ -170,6 +174,18 @@ async def _migrate_orphaned_threads(store, admin_user_id: str) -> int:
return migrated
async def _warm_memory_retrieval(manager) -> None:
"""Rebuild the derived retrieval index without delaying Gateway readiness."""
try:
rebuilt = await asyncio.to_thread(manager.warm_retrieval)
if rebuilt:
logger.info("Memory retrieval index rebuilt successfully")
else:
logger.warning("Memory retrieval index rebuild failed; scoped searches will retry lazily")
except Exception:
logger.warning("Memory retrieval index rebuild skipped", exc_info=True)
@asynccontextmanager
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
"""Application lifespan handler."""
@ -204,6 +220,26 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
except Exception: # observability must never break startup
logger.exception("Monocle tracing setup failed; continuing without it")
# Rebuild the derived memory retrieval index in the background. Scoped
# searches remain correct while this runs because DeerMem lazily rebuilds
# the requested scope when the full warm-up has not completed yet.
retrieval_warm_task: asyncio.Task[None] | None = None
try:
from deerflow.agents.memory import get_memory_manager
if startup_config.memory.enabled:
manager = get_memory_manager()
warm_retrieval = getattr(manager, "warm_retrieval", None)
if callable(warm_retrieval):
retrieval_warm_task = asyncio.create_task(
_warm_memory_retrieval(manager),
name="memory-retrieval-warm-up",
)
else:
logger.info("Memory is disabled; skipping retrieval index rebuild")
except Exception:
logger.warning("Memory retrieval index rebuild skipped", exc_info=True)
# Pre-warm tiktoken encoding cache so the first memory-injection request
# never blocks on the BPE data download (which hits an OpenAI/Azure URL
# that may be unreachable in restricted networks — see issue #3402).
@ -350,9 +386,26 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
#
# K8s caveat: ``shutdown_flush_timeout_seconds`` must fit inside the
# pod's ``terminationGracePeriodSeconds`` (channel stop + browser
# session close + this drain + buffer), set on the gateway Helm
# deployment -- or K8s SIGKILLs the drain mid-flight and the loss this
# is fixing is silently re-introduced.
# session close + the brief retrieval-warm wait + this drain + buffer),
# set on the gateway Helm deployment -- or K8s SIGKILLs the drain
# mid-flight and the loss this is fixing is silently re-introduced.
# The retrieval index is derived from canonical memory files, so its
# wait is independently capped and never consumes the flush budget.
retrieval_warm_finished = True
if retrieval_warm_task is not None and not retrieval_warm_task.done():
try:
await asyncio.wait_for(
asyncio.shield(retrieval_warm_task),
timeout=min(
_RETRIEVAL_WARM_SHUTDOWN_TIMEOUT_SECONDS,
startup_config.memory.shutdown_flush_timeout_seconds,
),
)
except TimeoutError:
retrieval_warm_finished = False
logger.warning("Memory retrieval index rebuild is still running; leaving its connection open during shutdown")
manager = None
try:
app_cfg = get_app_config()
if app_cfg.memory.enabled:
@ -370,6 +423,13 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
)
except Exception:
logger.exception("Failed to flush memory queue on shutdown")
finally:
close = getattr(manager, "close", None)
if callable(close) and retrieval_warm_finished:
try:
await asyncio.to_thread(close)
except Exception:
logger.exception("Failed to close memory backend on shutdown")
logger.info("Shutting down API Gateway")

View File

@ -24,7 +24,7 @@ from app.gateway.auth_disabled import (
get_auth_disabled_user,
is_auth_disabled,
)
from app.gateway.authz import _ALL_PERMISSIONS, AuthContext
from app.gateway.authz import AuthContext, resolve_route_permissions
from app.gateway.internal_auth import INTERNAL_AUTH_HEADER_NAME, get_internal_user, is_valid_internal_auth_token
from deerflow.runtime.user_context import reset_current_user, set_current_user
@ -151,7 +151,11 @@ class AuthMiddleware(BaseHTTPMiddleware):
# JWT-decode + DB-lookup pipeline a second time per request).
request.state.user = user
request.state.auth_source = auth_source
request.state.auth = AuthContext(user=user, permissions=_ALL_PERMISSIONS)
permissions = await resolve_route_permissions(
user,
is_internal=auth_source == AUTH_SOURCE_INTERNAL,
)
request.state.auth = AuthContext(user=user, permissions=permissions)
token = set_current_user(user)
try:
return await call_next(request)

View File

@ -29,19 +29,27 @@ Inspired by LangGraph Auth system: https://github.com/langchain-ai/langgraph/blo
from __future__ import annotations
import asyncio
import functools
import inspect
import logging
from collections.abc import Callable
from types import SimpleNamespace
from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar
from fastapi import HTTPException, Request
from deerflow.authz.principal import build_principal_from_context
from deerflow.authz.provider import AuthorizationProvider, AuthzDecision, AuthzRequest
from deerflow.authz.runtime import resolve_authorization_provider
from deerflow.config.authorization_config import AuthorizationConfig
if TYPE_CHECKING:
from app.gateway.auth.models import User
P = ParamSpec("P")
T = TypeVar("T")
logger = logging.getLogger(__name__)
# Permission constants
@ -128,6 +136,133 @@ def _make_test_request_stub() -> Any:
return SimpleNamespace(state=SimpleNamespace(), cookies={}, _deerflow_test_bypass_auth=True)
def _get_route_authorization_config() -> AuthorizationConfig:
"""Return the hot-reloaded authorization config for this request.
Falls back to a disabled config when AppConfig is not available (e.g. test
environments without a config.yaml), preserving legacy all-permissions behavior.
"""
from deerflow.config.app_config import get_app_config
try:
return get_app_config().authorization
except (FileNotFoundError, RuntimeError):
return AuthorizationConfig()
# --- Provider cache (W1/F1) ---
# Keyed by config object identity (id) so the expensive model_dump() signature
# is only recomputed when get_app_config() returns a new object (hot-reload).
_route_provider_cache: dict[str, AuthorizationProvider] = {}
_route_provider_config_id: int | None = None
_route_provider_config_sig: str | None = None
def _get_cached_route_provider(config: AuthorizationConfig) -> AuthorizationProvider | None:
"""Resolve (or reuse) the authorization provider for route permissions.
The provider is cached per config object identity. When ``get_app_config()``
returns a new object (hot-reload), the signature is recomputed and compared;
only an actual content change triggers re-resolution. This avoids calling
``model_dump()`` on every request the fast path is a single ``id()`` check.
"""
global _route_provider_config_id, _route_provider_config_sig, _route_provider_cache
config_id = id(config)
# Fast path: same config object as last time → return cached provider.
if config_id == _route_provider_config_id and _route_provider_cache:
return _route_provider_cache.get("provider")
# Config object changed (hot-reload): compute signature to check if
# content actually changed or just the wrapper object identity.
sig = repr(sorted(config.model_dump().items()))
if sig == _route_provider_config_sig and _route_provider_cache:
# Same content, different object — update id, reuse provider.
_route_provider_config_id = config_id
return _route_provider_cache.get("provider")
# Content changed (or first call): re-resolve into a local first,
# then publish id + sig + provider together to avoid a race window.
_route_provider_cache.clear()
provider = resolve_authorization_provider(config)
if provider is not None:
_route_provider_cache["provider"] = provider
_route_provider_config_id = config_id
_route_provider_config_sig = sig
return provider
async def resolve_route_permissions(user: User, *, is_internal: bool) -> list[str]:
"""Return the route permissions granted to an authenticated user.
Disabled authorization preserves the legacy all-permissions behavior.
When enabled, every registered ``resource:action`` permission is evaluated
independently so a provider failure affects only the route being checked.
Provider instances are cached per config signature (hot-reload safe).
"""
config = _get_route_authorization_config()
if config.enabled is not True:
return list(_ALL_PERMISSIONS)
try:
provider = _get_cached_route_provider(config)
if provider is None:
raise ValueError("authorization is enabled but provider resolution returned None")
except Exception:
logger.warning("Failed to resolve authorization provider for Gateway routes", exc_info=True)
return [] if config.fail_closed else list(_ALL_PERMISSIONS)
# Align with Phase 1B's tool path: internal callers (IM channel workers,
# scheduler) have system_role="internal", which is not a real RBAC role.
# Omit it so default_role applies, mirroring inject_authenticated_user_context
# which pops user_role for internal callers without a resolved owner.
from app.gateway.internal_auth import INTERNAL_SYSTEM_ROLE
user_role = getattr(user, "system_role", None)
if user_role == INTERNAL_SYSTEM_ROLE:
user_role = None
principal = build_principal_from_context(
{
"user_id": str(user.id),
"user_role": user_role,
"oauth_provider": getattr(user, "oauth_provider", None),
"oauth_id": getattr(user, "oauth_id", None),
"is_internal": is_internal,
},
default_role=config.default_role,
)
# Evaluate all permissions in parallel (W2).
async def _evaluate(permission: str) -> str | None:
_, action = permission.split(":", maxsplit=1)
request = AuthzRequest(
principal=principal,
resource="route",
action=action,
target=permission,
)
try:
decision = await provider.aauthorize(request)
if not isinstance(decision, AuthzDecision):
raise TypeError("AuthorizationProvider.aauthorize must return AuthzDecision")
except asyncio.CancelledError:
raise
except Exception:
logger.warning(
"Authorization provider failed while evaluating route permission %s",
permission,
exc_info=True,
)
return permission if not config.fail_closed else None
return permission if decision.allow else None
results = await asyncio.gather(*[_evaluate(p) for p in _ALL_PERMISSIONS])
return [p for p in results if p is not None]
async def _authenticate(request: Request) -> AuthContext:
"""Authenticate request and return AuthContext.
@ -140,8 +275,32 @@ async def _authenticate(request: Request) -> AuthContext:
if user is None:
return AuthContext(user=None, permissions=[])
# In future, permissions could be stored in user record
return AuthContext(user=user, permissions=_ALL_PERMISSIONS)
is_internal = _is_internal_caller(request, user)
permissions = await resolve_route_permissions(user, is_internal=is_internal)
return AuthContext(user=user, permissions=permissions)
def _is_internal_caller(request: Request, user: Any) -> bool:
"""Determine if the request originates from a trusted internal caller.
Checks three signals (any one suffices):
1. ``request.state.auth_source == AUTH_SOURCE_INTERNAL`` (set by AuthMiddleware).
2. ``user.system_role == INTERNAL_SYSTEM_ROLE`` (synthetic internal user).
3. The request carries a valid internal auth token header (decorator-only path
where AuthMiddleware may not have stamped ``auth_source`` yet).
"""
from app.gateway.auth_disabled import AUTH_SOURCE_INTERNAL
from app.gateway.internal_auth import INTERNAL_AUTH_HEADER_NAME, INTERNAL_SYSTEM_ROLE, is_valid_internal_auth_token
if getattr(getattr(request, "state", None), "auth_source", None) == AUTH_SOURCE_INTERNAL:
return True
if getattr(user, "system_role", None) == INTERNAL_SYSTEM_ROLE:
return True
# Decorator-only path: check the internal token header directly.
internal_token = request.headers.get(INTERNAL_AUTH_HEADER_NAME) if hasattr(request, "headers") else None
if internal_token and is_valid_internal_auth_token(internal_token):
return True
return False
def require_auth[**P, T](func: Callable[P, T]) -> Callable[P, T]:

View File

@ -13,6 +13,7 @@ from __future__ import annotations
import asyncio
import logging
import uuid
from copy import deepcopy
from datetime import UTC, datetime
from typing import Any, Literal
@ -50,6 +51,7 @@ REGENERATE_HISTORY_RAW_SCAN_LIMIT = REGENERATE_HISTORY_SCAN_LIMIT * 2
THREAD_MESSAGE_PAGE_SCAN_BATCH = 201
_MISSING_REGENERATE_BASE_DETAIL = "Could not find an addressable checkpoint before the target user message"
_UNSAFE_REGENERATE_LINEAGE_DETAIL = "Could not safely resolve the checkpoint before the target user message"
THREAD_MESSAGE_LEGACY_SCAN_BATCH = 201
def _is_duration_only_checkpoint(checkpoint_tuple: Any) -> bool:
@ -90,6 +92,16 @@ class RegeneratePrepareResponse(BaseModel):
target_run_id: str
class EditRegeneratePrepareRequest(BaseModel):
human_message_id: str = Field(..., min_length=1, description="Source human message id to edit and rerun")
replacement_text: str = Field(..., min_length=1, description="Replacement user-visible text")
class EditRegeneratePrepareResponse(RegeneratePrepareResponse):
replacement_human_message_id: str
source_message_ids: list[str]
class ThreadMessagesPageResponse(BaseModel):
data: list[dict[str, Any]]
has_more: bool
@ -263,6 +275,15 @@ def _message_additional_kwargs(message: Any) -> dict[str, Any]:
return dict(value or {}) if isinstance(value, dict) else {}
def _message_tool_calls(message: Any) -> list[Any]:
value = getattr(message, "tool_calls", None)
if value is None and isinstance(message, dict):
value = message.get("tool_calls")
if value is None:
value = _message_additional_kwargs(message).get("tool_calls")
return list(value) if isinstance(value, list) else []
def _is_hidden_or_control_message(message: Any) -> bool:
message_type = _message_type(message)
additional_kwargs = _message_additional_kwargs(message)
@ -286,6 +307,11 @@ def _checkpoint_messages(snapshot: Any) -> list[Any]:
return checkpoint_messages(snapshot)
def _checkpoint_values(snapshot: Any) -> dict[str, Any]:
values = getattr(snapshot, "values", None)
return dict(values) if isinstance(values, dict) else {}
def _checkpoint_configurable(checkpoint_tuple: Any) -> dict[str, Any]:
return checkpoint_configurable(checkpoint_tuple)
@ -322,6 +348,54 @@ def _clean_human_message_for_regenerate(message: Any) -> dict[str, Any]:
return clean_message
def _clean_human_message_for_edit(message: Any, *, replacement_id: str, replacement_text: str) -> dict[str, Any]:
source_kwargs = _message_additional_kwargs(message)
additional_kwargs: dict[str, Any] = {}
for key in ("files", "referenced_message_contexts"):
if key in source_kwargs:
additional_kwargs[key] = deepcopy(source_kwargs[key])
clean_message: dict[str, Any] = {
"type": "human",
"id": replacement_id,
"content": [{"type": "text", "text": replacement_text}],
"additional_kwargs": additional_kwargs,
}
name = _message_name(message)
if name:
clean_message["name"] = name
return clean_message
def _is_terminal_assistant_text_message(message: Any) -> bool:
return _is_visible_ai_message(message) and bool(_message_text(message).strip()) and not _message_tool_calls(message)
def _has_active_goal(snapshot: Any) -> bool:
goal = _checkpoint_values(snapshot).get("goal")
return isinstance(goal, dict) and goal.get("status") == "active"
def _latest_editable_turn(messages: list[Any], human_message_id: str) -> tuple[int, Any, int, Any, list[str]]:
latest_human_index = next((index for index in range(len(messages) - 1, -1, -1) if _is_visible_human_message(messages[index])), None)
if latest_human_index is None or _message_id(messages[latest_human_index]) != human_message_id:
raise HTTPException(status_code=409, detail="Only the latest completed user turn can be edited")
source_human = messages[latest_human_index]
last_ai_index: int | None = None
for index, message in enumerate(messages[latest_human_index + 1 :], start=latest_human_index + 1):
if _is_visible_human_message(message):
break
if _is_visible_ai_message(message):
last_ai_index = index
if last_ai_index is None or not _is_terminal_assistant_text_message(messages[last_ai_index]):
raise HTTPException(status_code=409, detail="Only completed assistant text turns can be edited")
source_message_ids = [message_id for message in messages[latest_human_index : last_ai_index + 1] if (message_id := _message_id(message))]
return latest_human_index, source_human, last_ai_index, messages[last_ai_index], source_message_ids
def _event_message_id(row: dict[str, Any]) -> str | None:
content = row.get("content")
if isinstance(content, BaseMessage):
@ -441,6 +515,32 @@ async def _find_base_checkpoint_before_human(
)
def _run_status_value(record: Any) -> str | None:
status = getattr(record, "status", None)
if isinstance(status, RunStatus):
return status.value
return str(status) if status is not None else None
async def _require_successful_source_run(thread_id: str, run_id: str, request: Request) -> RunRecord:
run_mgr = get_run_manager(request)
user_id = await get_current_user(request)
record = await run_mgr.get(run_id, user_id=user_id)
if record is None:
# The run-event journal is the authoritative lookup above. This fallback
# only covers recent in-memory/store hydration gaps for the latest turn.
records = await run_mgr.list_by_thread(thread_id, user_id=user_id, limit=20)
record = next((candidate for candidate in records if getattr(candidate, "run_id", None) == run_id), None)
if record is None:
raise HTTPException(status_code=409, detail="Could not find source run for assistant message")
record_thread_id = getattr(record, "thread_id", None)
if isinstance(record_thread_id, str) and record_thread_id and record_thread_id != thread_id:
raise HTTPException(status_code=409, detail="Could not find source run for assistant message")
if _run_status_value(record) != RunStatus.success.value:
raise HTTPException(status_code=409, detail="Only successful assistant runs can be edited and rerun")
return record
async def _prepare_regenerate_payload(thread_id: str, message_id: str, request: Request) -> RegeneratePrepareResponse:
accessor, latest_config = await build_thread_checkpoint_state_accessor(request, thread_id=thread_id)
try:
@ -490,14 +590,102 @@ async def _prepare_regenerate_payload(thread_id: str, message_id: str, request:
"regenerate_from_run_id": target_run_id,
"regenerate_checkpoint_id": checkpoint["checkpoint_id"],
}
regenerate_input: dict[str, Any] = {"messages": [_clean_human_message_for_regenerate(previous_human)]}
latest_values = latest_checkpoint.values if isinstance(latest_checkpoint.values, dict) else {}
latest_title = latest_values.get("title")
if isinstance(latest_title, str) and latest_title:
# Regenerate resumes from the checkpoint before the target human turn.
# That checkpoint can predate a manual rename, so replay the current
# title as graph input instead of letting checkpoint rollback restore
# the older automatically generated title (#4457).
regenerate_input["title"] = latest_title
return RegeneratePrepareResponse(
input={"messages": [_clean_human_message_for_regenerate(previous_human)]},
input=regenerate_input,
checkpoint=checkpoint,
metadata=metadata,
target_run_id=target_run_id,
)
async def _prepare_edit_regenerate_payload(
thread_id: str,
human_message_id: str,
replacement_text: str,
request: Request,
) -> EditRegeneratePrepareResponse:
normalized_text = replacement_text.strip()
if not normalized_text:
raise HTTPException(status_code=409, detail="Edited message cannot be empty")
accessor, latest_config = await build_thread_checkpoint_state_accessor(request, thread_id=thread_id)
try:
latest_checkpoint = await accessor.aget(latest_config)
except Exception as exc:
logger.exception("Failed to read latest checkpoint for edit replay thread %s", thread_id)
raise HTTPException(status_code=500, detail="Failed to read latest checkpoint") from exc
latest_checkpoint_id = _checkpoint_configurable(latest_checkpoint).get("checkpoint_id")
if not latest_checkpoint_id:
raise HTTPException(status_code=404, detail=f"Thread {thread_id} has no checkpoint")
messages = _checkpoint_messages(latest_checkpoint)
if _has_active_goal(latest_checkpoint):
raise HTTPException(status_code=409, detail="Cannot edit while a goal is active")
_, source_human, _, source_ai, source_message_ids = _latest_editable_turn(messages, human_message_id)
source_text = get_original_user_content_text(_message_content(source_human), _message_additional_kwargs(source_human)).strip()
if normalized_text == source_text:
raise HTTPException(status_code=409, detail="Edited message is unchanged")
source_human_id = _message_id(source_human)
source_ai_id = _message_id(source_ai)
if not source_human_id:
raise HTTPException(status_code=409, detail="The source user message is missing an id")
if not source_ai_id:
raise HTTPException(status_code=409, detail="The source assistant message is missing an id")
base_checkpoint_tuple = await _find_base_checkpoint_before_human(thread_id, source_human_id, request)
target_run_id = await _find_target_run_id(thread_id, source_ai_id, source_ai, source_human, request)
source_record = await _require_successful_source_run(thread_id, target_run_id, request)
checkpoint = _checkpoint_response(base_checkpoint_tuple)
replacement_human_message_id = str(uuid.uuid4())
source_metadata = getattr(source_record, "metadata", None) or {}
existing_group_id = source_metadata.get("edit_version_group_id") if isinstance(source_metadata, dict) else None
# Reserved for future edit-chain grouping across repeated edits of the same
# original prompt; current visibility still keys off regenerate_from_run_id.
edit_version_group_id = existing_group_id if isinstance(existing_group_id, str) and existing_group_id else source_human_id
metadata = {
"replay_kind": "edit",
"regenerate_from_message_id": source_ai_id,
"regenerate_from_run_id": target_run_id,
"regenerate_checkpoint_id": checkpoint["checkpoint_id"],
"edit_from_message_id": source_human_id,
"edit_message_id": replacement_human_message_id,
"edit_version_group_id": edit_version_group_id,
}
return EditRegeneratePrepareResponse(
input={
"messages": [
_clean_human_message_for_edit(
source_human,
replacement_id=replacement_human_message_id,
replacement_text=normalized_text,
)
]
},
checkpoint=checkpoint,
metadata=metadata,
target_run_id=target_run_id,
replacement_human_message_id=replacement_human_message_id,
source_message_ids=source_message_ids,
)
async def _default_history_hidden_run_ids(run_mgr: Any, thread_id: str, *, user_id: str | None) -> set[str]:
superseded_run_ids = await run_mgr.list_successful_regenerate_sources(thread_id, user_id=user_id)
edit_visibility = await run_mgr.list_edit_replay_visibility(thread_id, user_id=user_id)
return set(superseded_run_ids) | set(edit_visibility.hidden_source_run_ids) | set(edit_visibility.hidden_attempt_run_ids)
# ---------------------------------------------------------------------------
# Endpoints
# ---------------------------------------------------------------------------
@ -514,6 +702,17 @@ async def prepare_regenerate_run(
return await _prepare_regenerate_payload(thread_id, body.message_id, request)
@router.post("/{thread_id}/runs/edit-regenerate/prepare", response_model=EditRegeneratePrepareResponse)
@require_permission("runs", "create", owner_check=True, require_existing=True)
async def prepare_edit_regenerate_run(
thread_id: str,
body: EditRegeneratePrepareRequest,
request: Request,
) -> EditRegeneratePrepareResponse:
"""Prepare input and checkpoint for editing then rerunning the latest user turn."""
return await _prepare_edit_regenerate_payload(thread_id, body.human_message_id, body.replacement_text, request)
@router.post("/{thread_id}/runs", response_model=RunResponse)
@require_permission("runs", "create", owner_check=True, require_existing=True)
async def create_run(thread_id: str, body: RunCreateRequest, request: Request) -> RunResponse:
@ -744,12 +943,23 @@ async def list_thread_messages(
after_seq: int | None = Query(default=None, ge=1),
) -> list[dict]:
"""Return displayable messages for a thread (across all runs), with feedback attached."""
event_store = get_run_event_store(request)
messages = await event_store.list_messages(thread_id, limit=limit, before_seq=before_seq, after_seq=after_seq)
# Resolve the caller once; it is needed both to scope the feedback query
# below and to list the thread's runs for turn-duration injection.
user_id = await get_current_user(request)
run_mgr = get_run_manager(request)
hidden_run_ids = await _default_history_hidden_run_ids(run_mgr, thread_id, user_id=user_id)
messages, _ = await _scan_visible_thread_messages(
thread_id,
limit=limit,
before_seq=before_seq,
after_seq=after_seq,
request=request,
user_id=user_id,
hidden_run_ids=hidden_run_ids,
include_middleware=True,
include_extra=False,
batch_size=THREAD_MESSAGE_LEGACY_SCAN_BATCH,
)
# Find the last AI message per run_id. AI messages are persisted by
# RunJournal with event_type "llm.ai.response" (see runtime/journal.py);
@ -785,7 +995,6 @@ async def list_thread_messages(
else:
msg["feedback"] = None
run_mgr = get_run_manager(request)
runs = await run_mgr.list_by_thread(thread_id, user_id=user_id)
run_durations = compute_run_durations(runs)
@ -802,6 +1011,122 @@ async def list_thread_messages(
return messages
async def _scan_visible_thread_messages(
thread_id: str,
*,
limit: int,
before_seq: int | None,
after_seq: int | None,
request: Request,
user_id: str | None,
hidden_run_ids: set[str],
include_middleware: bool,
include_extra: bool,
batch_size: int,
) -> tuple[list[dict[str, Any]], bool]:
"""Scan raw message rows until ``limit`` visible rows survive filtering."""
event_store = get_run_event_store(request)
needed = limit + 1 if include_extra else limit
if after_seq is not None:
visible: list[dict[str, Any]] = []
scan_after = after_seq
while len(visible) < needed:
raw = await event_store.list_messages(
thread_id,
limit=batch_size,
after_seq=scan_after,
user_id=user_id,
)
if not raw:
break
_validate_message_scan_rows(raw, thread_id=thread_id, scan_before=None, scan_after=scan_after)
reached_before_bound = False
for row in raw:
if before_seq is not None and row["seq"] >= before_seq:
reached_before_bound = True
break
if (not include_middleware and _is_thread_history_hidden_message_row(row)) or row.get("run_id") in hidden_run_ids:
continue
visible.append(row)
if len(visible) == needed:
break
next_scan_after = max(row["seq"] for row in raw)
if next_scan_after <= scan_after:
_raise_non_advancing_message_scan(thread_id=thread_id, scan_before=None, scan_after=scan_after, next_cursor=next_scan_after, row_count=len(raw))
scan_after = next_scan_after
if reached_before_bound or len(raw) < batch_size:
break
has_more = len(visible) > limit
return visible[:limit], has_more
visible_desc: list[dict[str, Any]] = []
scan_before = before_seq
while len(visible_desc) < needed:
raw = await event_store.list_messages(
thread_id,
limit=batch_size,
before_seq=scan_before,
user_id=user_id,
)
if not raw:
break
_validate_message_scan_rows(raw, thread_id=thread_id, scan_before=scan_before, scan_after=None)
for row in reversed(raw):
if (not include_middleware and _is_thread_history_hidden_message_row(row)) or row.get("run_id") in hidden_run_ids:
continue
visible_desc.append(row)
if len(visible_desc) == needed:
break
next_scan_before = min(row["seq"] for row in raw)
if scan_before is not None and next_scan_before >= scan_before:
_raise_non_advancing_message_scan(thread_id=thread_id, scan_before=scan_before, scan_after=None, next_cursor=next_scan_before, row_count=len(raw))
scan_before = next_scan_before
if len(raw) < batch_size:
break
has_more = len(visible_desc) > limit
return list(reversed(visible_desc[:limit])), has_more
def _validate_message_scan_rows(
rows: list[dict[str, Any]],
*,
thread_id: str,
scan_before: int | None,
scan_after: int | None,
) -> None:
invalid_seq_rows = [row for row in rows if not isinstance(row.get("seq"), int)]
if invalid_seq_rows:
logger.error(
"Thread message scan found rows without sequence values: thread_id=%s scan_before=%s scan_after=%s row_count=%d invalid_count=%d",
thread_id,
scan_before,
scan_after,
len(rows),
len(invalid_seq_rows),
)
raise RuntimeError("Run event message rows are missing sequence values")
def _raise_non_advancing_message_scan(
*,
thread_id: str,
scan_before: int | None,
scan_after: int | None,
next_cursor: int,
row_count: int,
) -> None:
logger.error(
"Thread message scan cursor did not advance: thread_id=%s scan_before=%s scan_after=%s next_cursor=%s row_count=%d",
thread_id,
scan_before,
scan_after,
next_cursor,
row_count,
)
raise RuntimeError("Run event message scan did not advance its cursor")
async def _scan_thread_message_page(
thread_id: str,
*,
@ -811,57 +1136,20 @@ async def _scan_thread_message_page(
user_id: str | None,
) -> tuple[list[dict[str, Any]], bool]:
"""Select the newest ``limit + 1`` page-eligible rows before a cursor."""
event_store = get_run_event_store(request)
run_mgr = get_run_manager(request)
superseded_run_ids = await run_mgr.list_successful_regenerate_sources(thread_id, user_id=user_id)
visible_desc: list[dict[str, Any]] = []
scan_before = before_seq
while len(visible_desc) < limit + 1:
raw = await event_store.list_messages(
thread_id,
limit=THREAD_MESSAGE_PAGE_SCAN_BATCH,
before_seq=scan_before,
user_id=user_id,
)
if not raw:
break
invalid_seq_rows = [row for row in raw if not isinstance(row.get("seq"), int)]
if invalid_seq_rows:
logger.error(
"Thread message scan found rows without sequence values: thread_id=%s scan_before=%s row_count=%d invalid_count=%d",
thread_id,
scan_before,
len(raw),
len(invalid_seq_rows),
)
raise RuntimeError("Run event message rows are missing sequence values")
for row in reversed(raw):
if _is_thread_history_hidden_message_row(row) or row.get("run_id") in superseded_run_ids:
continue
visible_desc.append(row)
if len(visible_desc) == limit + 1:
break
raw_seqs = [row["seq"] for row in raw]
next_scan_before = min(raw_seqs)
if scan_before is not None and next_scan_before >= scan_before:
logger.error(
"Thread message scan cursor did not advance: thread_id=%s scan_before=%s next_scan_before=%s row_count=%d",
thread_id,
scan_before,
next_scan_before,
len(raw),
)
raise RuntimeError("Run event message scan did not advance its cursor")
scan_before = next_scan_before
if len(raw) < THREAD_MESSAGE_PAGE_SCAN_BATCH:
break
has_more = len(visible_desc) > limit
return list(reversed(visible_desc[:limit])), has_more
hidden_run_ids = await _default_history_hidden_run_ids(run_mgr, thread_id, user_id=user_id)
return await _scan_visible_thread_messages(
thread_id,
limit=limit,
before_seq=before_seq,
after_seq=None,
request=request,
user_id=user_id,
hidden_run_ids=hidden_run_ids,
include_middleware=False,
include_extra=True,
batch_size=THREAD_MESSAGE_PAGE_SCAN_BATCH,
)
async def _enrich_thread_message_page(

View File

@ -987,7 +987,6 @@ async def start_run(
body_context = getattr(body, "context", None) or {}
model_name = body_context.get("model_name")
# Coerce non-string model_name values to str before truncation.
if model_name is not None and not isinstance(model_name, str):
model_name = str(model_name)

View File

@ -109,10 +109,17 @@ Content-Type: application/json
**Run Option Compatibility:**
- Supported concurrency strategies: `reject`, `rollback`, and `interrupt`
- Compatibility default: `if_not_exists="create"`; this matches DeerFlow's current behavior
- Artifact delivery is enforced automatically when a run creates or modifies regular files under `/mnt/user-data/outputs`. `present_files` must present at least one path produced by the current run (or a directory containing it), and the terminal receipt must be persisted; presenting only an unrelated file does not satisfy delivery. Runs without changed outputs retain ordinary conversational behavior. `artifact_delivery` is not a client-settable run option.
- Unsupported options return `422`: `webhook`, `stream_resumable=true`, `after_seconds`, `feedback_keys`, any non-null `on_completion` value (including the SDK values `"complete"` and `"continue"`), `if_not_exists="reject"`, and `multitask_strategy="enqueue"`
- `stream_resumable=false` is accepted: it is the LangGraph SDK's default and requests the non-resumable stream DeerFlow already serves
- Undeclared SDK options, including `checkpoint_during` and `durability`, also return `422` instead of being silently discarded
When outputs changed during the run, `run.delivery` events retain the Slice 1
facts (`presented`, `paths`, and `by_tool`) and add `produced_paths`,
`presented_paths`, `matched_paths`, plus an explicit verdict: `verification`,
`stage` (`presented`, `mismatched`, or `not_started`), and `satisfied`. Receipts
for runs without changed outputs keep their existing shape.
**Recursion Limit:**
`config.recursion_limit` caps the number of graph steps LangGraph will execute

View File

@ -426,6 +426,15 @@ sandbox:
home_dir: /home/user # /mnt/user-data is remapped under this directory
idle_timeout: 600 # forwarded to e2b's server-side set_timeout()
replicas: 3 # max concurrent sandboxes per gateway process
ownership: # use Redis when more than one gateway shares E2B
type: redis
redis_url: $REDIS_URL
reconciliation_interval_seconds: 60
reconciliation_grace_seconds: 120
reconciliation_orphan_ttl_seconds: 3600
reconciliation_max_pages: 10
reconciliation_max_items: 200
reconciliation_max_seconds: 15
mounts: # one-shot upload of host files at sandbox start
- host_path: /path/on/host
container_path: /home/user/shared
@ -440,10 +449,19 @@ provider in `config.yaml`.
Notes specific to `E2BSandboxProvider`:
- Each DeerFlow thread is bound to its e2b sandbox via metadata
(`deer_flow_user`, `deer_flow_thread`), so the same thread reuses the same
sandbox across gateway restarts and across processes — no cross-process
file lock is needed because the e2b control plane is the source of truth.
- Each DeerFlow thread is bound to its E2B sandbox via metadata
(`deer_flow_user`, `deer_flow_thread`). Startup and periodic reconciliation
probe every bounded candidate, adopt one healthy canonical sandbox, and reap
duplicates after a grace period. Provider-tagged entries without a complete
user/thread identity are reaped only after the orphan TTL.
- Ownership leases prevent one gateway from adopting or destroying a sandbox
another live gateway is responsible for. The default in-memory store is safe
only for one gateway process. Multi-worker/load-balanced deployments must use
`sandbox.ownership.type: redis`; an existing Redis stream bridge configuration
is inferred automatically.
- Reconciliation is bounded by page, item, and wall-clock limits. Its summary log
exposes discovered, adopted, duplicate, deferred, killed, dead, and budget-exhausted
counts for operational monitoring.
- Idle expiry is enforced server-side by e2b's `set_timeout()`. The provider
refreshes the timeout on every release so warm sandboxes stay alive long
enough for the next acquire.

View File

@ -46,7 +46,7 @@ from deerflow.agents.middlewares.todo_middleware import TodoMiddleware
from deerflow.agents.middlewares.token_usage_middleware import TokenUsageMiddleware
from deerflow.agents.middlewares.tool_error_handling_middleware import build_lead_runtime_middlewares
from deerflow.agents.middlewares.view_image_middleware import ViewImageMiddleware
from deerflow.agents.thread_state import get_thread_state_schema, normalize_middleware_state_schemas
from deerflow.agents.thread_state import freeze_delta_snapshot_frequency, get_thread_state_schema, normalize_middleware_state_schemas
from deerflow.authz.tool_filter import apply_tool_authorization
from deerflow.config.agents_config import load_agent_config, validate_agent_name
from deerflow.config.app_config import AppConfig, get_app_config
@ -518,6 +518,7 @@ def make_lead_agent(config: RunnableConfig):
runtime_app_config.database.checkpoint_channel_mode,
)
mode = freeze_checkpoint_channel_mode(requested_mode)
freeze_delta_snapshot_frequency(runtime_app_config.database.checkpoint_delta_snapshot_frequency)
inject_checkpoint_mode(config, mode)
return _make_lead_agent(config, app_config=runtime_app_config)

View File

@ -24,6 +24,7 @@ from __future__ import annotations
import copy
import logging
import threading
from typing import Any, ClassVar, Literal
from pydantic import PrivateAttr
@ -132,6 +133,11 @@ class DeerMem(MemoryManager):
# mirroring pre-abstraction `model_name: null`. Standalone (no factory) -> None.
self._llm = self._config.host_llm if self._config.host_llm is not None else build_llm(self._config.model)
self._updater = MemoryUpdater(self._config, self._storage, self._llm, prompts_dir=self._config.prompts_dir, callbacks=self.callbacks)
# Retrieval is derived data. The first search for a scope lazily
# rebuilds it; Gateway warm-up performs the full rebuild off-loop.
self._retrieval_lock = threading.RLock()
self._retrieval_warmed_scopes: set[tuple[str | None, str | None]] = set()
self._retrieval_fully_warmed = False
# Validate the *global* explicit prompt templates at construction so a
# misconfigured prompts_dir surfaces at startup rather than as a silent
# dropped update. Per-agent overrides ({prompts_dir}/{agent}/*.yaml)
@ -319,40 +325,98 @@ class DeerMem(MemoryManager):
agent_name: str | None = None,
category: str | None = None,
) -> list[dict[str, Any]]:
"""Case-insensitive substring search over stored facts.
"""Search through the configured retrieval adapter.
Stand-in for the planned BM25+vector+MMR retrieval
(``core/retrieval.py``): returns facts whose ``content`` contains the
query, ranked by confidence desc, capped at ``top_k``. ``category``
filters BEFORE the ``top_k`` slice so a category-scoped search is not
starved by higher-confidence facts in other categories. Sufficient for
the tool-driven memory mode; upgrade to semantic retrieval later
without changing call sites.
Retrieval errors never make canonical memory unavailable: the existing
case-insensitive substring path remains the last-resort fallback.
"""
if not query or not query.strip() or top_k <= 0:
return []
query_lower = query.strip().lower()
search_facts = getattr(self._storage, "search_facts", None)
resolved_agent_name = _resolve_agent_name(agent_name)
scopes = [{"userId": user_id, "agentName": resolved_agent_name}]
indexed = (
search_facts(
query,
scopes=scopes,
top_k=top_k,
mode="hybrid",
filters={"category": category} if category else None,
indexed = self._fts5_search(query, top_k=top_k, user_id=user_id, agent_name=resolved_agent_name, category=category)
if indexed:
return indexed
return self._substring_search(query, top_k=top_k, user_id=user_id, agent_name=resolved_agent_name, category=category)
def _fts5_search(
self,
query: str,
*,
top_k: int,
user_id: str | None,
agent_name: str | None,
category: str | None,
) -> list[dict[str, Any]]:
"""Return adapter results in the public fact shape (compatibility helper)."""
agent_name = _resolve_agent_name(agent_name)
search_facts = getattr(self._storage, "search_facts", None)
scopes = [{"userId": user_id, "agentName": agent_name}]
try:
self._ensure_retrieval_scopes(scopes)
indexed = (
search_facts(
query,
scopes=scopes,
top_k=top_k,
mode="hybrid",
filters={"category": category} if category else None,
)
if callable(search_facts)
else []
)
if callable(search_facts)
else []
)
except Exception:
logger.exception("Memory retrieval adapter failed; using substring fallback")
indexed = []
if indexed:
return [_compat_document({"facts": [result.get("fact", result)]})["facts"][0] for result in indexed]
memory_data = _call_backend(lambda: self._updater.get_memory_data(agent_name=resolved_agent_name, user_id=user_id))
return []
def _substring_search(
self,
query: str,
*,
top_k: int,
user_id: str | None,
agent_name: str | None,
category: str | None,
) -> list[dict[str, Any]]:
query_lower = query.strip().lower()
memory_data = _call_backend(lambda: self._updater.get_memory_data(agent_name=agent_name, user_id=user_id))
matched = [fact for fact in memory_data.get("facts", []) if isinstance(fact.get("content"), str) and query_lower in fact["content"].lower() and (category is None or fact.get("category") == category)]
matched.sort(key=_coerce_source_confidence, reverse=True)
return _compat_document({"facts": matched[:top_k]})["facts"]
def _ensure_retrieval_scopes(self, scopes: list[dict[str, str | None]]) -> None:
"""Lazily rebuild every requested scope when warm-up was skipped."""
if not hasattr(self, "_retrieval_lock"):
self._retrieval_lock = threading.RLock()
if not hasattr(self, "_retrieval_warmed_scopes"):
self._retrieval_warmed_scopes = set()
if not hasattr(self, "_retrieval_fully_warmed"):
self._retrieval_fully_warmed = False
rebuild = getattr(self._storage, "rebuild_index", None)
if not callable(rebuild):
return
with self._retrieval_lock:
if self._retrieval_fully_warmed:
return
status = getattr(self._storage, "retrieval_status", lambda: {"configured": True})()
if not status.get("configured", True):
self._retrieval_warmed_scopes.update((scope.get("userId"), scope.get("agentName")) for scope in scopes)
return
for scope in scopes:
key = (scope.get("userId"), scope.get("agentName"))
if key in self._retrieval_warmed_scopes:
continue
try:
result = rebuild([scope])
except Exception:
logger.exception("Failed to lazily rebuild memory retrieval index for scope %r", key)
continue
if result.get("supported") and not result.get("fatal"):
self._retrieval_warmed_scopes.add(key)
# ── Manage ───────────────────────────────────────────────────────────
def get_memory(
self,
@ -408,9 +472,13 @@ class DeerMem(MemoryManager):
"""
return self._queue.flush_sync(timeout)
# ── Tier 3 hooks (override the base defaults; warm/reload/fact CRUD) ─
def close(self) -> None:
"""Close derived retrieval resources after pending updates drain."""
self._storage.close()
# ── Tier 3 hooks (override the base defaults; warm/reload/fact CRUD) ──
def warm(self) -> bool:
"""Pre-warm DeerMem-specific resources (the tiktoken encoding cache).
"""Pre-warm DeerMem's token-counting resources.
Overrides the base tier-3 hook (default None = nothing to warm). The
Gateway lifespan calls ``manager.warm()`` directly off the event loop;
@ -424,6 +492,29 @@ class DeerMem(MemoryManager):
return True
return warm_tiktoken_cache()
def warm_retrieval(self) -> bool:
"""Rebuild the complete derived retrieval index before serving traffic."""
rebuild = getattr(self._storage, "rebuild_index", None)
if not callable(rebuild):
return True
try:
result = rebuild()
index_ok = not bool(result.get("fatal"))
failed = int(result.get("failed") or 0)
if failed and index_ok:
logger.warning(
"Memory retrieval index rebuilt with %d fact(s) skipped",
failed,
)
if index_ok:
with self._retrieval_lock:
self._retrieval_fully_warmed = True
self._retrieval_warmed_scopes.clear()
return index_ok
except Exception:
logger.exception("Failed to rebuild memory retrieval index during warm-up")
return False
def reload_memory(
self,
*,

View File

@ -69,8 +69,8 @@ class DeerMemConfig(BaseModel):
description="Maximum wait for the per-scope cross-process advisory file lock.",
)
retrieval_adapter: str = Field(
default="",
description="Optional dotted retrieval-adapter factory. It receives DeerMemConfig and must implement RetrievalPort.",
default="fts5",
description="Retrieval adapter factory: 'fts5' (default), an empty string to disable, or a dotted factory receiving DeerMemConfig and implementing RetrievalPort.",
)
# ── Queue ────────────────────────────────────────────────────────────
debounce_seconds: int = Field(

View File

@ -0,0 +1,702 @@
"""FTS5-based retrieval engine for DeerMem.
Provides BM25 full-text search over stored facts with:
- jieba Chinese tokenization (optional, falls back to whitespace)
- FTS5 MATCH syntax support (AND/OR/NOT/phrase/prefix) with fallback
- Time-decay + confidence-weighted ranking
- Category filtering
- Scope (user_id) isolation
``FTS5Retrieval`` is the low-level SQLite engine. The storage integration is
owned by ``FTS5RetrievalAdapter``, which implements ``storage.RetrievalPort``
without importing storage and creating a circular dependency.
"""
from __future__ import annotations
import json
import logging
import math
import re
import sqlite3
import threading
from datetime import UTC, datetime
from pathlib import Path
from typing import Any
logger = logging.getLogger(__name__)
# ── Scoring weights ──────────────────────────────────────────────────
#
# SQLite FTS5 ``bm25(memory_fts)`` (no positional params → defaults K1=1.2,
# B=0.75) returns a negative value whose magnitude scales with document
# relevance. Critically the function takes positional params in order
# ``(table, k1, b, *column_weights)``: the original code passed
# ``bm25(..., 0.0, 0.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0)``, which set **k1 = 0**
# and silently zeroed out the entire BM25 score (disabled tf saturation),
# collapsing ranking to ``confidence * _CONFIDENCE_WEIGHT``. Use the
# no-arg form so SQLite defaults apply and BM25 is actually scored.
_CONFIDENCE_WEIGHT = 0.2
# ── jieba (optional) ──────────────────────────────────────────────────
try:
import jieba
_jieba_available = True
except ImportError:
_jieba_available = False
def _tokenize(text: str) -> list[str]:
"""Tokenize text: jieba for Chinese, whitespace split for English."""
if not text or not text.strip():
return []
if _jieba_available:
return [t for t in jieba.cut(text) if t.strip()]
return [t for t in text.split() if t.strip()]
# ── FTS5 query preprocessing ──────────────────────────────────────────
_FTS5_ADVANCED_RE = re.compile(
r"(\bAND\b|\bOR\b|\bNOT\b|\bNEAR\b"
r'|"\w.*?"' # phrase "..."
r"|\w+\*" # prefix prefix*
r"|\(.*?\))" # group (...)
)
def _is_advanced_query(query: str) -> bool:
"""Detect whether the query uses FTS5 advanced syntax."""
return bool(_FTS5_ADVANCED_RE.search(query))
def _build_fallback_query(query: str) -> str:
"""Convert natural-language query to FTS5 OR query (fallback strategy)."""
tokens = [token for token in _tokenize(query) if any(char.isalnum() for char in token)]
if not tokens:
return ""
# Quote each token so punctuation in natural-language input cannot become
# an FTS5 operator or syntax error. Double quotes inside a token are the
# FTS5 escape sequence for a literal quote.
return " OR ".join(f'"{token.replace(chr(34), chr(34) * 2)}"' for token in tokens)
# ── Core retrieval engine ─────────────────────────────────────────────
class FTS5Retrieval:
"""SQLite FTS5-based retrieval engine.
Query strategy:
1. Advanced FTS5 syntax -> pass through to MATCH
2. Natural language -> jieba tokenize + OR join
3. Syntax error -> fall back to tokenized OR query
4. Still fails -> return empty
Ranking:
BM25 score × time_decay + confidence × 0.2
"""
def __init__(self, db_path: str | Path = ":memory:"):
self._db_path = str(db_path)
# Gateway runs tool calls via asyncio.to_thread / ThreadPoolExecutor;
# SQLite connections are not safe to share across threads even when the
# top-level instance is single-process. We pick two defensive layers:
# 1. ``check_same_thread=False`` so a connection created in thread A
# is accessible from thread B (libsqlite itself is reentrant under
# a serialised wrapper, see #4208 hot-path discussion).
# 2. ``_lock`` guards all mutating sqlite calls so concurrent callers
# serialise through here (prevents interleaved writes / FTS5 index
# reorderings). Callers from outside instance methods MUST enter
# the lock via the public API and not bypass into ``self._conn``.
self._conn = sqlite3.connect(self._db_path, check_same_thread=False)
self._lock = threading.RLock()
try:
self._conn.execute("PRAGMA journal_mode=WAL")
self._conn.execute("PRAGMA busy_timeout=30000")
self._init_schema()
except Exception:
self._conn.close()
raise
def _init_schema(self) -> None:
conn = self._conn
conn.execute(
"""
CREATE VIRTUAL TABLE IF NOT EXISTS memory_fts USING fts5(
doc_id UNINDEXED,
content,
raw_content UNINDEXED,
category UNINDEXED,
scope_user UNINDEXED,
scope_agent UNINDEXED,
created_at UNINDEXED,
confidence UNINDEXED,
source UNINDEXED,
fact_json UNINDEXED,
tokenize='unicode61'
)
"""
)
conn.commit()
# ── Index operations ───────────────────────────────────────────────
def _preprocess_content(self, content: str) -> str:
"""Preprocess content for indexing: jieba tokenize for Chinese."""
if not content:
return ""
if _jieba_available:
tokens = _tokenize(content)
return " ".join(tokens)
return content
def _row_from_document(self, document: dict[str, Any]) -> tuple[Any, ...]:
now = datetime.now(UTC).isoformat().replace("+00:00", "Z")
return (
document["fact_id"],
self._preprocess_content(document["content"]),
document["content"],
document["category"],
document["scope_user"],
document["scope_agent"],
document.get("created_at") or now,
document.get("confidence", 0.5),
document.get("source"),
json.dumps(document.get("fact_data"), ensure_ascii=False, default=str),
)
def index_fact(
self,
fact_id: str,
content: str,
category: str = "context",
confidence: float = 0.5,
created_at: str | None = None,
scope_user: str | None = None,
scope_agent: str | None = None,
source: str | None = None,
fact_data: dict[str, Any] | None = None,
) -> None:
"""Insert or update a fact in the FTS5 index."""
with self._lock:
conn = self._conn
# Delete existing entry with same doc_id (INSERT OR REPLACE for FTS5)
conn.execute("DELETE FROM memory_fts WHERE doc_id = ?", (fact_id,))
conn.execute(
"""
INSERT INTO memory_fts(
doc_id, content, raw_content, category, scope_user, scope_agent,
created_at, confidence, source, fact_json
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
self._row_from_document(
{
"fact_id": fact_id,
"content": content,
"category": category,
"scope_user": scope_user or "",
"scope_agent": scope_agent or "",
"created_at": created_at,
"confidence": confidence,
"source": source,
"fact_data": fact_data,
}
),
)
conn.commit()
def replace_documents(self, documents: list[dict[str, Any]], *, scopes: list[tuple[str, str]] | None = None) -> None:
"""Atomically replace all or selected scope rows in one transaction."""
with self._lock:
conn = self._conn
try:
conn.execute("BEGIN IMMEDIATE")
if scopes is None:
conn.execute("DELETE FROM memory_fts")
else:
conn.executemany(
"DELETE FROM memory_fts WHERE scope_user = ? AND scope_agent = ?",
scopes,
)
conn.executemany(
"""
INSERT INTO memory_fts(
doc_id, content, raw_content, category, scope_user, scope_agent,
created_at, confidence, source, fact_json
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
[self._row_from_document(document) for document in documents],
)
conn.commit()
except Exception:
conn.rollback()
raise
def remove_fact(self, fact_id: str) -> None:
"""Remove a fact from the FTS5 index."""
with self._lock:
conn = self._conn
conn.execute("DELETE FROM memory_fts WHERE doc_id = ?", (fact_id,))
conn.commit()
def clear_index(self) -> None:
"""Clear the entire FTS5 index."""
with self._lock:
conn = self._conn
conn.execute("DELETE FROM memory_fts")
conn.commit()
def clear_scope(self, *, scope_user: str, scope_agent: str) -> None:
"""Clear one exact adapter scope without affecting other users."""
with self._lock:
self._conn.execute(
"DELETE FROM memory_fts WHERE scope_user = ? AND scope_agent = ?",
(scope_user, scope_agent),
)
self._conn.commit()
def rebuild_from_facts(
self,
facts: list[dict[str, Any]],
*,
scope_user: str | None = None,
scope_agent: str | None = None,
) -> None:
"""Rebuild the entire index from a list of fact dicts."""
with self._lock:
conn = self._conn
try:
conn.execute("BEGIN")
conn.execute("DELETE FROM memory_fts")
for fact in facts:
fact_id = fact.get("id", "")
content = fact.get("content", "")
if not fact_id or not isinstance(content, str) or not content:
continue
now = datetime.now(UTC).isoformat().replace("+00:00", "Z")
conn.execute(
"""
INSERT INTO memory_fts(
doc_id, content, raw_content, category, scope_user, scope_agent,
created_at, confidence, source, fact_json
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
fact_id,
self._preprocess_content(content),
content,
fact.get("category", "context"),
scope_user or "",
scope_agent or "",
fact.get("createdAt") or now,
fact.get("confidence", 0.5),
fact.get("source"),
json.dumps(fact, ensure_ascii=False, default=str),
),
)
conn.commit()
except Exception:
conn.rollback()
raise
# ── Search ─────────────────────────────────────────────────────────
def search(
self,
query: str,
*,
scope_user: str | None = None,
scope_agent: str | None = None,
category: str | None = None,
top_k: int = 5,
) -> list[dict[str, Any]]:
"""FTS5 BM25 search with category and scope filtering.
Returns list of fact dicts (id, content, category, confidence,
createdAt, source, score, bm25_score) sorted by relevance.
"""
import time
_t0 = time.perf_counter()
if not query or not query.strip() or top_k <= 0:
logger.debug("FTS5Retrieval.search: skipped (empty/invalid) query=%r top_k=%d", query, top_k)
return []
query = query.strip()
# Determine query strategy
if _is_advanced_query(query):
fts5_query = query
strategy = "advanced"
else:
fts5_query = _build_fallback_query(query)
strategy = "tokenized"
logger.debug(
"FTS5Retrieval.search: query=%r strategy=%r fts5_query=%r scope_user=%r scope_agent=%r category=%r top_k=%d",
query,
strategy,
fts5_query,
scope_user,
scope_agent,
category,
top_k,
)
# Try search
results = self._execute_search(fts5_query, scope_user, scope_agent, category, top_k)
if results is not None:
logger.debug(
"FTS5Retrieval.search: strategy=%r returned %d results in %.1fms",
strategy,
len(results),
(time.perf_counter() - _t0) * 1000,
)
return results
# Fallback: advanced syntax error -> tokenized OR
if strategy == "advanced":
fallback = _build_fallback_query(query)
if fallback and fallback != fts5_query:
logger.debug("FTS5Retrieval.search: advanced syntax failed, retrying with tokenized fts5_query=%r", fallback)
results = self._execute_search(fallback, scope_user, scope_agent, category, top_k)
if results is not None:
logger.debug(
"FTS5Retrieval.search: tokenized fallback returned %d results in %.1fms",
len(results),
(time.perf_counter() - _t0) * 1000,
)
return results
logger.debug(
"FTS5Retrieval.search: returning [] (no path produced results) in %.1fms",
(time.perf_counter() - _t0) * 1000,
)
return []
def _execute_search(
self,
fts5_query: str,
scope_user: str | None,
scope_agent: str | None,
category: str | None,
top_k: int,
) -> list[dict[str, Any]] | None:
"""Execute FTS5 query. Return None on syntax error."""
if not fts5_query:
return []
conditions = ["memory_fts MATCH ?"]
params: list[Any] = [fts5_query]
if scope_user:
conditions.append("scope_user = ?")
params.append(scope_user)
if scope_agent:
conditions.append("scope_agent = ?")
params.append(scope_agent)
if category:
conditions.append("category = ?")
params.append(category)
where_clause = " AND ".join(conditions)
sql = f"""
SELECT doc_id, content, raw_content, category, scope_user, scope_agent,
created_at, confidence, source, fact_json,
bm25(memory_fts) AS bm25_score
FROM memory_fts
WHERE {where_clause}
ORDER BY bm25_score
LIMIT ?
"""
params.append(top_k * 2)
with self._lock:
try:
conn = self._conn
rows = conn.execute(sql, params).fetchall()
except sqlite3.OperationalError as e:
logger.debug("FTS5 query syntax error: %s (query: %s)", e, fts5_query)
return None
results: list[dict[str, Any]] = []
for row in rows:
(
doc_id,
indexed_content,
raw_content,
cat,
s_user,
s_agent,
created_at,
confidence,
source,
fact_json,
bm25_score,
) = row
score = self._compute_final_score(
bm25_score=-bm25_score, # FTS5 returns negative
confidence=confidence,
created_at=created_at,
)
fact: dict[str, Any] = {}
if fact_json:
try:
decoded = json.loads(fact_json)
if isinstance(decoded, dict):
fact = decoded
except (TypeError, ValueError):
logger.debug("FTS5 fact metadata was not valid JSON for doc_id=%r", doc_id)
fact.setdefault("id", doc_id)
fact.setdefault("content", raw_content if raw_content is not None else indexed_content)
fact.setdefault("category", cat)
fact.setdefault("confidence", confidence)
fact.setdefault("createdAt", created_at)
fact.setdefault("source", source if source is not None else "fts5")
fact["score"] = score
fact["bm25_score"] = -bm25_score
results.append(fact)
logger.debug(
"FTS5 raw SQL: fts5_query=%r scope_user=%r scope_agent=%r category=%r -> %d rows. bm25 raw: %s",
fts5_query,
scope_user,
scope_agent,
category,
len(rows),
[(r["id"], r["bm25_score"]) for r in results[:10]],
)
results.sort(key=lambda r: r["score"], reverse=True)
return results[:top_k]
def _compute_final_score(
self,
bm25_score: float,
confidence: float,
created_at: str,
) -> float:
"""Combined score: BM25 × time_decay + confidence weight.
``bm25_score`` is negative for relevant docs (SQLite FTS5 convention).
The caller negates it (``-bm25_score``) before storing in the result
dict, so here we treat it as positive relevance magnitude.
"""
score = bm25_score
try:
dt = datetime.fromisoformat(created_at.replace("Z", "+00:00"))
age_days = (datetime.now(UTC) - dt).days
time_decay = 1.0 if age_days < 30 else math.exp(-0.01 * (age_days - 30))
score *= time_decay
except (AttributeError, ValueError, TypeError):
time_decay = -1.0 # sentinel for "unparseable" → skipped decay
age_days = -1
try:
normalized_confidence = float(confidence)
if not math.isfinite(normalized_confidence):
raise ValueError
except (TypeError, ValueError):
normalized_confidence = 0.5
score += normalized_confidence * _CONFIDENCE_WEIGHT
logger.debug(
"_compute_final_score: bm25_in=%.4f time_decay=%s age_days=%s conf=%.2f -> final=%.4f",
bm25_score,
time_decay,
age_days,
normalized_confidence,
score,
)
return score
# ── Stats ──────────────────────────────────────────────────────────
def stats(self) -> dict[str, Any]:
"""Index statistics."""
with self._lock:
conn = self._conn
total = conn.execute("SELECT COUNT(*) FROM memory_fts").fetchone()[0]
return {
"total_docs": total,
"jieba": _jieba_available,
"db_path": self._db_path,
}
def close(self) -> None:
with self._lock:
self._conn.close()
def _scope_value(value: str | None) -> str:
"""Encode a typed scope value so ``None`` cannot collide with a user id."""
return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
def _scope_key(scope: dict[str, str | None]) -> tuple[str, str]:
user_id = scope.get("userId")
agent_name = scope.get("agentName")
if user_id is not None and not isinstance(user_id, str):
raise ValueError("retrieval scope userId must be a string or null")
if agent_name is not None and not isinstance(agent_name, str):
raise ValueError("retrieval scope agentName must be a string or null")
return _scope_value(user_id), _scope_value(agent_name)
class FTS5RetrievalAdapter:
"""Scope-aware ``RetrievalPort`` adapter backed by one SQLite FTS5 DB.
The index is derived data. Canonical facts remain in Markdown and storage
notifications update only the addressed row. A deterministic composite
document id prevents equal fact ids in two user/agent scopes from
overwriting each other.
"""
def __init__(self, db_path: str | Path = ":memory:") -> None:
self._engine = FTS5Retrieval(db_path)
@staticmethod
def _document_id(fact_id: str, scope: dict[str, str | None]) -> str:
scope_user, scope_agent = _scope_key(scope)
return json.dumps([scope_user, scope_agent, fact_id], ensure_ascii=False, separators=(",", ":"))
def _document(self, fact: dict[str, Any], scope: dict[str, str | None]) -> dict[str, Any]:
fact_id = fact.get("id")
content = fact.get("content")
if not isinstance(fact_id, str) or not fact_id:
raise ValueError("retrieval fact.id must be a non-empty string")
if not isinstance(content, str) or not content.strip():
raise ValueError("retrieval fact.content must be a non-empty string")
scope_user, scope_agent = _scope_key(scope)
payload = dict(fact)
payload["scope"] = {"userId": scope.get("userId"), "agentName": scope.get("agentName")}
source = payload.get("source")
return {
"fact_id": self._document_id(fact_id, scope),
"content": content,
"category": str(payload.get("category") or "context"),
"confidence": float(payload.get("confidence") or 0.5),
"created_at": payload.get("createdAt") if isinstance(payload.get("createdAt"), str) else None,
"scope_user": scope_user,
"scope_agent": scope_agent,
"source": source if isinstance(source, str) else json.dumps(source, ensure_ascii=False, default=str),
"fact_data": payload,
}
def upsert(self, fact: dict[str, Any], *, scope: dict[str, str | None], path: str) -> None:
del path # Canonical location belongs to storage; the index is rebuildable.
document = self._document(fact, scope)
self._engine.index_fact(**document)
def rebuild(self, records: list[tuple[dict[str, Any], dict[str, str | None], str]], *, scopes: list[dict[str, str | None]] | None) -> None:
"""Atomically replace the records selected by a storage rebuild."""
documents = [self._document(fact, scope) for fact, scope, _path in records]
encoded_scopes = None
if scopes is not None:
encoded_scopes = [_scope_key(scope) for scope in scopes]
self._engine.replace_documents(documents, scopes=encoded_scopes)
def remove(self, fact_id: str, *, scope: dict[str, str | None]) -> None:
self._engine.remove_fact(self._document_id(fact_id, scope))
def search(
self,
query: str,
*,
scopes: list[dict[str, str | None]],
top_k: int,
mode: str,
filters: dict[str, Any] | None,
) -> list[dict[str, Any]]:
if not query.strip() or top_k <= 0:
return []
if mode not in {"hybrid", "fts5", "lexical"}:
raise ValueError(f"unsupported FTS5 retrieval mode: {mode}")
filters = filters or {}
category = filters.get("category")
if category is not None and not isinstance(category, str):
raise ValueError("retrieval category filter must be a string")
results: list[dict[str, Any]] = []
per_scope_limit = top_k * 4
for scope in scopes:
scope_user, scope_agent = _scope_key(scope)
for candidate in self._engine.search(
query,
scope_user=scope_user,
scope_agent=scope_agent,
category=category,
top_k=per_scope_limit,
):
fact = dict(candidate)
score = float(fact.pop("score", 0.0))
bm25_score = float(fact.pop("bm25_score", 0.0))
if any(fact.get(key) != value for key, value in filters.items()):
continue
results.append(
{
"fact": fact,
"score": score,
"matchType": "fts5",
"retrieval": {"bm25": bm25_score},
}
)
results.sort(key=lambda result: result["score"], reverse=True)
return results[:top_k]
def clear(self, *, scopes: list[dict[str, str | None]] | None = None) -> None:
if scopes is None:
self._engine.clear_index()
return
for scope in scopes:
scope_user, scope_agent = _scope_key(scope)
self._engine.clear_scope(scope_user=scope_user, scope_agent=scope_agent)
def stats(self) -> dict[str, Any]:
return self._engine.stats()
def close(self) -> None:
self._engine.close()
def create_fts5_retrieval(config: Any) -> FTS5RetrievalAdapter | None:
"""Build DeerMem's bundled adapter, using a persistent derived index.
Standalone ``DeerMem`` instances with no configured storage root use an
in-memory index. The host factory always injects an absolute storage root,
so normal Gateway instances persist the rebuildable index below it.
"""
storage_path = str(getattr(config, "storage_path", "") or "")
if not storage_path:
db_path: str | Path = ":memory:"
else:
index_dir = Path(storage_path) / ".retrieval"
index_dir.mkdir(parents=True, exist_ok=True)
db_path = index_dir / "memory-fts5.sqlite3"
try:
return FTS5RetrievalAdapter(db_path)
except sqlite3.DatabaseError as exc:
if db_path == ":memory:":
logger.warning("SQLite FTS5 is unavailable; DeerMem will use substring retrieval: %s", exc)
return None
logger.warning("Derived memory retrieval index is invalid; recreating it: %s", exc)
try:
for path in (Path(db_path), Path(f"{db_path}-wal"), Path(f"{db_path}-shm")):
path.unlink(missing_ok=True)
return FTS5RetrievalAdapter(db_path)
except (OSError, sqlite3.DatabaseError) as retry_exc:
logger.warning(
"SQLite FTS5 index recreation failed; DeerMem will use substring retrieval: %s",
retry_exc,
)
return None

View File

@ -68,6 +68,10 @@ class RetrievalPort(Protocol):
def search(self, query: str, *, scopes: list[dict[str, str | None]], top_k: int, mode: str, filters: dict[str, Any] | None) -> list[dict[str, Any]]: ...
def clear(self, *, scopes: list[dict[str, str | None]] | None = None) -> None: ...
def rebuild(self, records: list[tuple[dict[str, Any], dict[str, str | None], str]], *, scopes: list[dict[str, str | None]] | None) -> None: ...
RetrievalNotification = tuple[str, dict[str, Any] | str, str | None]
ScopedRetrievalNotifications = tuple[str, list[RetrievalNotification]]
@ -408,6 +412,9 @@ class MemoryStorage(abc.ABC):
"""Clear global summaries and every agent fact bucket for one user."""
raise NotImplementedError
def close(self) -> None:
"""Release optional storage resources."""
class FileMemoryStorage(MemoryStorage):
def __init__(self, config: DeerMemConfig, retrieval: RetrievalPort | None = None):
@ -416,6 +423,12 @@ class FileMemoryStorage(MemoryStorage):
self._memory_cache: dict[tuple[str | None, str | None], tuple[dict[str, Any], tuple[Any, ...]]] = {}
self._cache_lock = threading.Lock()
self._scope_locks: weakref.WeakValueDictionary[tuple[str | None, str | None], threading.RLock] = weakref.WeakValueDictionary()
self._retrieval_dirty_scopes: set[tuple[str | None, str | None]] = set()
def close(self) -> None:
"""Release the retrieval adapter owned by this storage instance."""
if self._retrieval is not None:
self._retrieval.close()
@staticmethod
def _cache_key(agent_name: str | None = None, *, user_id: str | None = None) -> tuple[str | None, str | None]:
@ -463,6 +476,8 @@ class FileMemoryStorage(MemoryStorage):
self._retrieval.remove(str(value), scope=scope)
except Exception:
logger.exception("Retrieval notification failed for %s", value)
with self._cache_lock:
self._retrieval_dirty_scopes.add(self._cache_key(agent_name, user_id=user_id))
@staticmethod
def _validate_loaded_fact(
@ -991,7 +1006,7 @@ class FileMemoryStorage(MemoryStorage):
self._memory_cache[key] = (copy.deepcopy(document), signature)
return copy.deepcopy(document)
def reload(self, agent_name: str | None = None, *, user_id: str | None = None) -> dict[str, Any]:
def reload(self, agent_name: str | None = None, *, user_id: str | None = None, _rebuild_retrieval: bool = True) -> dict[str, Any]:
path = self._get_memory_file_path(agent_name, user_id=user_id)
key = self._cache_key(agent_name, user_id=user_id)
legacy_path = self._legacy_agent_memory_path(path, agent_name) if agent_name is not None else None
@ -1012,6 +1027,8 @@ class FileMemoryStorage(MemoryStorage):
signature = self._scope_signature(path, agent_name)
with self._cache_lock:
self._memory_cache[key] = (copy.deepcopy(document), signature)
if _rebuild_retrieval and agent_name is not None and self._retrieval is not None:
self.rebuild_index([{"userId": user_id, "agentName": agent_name}])
return copy.deepcopy(document)
def migrate(
@ -1029,7 +1046,7 @@ class FileMemoryStorage(MemoryStorage):
self._recover_if_needed(path)
migrated, from_version, notifications = self._migrate_locked(path, agent_name, user_id=user_id, include_global=True)
self._dispatch_retrieval_notifications(notifications, user_id=user_id, agent_name=agent_name)
document = self.reload(agent_name, user_id=user_id)
document = self.reload(agent_name, user_id=user_id, _rebuild_retrieval=False)
return {
"migrated": migrated,
"fromVersion": from_version,
@ -1385,7 +1402,25 @@ class FileMemoryStorage(MemoryStorage):
filters: dict[str, Any] | None = None,
) -> list[dict[str, Any]]:
if self._retrieval is not None:
requested_keys = {self._cache_key(self._scope_kwargs(scope).get("agent_name"), user_id=self._scope_kwargs(scope).get("user_id")) for scope in scopes}
with self._cache_lock:
dirty = requested_keys & self._retrieval_dirty_scopes
if dirty:
dirty_scopes = [{"userId": user_id, "agentName": agent_name} for user_id, agent_name in dirty]
rebuild_result = self.rebuild_index(dirty_scopes)
if rebuild_result.get("failed"):
return self._search_substring(query, scopes=scopes, top_k=top_k, filters=filters)
return self._retrieval.search(query, scopes=scopes, top_k=top_k, mode=mode, filters=filters)
return self._search_substring(query, scopes=scopes, top_k=top_k, filters=filters)
def _search_substring(
self,
query: str,
*,
scopes: list[dict[str, str | None]],
top_k: int,
filters: dict[str, Any] | None,
) -> list[dict[str, Any]]:
query_lower = query.strip().lower()
if not query_lower or top_k <= 0:
return []
@ -1402,6 +1437,7 @@ class FileMemoryStorage(MemoryStorage):
def rebuild_index(self, scopes: list[dict[str, str | None]] | None = None) -> dict[str, Any]:
if self._retrieval is None:
return {"supported": False, "indexed": 0, "failed": 0, "reason": "retrieval_not_configured"}
records: list[tuple[dict[str, Any], dict[str, str | None], str]] = []
indexed = 0
failed = 0
if scopes is None:
@ -1422,8 +1458,7 @@ class FileMemoryStorage(MemoryStorage):
raise MemoryStorageCorruption(f"Fact user scope does not match directory for {path}")
validate_agent_name(expected_agent)
self._validate_loaded_fact(fact, path, user_id=original_user, agent_name=expected_agent)
self.notify_fact_upsert(fact, path=str(path))
indexed += 1
records.append((fact, _scope_dict(original_user, expected_agent), str(path)))
except Exception:
logger.exception("Failed to rebuild retrieval index for %s", path)
failed += 1
@ -1436,11 +1471,42 @@ class FileMemoryStorage(MemoryStorage):
continue
for fact in self.list_facts(**kwargs):
try:
self.notify_fact_upsert(fact, path=str(fact_file_path(memory_path, fact["id"], agent_name=agent_name)))
indexed += 1
records.append((fact, _scope_dict(kwargs.get("user_id"), agent_name), str(fact_file_path(memory_path, fact["id"], agent_name=agent_name))))
except Exception:
logger.exception("Failed to rebuild retrieval index for fact %s", fact.get("id"))
failed += 1
bulk_rebuild = getattr(self._retrieval, "rebuild", None)
if callable(bulk_rebuild):
try:
bulk_rebuild(records, scopes=scopes)
indexed = len(records)
with self._cache_lock:
if scopes is None:
self._retrieval_dirty_scopes.clear()
else:
self._retrieval_dirty_scopes.difference_update(self._cache_key(self._scope_kwargs(scope).get("agent_name"), user_id=self._scope_kwargs(scope).get("user_id")) for scope in scopes)
except Exception:
logger.exception("Failed to atomically rebuild retrieval index")
failed += len(records) or 1
return {"supported": True, "indexed": indexed, "failed": failed, "fatal": True}
else:
clear = getattr(self._retrieval, "clear", None)
if callable(clear):
clear(scopes=scopes)
for fact, scope, path in records:
try:
self.notify_fact_upsert(fact, path=path)
indexed += 1
except Exception:
logger.exception("Failed to rebuild retrieval index for fact %s", fact.get("id"))
failed += 1
if failed == 0:
with self._cache_lock:
if scopes is None:
self._retrieval_dirty_scopes.clear()
else:
self._retrieval_dirty_scopes.difference_update(self._cache_key(self._scope_kwargs(scope).get("agent_name"), user_id=self._scope_kwargs(scope).get("user_id")) for scope in scopes)
return {"supported": True, "indexed": indexed, "failed": failed}
def retrieval_status(self) -> dict[str, Any]:
@ -1459,8 +1525,13 @@ class FileMemoryStorage(MemoryStorage):
def create_storage(config: DeerMemConfig, retrieval: RetrievalPort | None = None) -> MemoryStorage:
if retrieval is None and config.retrieval_adapter:
try:
module_path, factory_name = config.retrieval_adapter.rsplit(".", 1)
factory = getattr(importlib.import_module(module_path), factory_name)
if config.retrieval_adapter == "fts5":
from .retrieval import create_fts5_retrieval
factory = create_fts5_retrieval
else:
module_path, factory_name = config.retrieval_adapter.rsplit(".", 1)
factory = getattr(importlib.import_module(module_path), factory_name)
retrieval = factory(config)
except Exception as exc:
raise ValueError(f"backend_config.retrieval_adapter={config.retrieval_adapter!r} failed to load: {exc}") from exc

View File

@ -481,6 +481,15 @@ class MemoryManager(BaseModel):
``cls(backend_config=backend_config, mode=mode)``.
"""
def close(self) -> None:
"""Release backend resources during graceful process shutdown.
Backends that own external resources may override this hook. The
default is intentionally a no-op for lightweight or third-party
implementations.
"""
return None
# ── Backend discovery (drop-in) ───────────────────────────────────────────
def _scan_backends() -> dict[str, type[MemoryManager]]:

View File

@ -15,6 +15,47 @@ from langgraph.types import Command
logger = logging.getLogger(__name__)
# Whitelisted form field types; anything else degrades to "text" so a bad
# model-provided type can never produce an unrenderable card.
FORM_FIELD_TYPES = frozenset({"text", "textarea", "number", "select", "multi_select", "checkbox", "date"})
_OPTION_FIELD_TYPES = frozenset({"select", "multi_select"})
# Field names that collide with JavaScript Object.prototype properties. The
# frontend stores form values in a plain object keyed by field name, so these
# would read inherited prototype members instead of user input.
_RESERVED_FIELD_NAMES = frozenset(
{
"__proto__",
"constructor",
"prototype",
"toString",
"toLocaleString",
"valueOf",
"hasOwnProperty",
"isPrototypeOf",
"propertyIsEnumerable",
"__defineGetter__",
"__defineSetter__",
"__lookupGetter__",
"__lookupSetter__",
}
)
# Hard caps so a runaway model cannot publish an unbounded form. Exceeding a
# cap is a structural error: the whole form degrades to the legacy modes
# instead of silently truncating business fields.
MAX_FORM_FIELDS = 16
MAX_FIELD_OPTIONS = 24
MAX_FIELD_TEXT_CHARS = 200
# Total budget over the serialized normalized fields, in UTF-8 bytes. The
# per-item caps alone still admit forms whose plain-text IM fallback exceeds
# channel delivery limits (Slack truncates at 40k chars per message; Feishu
# guides ~30KB per card), which would silently drop trailing fields — the very
# thing atomic validation exists to prevent. 16KB keeps the fallback text of
# any accepted form comfortably inside the strictest supported channel while
# leaving headroom for question/context.
MAX_FORM_SERIALIZED_BYTES = 16_384
class ClarificationMiddlewareState(AgentState):
"""Compatible with the `ThreadState` schema."""
@ -62,21 +103,141 @@ class ClarificationMiddleware(AgentMiddleware[ClarificationMiddlewareState]):
if not isinstance(options, list):
options = [options]
return [str(option) for option in options]
# Trim, drop blanks, and dedupe (order-preserving): the frontend parser
# rejects the whole payload on blank option labels, so they must never
# be emitted.
normalized: list[str] = []
seen: set[str] = set()
for option in options:
text = str(option).strip()
if not text or text in seen:
continue
seen.add(text)
normalized.append(text)
return normalized
def _build_human_input_payload(self, args: dict[str, Any], *, tool_call_id: str, request_id: str) -> dict[str, Any]:
"""Build the structured UI payload while keeping ToolMessage.content as fallback."""
@staticmethod
def _normalize_bool(raw: Any) -> bool:
"""Coerce a model-provided boolean; some models serialize booleans as strings or 1/0."""
if isinstance(raw, bool):
return raw
if isinstance(raw, int | float):
return bool(raw)
if isinstance(raw, str):
return raw.strip().lower() == "true"
return False
def _normalize_fields(self, raw_fields: Any) -> list[dict[str, Any]]:
"""Normalize tool-provided form fields into the validated v2 field schema.
Validation is atomic: any structurally broken entry (non-dict, bad or
reserved or duplicate name, over-cap counts/lengths) invalidates the
whole form so the card can never render "complete" while silently
missing a required business field. Benign issues keep their local
degradation: unknown types and option-less selects become ``text``.
"""
fields = raw_fields
if isinstance(fields, str):
try:
fields = json.loads(fields)
except (json.JSONDecodeError, TypeError):
return []
if not isinstance(fields, list):
return []
if len(fields) > MAX_FORM_FIELDS:
return []
normalized: list[dict[str, Any]] = []
seen_names: set[str] = set()
for entry in fields:
if not isinstance(entry, dict):
return []
raw_name = entry.get("name")
if not isinstance(raw_name, str) or not raw_name.strip():
return []
name = raw_name.strip()
if name in _RESERVED_FIELD_NAMES or name in seen_names or len(name) > MAX_FIELD_TEXT_CHARS:
return []
seen_names.add(name)
raw_label = entry.get("label")
label = raw_label.strip() if isinstance(raw_label, str) and raw_label.strip() else name
if len(label) > MAX_FIELD_TEXT_CHARS:
return []
field_type = entry.get("type")
# isinstance guard first: `type: []` / `type: {}` are legal JSON
# from a model, and an unhashable membership probe would raise
# TypeError instead of degrading.
if not isinstance(field_type, str) or field_type not in FORM_FIELD_TYPES:
field_type = "text"
options = self._normalize_options(entry.get("options")) if field_type in _OPTION_FIELD_TYPES else []
if len(options) > MAX_FIELD_OPTIONS or any(len(option) > MAX_FIELD_TEXT_CHARS for option in options):
return []
if field_type in _OPTION_FIELD_TYPES and not options:
field_type = "text"
field: dict[str, Any] = {
"name": name,
"label": label,
"type": field_type,
"required": self._normalize_bool(entry.get("required")),
}
if field_type in _OPTION_FIELD_TYPES:
field["options"] = [
{
"id": f"{name}-option-{index}",
"label": option,
"value": option,
}
for index, option in enumerate(options, 1)
]
placeholder = entry.get("placeholder")
if isinstance(placeholder, str) and placeholder.strip():
if len(placeholder.strip()) > MAX_FIELD_TEXT_CHARS:
return []
field["placeholder"] = placeholder.strip()
normalized.append(field)
if len(json.dumps(normalized, ensure_ascii=False).encode("utf-8")) > MAX_FORM_SERIALIZED_BYTES:
return []
return normalized
def _build_human_input_payload(self, args: dict[str, Any], *, tool_call_id: str, request_id: str, fields: list[dict[str, Any]] | None = None) -> dict[str, Any]:
"""Build the structured UI payload while keeping ToolMessage.content as fallback.
Protocol versioning: legacy modes (``free_text`` / ``choice_with_other``)
keep ``version: 1`` so their wire format is unchanged; the v2 ``form``
mode carries ``version: 2`` so older frontends reject the payload and
degrade to the plain-text ToolMessage content. Replies stay on the v1
response protocol (``text`` / ``option``) the form card submits a
readable ``value`` summary, so no new response kind is introduced.
``fields`` accepts an already-normalized list so callers rendering both
the payload and the text fallback normalize only once.
"""
if fields is None:
fields = self._normalize_fields(args.get("fields"))
options = self._normalize_options(args.get("options", []))
clarification_type = str(args.get("clarification_type", "missing_info"))
if fields:
version, input_mode = 2, "form"
elif options:
version, input_mode = 1, "choice_with_other"
else:
version, input_mode = 1, "free_text"
payload: dict[str, Any] = {
"version": 1,
"version": version,
"kind": "human_input_request",
"source": "ask_clarification",
"request_id": request_id,
"clarification_type": clarification_type,
"question": str(args.get("question") or ""),
"input_mode": "choice_with_other" if options else "free_text",
"input_mode": input_mode,
}
if tool_call_id:
@ -86,7 +247,9 @@ class ClarificationMiddleware(AgentMiddleware[ClarificationMiddlewareState]):
context = args.get("context")
payload["context"] = None if context is None else str(context)
if options:
if input_mode == "form":
payload["fields"] = fields
elif options:
payload["options"] = [
{
"id": f"option-{index}",
@ -109,18 +272,24 @@ class ClarificationMiddleware(AgentMiddleware[ClarificationMiddlewareState]):
"""
return any("\u4e00" <= char <= "\u9fff" for char in text)
def _format_clarification_message(self, args: dict) -> str:
def _format_clarification_message(self, args: dict, fields: list[dict[str, Any]] | None = None) -> str:
"""Format the clarification arguments into a user-friendly message.
Args:
args: The tool call arguments containing clarification details
fields: Already-normalized form fields, so callers rendering both
the payload and this fallback normalize only once
Returns:
Formatted message string
"""
question = args.get("question", "")
clarification_type = args.get("clarification_type", "missing_info")
# str() coercion keeps the icon lookup hashable — `clarification_type:
# []` is legal JSON from a model and would raise TypeError as a dict key.
clarification_type = str(args.get("clarification_type", "missing_info"))
context = args.get("context")
if fields is None:
fields = self._normalize_fields(args.get("fields"))
options = self._normalize_options(args.get("options", []))
# Type-specific icons
@ -146,8 +315,22 @@ class ClarificationMiddleware(AgentMiddleware[ClarificationMiddlewareState]):
# Just the question with icon
message_parts.append(f"{icon} {question}")
# Add options in a cleaner format
if options and len(options) > 0:
# Form fields take precedence over options, mirroring the payload logic.
if fields:
message_parts.append("") # blank line for spacing
for i, field in enumerate(fields, 1):
line = f" {i}. {field['label']}"
if field["required"]:
line += " (required)"
field_options = field.get("options")
if field_options:
line += " — options: " + " / ".join(option["label"] for option in field_options)
if field["type"] == "multi_select":
line += " (multiple allowed)"
message_parts.append(line)
message_parts.append("")
message_parts.append("Please reply with a value for each field.")
elif options and len(options) > 0:
message_parts.append("") # blank line for spacing
for i, option in enumerate(options, 1):
message_parts.append(f" {i}. {option}")
@ -208,14 +391,18 @@ class ClarificationMiddleware(AgentMiddleware[ClarificationMiddlewareState]):
logger.info("Intercepted clarification request")
logger.debug("Clarification question: %s", question)
# Normalize form fields once; both the text fallback and the payload
# consume the same result.
fields = self._normalize_fields(args.get("fields"))
# Format the clarification message
formatted_message = self._format_clarification_message(args)
formatted_message = self._format_clarification_message(args, fields=fields)
# Get the tool call ID
tool_call_id = request.tool_call.get("id", "")
request_id = self._stable_message_id(tool_call_id, formatted_message)
human_input_payload = self._build_human_input_payload(args, tool_call_id=tool_call_id, request_id=request_id)
human_input_payload = self._build_human_input_payload(args, tool_call_id=tool_call_id, request_id=request_id, fields=fields)
# Create a ToolMessage with the formatted question
# This will be added to the message history

View File

@ -352,16 +352,61 @@ def merge_message_writes(state: list[AnyMessage], writes: Sequence[Any]) -> list
return [message for message in messages if message is not None]
DELTA_MESSAGES_FIELD = Annotated[
list[AnyMessage],
DeltaChannel(merge_message_writes, snapshot_frequency=1000),
]
DEFAULT_DELTA_SNAPSHOT_FREQUENCY = 1000
_frozen_delta_snapshot_frequency: int | None = None
def delta_messages_field(snapshot_frequency: int) -> Any:
return Annotated[
list[AnyMessage],
DeltaChannel(merge_message_writes, snapshot_frequency=snapshot_frequency),
]
DELTA_MESSAGES_FIELD = delta_messages_field(DEFAULT_DELTA_SNAPSHOT_FREQUENCY)
class DeltaThreadState(ThreadState):
messages: DELTA_MESSAGES_FIELD
def freeze_delta_snapshot_frequency(frequency: int) -> int:
"""Freeze the process-wide delta snapshot frequency (restart-required)."""
# Lazy import: deerflow.runtime.__init__ imports checkpoint_state, which
# imports this module — a top-level import would be circular.
from deerflow.runtime.checkpoint_mode import CheckpointModeReconfigurationError
global _frozen_delta_snapshot_frequency
if frequency <= 0:
raise ValueError("snapshot frequency must be positive")
if _frozen_delta_snapshot_frequency is None:
_frozen_delta_snapshot_frequency = frequency
elif _frozen_delta_snapshot_frequency != frequency:
raise CheckpointModeReconfigurationError(f"checkpoint delta snapshot frequency is frozen at {_frozen_delta_snapshot_frequency}; refusing to reconfigure to {frequency} — restart the process to change it")
return _frozen_delta_snapshot_frequency
def resolved_delta_snapshot_frequency() -> int:
return _frozen_delta_snapshot_frequency or DEFAULT_DELTA_SNAPSHOT_FREQUENCY
def reset_delta_snapshot_frequency_for_tests() -> None:
global _frozen_delta_snapshot_frequency
_frozen_delta_snapshot_frequency = None
@cache
def _delta_thread_state_for_frequency(snapshot_frequency: int) -> type:
if snapshot_frequency == DEFAULT_DELTA_SNAPSHOT_FREQUENCY:
# Preserve the canonical class so identity checks (and its public
# re-export) keep holding at the default frequency.
return DeltaThreadState
annotations = get_type_hints(ThreadState, include_extras=True)
annotations["messages"] = delta_messages_field(snapshot_frequency)
return TypedDict(f"DeltaThreadState_f{snapshot_frequency}", annotations, total=True)
THREAD_STATE_REDUCER_FIELDS = frozenset(
{
"messages",
@ -378,17 +423,23 @@ THREAD_STATE_REDUCER_FIELDS = frozenset(
def get_thread_state_schema(mode: CheckpointChannelMode) -> type:
return DeltaThreadState if mode == "delta" else ThreadState
if mode == "delta":
return _delta_thread_state_for_frequency(resolved_delta_snapshot_frequency())
return ThreadState
@cache
def adapt_state_schema_for_mode(schema: type, mode: CheckpointChannelMode) -> type:
if mode == "full":
return schema
return _adapt_state_schema_for_mode(schema, mode, resolved_delta_snapshot_frequency())
@cache
def _adapt_state_schema_for_mode(schema: type, mode: CheckpointChannelMode, snapshot_frequency: int) -> type:
annotations = get_type_hints(schema, include_extras=True)
annotations["messages"] = DELTA_MESSAGES_FIELD
annotations["messages"] = delta_messages_field(snapshot_frequency)
return TypedDict(
f"Delta{schema.__module__.replace('.', '_')}_{schema.__name__}",
f"Delta{schema.__module__.replace('.', '_')}_{schema.__name__}_f{snapshot_frequency}",
annotations,
total=getattr(schema, "__total__", True),
)

View File

@ -37,7 +37,7 @@ from langchain_core.runnables import RunnableConfig
from deerflow.agents.lead_agent.agent import build_middlewares
from deerflow.agents.lead_agent.prompt import apply_prompt_template, get_enabled_skills_for_config
from deerflow.agents.thread_state import get_thread_state_schema, normalize_middleware_state_schemas
from deerflow.agents.thread_state import freeze_delta_snapshot_frequency, get_thread_state_schema, normalize_middleware_state_schemas
from deerflow.authz.principal import build_principal_from_context
from deerflow.config.agents_config import AGENT_NAME_PATTERN
from deerflow.config.app_config import get_app_config, is_trace_correlation_enabled, reload_app_config
@ -192,6 +192,7 @@ class DeerFlowClient:
reload_app_config(config_path)
self._app_config = get_app_config()
self._checkpoint_channel_mode = freeze_checkpoint_channel_mode(self._app_config.database.checkpoint_channel_mode)
freeze_delta_snapshot_frequency(self._app_config.database.checkpoint_delta_snapshot_frequency)
if agent_name is not None and not AGENT_NAME_PATTERN.match(agent_name):
raise ValueError(f"Invalid agent name '{agent_name}'. Must match pattern: {AGENT_NAME_PATTERN.pattern}")

View File

@ -13,6 +13,15 @@ Configuration example (``config.yaml``)::
domain: e2b.dev # optional e2b domain (e.g. self-hosted)
idle_timeout: 600 # forwarded to e2b ``set_timeout`` (seconds)
replicas: 3 # max concurrent sandboxes (LRU eviction beyond)
ownership: # required for safe multi-worker reconciliation
type: redis
redis_url: $REDIS_URL
reconciliation_interval_seconds: 60
reconciliation_grace_seconds: 120
reconciliation_orphan_ttl_seconds: 3600
reconciliation_max_pages: 10
reconciliation_max_items: 200
reconciliation_max_seconds: 15
mounts: # one-shot upload of host files into the sandbox
- host_path: /path/on/host
container_path: /path/in/sandbox

View File

@ -38,6 +38,7 @@ import time
import uuid
from collections import OrderedDict
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from decimal import Decimal, InvalidOperation
from functools import partial
from pathlib import Path
@ -51,6 +52,14 @@ from deerflow.sandbox.exceptions import SandboxCapacityExceededError
from deerflow.sandbox.sandbox import Sandbox
from deerflow.sandbox.sandbox_provider import SandboxProvider
from ..aio_sandbox.ownership import (
OwnershipBackendError,
RenewOutcome,
SandboxOwnershipStore,
generate_owner_id,
make_sandbox_ownership_store,
resolve_ownership_config,
)
from .e2b_sandbox import DEFAULT_E2B_HOME_DIR, E2BSandbox, _is_sandbox_gone_error
logger = logging.getLogger(__name__)
@ -62,6 +71,12 @@ DEFAULT_IDLE_TIMEOUT = 1800 # 30 minutes; passed to ``Sandbox.set_timeout``.
DEFAULT_REPLICAS = 3
DEFAULT_OVERFLOW_POLICY = "wait" # wait | reject | burst
DEFAULT_ACQUIRE_TIMEOUT = 30 # seconds for wait policy
DEFAULT_RECONCILIATION_INTERVAL_SECONDS = 60.0
DEFAULT_RECONCILIATION_GRACE_SECONDS = 120.0
DEFAULT_RECONCILIATION_ORPHAN_TTL_SECONDS = 3600.0
DEFAULT_RECONCILIATION_MAX_PAGES = 10
DEFAULT_RECONCILIATION_MAX_ITEMS = 200
DEFAULT_RECONCILIATION_MAX_SECONDS = 15.0
# Hard upper bound for ``set_timeout`` (e2b currently caps at 24h on the
# free plan; passing an excessive value is rejected by the control-plane).
MAX_E2B_TIMEOUT = 24 * 60 * 60
@ -71,10 +86,25 @@ MAX_E2B_TIMEOUT = 24 * 60 * 60
META_KEY_USER = "deer_flow_user"
META_KEY_THREAD = "deer_flow_thread"
META_KEY_PROVIDER = "deer_flow_provider"
META_KEY_GATEWAY = "deer_flow_gateway"
META_KEY_CREATED_AT = "deer_flow_created_at"
META_VAL_PROVIDER = "e2b_sandbox_provider"
E2B_EXTRA_CONFIG_KEYS = frozenset({"api_key", "domain", "home_dir", "template"})
@dataclass
class ReconciliationStats:
"""Bounded reconciliation outcome, also suitable for metrics/logging."""
discovered: int = 0
adopted: int = 0
duplicates: int = 0
deferred: int = 0
killed: int = 0
dead: int = 0
budget_exhausted: bool = False
class E2BSandboxProvider(SandboxProvider):
"""Sandbox provider backed by the e2b code-interpreter cloud SDK."""
@ -116,6 +146,13 @@ class E2BSandboxProvider(SandboxProvider):
self._transitioning_slots = 0
self._capacity_cond = threading.Condition(self._lock)
self._shutdown_called = False
self._owned_sandbox_ids: set[str] = set()
self._acquire_inflight: set[str] = set()
self._orphan_first_seen: dict[str, float] = {}
self._maintenance_stop = threading.Event()
self._lease_thread: threading.Thread | None = None
self._reconcile_thread: threading.Thread | None = None
self._owner_id = generate_owner_id()
self._config = self._load_config()
acquire_workers = max(4, min(32, self._capacity_limit() + 1))
@ -123,9 +160,20 @@ class E2BSandboxProvider(SandboxProvider):
max_workers=acquire_workers,
thread_name_prefix="e2b-sandbox-acquire",
)
self._ownership_config = resolve_ownership_config(
self._config.get("ownership"),
stream_bridge=self._config.get("stream_bridge"),
)
self._ownership: SandboxOwnershipStore = make_sandbox_ownership_store(
self._ownership_config,
owner_id=self._owner_id,
)
if not self._ownership.supports_cross_process:
logger.warning("E2B sandbox ownership is process-local. Multi-worker gateways must configure sandbox.ownership.type: redis for safe reconciliation.")
atexit.register(self.shutdown)
self._register_signal_handlers()
self._start_maintenance_threads()
def _load_config(self) -> dict[str, Any]:
"""Read e2b options off ``SandboxConfig`` (``extra="allow"``)."""
@ -181,6 +229,32 @@ class E2BSandboxProvider(SandboxProvider):
"burst_limit": burst_limit,
"mounts": _opt("mounts") or [],
"environment": self._resolve_env_vars(_opt("environment") or {}),
"ownership": _opt("ownership"),
"stream_bridge": getattr(get_app_config(), "stream_bridge", None),
"reconciliation_interval_seconds": max(
1.0,
float(_opt("reconciliation_interval_seconds", DEFAULT_RECONCILIATION_INTERVAL_SECONDS)),
),
"reconciliation_grace_seconds": max(
0.0,
float(_opt("reconciliation_grace_seconds", DEFAULT_RECONCILIATION_GRACE_SECONDS)),
),
"reconciliation_orphan_ttl_seconds": max(
0.0,
float(_opt("reconciliation_orphan_ttl_seconds", DEFAULT_RECONCILIATION_ORPHAN_TTL_SECONDS)),
),
"reconciliation_max_pages": max(
1,
int(_opt("reconciliation_max_pages", DEFAULT_RECONCILIATION_MAX_PAGES)),
),
"reconciliation_max_items": max(
1,
int(_opt("reconciliation_max_items", DEFAULT_RECONCILIATION_MAX_ITEMS)),
),
"reconciliation_max_seconds": max(
0.1,
float(_opt("reconciliation_max_seconds", DEFAULT_RECONCILIATION_MAX_SECONDS)),
),
}
@staticmethod
@ -322,6 +396,9 @@ class E2BSandboxProvider(SandboxProvider):
self._refresh_remote_timeout(sandbox.client)
except Exception as e: # pragma: no cover - defensive
logger.debug("Failed to refresh timeout on reuse: %s", e)
self._publish_ownership(sid)
with self._lock:
self._acquire_inflight.discard(sid)
logger.info(
"Reusing in-process e2b sandbox %s for user/thread %s/%s",
@ -370,13 +447,17 @@ class E2BSandboxProvider(SandboxProvider):
self._complete_transition_remote_op(target_id, remote_destroyed=True)
return None
try:
self._publish_ownership(target_id)
except Exception:
self._complete_transition_remote_op(target_id, remote_destroyed=False)
self._safe_close_client(client)
raise
self._refresh_remote_timeout(client)
bootstrap_error, remote_destroyed = self._bootstrap_or_discard(client, target_id)
if bootstrap_error is not None:
if remote_destroyed:
self._complete_transition_remote_op(target_id, remote_destroyed=True)
else:
self._complete_transition_remote_op(target_id, remote_destroyed=False)
self._complete_transition_remote_op(target_id, remote_destroyed=remote_destroyed)
return None
discard_after_shutdown = False
@ -393,10 +474,11 @@ class E2BSandboxProvider(SandboxProvider):
self._end_transition_locked()
if discard_after_shutdown:
self._kill_client(client)
if self._claim_ownership(target_id, for_destroy=True):
self._kill_client(client)
self._release_ownership(target_id)
self._safe_close_client(client)
return None
logger.info(
"Reclaimed warm-pool e2b sandbox %s for user/thread %s/%s",
target_id,
@ -414,81 +496,39 @@ class E2BSandboxProvider(SandboxProvider):
"""
sandbox_cls = self._get_sandbox_cls()
seed = self._stable_seed(thread_id, user_id)
list_kwargs = self._common_kwargs()
try:
running = sandbox_cls.list( # type: ignore[attr-defined]
query={
"metadata": {
META_KEY_PROVIDER: META_VAL_PROVIDER,
META_KEY_USER: user_id,
META_KEY_THREAD: thread_id,
}
},
**list_kwargs,
entries, _ = self._list_remote_entries(
{
META_KEY_PROVIDER: META_VAL_PROVIDER,
META_KEY_USER: user_id,
META_KEY_THREAD: thread_id,
}
)
candidates = sorted(
((sandbox_id, metadata) for entry in entries if (sandbox_id := self._entry_id(entry)) and (metadata := self._entry_metadata(entry)).get(META_KEY_USER) == user_id and metadata.get(META_KEY_THREAD) == thread_id),
key=lambda item: (item[1].get(META_KEY_CREATED_AT, ""), item[0]),
)
for target_id, _metadata in candidates:
adopted = self._adopt_remote_candidate(
sandbox_cls,
target_id,
thread_id=thread_id,
user_id=user_id,
seed=seed,
)
except TypeError:
try:
running = sandbox_cls.list(
metadata={
META_KEY_PROVIDER: META_VAL_PROVIDER,
META_KEY_USER: user_id,
META_KEY_THREAD: thread_id,
},
**list_kwargs,
)
except Exception as e:
logger.debug("e2b Sandbox.list() unavailable, skipping discovery: %s", e)
return None
except Exception as e:
logger.debug(
"e2b Sandbox.list() raised while discovering thread %s: %s",
thread_id,
e,
)
return None
if adopted is not None:
return adopted
return None
# Pick the first matching candidate; tolerate either ``SandboxInfo``
# objects with ``sandbox_id`` or plain dicts.
# Normalise the return value of ``Sandbox.list()``:
# * Older SDKs (<= 1.x) returned a plain ``list[SandboxInfo]`` — directly iterable.
# * e2b-code-interpreter >= 2.x returns a ``SandboxPaginator`` exposing
# ``has_next: bool`` and ``next_items() -> list[SandboxInfo]`` instead
# of being iterable. Walking pages keeps discovery correct when the
# org has more sandboxes than fit in a single page.
def _iter_running(obj):
if obj is None:
return
if hasattr(obj, "next_items") and hasattr(obj, "has_next"):
for _ in range(50):
try:
page = obj.next_items()
except Exception as exc:
logger.debug("SandboxPaginator.next_items() failed: %s", exc)
return
if not page:
return
yield from page
if not getattr(obj, "has_next", False):
return
return
try:
yield from obj
except TypeError:
logger.debug("Sandbox.list() returned non-iterable %s; ignoring", type(obj).__name__)
target_id: str | None = None
for entry in _iter_running(running):
sid = getattr(entry, "sandbox_id", None) or (entry.get("sandbox_id") if isinstance(entry, dict) else None)
metadata = getattr(entry, "metadata", None) or (entry.get("metadata") if isinstance(entry, dict) else {}) or {}
if metadata.get(META_KEY_USER) != user_id:
continue
if metadata.get(META_KEY_THREAD) != thread_id:
continue
target_id = sid
break
if not target_id:
return None
def _adopt_remote_candidate(
self,
sandbox_cls: type[E2BClientSandbox],
target_id: str,
*,
thread_id: str,
user_id: str,
seed: str,
) -> str | None:
"""Try to adopt one discovered candidate without harming peer-owned VMs."""
try:
client = self._reconnect_live_client(sandbox_cls, target_id)
@ -528,6 +568,20 @@ class E2BSandboxProvider(SandboxProvider):
self._safe_close_client(client)
raise
with self._lock:
shutdown_before_ownership = self._shutdown_called
if shutdown_before_ownership:
self._complete_reserved_remote_op(target_id, remote_destroyed=False)
self._safe_close_client(client)
return None
try:
self._publish_ownership(target_id)
except Exception:
self._complete_reserved_remote_op(target_id, remote_destroyed=False)
self._safe_close_client(client)
raise
self._refresh_remote_timeout(client)
bootstrap_error, remote_destroyed = self._bootstrap_or_discard(client, target_id)
if bootstrap_error is not None:
@ -543,6 +597,9 @@ class E2BSandboxProvider(SandboxProvider):
self._register_connected_sandbox(target_id, client, thread_id=thread_id, user_id=user_id)
self._commit_capacity()
if discard_after_shutdown:
if self._claim_ownership(target_id, for_destroy=True):
self._kill_client(client)
self._release_ownership(target_id)
self._safe_close_client(client)
return None
logger.info(
@ -787,6 +844,8 @@ class E2BSandboxProvider(SandboxProvider):
sandbox_cls = self._get_sandbox_cls()
metadata: dict[str, str] = {
META_KEY_PROVIDER: META_VAL_PROVIDER,
META_KEY_GATEWAY: self._owner_id,
META_KEY_CREATED_AT: str(time.time()),
}
if thread_id:
metadata[META_KEY_USER] = user_id
@ -849,6 +908,17 @@ class E2BSandboxProvider(SandboxProvider):
reason="shutdown",
)
try:
self._publish_ownership(sandbox_id)
except Exception:
remote_destroyed = False
if self._claim_ownership(sandbox_id, for_destroy=True):
remote_destroyed = self._kill_client(client) is None
self._release_ownership(sandbox_id)
self._safe_close_client(client)
self._complete_reserved_remote_op(sandbox_id, remote_destroyed=remote_destroyed)
raise
# Materialise DeerFlow's virtual path layout (/mnt/user-data/...) inside
# the e2b VM. Without this step shell commands the agent emits — which
# use the same /mnt/user-data prefix as LocalSandbox / AioSandbox — fail
@ -883,7 +953,9 @@ class E2BSandboxProvider(SandboxProvider):
self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id
if should_kill:
self._kill_client(client)
if self._claim_ownership(sandbox_id, for_destroy=True):
self._kill_client(client)
self._release_ownership(sandbox_id)
self._safe_close_client(client)
raise SandboxCapacityExceededError(
f"Sandbox provider shut down during sandbox creation; killed remote sandbox {sandbox_id}",
@ -912,6 +984,339 @@ class E2BSandboxProvider(SandboxProvider):
kwargs["domain"] = self._config["domain"]
return kwargs
@staticmethod
def _entry_id(entry: Any) -> str | None:
value = getattr(entry, "sandbox_id", None)
if value is None and isinstance(entry, dict):
value = entry.get("sandbox_id")
return str(value) if value else None
@staticmethod
def _entry_metadata(entry: Any) -> dict[str, Any]:
value = getattr(entry, "metadata", None)
if value is None and isinstance(entry, dict):
value = entry.get("metadata")
return value if isinstance(value, dict) else {}
def _list_remote_entries(self, metadata: dict[str, str]) -> tuple[list[Any], bool]:
"""List matching E2B entries within configured page/item/time budgets."""
sandbox_cls = self._get_sandbox_cls()
try:
result = sandbox_cls.list(query={"metadata": metadata}, **self._common_kwargs()) # type: ignore[attr-defined]
except TypeError:
try:
result = sandbox_cls.list(metadata=metadata, **self._common_kwargs()) # type: ignore[attr-defined]
except Exception as e:
logger.warning("E2B reconciliation list failed: %s", e)
return [], False
except Exception as e:
logger.warning("E2B reconciliation list failed: %s", e)
return [], False
max_pages = int(self._config["reconciliation_max_pages"])
max_items = int(self._config["reconciliation_max_items"])
deadline = time.monotonic() + float(self._config["reconciliation_max_seconds"])
entries: list[Any] = []
exhausted = False
if hasattr(result, "next_items") and hasattr(result, "has_next"):
for page_number in range(max_pages):
if time.monotonic() >= deadline or len(entries) >= max_items:
exhausted = True
break
try:
page = result.next_items()
except Exception as e:
logger.warning("E2B reconciliation paginator failed: %s", e)
break
if not page:
break
room = max_items - len(entries)
entries.extend(list(page)[:room])
if len(page) > room:
exhausted = True
break
if not getattr(result, "has_next", False):
break
if page_number + 1 >= max_pages:
exhausted = True
else:
try:
all_entries = list(result or [])
except TypeError:
logger.warning("E2B Sandbox.list returned non-iterable %s", type(result).__name__)
return [], False
entries = all_entries[:max_items]
exhausted = len(all_entries) > max_items
if time.monotonic() >= deadline:
exhausted = True
return entries, exhausted
def _publish_ownership(self, sandbox_id: str) -> None:
"""Publish acquire-side ownership before exposing a sandbox locally."""
with self._lock:
self._acquire_inflight.add(sandbox_id)
try:
if not self._ownership.take(sandbox_id):
raise RuntimeError(f"E2B sandbox {sandbox_id} is being destroyed")
except Exception:
with self._lock:
self._acquire_inflight.discard(sandbox_id)
raise
with self._lock:
self._owned_sandbox_ids.add(sandbox_id)
def _claim_ownership(self, sandbox_id: str, *, for_destroy: bool = False) -> bool:
"""Exclusively claim an unowned sandbox, failing closed on store errors."""
try:
claimed = self._ownership.claim(sandbox_id, for_destroy=for_destroy)
except OwnershipBackendError as e:
logger.warning("E2B ownership claim failed for %s: %s", sandbox_id, e)
return False
if claimed:
with self._lock:
self._owned_sandbox_ids.add(sandbox_id)
return claimed
def _release_ownership(self, sandbox_id: str) -> None:
try:
self._ownership.release(sandbox_id)
except OwnershipBackendError as e:
logger.warning("Failed to release E2B ownership for %s: %s", sandbox_id, e)
with self._lock:
self._owned_sandbox_ids.discard(sandbox_id)
self._acquire_inflight.discard(sandbox_id)
def _forget_local_sandbox(self, sandbox_id: str) -> None:
"""Forget a lease taken by a peer without touching the remote VM."""
with self._lock:
if sandbox_id in self._acquire_inflight:
return
sandbox = self._sandboxes.pop(sandbox_id, None)
self._warm_pool.pop(sandbox_id, None)
self._owned_sandbox_ids.discard(sandbox_id)
for key, sid in list(self._thread_sandboxes.items()):
if sid == sandbox_id:
self._thread_sandboxes.pop(key, None)
if sandbox is not None:
try:
sandbox.close()
except Exception:
pass
def _refresh_owned_leases(self) -> None:
with self._lock:
sandbox_ids = list(self._owned_sandbox_ids)
for sandbox_id in sandbox_ids:
try:
outcome = self._ownership.renew(sandbox_id)
except OwnershipBackendError as e:
logger.warning("Could not renew E2B ownership for %s; will retry: %s", sandbox_id, e)
continue
if outcome is RenewOutcome.RENEWED:
continue
if outcome is RenewOutcome.LAPSED:
try:
if self._ownership.claim(sandbox_id):
continue
except OwnershipBackendError as e:
logger.warning("Could not re-establish E2B ownership for %s: %s", sandbox_id, e)
continue
logger.info("E2B sandbox %s ownership moved to a peer; forgetting local client", sandbox_id)
self._forget_local_sandbox(sandbox_id)
def _start_maintenance_threads(self) -> None:
def renew() -> None:
interval = self._ownership_config.renewal_interval_seconds
while not self._maintenance_stop.wait(interval):
self._refresh_owned_leases()
def reconcile() -> None:
interval = float(self._config["reconciliation_interval_seconds"])
while not self._maintenance_stop.is_set():
try:
self._reconcile_remote_sandboxes()
except Exception:
logger.exception("Periodic E2B sandbox reconciliation failed")
if self._maintenance_stop.wait(interval):
break
self._lease_thread = threading.Thread(target=renew, name="e2b-lease-renewal", daemon=True)
self._reconcile_thread = threading.Thread(target=reconcile, name="e2b-reconciliation", daemon=True)
self._lease_thread.start()
self._reconcile_thread.start()
def _reserve_reconciliation_capacity(self, sandbox_id: str) -> bool:
"""Reserve one local slot for adoption without blocking maintenance."""
with self._lock:
if self._shutdown_called or self._total_capacity_used_locked() >= self._capacity_limit():
return False
self._reserved_slots += 1
self._unowned_remote_ops_in_progress.add(sandbox_id)
self._acquire_inflight.add(sandbox_id)
return True
def _reconcile_remote_sandboxes(self, *, now: float | None = None) -> ReconciliationStats:
"""Adopt canonical E2B sandboxes and safely reap duplicates/orphans."""
observed_at = time.monotonic() if now is None else now
deadline = time.monotonic() + float(self._config["reconciliation_max_seconds"])
stats = ReconciliationStats()
entries, stats.budget_exhausted = self._list_remote_entries({META_KEY_PROVIDER: META_VAL_PROVIDER})
stats.discovered = len(entries)
groups: dict[tuple[str, str], list[tuple[str, dict[str, Any]]]] = {}
orphans: list[tuple[str, dict[str, Any]]] = []
present_ids: set[str] = set()
for entry in entries:
sandbox_id = self._entry_id(entry)
if not sandbox_id:
continue
present_ids.add(sandbox_id)
metadata = self._entry_metadata(entry)
user_id = metadata.get(META_KEY_USER)
thread_id = metadata.get(META_KEY_THREAD)
if isinstance(user_id, str) and user_id and isinstance(thread_id, str) and thread_id:
groups.setdefault((user_id, thread_id), []).append((sandbox_id, metadata))
else:
orphans.append((sandbox_id, metadata))
for (user_id, thread_id), candidates in groups.items():
candidates.sort(key=lambda item: (item[1].get(META_KEY_CREATED_AT, ""), item[0]))
with self._lock:
local_id = self._thread_sandboxes.get((user_id, thread_id))
if local_id:
candidates.sort(key=lambda item: item[0] != local_id)
live: list[tuple[str, dict[str, Any], E2BClientSandbox]] = []
for sandbox_id, metadata in candidates:
if time.monotonic() >= deadline:
stats.budget_exhausted = True
break
try:
client = self._reconnect_live_client(self._get_sandbox_cls(), sandbox_id)
except Exception as e:
logger.debug("Could not probe E2B sandbox %s during reconciliation: %s", sandbox_id, e)
continue
if client is None:
stats.dead += 1
continue
live.append((sandbox_id, metadata, client))
if not live:
continue
stats.duplicates += max(0, len(live) - 1)
canonical_id, _metadata, canonical_client = live[0]
with self._lock:
already_local = canonical_id in self._sandboxes
if already_local:
self._safe_close_client(canonical_client)
elif not self._reserve_reconciliation_capacity(canonical_id):
self._safe_close_client(canonical_client)
stats.deferred += 1
elif not self._claim_ownership(canonical_id):
self._complete_reserved_remote_op(canonical_id, remote_destroyed=False)
with self._lock:
self._acquire_inflight.discard(canonical_id)
self._safe_close_client(canonical_client)
stats.deferred += 1
else:
bootstrap_error, remote_destroyed = self._bootstrap_or_discard(canonical_client, canonical_id)
if bootstrap_error is not None:
self._complete_reserved_remote_op(canonical_id, remote_destroyed=remote_destroyed)
with self._lock:
self._acquire_inflight.discard(canonical_id)
stats.deferred += 1
else:
discard_after_shutdown = False
with self._lock:
if self._shutdown_called:
discard_after_shutdown = True
else:
self._owned_sandbox_ids.add(canonical_id)
self._unowned_remote_ops_in_progress.discard(canonical_id)
self._register_connected_sandbox(
canonical_id,
canonical_client,
thread_id=thread_id,
user_id=user_id,
)
self._commit_capacity()
if discard_after_shutdown:
if self._claim_ownership(canonical_id, for_destroy=True):
self._kill_client(canonical_client)
self._release_ownership(canonical_id)
self._safe_close_client(canonical_client)
else:
stats.adopted += 1
for sandbox_id, _metadata, client in live[1:]:
first_seen = self._orphan_first_seen.setdefault(sandbox_id, observed_at)
if observed_at - first_seen < float(self._config["reconciliation_grace_seconds"]):
self._safe_close_client(client)
stats.deferred += 1
continue
if not self._claim_ownership(sandbox_id, for_destroy=True):
self._safe_close_client(client)
stats.deferred += 1
continue
if error := self._kill_client(client):
logger.warning("Failed to kill duplicate E2B sandbox %s: %s", sandbox_id, error)
self._release_ownership(sandbox_id)
stats.deferred += 1
else:
stats.killed += 1
self._orphan_first_seen.pop(sandbox_id, None)
self._forget_local_sandbox(sandbox_id)
self._release_ownership(sandbox_id)
self._safe_close_client(client)
for sandbox_id, metadata in orphans:
if time.monotonic() >= deadline:
stats.budget_exhausted = True
break
first_seen = self._orphan_first_seen.setdefault(sandbox_id, observed_at)
created_at = metadata.get(META_KEY_CREATED_AT)
try:
age = time.time() - float(created_at) if created_at is not None else observed_at - first_seen
except (TypeError, ValueError):
age = observed_at - first_seen
if age < float(self._config["reconciliation_orphan_ttl_seconds"]):
stats.deferred += 1
continue
if not self._claim_ownership(sandbox_id, for_destroy=True):
stats.deferred += 1
continue
try:
client = self._reconnect_client(self._get_sandbox_cls(), sandbox_id)
except Exception:
self._release_ownership(sandbox_id)
continue
if error := self._kill_client(client):
logger.warning("Failed to kill orphan E2B sandbox %s: %s", sandbox_id, error)
self._release_ownership(sandbox_id)
stats.deferred += 1
else:
stats.killed += 1
self._orphan_first_seen.pop(sandbox_id, None)
self._forget_local_sandbox(sandbox_id)
self._release_ownership(sandbox_id)
self._safe_close_client(client)
for sandbox_id in set(self._orphan_first_seen) - present_ids:
self._orphan_first_seen.pop(sandbox_id, None)
logger.info(
"E2B reconciliation: discovered=%d adopted=%d duplicates=%d deferred=%d killed=%d dead=%d budget_exhausted=%s",
stats.discovered,
stats.adopted,
stats.duplicates,
stats.deferred,
stats.killed,
stats.dead,
stats.budget_exhausted,
)
return stats
def _reconnect_client(self, sandbox_cls: type[E2BClientSandbox], sandbox_id: str) -> E2BClientSandbox:
"""Connect to an existing e2b sandbox by id, with consistent kwargs."""
return sandbox_cls.connect(sandbox_id, **self._common_kwargs()) # type: ignore[attr-defined]
@ -938,7 +1343,7 @@ class E2BSandboxProvider(SandboxProvider):
sandbox_id: str,
client: E2BClientSandbox,
*,
thread_id: str,
thread_id: str | None,
user_id: str,
) -> None:
"""Track a live reconnected sandbox under its thread ownership.
@ -947,7 +1352,10 @@ class E2BSandboxProvider(SandboxProvider):
"""
sandbox = E2BSandbox(id=sandbox_id, client=client, home_dir=self._config["home_dir"])
self._sandboxes[sandbox_id] = sandbox
self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id
self._warm_pool.pop(sandbox_id, None)
if thread_id:
self._thread_sandboxes[self._thread_key(thread_id, user_id)] = sandbox_id
self._acquire_inflight.discard(sandbox_id)
def _refresh_remote_timeout(self, client: E2BClientSandbox) -> None:
"""Push the configured idle timeout to the e2b control plane."""
@ -1016,11 +1424,20 @@ class E2BSandboxProvider(SandboxProvider):
self._bootstrap_sandbox_paths(client)
except Exception as e:
logger.exception("Failed to bootstrap e2b sandbox %s. Discarding the unusable sandbox.", sandbox_id)
kill_error = self._kill_client(client)
if kill_error:
logger.warning("Failed to kill e2b sandbox %s after bootstrap failure: %s", sandbox_id, kill_error)
remote_destroyed = False
if self._claim_ownership(sandbox_id, for_destroy=True):
kill_error = self._kill_client(client)
remote_destroyed = kill_error is None
if kill_error:
logger.warning("Failed to kill e2b sandbox %s after bootstrap failure: %s", sandbox_id, kill_error)
self._release_ownership(sandbox_id)
else:
logger.info(
"Not killing E2B sandbox %s after bootstrap failure because a peer owns it",
sandbox_id,
)
self._safe_close_client(client)
return e, kill_error is None
return e, remote_destroyed
return None, True
def _bootstrap_sandbox_paths(self, client: E2BClientSandbox) -> None:
@ -1459,6 +1876,15 @@ class E2BSandboxProvider(SandboxProvider):
else:
return None
if not self._claim_ownership(evict_id, for_destroy=True):
logger.info("Deferring warm-pool eviction for %s because a peer owns it", evict_id)
self._forget_local_sandbox(evict_id)
with self._lock:
if evict_id in self._evictions_in_progress:
self._evictions_in_progress.discard(evict_id)
self._eviction_tombstones.discard(evict_id)
self._end_transition_locked()
return evict_id
try:
client = self._reconnect_live_client(self._get_sandbox_cls(), evict_id)
except Exception as e:
@ -1472,6 +1898,7 @@ class E2BSandboxProvider(SandboxProvider):
self._evictions_in_progress.discard(evict_id)
if not self._shutdown_called:
self._eviction_tombstones.add(evict_id)
self._release_ownership(evict_id)
return None
if client is None:
@ -1480,6 +1907,7 @@ class E2BSandboxProvider(SandboxProvider):
self._evictions_in_progress.discard(evict_id)
self._eviction_tombstones.discard(evict_id)
self._end_transition_locked()
self._release_ownership(evict_id)
logger.info("Evicted warm-pool e2b sandbox %s was already gone", evict_id)
return evict_id
@ -1491,6 +1919,7 @@ class E2BSandboxProvider(SandboxProvider):
self._evictions_in_progress.discard(evict_id)
if not self._shutdown_called:
self._eviction_tombstones.add(evict_id)
self._release_ownership(evict_id)
return None
self._safe_close_client(client)
@ -1499,6 +1928,7 @@ class E2BSandboxProvider(SandboxProvider):
self._evictions_in_progress.discard(evict_id)
self._eviction_tombstones.discard(evict_id)
self._end_transition_locked()
self._release_ownership(evict_id)
logger.info("Evicted warm-pool e2b sandbox %s", evict_id)
return evict_id
@ -1608,8 +2038,10 @@ class E2BSandboxProvider(SandboxProvider):
"Provider shut down during release of sandbox %s; killing instead of parking in warm pool",
sandbox_id,
)
if error := self._kill_client(client):
logger.debug("Failed to kill e2b sandbox %s during release: %s", sandbox_id, error)
if self._claim_ownership(sandbox_id, for_destroy=True):
if error := self._kill_client(client):
logger.debug("Failed to kill e2b sandbox %s during release: %s", sandbox_id, error)
self._release_ownership(sandbox_id)
self._safe_close_client(client)
return
@ -1622,12 +2054,20 @@ class E2BSandboxProvider(SandboxProvider):
self._free_transitioning_slot()
def _kill_and_close(self, sandbox: E2BSandbox) -> None:
if not self._claim_ownership(sandbox.id, for_destroy=True):
logger.info("Not killing E2B sandbox %s because a peer owns it", sandbox.id)
try:
sandbox.close()
except Exception:
pass
return
if error := self._kill_client(getattr(sandbox, "_client", None)):
logger.debug(
"kill() on e2b sandbox %s raised (probably already gone): %s",
sandbox.id,
error,
)
self._release_ownership(sandbox.id)
try:
sandbox.close()
except Exception:
@ -1658,8 +2098,14 @@ class E2BSandboxProvider(SandboxProvider):
if self._shutdown_called:
return
self._shutdown_called = True
self._maintenance_stop.set()
for thread in (self._lease_thread, self._reconcile_thread):
if thread is not None and thread is not threading.current_thread():
thread.join(timeout=max(5.0, float(self._config["reconciliation_max_seconds"]) + 1.0))
with self._lock:
active = list(self._sandboxes.items())
warm_ids = list(self._warm_pool.keys() | self._eviction_tombstones | self._remote_ops_in_progress)
owned_ids = set(self._owned_sandbox_ids)
self._sandboxes.clear()
self._warm_pool.clear()
self._eviction_tombstones.clear()
@ -1668,6 +2114,9 @@ class E2BSandboxProvider(SandboxProvider):
self._unowned_remote_ops_in_progress.clear()
self._thread_sandboxes.clear()
self._thread_locks.clear()
self._owned_sandbox_ids.clear()
self._acquire_inflight.clear()
self._orphan_first_seen.clear()
self._reserved_slots = 0
self._transitioning_slots = 0
self._capacity_cond.notify_all()
@ -1683,12 +2132,14 @@ class E2BSandboxProvider(SandboxProvider):
)
for sandbox_id, sandbox in active:
if error := self._kill_client(sandbox.client):
logger.warning(
"Failed to kill active e2b sandbox %s during shutdown: %s",
sandbox_id,
error,
)
if sandbox_id in owned_ids and self._claim_ownership(sandbox_id, for_destroy=True):
if error := self._kill_client(sandbox.client):
logger.warning(
"Failed to kill active e2b sandbox %s during shutdown: %s",
sandbox_id,
error,
)
self._release_ownership(sandbox_id)
try:
sandbox.close()
except Exception:
@ -1696,6 +2147,10 @@ class E2BSandboxProvider(SandboxProvider):
sandbox_cls = self._get_sandbox_cls()
for sandbox_id in warm_ids:
if sandbox_id not in owned_ids:
continue
if not self._claim_ownership(sandbox_id, for_destroy=True):
continue
try:
client = self._reconnect_client(sandbox_cls, sandbox_id)
except Exception as e:
@ -1704,6 +2159,7 @@ class E2BSandboxProvider(SandboxProvider):
sandbox_id,
e,
)
self._release_ownership(sandbox_id)
continue
if error := self._kill_client(client):
logger.warning(
@ -1711,9 +2167,14 @@ class E2BSandboxProvider(SandboxProvider):
sandbox_id,
error,
)
self._release_ownership(sandbox_id)
close = getattr(client, "close", None)
if callable(close):
try:
close()
except Exception:
pass
try:
self._ownership.close()
except Exception as e:
logger.warning("Failed to close E2B ownership store: %s", e)

View File

@ -0,0 +1,97 @@
# Tenki backend
Runs each DeerFlow sandbox as a [Tenki](https://tenki.cloud) cloud sandbox — an
isolated microVM created from a stock base image, with no daemon or local
virtualization to manage. A cloud-hosted alternative to the container-based AIO
sandbox and the local-virtualization BoxLite backend.
## Configuration
```yaml
sandbox:
use: deerflow.community.tenki:TenkiSandboxProvider
api_key: $TENKI_API_KEY # falls back to TENKI_API_KEY / TENKI_AUTH_TOKEN env var
base_url: https://tenki.cloud # optional; SDK default when omitted
image: my-base-image # optional; Tenki account default base image when omitted
project_id: proj_... # optional; auto-selected if the account has exactly one
workspace_id: ws_... # optional; auto-selected if the account has exactly one
cpu_cores: 2 # optional per-sandbox vCPUs
memory_mb: 2048 # optional per-sandbox memory
replicas: 3 # active + warm microVM cap per gateway process (default: 3)
idle_timeout: 600 # warm microVM idle seconds before terminate; 0 disables
max_duration: 14400 # Tenki sandbox lifetime in seconds (default: 4h); 0 uses the account default
sticky: false # pin the microVM to its host (only matters with pause/resume)
home_dir: /home/tenki # writable dir backing /mnt/user-data (default: /home/tenki)
environment: # injected into every command (and as create-time env)
PYTHONUNBUFFERED: "1"
```
Install the optional SDK before selecting this provider:
```bash
pip install "deerflow-harness[tenki]"
```
The `tenki-sandbox` package is an optional DeerFlow harness extra, not part of
the default install. Get an API key from <https://tenki.cloud/docs/sandbox/sdk>.
## Design
Tenki's Python SDK is synchronous, so — unlike BoxLite — the adapter calls the
SDK directly with no event-loop bridge. Sandboxes are named deterministically
from `user_id:thread_id`, released into an in-process warm pool after each agent
turn, and reclaimed (after a liveness health check) by the same thread on the
next acquire. A terminal session error evicts the sandbox and the next acquire
rebuilds it; other transport errors surface to the caller (`exec` is never
auto-retried — it is not idempotent, so re-running could double a command's side
effects). Sandboxes are created with `wait=False` and awaited via `wait_ready()`
so a readiness failure still leaves this provider holding the handle to
terminate, and with an explicit `max_duration` so a long-lived thread does not
lose its sandbox to Tenki's default lifetime mid-conversation.
| File | Role |
| --- | --- |
| `provider.py` | `SandboxProvider` lifecycle, warm pool, scope resolution |
| `sandbox.py` | `Sandbox` adapter; `execute_command` + file ops |
## Contract coverage
The full `Sandbox` surface is implemented. File transport uses Tenki's native `sandbox.fs`
API; directory and content search shell out and reuse `deerflow.sandbox.search`,
mirroring `e2b_sandbox`:
- `execute_command``sh -lc`, with per-call env and timeout.
- `read_file` / `write_file` / `update_file` — native `fs.read_text` / `fs.mkdir` / `fs.write_stream` (binary-safe, streamed).
- `download_file` — native `fs.read_stream`, restricted to the `/mnt/user-data` prefix; the 100 MB cap is enforced on bytes actually received, so a file growing mid-transfer cannot slip past it.
- `list_dir` / `glob` / `grep``find` / `grep` with busybox-portable flags (the fs API is single-level and has no content search); results filtered/capped in Python and reported back under `/mnt/user-data`.
Tenki sandboxes run as the unprivileged `tenki` user with `/mnt` root-owned, so
DeerFlow's `/mnt/user-data` virtual prefix is remapped under the writable
`home_dir` (like `e2b_sandbox`). The provider also best-effort `sudo`-symlinks
`/mnt/user-data``home_dir` at create time so agent shell commands using the
literal `/mnt/...` path still work; if `sudo` is unavailable the file APIs keep
working via the remap.
Warm-pool capacity is governed by `sandbox.replicas` across active + warm
sandboxes. `sandbox.idle_timeout` controls how long released warm sandboxes stay
running; `0` disables idle reaping. Active sandboxes are never evicted to satisfy
the cap.
## Scope: stable features only
Only the stable Tenki surface is used — sandbox create/terminate plus
exec/shell/filesystem over shell commands. Volumes, snapshots, and template
builds are intentionally **not** used, so no prebaked image or unstable Tenki
feature is required and any stock base image works.
Not yet implemented (follow-ups): cross-process orphan reconciliation (adopting
sandboxes left by a previous gateway process) and a preview-URL surface.
## Status
Verified end-to-end against live Tenki sandboxes: provider resolution →
`execute_command` → full file-op surface (`read`/`write`/`update`/`download`,
`list_dir`/`glob`/`grep`) → warm-pool reclaim → terminate, plus the
`/mnt/user-data` sudo symlink. Unit tests run in CI without `tenki-sandbox`
installed; `test_integration_real_sandbox` exercises a real microVM when
`TENKI_API_KEY` is set.

View File

@ -0,0 +1,46 @@
"""Tenki cloud sandbox provider for DeerFlow.
Integrates `Tenki <https://tenki.cloud>`_ cloud sandboxes behind DeerFlow's
:class:`Sandbox` / :class:`SandboxProvider` contract. Each sandbox is an
isolated cloud microVM created from a stock base image; the full contract is
implemented ``execute_command`` plus ``read_file`` / ``write_file`` /
``update_file`` / ``download_file`` / ``list_dir`` / ``glob`` / ``grep`` (file
transport uses Tenki's native ``sandbox.fs`` API; search shells out to
``find`` / ``grep``).
Configuration example (``config.yaml``)::
sandbox:
use: deerflow.community.tenki:TenkiSandboxProvider
# Tenki-specific options (read via SandboxConfig's ``extra="allow"``):
api_key: $TENKI_API_KEY # falls back to TENKI_API_KEY / TENKI_AUTH_TOKEN env var
base_url: https://tenki.cloud # optional; SDK default when omitted
image: my-base-image # optional; Tenki account default base image when omitted
project_id: proj_... # optional; auto-selected if the account has exactly one
workspace_id: ws_... # optional; auto-selected if the account has exactly one
cpu_cores: 2 # optional per-sandbox vCPUs
memory_mb: 2048 # optional per-sandbox memory
replicas: 3 # active + warm microVM cap per gateway process
idle_timeout: 600 # warm microVM idle seconds before terminate; 0 disables
max_duration: 14400 # Tenki sandbox lifetime in seconds; 0 uses the account default
sticky: false # pin the microVM to its host (only matters with pause/resume)
home_dir: /home/tenki # writable dir backing /mnt/user-data
environment: # injected into every command (and as create-time env)
PYTHONUNBUFFERED: "1"
Install the optional SDK before selecting this provider::
pip install "deerflow-harness[tenki]"
Only the stable Tenki surface is used sandbox create/terminate plus
exec/shell/filesystem. Volumes, snapshots, and template builds are intentionally
not used, so no prebaked image or unstable Tenki feature is required.
"""
from .provider import TenkiSandboxProvider
from .sandbox import TenkiSandbox
__all__ = [
"TenkiSandbox",
"TenkiSandboxProvider",
]

View File

@ -0,0 +1,450 @@
"""``TenkiSandboxProvider`` — DeerFlow :class:`SandboxProvider` backed by Tenki.
Integrates `Tenki <https://tenki.cloud>`_ cloud sandboxes as a DeerFlow sandbox
backend. Each sandbox is an isolated cloud microVM created from a stock base
image; the provider creates one per ``(user, thread)`` and reuses it within the
process, parking released sandboxes in a warm pool (shared
:class:`WarmPoolLifecycleMixin` machinery) for fast reclaim.
Config is read off :class:`SandboxConfig` (``extra="allow"``), so Tenki keys may
appear under ``sandbox:`` in ``config.yaml`` even though they are not declared on
the model see this package's ``__init__`` docstring for the full set.
The Tenki SDK is imported lazily (``_import_client``) so the harness and every
other provider installs without ``tenki-sandbox``; the dependency is only
needed once this provider is selected.
"""
from __future__ import annotations
import atexit
import hashlib
import logging
import shlex
import threading
import time
import uuid
from typing import TYPE_CHECKING, Any
from deerflow.config import get_app_config
from deerflow.sandbox.sandbox import Sandbox, _validate_extra_env
from deerflow.sandbox.sandbox_provider import SandboxProvider
from ..warm_pool_lifecycle import WarmPoolLifecycleMixin
from .sandbox import DEFAULT_TENKI_HOME_DIR, TenkiSandbox
if TYPE_CHECKING:
from tenki_sandbox import Client
logger = logging.getLogger(__name__)
_SANDBOX_NAME_PREFIX = "deer-flow-tenki-"
# Tenki terminates a sandbox at its max lifetime (~30 min by default), which
# would silently drop a long-running thread's state mid-conversation. DeerFlow
# owns the lifecycle here — the warm pool's idle_timeout reaps unused sandboxes
# — so ask for a lifetime that comfortably outlives a research run. Override
# with sandbox.max_duration (seconds); 0 falls back to the Tenki account default.
DEFAULT_MAX_DURATION = 4 * 60 * 60
# Bootstrap is best-effort and holds the per-scope acquire lock, so bound its
# exec: a hang here (e.g. a stuck sudo) would otherwise stall acquire for that
# scope indefinitely and queue later acquires behind it. A timeout drops it to
# the existing warning path instead.
_BOOTSTRAP_TIMEOUT = 30.0
def _bootstrap_script(home_dir: str) -> str:
"""Materialise DeerFlow's virtual path layout inside the sandbox.
Tenki sandboxes run as the unprivileged ``tenki`` user with ``/mnt``
root-owned, so ``mkdir -p /mnt/user-data/...`` fails with Permission denied.
Mirroring ``community/e2b_sandbox``, we create the real backing directories
under the writable HOME and best-effort ``sudo``-symlink ``/mnt/user-data``
to it so commands using the documented ``/mnt/...`` paths still work. If
sudo is unavailable the symlink step is skipped; the file APIs keep working
via :meth:`TenkiSandbox._resolve_path`'s home remap.
``sudo -n`` (non-interactive) is deliberate: if the tenki user's sudoers
entry needs a password, a bare ``sudo`` on a tty-allocating exec would block
on the password prompt and with this call's exec timeout — stall the whole
acquire. ``-n`` fails fast instead, and the existing ``|| true`` swallows it.
"""
home = shlex.quote(home_dir)
return (
f"mkdir -p {home}/workspace {home}/uploads {home}/outputs; "
f"if command -v sudo >/dev/null 2>&1; then "
f" if [ ! -e /mnt/user-data ] || [ -L /mnt/user-data ]; then "
f" sudo -n ln -sfn {home} /mnt/user-data 2>/dev/null || true; "
f" fi; "
f"fi; "
f"echo BOOTSTRAP_OK"
)
def _import_client() -> type[Client]:
"""Import Tenki's ``Client`` lazily.
Kept out of module import so the harness (and every other provider) installs
without Tenki; the dependency is only needed once this provider is selected.
"""
try:
from tenki_sandbox import Client
except ImportError as e: # pragma: no cover - depends on the optional dependency
raise ImportError("TenkiSandboxProvider requires the optional 'tenki-sandbox' dependency. Install it with: pip install 'deerflow-harness[tenki]' or pip install tenki-sandbox.") from e
return Client
class TenkiSandboxProvider(WarmPoolLifecycleMixin[TenkiSandbox], SandboxProvider):
"""Run each DeerFlow sandbox as a Tenki cloud microVM."""
uses_thread_data_mounts = False
needs_upload_permission_adjustment = True
_idle_checker_thread_name = "tenki-idle-reaper"
@staticmethod
def _sandbox_id(thread_id: str, user_id: str) -> str:
"""Deterministic sandbox ID from user/thread scope.
Includes user_id so a sandbox created for one user's bucket cannot be
reclaimed by another user's thread with the same thread_id. The warm
pool is keyed by this id alone (``_reclaim_warm_pool`` looks it up
directly, with no full-seed fallback), so on a hosted multi-tenant
gateway a hash collision would let one user reclaim another's parked
sandbox. 64 bits, like community/e2b_sandbox, keeps that negligible.
"""
return hashlib.sha256(f"{user_id}:{thread_id}".encode()).hexdigest()[:16]
# ── Provider lifecycle ───────────────────────────────────────────────
def __init__(self) -> None:
self._lock = threading.Lock()
self._sandboxes: dict[str, TenkiSandbox] = {}
self._thread_sandboxes: dict[tuple[str, str], str] = {}
self._warm_pool: dict[str, tuple[TenkiSandbox, float]] = {}
self._acquire_locks: dict[str, threading.Lock] = {}
self._idle_checker_stop = threading.Event()
self._idle_checker_thread: threading.Thread | None = None
self._shutdown_called = False
self._client: Client | None = None
self._config = self._load_config()
atexit.register(self.shutdown)
self._start_idle_checker()
def _load_config(self) -> dict[str, Any]:
sandbox_config = get_app_config().sandbox
def _opt(name: str, default: Any = None) -> Any:
return getattr(sandbox_config, name, default)
api_key = _opt("api_key")
replicas = _opt("replicas")
idle_timeout = _opt("idle_timeout")
max_duration = _opt("max_duration")
environment = dict(_opt("environment") or {})
# Fail fast on a misconfigured key (e.g. "bad-key"): the per-call env goes
# through the same POSIX-name check in execute_command, but this static
# config env is merged into every command and would otherwise only surface
# as a confusing SDK error at create/exec time.
_validate_extra_env(environment)
return {
"max_duration": float(max_duration if max_duration is not None else DEFAULT_MAX_DURATION),
# Off by default (the SDK default). Warm-pool sandboxes stay running
# between turns, so host pinning only matters to deployments that
# also pause/resume; expose it rather than decide for them.
"sticky": bool(_opt("sticky", False)),
"api_key": api_key, # None → SDK falls back to TENKI_API_KEY / TENKI_AUTH_TOKEN
"base_url": _opt("base_url"),
"image": _opt("image"), # None → Tenki account default base image
"home_dir": _opt("home_dir") or DEFAULT_TENKI_HOME_DIR,
"project_id": _opt("project_id"),
"workspace_id": _opt("workspace_id"),
"cpu_cores": _opt("cpu_cores"),
"memory_mb": _opt("memory_mb"),
"environment": environment,
"replicas": replicas if replicas is not None else self.DEFAULT_REPLICAS,
"idle_timeout": idle_timeout if idle_timeout is not None else self.DEFAULT_IDLE_TIMEOUT,
}
def _get_client(self) -> Client:
with self._lock:
if self._client is not None:
return self._client
client_cls = _import_client()
client = client_cls(auth_token=self._config["api_key"], base_url=self._config["base_url"])
with self._lock:
if self._client is None:
self._client = client
return self._client
def _resolve_scope(self) -> tuple[str | None, str | None]:
"""Return (project_id, workspace_id), auto-selecting when unambiguous.
Tenki's ``create`` needs a project scope. When the caller didn't set one
in config, pick it if the account has exactly one workspace and project;
otherwise raise with the choices so the operator can set ``project_id``.
"""
project_id = self._config["project_id"]
workspace_id = self._config["workspace_id"]
if project_id is not None:
return project_id, workspace_id
identity = self._get_client().who_am_i()
workspaces = list(identity.workspaces or [])
if workspace_id is not None:
workspaces = [w for w in workspaces if w.id == workspace_id]
workspace = self._require_single(workspaces, "workspace", "workspace_id")
project = self._require_single(list(workspace.projects or []), "project", "project_id")
return project.id, workspace.id
@staticmethod
def _require_single(items: list[Any], kind: str, param: str) -> Any:
if len(items) != 1:
raise ValueError(f"Could not auto-select a Tenki {kind}; set sandbox.{param} in config.yaml. Found: {[(getattr(i, 'id', None), getattr(i, 'name', None)) for i in items]}")
return items[0]
@staticmethod
def _thread_key(thread_id: str, user_id: str | None) -> tuple[str, str]:
return (user_id or "", thread_id)
@classmethod
def _sandbox_name(cls, sandbox_id: str) -> str:
return f"{_SANDBOX_NAME_PREFIX}{sandbox_id}"
def _lock_for_sandbox(self, sandbox_id: str) -> threading.Lock:
with self._lock:
lock = self._acquire_locks.get(sandbox_id)
if lock is None:
lock = threading.Lock()
self._acquire_locks[sandbox_id] = lock
return lock
def _start_idle_checker(self) -> None:
"""Start idle cleanup when enabled; idle_timeout=0 keeps it disabled."""
if self._config["idle_timeout"] <= 0:
return
super()._start_idle_checker()
def _active_count_locked(self) -> int:
return len(self._sandboxes)
def _destroy_warm_entry(self, sandbox_id: str, entry: TenkiSandbox, *, reason: str) -> None:
self._close_quietly(entry, context=f"warm pool, reason={reason}")
def _invalidate_sandbox(self, sandbox_id: str, reason: str) -> None:
"""Destroy and deregister a sandbox after a terminal command-path failure."""
to_close: TenkiSandbox | None = None
with self._lock:
active = self._sandboxes.pop(sandbox_id, None)
warm_entry = self._warm_pool.pop(sandbox_id, None)
for key in [k for k, sid in self._thread_sandboxes.items() if sid == sandbox_id]:
self._thread_sandboxes.pop(key, None)
to_close = active or (warm_entry[0] if warm_entry is not None else None)
if to_close is None:
logger.warning("Tenki sandbox %s failed terminally but was not tracked: %s", sandbox_id, reason)
return
logger.warning("Invalidating Tenki sandbox %s after terminal failure: %s", sandbox_id, reason)
self._close_quietly(to_close, context="terminal failure")
# ── Acquire / release ────────────────────────────────────────────────
def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str:
if thread_id is None:
sandbox_id = str(uuid.uuid4())[:8]
sandbox = self._create_sandbox(sandbox_id)
with self._lock:
self._sandboxes[sandbox.id] = sandbox
return sandbox.id
key = self._thread_key(thread_id, user_id)
sandbox_id = self._sandbox_id(thread_id, user_id or "")
acquire_lock = self._lock_for_sandbox(sandbox_id)
with acquire_lock:
with self._lock:
existing = self._thread_sandboxes.get(key)
if existing is not None and existing in self._sandboxes:
return existing
reclaimed = self._reclaim_warm_pool(sandbox_id)
if reclaimed is not None:
with self._lock:
self._thread_sandboxes[key] = reclaimed
return reclaimed
sandbox = self._create_sandbox(sandbox_id)
with self._lock:
self._sandboxes[sandbox.id] = sandbox
self._thread_sandboxes[key] = sandbox.id
return sandbox.id
def _create_sandbox(self, sandbox_id: str) -> TenkiSandbox:
# Enforce replica soft cap: evict the oldest warm sandbox if active + warm
# are at capacity.
replicas, total = self._replica_count()
if total >= replicas:
evicted = self._evict_oldest_warm()
self._log_replicas_soft_cap(replicas, sandbox_id, evicted)
client = self._get_client()
project_id, workspace_id = self._resolve_scope()
create_kwargs: dict[str, Any] = {
"name": self._sandbox_name(sandbox_id),
"project_id": project_id,
"workspace_id": workspace_id,
"sticky": self._config["sticky"],
# Wait for readiness ourselves (below) instead of inside create():
# create(wait=True) raises with the session handle still local to the
# SDK, so a readiness failure would leak a running, billed microVM
# that this provider never sees and can never terminate.
"wait": False,
}
if self._config["max_duration"] > 0:
create_kwargs["max_duration"] = self._config["max_duration"]
for key in ("image", "cpu_cores", "memory_mb"):
if self._config[key] is not None:
create_kwargs[key] = self._config[key]
if self._config["environment"]:
create_kwargs["env"] = self._config["environment"]
remote = client.create(**create_kwargs)
try:
remote.wait_ready()
except Exception:
self._terminate_orphan(sandbox_id, remote)
raise
# Materialise DeerFlow's virtual path layout under the writable HOME.
# Best-effort: on failure the file APIs still work via the home remap.
try:
result = remote.exec("sh", "-lc", _bootstrap_script(self._config["home_dir"]), timeout=_BOOTSTRAP_TIMEOUT)
if result.exit_code not in (0, None) or "BOOTSTRAP_OK" not in (result.stdout_text or ""):
logger.warning(
"Tenki bootstrap for %s exited code=%s stderr=%s",
sandbox_id,
result.exit_code,
(result.stderr_text or "").strip(),
)
except Exception as e:
logger.warning("Tenki bootstrap for %s raised: %s", sandbox_id, e)
logger.info("Created Tenki sandbox %s (name=%s)", sandbox_id, self._sandbox_name(sandbox_id))
return TenkiSandbox(
sandbox_id,
remote,
default_env=self._config["environment"],
home_dir=self._config["home_dir"],
on_terminal_failure=self._invalidate_sandbox,
)
@staticmethod
def _terminate_orphan(sandbox_id: str, remote: Any) -> None:
"""Terminate a microVM that was created but never handed to the adapter."""
try:
remote.close()
logger.warning("Terminated Tenki sandbox %s after it failed to become ready", sandbox_id)
except Exception as e:
logger.error("Leaked Tenki sandbox %s (id=%s): could not terminate after readiness failure: %s", sandbox_id, getattr(remote, "id", "?"), e)
@staticmethod
def _close_quietly(sandbox: TenkiSandbox, *, context: str) -> None:
"""Close a sandbox where the caller has no way to act on a failure."""
try:
sandbox.close()
except Exception as e:
logger.warning("Error closing Tenki sandbox %s (%s): %s", sandbox.id, context, e)
def get(self, sandbox_id: str) -> Sandbox | None:
with self._lock:
return self._sandboxes.get(sandbox_id)
def release(self, sandbox_id: str) -> None:
"""Release a sandbox into the warm pool — the microVM stays running.
The sandbox moves from ``_sandboxes`` to ``_warm_pool`` and its
``_thread_sandboxes`` entries are cleared. It is NOT terminated unless
shutdown has already begun.
"""
close_sandbox: TenkiSandbox | None = None
with self._lock:
sandbox = self._sandboxes.pop(sandbox_id, None)
for key in [k for k, sid in self._thread_sandboxes.items() if sid == sandbox_id]:
self._thread_sandboxes.pop(key, None)
if sandbox is None:
return
if self._shutdown_called:
close_sandbox = sandbox
else:
self._warm_pool[sandbox_id] = (sandbox, time.time())
if close_sandbox is not None:
self._close_quietly(close_sandbox, context="released during shutdown")
logger.info("Closed released Tenki sandbox %s because shutdown is in progress", sandbox_id)
else:
logger.info("Released Tenki sandbox %s to warm pool (microVM still running)", sandbox_id)
def _reclaim_warm_pool(self, sandbox_id: str) -> str | None:
"""Reclaim a warm-pool sandbox by id after a liveness health check.
Returns sandbox_id on success, None if not present or the health check
fails (the dead entry is destroyed).
"""
with self._lock:
if sandbox_id not in self._warm_pool:
return None
sandbox, _ = self._warm_pool[sandbox_id]
try:
result = sandbox.execute_command("echo ok", timeout=10)
healthy = "ok" in result
except Exception as e:
logger.warning("Tenki warm-pool sandbox %s health check error: %s", sandbox_id, e)
healthy = False
if not healthy:
with self._lock:
warm_entry = self._warm_pool.pop(sandbox_id, None)
if warm_entry is not None:
self._destroy_warm_entry(sandbox_id, warm_entry[0], reason="health_check_failed")
return None
with self._lock:
warm_entry = self._warm_pool.pop(sandbox_id, None)
if warm_entry is None:
return None # Raced with another thread
self._sandboxes[sandbox_id] = warm_entry[0]
logger.info("Reclaimed warm-pool Tenki sandbox %s", sandbox_id)
return sandbox_id
def reset(self) -> None:
"""Park tracked sandboxes into this instance's warm pool for later cleanup.
``reset_sandbox_provider()`` drops the singleton and calls this so config
changes take effect on the next construction. Teardown belongs to
``shutdown()``; reset leaves running microVMs alive but visible to this
instance's idle reaper and atexit shutdown instead of orphaning them.
"""
with self._lock:
now = time.time()
for sandbox_id, sandbox in self._sandboxes.items():
self._warm_pool.setdefault(sandbox_id, (sandbox, now))
self._sandboxes.clear()
self._thread_sandboxes.clear()
self._acquire_locks.clear()
def shutdown(self) -> None:
with self._lock:
if self._shutdown_called:
return
self._shutdown_called = True
self._stop_idle_checker()
with self._lock:
active = list(self._sandboxes.values())
warm = [sandbox for sandbox, _ in self._warm_pool.values()]
self._sandboxes.clear()
self._warm_pool.clear()
self._thread_sandboxes.clear()
self._acquire_locks.clear()
for sandbox in active + warm:
self._close_quietly(sandbox, context="shutdown")

View File

@ -0,0 +1,464 @@
"""``TenkiSandbox`` — DeerFlow :class:`Sandbox` backed by a Tenki cloud sandbox.
Tenki's Python SDK (``tenki-sandbox``) is synchronous, so — unlike
``community/boxlite`` this adapter calls the SDK directly with no event-loop
bridge. File transport uses Tenki's native ``sandbox.fs`` API (``read_text`` /
``read_stream`` / ``write_stream`` / ``mkdir`` / ``stat``), which is binary-safe
and streams, so no base64/shell encoding is involved. Directory and content
*search* (``list_dir`` / ``glob`` / ``grep``) still shells out to ``find`` /
``grep`` the fs API is single-level and has no content search and is parsed
with the shared ``deerflow.sandbox.search`` helpers, the same approach as
``community/e2b_sandbox``. Those commands use only busybox-portable flags so any
Tenki base image works.
The Tenki SDK is not imported at module load (only its exception *class names*
are matched, as strings), so importing this package never requires
``tenki-sandbox`` to be installed it is needed only once the provider is
selected and a sandbox is actually created.
"""
from __future__ import annotations
import errno
import logging
import posixpath
import re
import shlex
import threading
from typing import TYPE_CHECKING, Any, TypeVar
from deerflow.config.paths import VIRTUAL_PATH_PREFIX
from deerflow.sandbox.sandbox import Sandbox, _validate_extra_env
from deerflow.sandbox.search import GrepMatch, path_matches, should_ignore_path, truncate_line
if TYPE_CHECKING:
from collections.abc import Callable, Iterator
from tenki_sandbox import Sandbox as TenkiClientSandbox
from tenki_sandbox.fs import SandboxFS
T = TypeVar("T")
logger = logging.getLogger(__name__)
_MAX_DOWNLOAD_SIZE = 100 * 1024 * 1024 # 100 MB
# Tenki sandboxes run as the unprivileged ``tenki`` user (HOME=/home/tenki) and
# ``/mnt`` is root-owned, so DeerFlow's ``/mnt/user-data`` virtual prefix is not
# writable directly. Like ``community/e2b_sandbox``, file ops are remapped under
# this home dir (the provider also best-effort symlinks /mnt/user-data → here so
# agent shell commands using the literal path still work).
DEFAULT_TENKI_HOME_DIR = "/home/tenki"
# Frame size for fs.write_stream uploads.
_STREAM_CHUNK = 1024 * 1024
# Tenki SDK exception *class names* that mean the remote session is gone for
# good — matched as strings so this module imports without ``tenki-sandbox``.
# A terminated/not-found/closed session is unrecoverable; the provider drops it
# and rebuilds on the next call. This is only the named-error half of the rule:
# _is_terminal_failure ALSO treats the builtin ConnectionError / BrokenPipeError
# / EOFError as terminal via isinstance, so a transport reset evicts the sandbox
# and cold-starts the next acquire too. That is a deliberate fail-safe (a reset
# often means the microVM is gone); the cost is churning a warm sandbox on a
# one-off flaky-network blip.
_TERMINAL_ERROR_NAMES = frozenset(
{
"SessionTerminatedError",
"SessionNotFoundError",
"InvalidStateError",
"StreamClosedError",
}
)
class TenkiSandbox(Sandbox):
"""DeerFlow Sandbox adapter that delegates to a live Tenki cloud sandbox.
Args:
id: DeerFlow-side sandbox id (the provider's cache key).
sandbox: A live, started ``tenki_sandbox.Sandbox``. The provider owns
its lifecycle; this adapter terminates it on :meth:`close`.
default_env: Static environment merged into every command, overridden
per-call by the ``env`` passed to :meth:`execute_command`
(request-scoped secrets).
home_dir: Writable directory that backs the ``VIRTUAL_PATH_PREFIX``
(``/mnt/user-data``) prefix inside the sandbox. Defaults to
:data:`DEFAULT_TENKI_HOME_DIR`.
on_terminal_failure: Optional callback ``(sandbox_id, reason)`` invoked
when an operation fails with a terminal Tenki error, so the provider
can evict the dead sandbox.
"""
def __init__(
self,
id: str,
sandbox: TenkiClientSandbox,
*,
default_env: dict[str, str] | None = None,
home_dir: str = DEFAULT_TENKI_HOME_DIR,
on_terminal_failure: Callable[[str, str], None] | None = None,
) -> None:
super().__init__(id)
self._sandbox = sandbox
self._default_env = dict(default_env or {})
self._home_dir = home_dir.rstrip("/") or "/"
self._on_terminal_failure = on_terminal_failure
self._lock = threading.Lock()
# Serialises the append read-modify-write across its three fs ops. A
# lock distinct from _lock, so it can wrap the whole sequence without the
# per-op eviction callback (which reaches back into the provider) ever
# running under it.
self._write_lock = threading.Lock()
self._closed = False
@property
def is_closed(self) -> bool:
with self._lock:
return self._closed
@staticmethod
def _is_terminal_failure(error: Exception) -> bool:
if isinstance(error, (BrokenPipeError, ConnectionError, EOFError)):
return True
return type(error).__name__ in _TERMINAL_ERROR_NAMES
def close(self) -> None:
"""Terminate the underlying Tenki session (idempotent).
The microVM is terminated *first*; the adapter is only marked closed once
the session is actually gone, so a failed termination stays retryable
instead of silently leaking a running (billed) sandbox. A terminal
session error means it is already gone, which counts as closed; anything
else is raised so the caller can retry or alert.
"""
with self._lock:
if self._closed:
return
sandbox = self._sandbox
try:
sandbox.close()
except Exception as e:
if not self._is_terminal_failure(e):
logger.error("Error terminating Tenki sandbox %s: %s", self.id, e)
raise
logger.info("Tenki sandbox %s was already gone at close: %s", self.id, e)
with self._lock:
self._closed = True
# ── bridge helpers ──────────────────────────────────────────────────
def _note_failure(self, error: Exception) -> None:
"""Evict this sandbox when an operation failed with a terminal error."""
if self._on_terminal_failure is None or not self._is_terminal_failure(error):
return
try:
self._on_terminal_failure(self.id, str(error))
except Exception:
logger.exception("Terminal Tenki failure callback errored for %s", self.id)
def _fs_op(self, op: Callable[[SandboxFS], T]) -> T:
"""Run a native ``sandbox.fs`` call, evicting the sandbox on terminal errors.
The lock is held across ``op`` (not just the fs lookup) so concurrent
calls on the same sandbox serialise: the Tenki SDK shares one connection
per instance, like community/e2b_sandbox. ``_note_failure`` runs *after*
the lock is released it reaches back into the provider, which locks in
the opposite order (provider then sandbox), so holding both at once could
deadlock.
"""
with self._lock:
if self._closed:
raise RuntimeError("sandbox has been closed")
fs = self._sandbox.fs
try:
return op(fs)
except Exception as e:
failure = e
self._note_failure(failure)
raise failure
def _exec(self, *argv: str, env: dict[str, str] | None = None, timeout: float | None = None) -> Any:
# No forced cwd: commands run in the sandbox default working directory
# (like community/e2b_sandbox and community/boxlite); file ops address
# absolute, home-remapped paths, so cwd is irrelevant to them.
#
# No auto-retry: exec is not idempotent (the command may have run
# server-side before a transport ack dropped), so re-running it risks
# double side effects. Like boxlite, a transient error is surfaced to the
# caller (returned as text by execute_command); a terminal session error
# additionally evicts the sandbox so the next acquire rebuilds it.
with self._lock:
if self._closed:
raise RuntimeError("sandbox has been closed")
sandbox = self._sandbox
try:
return sandbox.exec(*argv, env=env, timeout=timeout)
except Exception as e:
self._note_failure(e)
raise
def _sh(self, script: str, env: dict[str, str] | None = None, timeout: float | None = None) -> Any:
return self._exec("sh", "-lc", script, env=env, timeout=timeout)
# ── path safety (mirrors community/e2b_sandbox) ──────────────────────
@staticmethod
def _guard_traversal(path: str) -> str:
if not path:
raise ValueError("path must be a non-empty string")
normalized = path.replace("\\", "/")
for segment in normalized.split("/"):
if segment == "..":
raise PermissionError(f"Access denied: path traversal detected in '{path}'")
return normalized
def _resolve_path(self, path: str) -> str:
"""Map DeerFlow virtual paths into the writable sandbox home dir.
``VIRTUAL_PATH_PREFIX`` (``/mnt/user-data``) is rewritten under
:attr:`_home_dir`; other absolute paths pass through so the sandbox can
reach system directories when needed. Traversal is always rejected.
"""
normalized = self._guard_traversal(path)
if normalized == VIRTUAL_PATH_PREFIX or normalized.startswith(f"{VIRTUAL_PATH_PREFIX}/"):
tail = normalized[len(VIRTUAL_PATH_PREFIX) :].lstrip("/")
return f"{self._home_dir}/{tail}".rstrip("/") if tail else self._home_dir
return normalized
def _virtual_path(self, resolved: str) -> str:
"""Inverse of :meth:`_resolve_path` — the form callers gave us.
Everything that *returns* paths (``list_dir``/``glob``/``grep``) reports
them under ``VIRTUAL_PATH_PREFIX``, not the sandbox-internal home dir, so
results can be fed straight back into the other file APIs.
"""
if resolved == self._home_dir:
return VIRTUAL_PATH_PREFIX
if resolved.startswith(f"{self._home_dir}/"):
return f"{VIRTUAL_PATH_PREFIX}/{resolved[len(self._home_dir) :].lstrip('/')}"
return resolved
# ── command execution ───────────────────────────────────────────────
def execute_command(
self,
command: str,
env: dict[str, str] | None = None,
timeout: float | None = None,
) -> str:
"""Run ``command`` through a shell in the Tenki sandbox and return output.
DeerFlow passes a bash command *string*; it runs through ``sh -lc``.
Per-call ``env`` is layered over the static config environment and
scoped to this command only (request-scoped secrets, issue #3861).
"""
_validate_extra_env(env) # POSIX env-var key rule; raises ValueError on a bad key
if self.is_closed:
return "Error: sandbox has been closed"
merged_env = {**self._default_env, **(env or {})} or None
try:
result = self._sh(command, env=merged_env, timeout=timeout)
except Exception as e:
logger.error("Failed to execute command in Tenki sandbox %s: %s", self.id, e)
return f"Error: {e}"
stdout = result.stdout_text or ""
stderr = result.stderr_text or ""
if stdout and stderr:
output = f"{stdout}\n{stderr}"
else:
output = stdout or stderr
if result.exit_code not in (0, None) and not output:
output = f"Command exited with code {result.exit_code}"
return output if output else "(no output)"
# ── file operations ─────────────────────────────────────────────────
def read_file(self, path: str) -> str:
resolved = self._resolve_path(path)
try:
return self._fs_op(lambda fs: fs.read_text(resolved))
except Exception as e:
logger.error("read_file %s failed: %s", resolved, e)
return f"Error: {e}"
def write_file(self, path: str, content: str, append: bool = False) -> None:
self._write_bytes(self._resolve_path(path), content.encode("utf-8"), append=append)
def update_file(self, path: str, content: bytes) -> None:
self._write_bytes(self._resolve_path(path), content, append=False)
def _write_bytes(self, resolved: str, data: bytes, *, append: bool) -> None:
parent = posixpath.dirname(resolved)
if not append:
if parent:
self._fs_op(lambda fs: fs.mkdir(parent))
self._fs_op(lambda fs: fs.write_stream(resolved, _frames(data)))
return
# Tenki's write stream has no append mode (it starts at offset 0), so we
# read-modify-write like community/e2b_sandbox. The read and the write are
# separate fs ops, so two concurrent appends could both read the same
# pre-image and the second would clobber the first; _write_lock makes the
# whole sequence atomic.
with self._write_lock:
if parent:
self._fs_op(lambda fs: fs.mkdir(parent))
try:
data = self._fs_op(lambda fs: fs.read_bytes(resolved)) + data
except Exception as e:
if type(e).__name__ != "FileNotFoundError":
raise
self._fs_op(lambda fs: fs.write_stream(resolved, _frames(data)))
def download_file(self, path: str) -> bytes:
normalized = self._guard_traversal(path)
stripped = normalized.lstrip("/")
allowed = VIRTUAL_PATH_PREFIX.lstrip("/")
if stripped != allowed and not stripped.startswith(f"{allowed}/"):
raise PermissionError(f"Access denied: path must be under '{VIRTUAL_PATH_PREFIX}': '{path}'")
resolved = self._resolve_path(path)
with self._lock:
if self._closed:
raise RuntimeError("sandbox has been closed")
fs = self._sandbox.fs
# Deliberate: the lock is dropped before streaming, unlike _fs_op which
# holds it across its op. _fs_op's serialization guards short, bounded
# calls; a download can be up to _MAX_DOWNLOAD_SIZE (100 MB), and holding
# the instance lock across it would block every other tool on this
# sandbox for the whole transfer. The Tenki read stream is safe to run
# alongside other ops (the SDK multiplexes over its connection), so we
# accept the interleave here for latency and still evict on a terminal
# transport error via _note_failure below.
#
# The cap is enforced on bytes actually received, so a file that grows
# mid-transfer still can't exceed it (a stat-then-read check could).
chunks: list[bytes] = []
total = 0
try:
for chunk in fs.read_stream(resolved):
total += len(chunk)
if total > _MAX_DOWNLOAD_SIZE:
raise OSError(errno.EFBIG, f"File exceeds maximum download size of {_MAX_DOWNLOAD_SIZE} bytes", path)
chunks.append(chunk)
except OSError as e:
# Our own EFBIG size-cap is not a session death — let it pass through
# without evicting. Every other OSError is a real transport failure:
# ConnectionError / BrokenPipeError / EOFError are OSError subclasses
# that _is_terminal_failure treats as terminal, so they must route
# through _note_failure like _fs_op/_exec do. Without this, a session
# that dies mid-download is never evicted and the agent keeps hitting
# OSErrors until some other op happens to reap it.
if e.errno == errno.EFBIG:
raise
self._note_failure(e)
raise
except Exception as e:
self._note_failure(e)
raise OSError(f"cannot read '{path}' from sandbox: {e}") from e
return b"".join(chunks)
def list_dir(self, path: str, max_depth: int = 2) -> list[str]:
resolved = self._resolve_path(path)
r = self._sh(f"find {shlex.quote(resolved)} -maxdepth {int(max_depth)} \\( -type f -o -type d \\) 2>/dev/null | head -500")
return [self._virtual_path(line.strip()) for line in (r.stdout_text or "").splitlines() if line.strip()]
def glob(
self,
path: str,
pattern: str,
*,
include_dirs: bool = False,
max_results: int = 200,
) -> tuple[list[str], bool]:
resolved = self._resolve_path(path)
types = ("f", "d") if include_dirs else ("f",)
type_expr = " -o ".join(f"-type {t}" for t in types)
hard_limit = max(max_results * 4, max_results + 50)
r = self._sh(f"find {shlex.quote(resolved)} \\( {type_expr} \\) -print 2>/dev/null | head -{hard_limit}")
matches: list[str] = []
root = resolved.rstrip("/") or "/"
root_prefix = root if root == "/" else f"{root}/"
for entry in (r.stdout_text or "").splitlines():
entry = entry.strip()
if not entry or (entry != root and not entry.startswith(root_prefix)):
continue
if should_ignore_path(entry):
continue
rel_path = entry[len(root) :].lstrip("/")
if not rel_path:
continue
if path_matches(pattern, rel_path):
matches.append(self._virtual_path(entry))
if len(matches) >= max_results:
return matches, True
return matches, False
def grep(
self,
path: str,
pattern: str,
*,
glob: str | None = None,
literal: bool = False,
case_sensitive: bool = False,
max_results: int = 100,
) -> tuple[list[GrepMatch], bool]:
# Validate a regex pattern at the boundary (grep uses POSIX ERE, but this
# catches gross errors); a literal needs none. grep receives the RAW
# pattern: -F matches it literally, -E as a regex.
if not literal:
re.compile(pattern, 0 if case_sensitive else re.IGNORECASE)
resolved = self._resolve_path(path)
# busybox+GNU-portable flags: -r recursive, -H always print the filename
# (without it, grep -r on a path that resolves to a single file prints
# "line:text" and the file:line:text unpack below drops every match), -n
# line numbers, -I skip binary, -E/-F regex vs fixed. --include and -m are
# omitted for busybox portability; glob-scoping and the result cap are
# applied in Python below.
flags = ["-r", "-H", "-n", "-I"]
if not case_sensitive:
flags.append("-i")
flags.append("-F" if literal else "-E")
total_cap = max(max_results * 4, max_results + 50)
cmd = "grep " + " ".join(flags) + f" -e {shlex.quote(pattern)} {shlex.quote(resolved)} 2>/dev/null | head -{total_cap}"
r = self._sh(cmd)
root = resolved.rstrip("/") or "/"
root_prefix = root if root == "/" else f"{root}/"
matches: list[GrepMatch] = []
truncated = False
for raw in (r.stdout_text or "").splitlines():
try:
file_path, line_no_str, line_text = raw.split(":", 2)
except ValueError:
continue
try:
line_number = int(line_no_str)
except ValueError:
continue
if should_ignore_path(file_path):
continue
if glob is not None:
# Match the caller's real directory scope: a pattern like
# "src/*.js" must not broaden to every *.js in the tree. Same
# helper, same relative-to-root semantics as glob() above.
if file_path != root and not file_path.startswith(root_prefix):
continue
rel_path = file_path[len(root) :].lstrip("/")
if not rel_path or not path_matches(glob, rel_path):
continue
matches.append(GrepMatch(path=self._virtual_path(file_path), line_number=line_number, line=truncate_line(line_text)))
if len(matches) >= max_results:
truncated = True
break
return matches, truncated
def _frames(data: bytes) -> Iterator[bytes]:
"""Slice ``data`` into upload frames for ``fs.write_stream``."""
for i in range(0, len(data), _STREAM_CHUNK):
yield data[i : i + _STREAM_CHUNK]

View File

@ -17,6 +17,7 @@ from deerflow.config.authorization_config import AuthorizationConfig, load_autho
from deerflow.config.channel_connections_config import ChannelConnectionsConfig
from deerflow.config.checkpointer_config import CheckpointerConfig, load_checkpointer_config_from_dict
from deerflow.config.database_config import DatabaseConfig
from deerflow.config.dedupe_storage_config import DedupeStorageConfig
from deerflow.config.extensions_config import ExtensionsConfig
from deerflow.config.file_signature import ConfigSignature as _ConfigSignature
from deerflow.config.file_signature import get_config_signature as _get_config_signature
@ -296,6 +297,13 @@ class AppConfig(BaseModel):
field_doc="Run ownership and lease configuration for multi-worker deployments.",
),
)
dedupe_storage: DedupeStorageConfig = Field(
default_factory=DedupeStorageConfig,
description=format_field_description(
"dedupe_storage",
field_doc="Inbound webhook dedupe storage backend (memory / postgres / auto) for cross-pod redelivery dedup. See issue #4120.",
),
)
# Name -> config lookup tables, (re)built after validation by
# ``_build_name_indexes``. They make ``get_model_config`` / ``get_tool_config``

View File

@ -53,6 +53,16 @@ class DatabaseConfig(BaseModel):
"processes sharing one checkpoint database must use the same value."
),
)
checkpoint_delta_snapshot_frequency: int = Field(
default=1000,
ge=1,
description=(
"DeltaChannel snapshot frequency used in checkpoint 'delta' mode: "
"every N message-channel writes the saver stores a full snapshot "
"instead of a delta sentinel. Restart-required, and all processes "
"sharing one checkpoint database must use the same value."
),
)
sqlite_dir: str = Field(
default=".deer-flow/data",
description=("Directory for the SQLite database file. Both checkpointer and application data share {sqlite_dir}/deerflow.db."),

View File

@ -0,0 +1,41 @@
"""Configuration for inbound webhook dedupe storage.
Controls where the ChannelManager's inbound dedupe state lives. See issue #4120
(cross-pod webhook dedupe). The default ``auto`` reuses the Postgres application
database whenever database.backend='postgres', otherwise an in-process memory
store. ``memory`` is per-pod and not shared across replicas; ``postgres`` shares
state across pods.
"""
from __future__ import annotations
from enum import StrEnum
from pydantic import BaseModel, Field
from deerflow.config.reload_boundary import format_field_description
class DedupeStorageBackend(StrEnum):
AUTO = "auto"
MEMORY = "memory"
POSTGRES = "postgres"
class DedupeStorageConfig(BaseModel):
"""Where inbound webhook dedupe state lives."""
backend: DedupeStorageBackend = Field(
default=DedupeStorageBackend.AUTO,
description=format_field_description(
"dedupe_storage",
field_doc=(
"Storage backend for inbound webhook dedupe state. "
"'auto' uses the Postgres application database whenever database.backend='postgres', "
"otherwise an in-process memory store (single-pod). "
"'memory' forces the in-process store (per-pod; not shared across replicas). "
"'postgres' shares dedupe state across pods via the application database. "
"See issue #4120."
),
),
)

View File

@ -73,6 +73,10 @@ STARTUP_ONLY_FIELDS: dict[str, str] = {
"RunOwnershipConfig is captured once into RunManager at langgraph_runtime() startup; the lease heartbeat background task is created and "
"started there, and heartbeat_enabled / lease_seconds / grace_seconds are not re-read on config.yaml edits."
),
"dedupe_storage": (
"make_inbound_dedupe_store() resolves the inbound dedupe store once when ChannelService is constructed at startup; the store "
"(in-process memory or shared Postgres) is captured onto ChannelManager and is not rebuilt on config.yaml edits."
),
}

View File

@ -88,8 +88,10 @@ class SandboxConfig(BaseModel):
port: Base port for sandbox containers (default: 8080)
container_prefix: Prefix for container names (default: deer-flow-sandbox)
mounts: List of volume mounts to share directories with the container
ownership: Cross-instance container ownership store (memory | redis). Multi-instance
deployments sharing a container backend need redis; see SandboxOwnershipConfig.
AioSandboxProvider and E2BSandboxProvider shared options:
ownership: Cross-instance sandbox ownership store (memory | redis). Multi-instance
deployments sharing a sandbox backend need redis; see SandboxOwnershipConfig.
"""
use: str = Field(
@ -143,8 +145,8 @@ class SandboxConfig(BaseModel):
ownership: SandboxOwnershipConfig | None = Field(
default=None,
description=(
"AioSandboxProvider-only: where cross-instance container ownership is tracked (#4206). Omitted = memory (single-instance). "
"Multi-worker / load-balanced gateways sharing one container backend must set type: redis, or peers will adopt and idle-destroy each other's live sandboxes."
"AioSandboxProvider/E2BSandboxProvider: where cross-instance sandbox ownership is tracked (#4206, #4341). Omitted = memory (single-instance). "
"Multi-worker / load-balanced gateways sharing one sandbox backend must set type: redis, or peers can adopt and destroy each other's live sandboxes."
),
)
mounts: list[VolumeMountConfig] = Field(

View File

@ -0,0 +1,61 @@
"""Cross-pod inbound webhook dedupe table (issue #4120).
Revision ID: 0009_webhook_dedupe
Revises: 0008_thread_operation_kind
Create Date: 2026-07-15
"""
from __future__ import annotations
from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
revision: str = "0009_webhook_dedupe"
down_revision: str | Sequence[str] | None = "0008_thread_operation_kind"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None
def upgrade() -> None:
bind = op.get_bind()
inspector = sa.inspect(bind)
if inspector.has_table("webhook_deliveries"):
# Idempotent: a DB whose full-metadata create_all already provisioned
# the table (e.g. a fresh DB, or a legacy test seed) must not have it
# re-created here.
return
op.create_table(
"webhook_deliveries",
sa.Column("channel", sa.String(length=64), nullable=False),
sa.Column("workspace_id", sa.String(length=512), nullable=False),
sa.Column("chat_id", sa.String(length=512), nullable=False),
sa.Column("message_id", sa.String(length=1024), nullable=False),
sa.Column(
"first_seen",
sa.DateTime(timezone=True),
server_default=sa.func.now(),
nullable=False,
),
# Composite PK mirrors ChannelManager._inbound_dedupe_key exactly. A
# single joined surrogate key is deliberately avoided: the components
# can contain characters (e.g. NUL) that are illegal in a Postgres TEXT
# column, and a joined key would also risk length overflow.
sa.PrimaryKeyConstraint(
"channel",
"workspace_id",
"chat_id",
"message_id",
name="pk_webhook_deliveries",
),
)
op.create_index("ix_webhook_deliveries_first_seen", "webhook_deliveries", ["first_seen"], unique=False)
def downgrade() -> None:
bind = op.get_bind()
inspector = sa.inspect(bind)
if inspector.has_table("webhook_deliveries"):
op.drop_index("ix_webhook_deliveries_first_seen", table_name="webhook_deliveries")
op.drop_table("webhook_deliveries")

View File

@ -1,7 +1,7 @@
"""feedback tags.
Revision ID: 0009_feedback_tags
Revises: 0008_thread_operation_kind
Revision ID: 0010_feedback_tags
Revises: 0009_webhook_dedupe
Create Date: 2026-07-23
"""
@ -12,8 +12,8 @@ from collections.abc import Sequence
import sqlalchemy as sa
from alembic import op
revision: str = "0009_feedback_tags"
down_revision: str | Sequence[str] | None = "0008_thread_operation_kind"
revision: str = "0010_feedback_tags"
down_revision: str | Sequence[str] | None = "0009_webhook_dedupe"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None

View File

@ -28,6 +28,7 @@ from deerflow.persistence.scheduled_task_runs.model import ScheduledTaskRunRow
from deerflow.persistence.scheduled_tasks.model import ScheduledTaskRow
from deerflow.persistence.thread_meta.model import ThreadMetaRow
from deerflow.persistence.user.model import UserRow
from deerflow.persistence.webhook_delivery.model import WebhookDeliveryRow
__all__ = [
"AgentRow",
@ -42,4 +43,5 @@ __all__ = [
"ScheduledTaskRunRow",
"ThreadMetaRow",
"UserRow",
"WebhookDeliveryRow",
]

View File

@ -191,6 +191,28 @@ class RunRepository(RunStore):
result = await session.execute(stmt)
return {value for value in result.scalars() if isinstance(value, str) and value}
async def list_edit_regenerate_runs(
self,
thread_id,
*,
user_id: str | None | _AutoSentinel = AUTO,
):
resolved_user_id = resolve_user_id(user_id, method_name="RunRepository.list_edit_regenerate_runs")
replay_kind = RunRow.metadata_json["replay_kind"].as_string()
source = RunRow.metadata_json["regenerate_from_run_id"].as_string()
stmt = select(RunRow).where(
RunRow.thread_id == thread_id,
replay_kind == "edit",
source.is_not(None),
source != "",
)
if resolved_user_id is not None:
stmt = stmt.where(RunRow.user_id == resolved_user_id)
stmt = stmt.order_by(RunRow.created_at.asc())
async with self._sf() as session:
result = await session.execute(stmt)
return [self._row_to_dict(row) for row in result.scalars()]
async def get_many_by_thread(
self,
thread_id,

View File

@ -6,7 +6,7 @@ import logging
from datetime import UTC, datetime
from typing import Any
from sqlalchemy import case, select, update
from sqlalchemy import case, select, text, update
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from sqlalchemy.orm.attributes import flag_modified
@ -204,16 +204,27 @@ class ThreadMetaRepository(ThreadMetaStore):
) -> None:
"""Merge ``metadata`` into ``metadata_json``.
Read-modify-write inside a single session/transaction so concurrent
callers see consistent state. No-op if the row does not exist or
the user_id check fails.
The row is locked before the read-modify-write merge so concurrent
callers cannot replace each other's keys. SQLite acquires its write
transaction before reading; databases with row-level locking use
``SELECT ... FOR UPDATE``. No-op if the row does not exist or the
user_id check fails.
``touch`` refreshes ``updated_at`` (default); pass ``touch=False`` to
preserve recency ordering for metadata-only changes such as pin/unpin.
"""
resolved_user_id = resolve_user_id(user_id, method_name="ThreadMetaRepository.update_metadata")
async with self._sf() as session:
row = await session.get(ThreadMetaRow, thread_id)
if session.get_bind().dialect.name == "sqlite":
# A deferred SQLite transaction does not reserve the writer
# until the UPDATE, which is too late for a read-modify-write
# merge. BEGIN IMMEDIATE serializes writers before the read,
# including writers in other processes using the same file.
await session.execute(text("BEGIN IMMEDIATE"))
row = await session.get(ThreadMetaRow, thread_id)
else:
result = await session.execute(select(ThreadMetaRow).where(ThreadMetaRow.thread_id == thread_id).with_for_update())
row = result.scalar_one_or_none()
if row is None:
return
if resolved_user_id is not None and row.user_id != resolved_user_id:

View File

@ -0,0 +1,5 @@
"""Shared inbound webhook dedupe table (issue #4120).
ORM model only; row writes/reads go through raw SQL in
``app.channels.dedupe_store.PostgresInboundDedupeStore``.
"""

View File

@ -0,0 +1,37 @@
"""ORM model for the shared inbound webhook dedupe table (issue #4120).
A single row records that a particular inbound webhook (identified by the
``_inbound_dedupe_key`` 4-tuple) has already been dispatched, so a redelivery
routed to a different gateway pod is still dropped as a duplicate. Rows expire
via lazy cleanup (see ``PostgresInboundDedupeStore``), which deletes rows older
than ``INBOUND_DEDUPE_TTL_SECONDS`` using the ``first_seen`` column.
"""
from __future__ import annotations
from datetime import datetime
from sqlalchemy import DateTime, Index, PrimaryKeyConstraint, String, func
from sqlalchemy.orm import Mapped, mapped_column
from deerflow.persistence.base import Base
class WebhookDeliveryRow(Base):
__tablename__ = "webhook_deliveries"
# Composite primary key mirrors ChannelManager._inbound_dedupe_key exactly:
# (channel, workspace_id, chat_id, message_id). Using the four columns
# directly avoids any string-joined surrogate key — important because the
# components can contain characters (e.g. NUL) that are illegal in a single
# Postgres TEXT column, and keeps the ON CONFLICT target natural.
channel: Mapped[str] = mapped_column(String(64), nullable=False)
workspace_id: Mapped[str] = mapped_column(String(512), nullable=False)
chat_id: Mapped[str] = mapped_column(String(512), nullable=False)
message_id: Mapped[str] = mapped_column(String(1024), nullable=False)
first_seen: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False, server_default=func.now())
__table_args__ = (
PrimaryKeyConstraint("channel", "workspace_id", "chat_id", "message_id", name="pk_webhook_deliveries"),
Index("ix_webhook_deliveries_first_seen", "first_seen"),
)

View File

@ -21,6 +21,7 @@ from deerflow.utils.time import is_lease_expired
from deerflow.utils.time import now_iso as _now_iso
from .schemas import DisconnectMode, RunStatus, ThreadOperationKind
from .store.base import EditReplayVisibility
if TYPE_CHECKING:
from deerflow.config.run_ownership_config import RunOwnershipConfig
@ -706,6 +707,67 @@ class RunManager:
sources.add(source)
return sources
@staticmethod
def _record_status_value(record: RunRecord) -> str:
status = record.status
return status.value if isinstance(status, RunStatus) else str(status)
@staticmethod
def _compute_edit_replay_visibility(records: list[RunRecord]) -> EditReplayVisibility:
latest_attempt_by_source: dict[str, tuple[str, str]] = {}
failed_attempts: set[str] = set()
for record in sorted(records, key=lambda item: item.created_at):
metadata = record.metadata or {}
if metadata.get("replay_kind") != "edit":
continue
source = metadata.get("regenerate_from_run_id")
if not isinstance(source, str) or not source:
continue
status = RunManager._record_status_value(record)
latest_attempt_by_source[source] = (record.run_id, status)
if status in {RunStatus.error.value, RunStatus.timeout.value, RunStatus.interrupted.value}:
failed_attempts.add(record.run_id)
hidden_sources: set[str] = set()
for source, (_, status) in latest_attempt_by_source.items():
if status in {RunStatus.pending.value, RunStatus.running.value, RunStatus.success.value}:
hidden_sources.add(source)
return EditReplayVisibility(
hidden_source_run_ids=hidden_sources,
hidden_attempt_run_ids=failed_attempts,
)
async def list_edit_replay_visibility(
self,
thread_id: str,
*,
user_id: str | None | _AutoSentinel = AUTO,
) -> EditReplayVisibility:
"""Return run-id visibility rules for edit-and-rerun attempts.
Store rows cover reload/multi-worker history. Current-process records
override the same run ids because they may have newer terminal status
than the persisted snapshot visible when this query started.
"""
resolved_user_id = resolve_user_id(user_id, method_name="RunManager.list_edit_replay_visibility")
records_by_id: dict[str, RunRecord] = {}
if self._store is not None:
rows = await self._store.list_edit_regenerate_runs(thread_id, user_id=resolved_user_id)
for row in rows:
try:
record = self._record_from_store(row)
except Exception:
logger.warning("Failed to map edit replay run row for %s", row.get("run_id"), exc_info=True)
continue
records_by_id[record.run_id] = record
async with self._lock:
memory_records = [record for record in self._thread_records_locked(thread_id) if resolved_user_id is None or record.user_id == resolved_user_id]
for record in memory_records:
records_by_id[record.run_id] = record
return self._compute_edit_replay_visibility(list(records_by_id.values()))
async def try_start(self, run_id: str) -> RunStartOutcome:
"""Transition an uncancelled pending run to running before building the agent."""
async with self._lock:

View File

@ -11,9 +11,16 @@ When user_id is None, no user filtering is applied (single-user mode).
from __future__ import annotations
import abc
from dataclasses import dataclass, field
from typing import Any
@dataclass(frozen=True)
class EditReplayVisibility:
hidden_source_run_ids: set[str] = field(default_factory=set)
hidden_attempt_run_ids: set[str] = field(default_factory=set)
class RunStore(abc.ABC):
@abc.abstractmethod
async def put(
@ -69,6 +76,15 @@ class RunStore(abc.ABC):
"""
raise NotImplementedError
async def list_edit_regenerate_runs(
self,
thread_id: str,
*,
user_id: str | None = None,
) -> list[dict[str, Any]]:
"""Return all edit-regenerate attempt runs for one thread, oldest first."""
raise NotImplementedError
async def get_many_by_thread(
self,
thread_id: str,

View File

@ -105,6 +105,22 @@ class MemoryRunStore(RunStore):
sources.add(source)
return sources
async def list_edit_regenerate_runs(self, thread_id, *, user_id=None):
run_ids = self._runs_by_thread.get(thread_id) or ()
results = []
for run_id in run_ids:
run = self._runs.get(run_id)
if run is None:
continue
if user_id is not None and run.get("user_id") != user_id:
continue
metadata = run.get("metadata") or {}
source = metadata.get("regenerate_from_run_id")
if metadata.get("replay_kind") == "edit" and isinstance(source, str) and source:
results.append(run)
results.sort(key=lambda r: r["created_at"])
return results
async def get_many_by_thread(self, thread_id, run_ids, *, user_id=None):
thread_run_ids = self._runs_by_thread.get(thread_id) or ()
return {run_id: run for run_id in thread_run_ids if run_id in run_ids and (run := self._runs.get(run_id)) is not None and run.get("operation_kind", "run") == "run" and (user_id is None or run.get("user_id") == user_id)}

View File

@ -78,7 +78,7 @@ from deerflow.trace_context import (
)
from deerflow.tracing import inject_langfuse_metadata
from deerflow.utils.messages import message_to_text
from deerflow.workspace_changes import capture_workspace_snapshot, record_workspace_changes
from deerflow.workspace_changes import capture_workspace_snapshot, get_changed_output_paths, record_workspace_changes
from deerflow.workspace_changes.types import WorkspaceSnapshot
from .manager import RunManager, RunRecord, RunStartOutcome
@ -140,7 +140,7 @@ async def _persist_delivery_receipt(
except Exception:
if attempt == attempts - 1:
logger.warning(
"Failed to persist delivery receipt for run %s after %d attempts; preserving the real terminal status without a receipt",
"Failed to persist delivery receipt for run %s after %d attempts; applying terminal delivery semantics without a receipt",
run_id,
attempts,
exc_info=True,
@ -160,6 +160,68 @@ async def _persist_delivery_receipt(
return False # pragma: no cover - loop always returns
_DELIVERY_INCOMPLETE_ERROR = "Artifact delivery incomplete: no produced output artifact was presented"
_DELIVERY_RECEIPT_FAILED_ERROR = "Artifact delivery verification failed: terminal delivery receipt could not be persisted"
def _empty_delivery_content() -> dict[str, Any]:
return {"presented": 0, "paths": [], "by_tool": {}}
def _presented_path_covers_output(presented_path: str, produced_path: str) -> bool:
presented_path = presented_path.rstrip("/")
return bool(presented_path) and (produced_path == presented_path or produced_path.startswith(f"{presented_path}/"))
def _delivery_content_with_outputs(
content: dict[str, Any],
produced_paths: list[str],
) -> dict[str, Any]:
"""Attach a delivery verdict when this run created or modified outputs."""
if not produced_paths:
return content
presented_paths = content.get("by_tool", {}).get("present_files", [])
matched_paths = [produced_path for produced_path in produced_paths if any(_presented_path_covers_output(presented_path, produced_path) for presented_path in presented_paths)]
satisfied = bool(matched_paths)
return {
**content,
"verification": {
"source": "outputs_changed",
"requirement": "present_files_matches_produced_output",
},
"produced_paths": produced_paths,
"presented_paths": presented_paths,
"matched_paths": matched_paths,
"stage": "presented" if satisfied else ("mismatched" if presented_paths else "not_started"),
"satisfied": satisfied,
}
def _delivery_error(content: dict[str, Any]) -> str | None:
"""Return the terminal error when no changed output was presented."""
if not content.get("produced_paths") or content.get("satisfied") is True:
return None
return _DELIVERY_INCOMPLETE_ERROR
async def _produced_output_paths(
before: WorkspaceSnapshot | None,
*,
thread_id: str,
user_id: str | None,
) -> list[str]:
"""Detect regular output files created or modified by this run."""
if before is None:
return []
try:
after = await capture_workspace_snapshot(thread_id, user_id=user_id, include_text=False)
return get_changed_output_paths(before, after)
except Exception:
logger.warning("Could not detect produced output artifacts for run thread %s", thread_id, exc_info=True)
return []
# Keep this streaming policy separate from middleware write-authorization sets.
_LARGE_FILE_TOOL_NAMES = frozenset({"str_replace", "write_file"})
_LARGE_FILE_TOOL_BATCH_SIZE = 32
@ -459,6 +521,7 @@ async def run_agent(
workspace_changes_user_id: str | None = None
snapshot_capture_failed = False
llm_error_fallback_message: str | None = None
checkpoint_rollback_completed = False
# Message ids checkpointed *before* this run started. The stream loop uses
# this set to mask out ``deerflow_error_fallback`` markers that belong to
# earlier runs on the same thread — without it, one stale fallback in
@ -471,6 +534,8 @@ async def run_agent(
accessor: CheckpointStateAccessor | None = None
rollback_point: RollbackPoint | None = None
journal = None
delivery_content: dict[str, Any] | None = None
produced_output_paths: list[str] | None = None
# Journal construction moved ahead of preflight so every terminal run can
# emit a receipt. Completion persistence keeps its prior boundary: before
# #4272 the journal did not exist until preflight had succeeded, so early
@ -826,7 +891,7 @@ async def run_agent(
**terminal_status_kwargs,
)
try:
await _rollback_to_pre_run_checkpoint(
checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint(
accessor=accessor,
checkpointer=checkpointer,
thread_id=thread_id,
@ -848,6 +913,7 @@ async def run_agent(
if error_msg is None and journal is not None:
error_msg = journal.llm_error_fallback_message
error_msg = error_msg or "LLM provider failed after retries"
await _ensure_finalizing_before_edit_failure(run_manager, record)
await run_manager.set_status(
run_id,
RunStatus.error,
@ -871,9 +937,20 @@ async def run_agent(
# collects the most severe / first / all reasons) instead of each
# guard writing directly to the same key.
stop_reason = runtime_context.get("stop_reason") if runtime_context is not None else None
produced_output_paths = await _produced_output_paths(
pre_run_workspace_snapshot,
thread_id=thread_id,
user_id=workspace_changes_user_id,
)
delivery_content = _delivery_content_with_outputs(
journal.get_delivery_content() if journal is not None else _empty_delivery_content(),
produced_output_paths,
)
delivery_error = _delivery_error(delivery_content)
await run_manager.set_status(
run_id,
RunStatus.success,
RunStatus.error if delivery_error else RunStatus.success,
error=delivery_error,
stop_reason=stop_reason,
**terminal_status_kwargs,
)
@ -889,7 +966,7 @@ async def run_agent(
**terminal_status_kwargs,
)
try:
await _rollback_to_pre_run_checkpoint(
checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint(
accessor=accessor,
checkpointer=checkpointer,
thread_id=thread_id,
@ -901,6 +978,7 @@ async def run_agent(
except Exception:
logger.warning("Run %s cancellation rollback failed", run_id, exc_info=True)
else:
await _ensure_finalizing_before_edit_failure(run_manager, record)
await run_manager.set_status(
run_id,
RunStatus.interrupted,
@ -911,6 +989,7 @@ async def run_agent(
except Exception as exc:
error_msg = f"{exc}"
logger.exception("Run %s failed: %s", run_id, error_msg)
await _ensure_finalizing_before_edit_failure(run_manager, record)
await run_manager.set_status(
run_id,
RunStatus.error,
@ -933,6 +1012,30 @@ async def run_agent(
run_id,
)
if not record.ownership_lost and _is_edit_replay_run(record) and record.status != RunStatus.success:
if not record.finalizing:
await run_manager.set_finalizing(run_id, True)
try:
if not checkpoint_rollback_completed:
checkpoint_rollback_completed = await _rollback_to_pre_run_checkpoint(
accessor=accessor,
checkpointer=checkpointer,
thread_id=thread_id,
run_id=run_id,
rollback_point=rollback_point,
snapshot_capture_failed=snapshot_capture_failed,
)
if checkpoint_rollback_completed:
await _publish_restored_checkpoint_values(
bridge=bridge,
run_id=run_id,
accessor=accessor,
thread_id=thread_id,
)
logger.info("Run %s edit replay restored pre-run checkpoint %s", run_id, pre_run_checkpoint_id)
except Exception:
logger.warning("Run %s edit replay rollback failed", run_id, exc_info=True)
# Persist any subagent step events still buffered (#3779) — including on
# abort/exception paths, where the stream loop broke before its own flush.
if not record.ownership_lost and subagent_events is not None:
@ -961,12 +1064,27 @@ async def run_agent(
except Exception:
logger.warning("Failed to flush journal for run %s", run_id, exc_info=True)
await _persist_delivery_receipt(
if delivery_content is None:
if produced_output_paths is None:
produced_output_paths = await _produced_output_paths(
pre_run_workspace_snapshot,
thread_id=thread_id,
user_id=workspace_changes_user_id,
)
delivery_content = _delivery_content_with_outputs(journal.get_delivery_content(), produced_output_paths)
receipt_persisted = await _persist_delivery_receipt(
event_store,
thread_id=thread_id,
run_id=run_id,
content=journal.get_delivery_content(),
content=delivery_content,
)
if produced_output_paths and record.status == RunStatus.success and not receipt_persisted:
await run_manager.set_status(
run_id,
RunStatus.error,
error=_DELIVERY_RECEIPT_FAILED_ERROR,
persist=False,
)
if not record.ownership_lost and event_store is not None:
try:
@ -986,7 +1104,7 @@ async def run_agent(
except Exception:
logger.warning("Failed to persist run completion for %s (non-fatal)", run_id, exc_info=True)
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.interrupted:
if started and not record.ownership_lost and checkpointer is not None and record.status == RunStatus.interrupted and not _is_edit_replay_run(record):
try:
await run_manager.wait_for_prior_finalizing(thread_id, run_id)
if not await run_manager.has_later_started_run(thread_id, run_id):
@ -1389,6 +1507,31 @@ async def _prepare_goal_continuation_input(
return {"messages": [make_goal_continuation_message(updated_goal, evaluation)]}
def _is_edit_replay_run(record: RunRecord) -> bool:
metadata = record.metadata or {}
return metadata.get("replay_kind") == "edit"
async def _ensure_finalizing_before_edit_failure(run_manager: RunManager, record: RunRecord) -> None:
if _is_edit_replay_run(record) and not record.finalizing:
await run_manager.set_finalizing(record.run_id, True)
async def _publish_restored_checkpoint_values(
*,
bridge: StreamBridge,
run_id: str,
accessor: CheckpointStateAccessor | None,
thread_id: str,
) -> None:
if accessor is None:
return
snapshot = await accessor.aget({"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}})
values = getattr(snapshot, "values", None)
if isinstance(values, dict):
await bridge.publish(run_id, "values", serialize(values, mode="values"))
@dataclass(frozen=True)
class RollbackPoint:
"""Materialized pre-run state used to restore the thread after cancellation.
@ -1567,8 +1710,8 @@ async def _rollback_to_pre_run_checkpoint(
run_id: str,
rollback_point: RollbackPoint | None,
snapshot_capture_failed: bool,
) -> None:
"""Restore the complete pre-run state after a cancelled run.
) -> bool:
"""Restore the complete pre-run state and report whether it completed.
Full mode forks the captured pre-run checkpoint and overwrites messages;
all other channels inherit from that parent. Delta mode cannot safely fork
@ -1579,27 +1722,27 @@ async def _rollback_to_pre_run_checkpoint(
"""
if checkpointer is None:
logger.info("Run %s rollback requested but no checkpointer is configured", run_id)
return
return False
if snapshot_capture_failed:
logger.warning("Run %s rollback skipped: pre-run checkpoint capture failed", run_id)
return
return False
if rollback_point is None:
await _call_checkpointer_method(checkpointer, "adelete_thread", "delete_thread", thread_id)
logger.info("Run %s rollback reset thread %s to empty state", run_id, thread_id)
return
return True
configurable = rollback_point.config.get("configurable", {})
if not configurable.get("checkpoint_id"):
logger.warning("Run %s rollback skipped: pre-run checkpoint has no checkpoint id", run_id)
return
return False
if accessor is None:
# Unreachable in practice: a rollback point can only be captured
# through the bound accessor. Stay fail-closed.
logger.warning("Run %s rollback skipped: agent accessor unavailable", run_id)
return
return False
# Compile with the thread's effective schema so middleware-contributed
# channels survive (the base ThreadState fallback would silently drop
@ -1643,7 +1786,7 @@ async def _rollback_to_pre_run_checkpoint(
pending_writes = rollback_point.pending_writes
if not pending_writes:
return
return True
writes_by_task: dict[str, list[tuple[str, Any]]] = {}
for item in pending_writes:
@ -1663,6 +1806,7 @@ async def _rollback_to_pre_run_checkpoint(
writes,
task_id=task_id,
)
return True
def _new_checkpoint_marker() -> dict[str, str]:

View File

@ -1,8 +1,24 @@
from typing import Literal
from typing import Literal, Required, TypedDict
from langchain.tools import tool
class ClarificationFormField(TypedDict, total=False):
"""One form field definition for a structured clarification card.
The model-visible schema documents the item shape; runtime validation
still happens defensively in ``ClarificationMiddleware`` because the
middleware intercepts the call before tool execution.
"""
name: Required[str]
label: str
type: Literal["text", "textarea", "number", "select", "multi_select", "checkbox", "date"]
required: bool
options: list[str]
placeholder: str
@tool("ask_clarification", parse_docstring=True, return_direct=True)
def ask_clarification_tool(
question: str,
@ -15,6 +31,7 @@ def ask_clarification_tool(
],
context: str | None = None,
options: list[str] | None = None,
fields: list[ClarificationFormField] | None = None,
) -> str:
"""Ask the user for clarification when you need more information to proceed.
@ -36,11 +53,22 @@ def ask_clarification_tool(
- You're about to perform a potentially dangerous operation
- You have a recommendation but need user approval
Choosing the interaction shape:
- One open question -> just `question` (free text input)
- Pick exactly one option -> `options`
- Pick several options -> a single `fields` entry of type `multi_select`
- Collect several values at once (e.g. a set of parameters for one action) ->
`fields`, which renders a single structured form instead of several
sequential questions. Prefer one form over asking field-by-field.
Best practices:
- Ask ONE clarification at a time for clarity
- Ask ONE clarification at a time for clarity; a form with several fields
still counts as one clarification
- Be specific and clear in your question
- Don't make assumptions when clarification is needed
- For risky operations, ALWAYS ask for confirmation
- If a skill provides a predefined field template, pass it through `fields`
unchanged instead of redesigning it
- After calling this tool, execution will be interrupted automatically
Args:
@ -48,6 +76,15 @@ def ask_clarification_tool(
clarification_type: The type of clarification needed (missing_info, ambiguous_requirement, approach_choice, risk_confirmation, suggestion).
context: Optional context explaining why clarification is needed. Helps the user understand the situation.
options: Optional list of choices (for approach_choice or suggestion types). Present clear options for the user to choose from.
fields: Optional form field definitions for collecting multiple values in one card; takes precedence over `options`.
Each field is an object with `name` (unique identifier, required; avoid JavaScript prototype names like
`constructor` or `toString`), `label` (display text, defaults to name), `type` (one of: text, textarea,
number, select, multi_select, checkbox, date; defaults to text), `required` (boolean, defaults to false),
`options` (list of strings, required for select/multi_select types), and `placeholder` (optional hint text).
A `checkbox` field is a boolean that defaults to "no"; set `required` on a checkbox only for
must-agree/consent semantics (the user has to tick it to submit). Keep forms bounded: at most 16 fields,
24 options per field, and 200 characters per name/label/option/placeholder exceeding a limit degrades
the whole request to a plain-text question.
"""
# This is a placeholder implementation
# The actual logic is handled by ClarificationMiddleware which intercepts this tool call

View File

@ -1,5 +1,5 @@
from .api import get_workspace_changes_response
from .diff import compare_snapshots, get_changed_paths
from .diff import compare_snapshots, get_changed_output_paths, get_changed_paths
from .recorder import capture_workspace_snapshot, record_workspace_changes
from .scanner import scan_workspace_roots
from .types import (
@ -26,6 +26,7 @@ __all__ = [
"WorkspaceSnapshot",
"capture_workspace_snapshot",
"compare_snapshots",
"get_changed_output_paths",
"get_changed_paths",
"get_workspace_changes_response",
"record_workspace_changes",

View File

@ -109,6 +109,16 @@ def get_changed_paths(before: WorkspaceSnapshot, after: WorkspaceSnapshot) -> se
return changed
def get_changed_output_paths(before: WorkspaceSnapshot, after: WorkspaceSnapshot) -> list[str]:
"""Return created or modified regular files under the outputs root."""
paths: list[str] = []
for path in sorted(get_changed_paths(before, after)):
snapshot = after.files.get(path)
if snapshot is not None and snapshot.root == "outputs" and not snapshot.symlink:
paths.append(path)
return paths
def _status(
before_file: FileSnapshot | None,
after_file: FileSnapshot | None,

View File

@ -67,6 +67,9 @@ postgres = [
redis = ["redis>=5.0.0"]
pymupdf = ["pymupdf4llm>=0.0.17"]
boxlite = ["boxlite>=0.9.7"]
# Tenki cloud sandbox provider (deerflow.community.tenki). Optional so a default
# install stays free of the Tenki SDK; only pulled in when the provider is used.
tenki = ["tenki-sandbox>=0.4.0"]
# Agent observability (Monocle). Optional so a default install stays free of the
# OpenTelemetry stack; only pulled in when MONOCLE_TRACING is used.
monocle = ["monocle_apptrace>=0.8.8"]
@ -74,6 +77,8 @@ monocle = ["monocle_apptrace>=0.8.8"]
# so the core harness install stays lean; import is lazy inside the private
# Playwright loop. After install, run `playwright install chromium` once.
browser = ["playwright>=1.40"]
# Optional Chinese tokenization for the FTS5 memory retrieval adapter.
memory-zh = ["jieba>=0.42.1"]
[build-system]
requires = ["hatchling"]

View File

@ -30,6 +30,7 @@ redis = ["deerflow-harness[redis]"]
discord = ["discord.py>=2.7.0"]
monocle = ["deerflow-harness[monocle]"]
browser = ["deerflow-harness[browser]"]
memory-zh = ["deerflow-harness[memory-zh]"]
[dependency-groups]
dev = [

View File

@ -8,12 +8,12 @@ warming the other.
Examples::
PYTHONPATH=. uv run python scripts/benchmark/bench_checkpoint_channels.py \
PYTHONPATH=. uv run python scripts/benchmark/checkpoint/bench_channels.py \
--updates 10,100,500,999,1000,1001,2000 \
--payload-bytes 128,4096 \
--output checkpoint-bench.jsonl
PYTHONPATH=. uv run python scripts/benchmark/bench_checkpoint_channels.py \
PYTHONPATH=. uv run python scripts/benchmark/checkpoint/bench_channels.py \
--backends sqlite --updates 1000 --payload-bytes 128 \
--repetitions 7 --output snapshot-boundary.jsonl
@ -27,15 +27,11 @@ and memory use.
from __future__ import annotations
import argparse
import cProfile
import gc
import hashlib
import importlib.metadata
import json
import os
import platform
import statistics
import subprocess
import sys
import tempfile
import time
@ -55,10 +51,19 @@ from deerflow.agents.thread_state import merge_message_writes
from deerflow.runtime.checkpoint_mode import inject_checkpoint_mode
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor
try:
import resource
except ImportError: # pragma: no cover - Windows only
resource = None # type: ignore[assignment]
sys.path.insert(0, str(Path(__file__).resolve().parent))
import checkpoint_bench_common as _common # noqa: E402
_GIT_SHA_ENV = _common.GIT_SHA_ENV
_parse_positive_int_csv = _common.parse_positive_int_csv
_parse_choice_csv = _common.parse_choice_csv
_canonical_messages_digest = _common.canonical_messages_digest
_percentile = _common.percentile
_window_median = _common.window_median
_resolve_git_sha = _common.resolve_git_sha
_safe_error = _common.safe_error
_file_size = _common.file_size
_peak_rss_bytes = _common.peak_rss_bytes
Mode = Literal["full", "delta"]
Backend = Literal["memory", "sqlite"]
@ -67,7 +72,6 @@ SCHEMA_VERSION = 1
BENCHMARK_VERSION = 1
PRODUCTION_SNAPSHOT_FREQUENCY = 1000
DEFAULT_MAX_ESTIMATED_FULL_BYTES = 1024**3
_GIT_SHA_ENV = "DEERFLOW_CHECKPOINT_BENCH_GIT_SHA"
_MODES: tuple[Mode, ...] = ("full", "delta")
_BACKENDS: tuple[Backend, ...] = ("memory", "sqlite")
_STORAGE_STAT_FIELDS = (
@ -117,53 +121,6 @@ class BenchmarkCase:
raise ValueError(f"unsupported scenario: {self.scenario!r}")
def _parse_positive_int_csv(value: str, *, option: str) -> list[int]:
if not value or value.startswith(",") or value.endswith(",") or ",," in value:
raise ValueError(f"{option} must be a comma-separated list of positive integers")
result: list[int] = []
seen: set[int] = set()
duplicates: list[int] = []
try:
parsed = [int(part.strip()) for part in value.split(",")]
except ValueError as exc:
raise ValueError(f"{option} must be a comma-separated list of positive integers") from exc
if any(item <= 0 for item in parsed):
raise ValueError(f"{option} values must be positive integers")
for item in parsed:
if item not in seen:
result.append(item)
seen.add(item)
elif item not in duplicates:
duplicates.append(item)
if duplicates:
print(
f"{option}: ignored duplicate value(s): {', '.join(str(item) for item in duplicates)}; use --repetitions for repeated samples.",
file=sys.stderr,
)
return result
def _parse_choice_csv(value: str, *, option: str, choices: tuple[str, ...]) -> list[str]:
if not value or value.startswith(",") or value.endswith(",") or ",," in value:
raise ValueError(f"{option} must contain one or more of: {', '.join(choices)}")
result: list[str] = []
duplicates: list[str] = []
for raw in value.split(","):
item = raw.strip()
if item not in choices:
raise ValueError(f"{option} contains unsupported value {item!r}; expected: {', '.join(choices)}")
if item not in result:
result.append(item)
elif item not in duplicates:
duplicates.append(item)
if duplicates:
print(
f"{option}: ignored duplicate value(s): {', '.join(duplicates)}; use --repetitions for repeated samples.",
file=sys.stderr,
)
return result
def _expand_cases(
*,
modes: list[str],
@ -219,47 +176,6 @@ def _message_for_update(index: int, payload_bytes: int) -> BaseMessage:
return AIMessage(id=message_id, content=content)
def _canonical_messages_digest(messages: list[AnyMessage]) -> str:
canonical = [
{
"id": message.id,
"type": message.type,
"content": message.content,
}
for message in messages
]
payload = json.dumps(canonical, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode("utf-8")
return hashlib.sha256(payload).hexdigest()
def _percentile(values: list[float], percentile: float) -> float:
if not values:
return 0.0
ordered = sorted(values)
if len(ordered) == 1:
return ordered[0]
rank = (len(ordered) - 1) * percentile / 100
lower = int(rank)
upper = min(lower + 1, len(ordered) - 1)
fraction = rank - lower
return ordered[lower] * (1 - fraction) + ordered[upper] * fraction
def _window_median(values: list[float], window: Literal["first", "middle", "last"]) -> float:
if not values:
return 0.0
width = max(1, len(values) // 10)
if window == "first":
selected = values[:width]
elif window == "last":
selected = values[-width:]
else:
center = len(values) // 2
start = max(0, center - width // 2)
selected = values[start : start + width]
return statistics.median(selected)
def _noop(_state: dict[str, Any]) -> dict[str, Any]:
return {}
@ -283,20 +199,6 @@ def _config(case: BenchmarkCase) -> dict[str, Any]:
return config
def _resolve_git_sha() -> str:
try:
return subprocess.run(
["git", "rev-parse", "HEAD"],
cwd=Path(__file__).resolve().parents[3],
check=True,
capture_output=True,
text=True,
timeout=5,
).stdout.strip()
except (OSError, subprocess.SubprocessError):
return "unknown"
def _base_row(case: BenchmarkCase) -> dict[str, Any]:
git_sha = os.environ.get(_GIT_SHA_ENV) or _resolve_git_sha()
try:
@ -324,13 +226,6 @@ def _base_row(case: BenchmarkCase) -> dict[str, Any]:
}
def _safe_error(error: BaseException | str, *, work_dir: Path | None = None) -> str:
message = str(error).replace(str(Path.home()), "<home>")
if work_dir is not None:
message = message.replace(str(work_dir), "<work-dir>")
return message[:2000]
def _collect_storage_stats(collector: Callable[[], dict[str, int]]) -> dict[str, Any]:
"""Keep timing data usable when a saver's diagnostic layout changes."""
try:
@ -387,20 +282,6 @@ def _sqlite_storage_stats(saver: SqliteSaver, thread_id: str) -> dict[str, int]:
}
def _file_size(path: Path) -> int:
try:
return path.stat().st_size
except FileNotFoundError:
return 0
def _peak_rss_bytes() -> int | None:
if resource is None:
return None
peak_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
return int(peak_rss) if sys.platform == "darwin" else int(peak_rss * 1024)
def _write_and_read(case: BenchmarkCase, saver: Any, messages: list[BaseMessage]) -> tuple[dict[str, Any], list[AnyMessage]]:
graph = _build_graph(case.mode, saver)
accessor = CheckpointStateAccessor.bind(graph, saver, mode=case.mode)
@ -532,12 +413,7 @@ def _run_case(case: BenchmarkCase, *, work_dir: Path) -> dict[str, Any]:
def _run_profiled_case(case: BenchmarkCase, *, work_dir: Path, profile_path: Path) -> dict[str, Any]:
profile_path.parent.mkdir(parents=True, exist_ok=True)
profiler = cProfile.Profile()
row = profiler.runcall(_run_case, case, work_dir=work_dir)
row["profiled"] = True
profiler.dump_stats(profile_path)
return row
return _common.run_profiled(_run_case, case, work_dir=work_dir, profile_path=profile_path)
def _comparison_key(row: dict[str, Any]) -> tuple[Any, ...]:
@ -583,37 +459,16 @@ def _profile_filename(case: BenchmarkCase) -> str:
def _run_child_case(case: BenchmarkCase, *, timeout_seconds: float, git_sha: str, profile_dir: Path | None = None) -> dict[str, Any]:
encoded_case = json.dumps(asdict(case), separators=(",", ":"))
command = [sys.executable, str(Path(__file__).resolve()), "--worker-case", encoded_case]
worker_args = ["--worker-case", json.dumps(asdict(case), separators=(",", ":"))]
if profile_dir is not None:
command.extend(["--worker-profile", str(profile_dir / _profile_filename(case))])
started = time.perf_counter()
child_env = os.environ.copy()
child_env[_GIT_SHA_ENV] = git_sha
try:
completed = subprocess.run(
command,
check=False,
capture_output=True,
text=True,
timeout=timeout_seconds,
env=child_env,
)
except subprocess.TimeoutExpired:
return _failure_row(case, f"child process timed out after {timeout_seconds:g} seconds")
child_process_ms = (time.perf_counter() - started) * 1000
output_lines = [line for line in completed.stdout.splitlines() if line.strip()]
if not output_lines:
return _failure_row(case, f"child process returned {completed.returncode} without a result")
try:
row = json.loads(output_lines[-1])
except json.JSONDecodeError:
return _failure_row(case, f"child process returned {completed.returncode} with malformed JSON")
row["child_process_ms"] = child_process_ms
if completed.returncode != 0 and row.get("success"):
row["success"] = False
row["error"] = f"child process exited with status {completed.returncode}"
return row
worker_args.extend(["--worker-profile", str(profile_dir / _profile_filename(case))])
return _common.run_child_case(
script=Path(__file__).resolve(),
worker_args=worker_args,
failure_row=lambda error: _failure_row(case, error),
timeout_seconds=timeout_seconds,
git_sha=git_sha,
)
def _build_parser() -> argparse.ArgumentParser:

View File

@ -0,0 +1,716 @@
#!/usr/bin/env python3
"""Production-shaped full/delta checkpoint benchmark.
Measures user-visible run, state, and history latency on the production path:
the real lead-agent graph (scripted deterministic model), a real
AsyncSqliteSaver, and the real Gateway route stack for reads. Every case runs
in a fresh child process with a fresh database, mirroring the
restart-required checkpoint mode boundary.
Examples::
PYTHONPATH=. uv run python scripts/benchmark/checkpoint/bench_production.py \
--turns 100,1000 --payload-bytes 128 \
--snapshot-frequencies 10,100,1000 \
--repetitions 7 --output production-bench.jsonl
PYTHONPATH=. uv run python scripts/benchmark/checkpoint/summarize_production.py \
production-bench.jsonl
"""
from __future__ import annotations
import argparse
import json
import sys
import tempfile
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Literal
sys.path.insert(0, str(Path(__file__).resolve().parent))
import checkpoint_bench_common as common # noqa: E402
Mode = Literal["full", "delta"]
SCHEMA_VERSION = 1
BENCHMARK_VERSION = 1
BENCHMARK_NAME = "checkpoint_production"
DEFAULT_HISTORY_LIMITS = (10, 50, 200)
DEFAULT_READ_REPETITIONS = 25
WARMUP_TURNS = 2
_MODES: tuple[Mode, ...] = ("full", "delta")
@dataclass(frozen=True)
class ProductionCase:
mode: Mode
turns: int
payload_bytes: int
snapshot_frequency: int | None
history_limits: tuple[int, ...]
read_repetitions: int
repetition: int
seed: int
def __post_init__(self) -> None:
if self.mode not in _MODES:
raise ValueError(f"unsupported mode: {self.mode!r}")
if self.turns <= WARMUP_TURNS:
raise ValueError(f"turns must exceed the {WARMUP_TURNS}-turn warm-up")
if self.payload_bytes <= 0:
raise ValueError("payload_bytes must be positive")
if self.mode == "full" and self.snapshot_frequency is not None:
raise ValueError("snapshot_frequency is a delta-only axis")
if self.mode == "delta" and (self.snapshot_frequency is None or self.snapshot_frequency <= 0):
raise ValueError("delta mode requires a positive snapshot_frequency")
if not self.history_limits or any(limit <= 0 for limit in self.history_limits):
raise ValueError("history_limits must be positive")
if self.read_repetitions <= 0:
raise ValueError("read_repetitions must be positive")
if self.repetition < 0:
raise ValueError("repetition must be non-negative")
def expand_cases(
*,
modes: list[str],
turn_counts: list[int],
payload_bytes: list[int],
snapshot_frequencies: list[int],
repetitions: int,
seed: int = 1,
history_limits: tuple[int, ...] = DEFAULT_HISTORY_LIMITS,
read_repetitions: int = DEFAULT_READ_REPETITIONS,
) -> list[ProductionCase]:
"""Build the matrix with alternating mode order to reduce order bias."""
cases: list[ProductionCase] = []
group_index = 0
for repetition in range(repetitions):
for payload in payload_bytes:
for turns in turn_counts:
ordered_modes = list(modes)
if group_index % 2 == 1:
ordered_modes.reverse()
for mode in ordered_modes:
frequencies: list[int | None] = snapshot_frequencies if mode == "delta" else [None]
for frequency in frequencies:
cases.append(
ProductionCase(
mode=mode, # type: ignore[arg-type]
turns=turns,
payload_bytes=payload,
snapshot_frequency=frequency,
history_limits=history_limits,
read_repetitions=read_repetitions,
repetition=repetition,
seed=seed,
)
)
group_index += 1
return cases
def _base_row(case: ProductionCase) -> dict:
import importlib.metadata
import os
import platform
git_sha = os.environ.get(common.GIT_SHA_ENV) or common.resolve_git_sha()
try:
langgraph_version = importlib.metadata.version("langgraph")
except importlib.metadata.PackageNotFoundError:
langgraph_version = "unknown"
return {
"schema_version": SCHEMA_VERSION,
"benchmark_version": BENCHMARK_VERSION,
"benchmark": BENCHMARK_NAME,
"success": True,
"error": None,
"profiled": False,
"git_sha": git_sha,
"python_version": platform.python_version(),
"langgraph_version": langgraph_version,
"platform": platform.platform(),
"mode": case.mode,
"snapshot_frequency": case.snapshot_frequency,
"turns": case.turns,
"payload_bytes": case.payload_bytes,
"history_limits": list(case.history_limits),
"read_repetitions": case.read_repetitions,
"repetition": case.repetition,
"seed": case.seed,
}
def _failure_row(case: ProductionCase, error: str) -> dict:
row = _base_row(case)
row["success"] = False
row["error"] = common.safe_error(error)
return row
def _comparison_key(row: dict) -> tuple:
# snapshot_frequency is intentionally excluded: every delta frequency
# pairs against the same full row, and all frequencies must agree on
# the materialized digests.
return (row.get("turns"), row.get("payload_bytes"), row.get("repetition"))
def validate_cross_mode_rows(rows: list[dict]) -> None:
grouped: dict[tuple, list[dict]] = {}
for row in rows:
grouped.setdefault(_comparison_key(row), []).append(row)
for group in grouped.values():
successful = [row for row in group if row.get("success")]
modes = {row.get("mode") for row in successful}
if not {"full", "delta"}.issubset(modes):
continue
signatures = {(row.get("message_count"), row.get("seed_content_sha256"), row.get("state_wire_sha256"), row.get("history_wire_sha256")) for row in successful}
if len(signatures) == 1:
continue
for row in successful:
row["success"] = False
row["error"] = "cross-mode materialized state mismatch"
def _profile_filename(case: ProductionCase) -> str:
return f"production-{case.mode}-turns-{case.turns}-payload-{case.payload_bytes}-freq-{case.snapshot_frequency}-rep-{case.repetition}.prof"
def _worker_main(encoded_case: str, *, profile_path: Path | None = None) -> int:
try:
raw = json.loads(encoded_case)
history_limits = raw.get("history_limits")
if history_limits is not None:
raw["history_limits"] = tuple(history_limits)
case = ProductionCase(**raw)
except (TypeError, ValueError) as exc:
print(json.dumps({"schema_version": SCHEMA_VERSION, "benchmark_version": BENCHMARK_VERSION, "benchmark": BENCHMARK_NAME, "success": False, "error": common.safe_error(exc)}, separators=(",", ":")))
return 2
with tempfile.TemporaryDirectory(prefix="deerflow-checkpoint-production-") as temp_dir:
if profile_path is None:
row = _run_case(case, work_dir=Path(temp_dir))
else:
row = common.run_profiled(_run_case, case, work_dir=Path(temp_dir), profile_path=profile_path)
print(json.dumps(row, ensure_ascii=False, separators=(",", ":")))
return 0 if row.get("success") else 1
def _build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Production-shaped full/delta checkpoint benchmark")
parser.add_argument("--modes", default="full,delta")
parser.add_argument("--turns", default="10,100,500,1000,2000")
parser.add_argument("--payload-bytes", default="128")
parser.add_argument("--snapshot-frequencies", default="10,50,100,500,1000")
parser.add_argument("--history-limits", default=",".join(str(v) for v in DEFAULT_HISTORY_LIMITS))
parser.add_argument("--read-repetitions", type=int, default=DEFAULT_READ_REPETITIONS)
parser.add_argument("--repetitions", type=int, default=3)
parser.add_argument("--seed", type=int, default=1)
parser.add_argument("--timeout-seconds", type=float, default=900)
parser.add_argument("--profile-dir", type=Path)
parser.add_argument("--output", type=Path)
parser.add_argument("--worker-case", help=argparse.SUPPRESS)
parser.add_argument("--worker-profile", type=Path, help=argparse.SUPPRESS)
return parser
def main(argv: list[str] | None = None) -> int:
parser = _build_parser()
args = parser.parse_args(argv)
if args.worker_case is not None:
return _worker_main(args.worker_case, profile_path=args.worker_profile)
if args.output is None:
parser.error("--output is required")
try:
modes = common.parse_choice_csv(args.modes, option="--modes", choices=_MODES)
turn_counts = common.parse_positive_int_csv(args.turns, option="--turns")
payload_bytes = common.parse_positive_int_csv(args.payload_bytes, option="--payload-bytes")
frequencies = common.parse_positive_int_csv(args.snapshot_frequencies, option="--snapshot-frequencies")
history_limits = tuple(common.parse_positive_int_csv(args.history_limits, option="--history-limits"))
if args.repetitions <= 0 or args.read_repetitions <= 0 or args.timeout_seconds <= 0:
raise ValueError("repetition and timeout values must be positive")
except ValueError as exc:
parser.error(str(exc))
cases = expand_cases(
modes=modes,
turn_counts=turn_counts,
payload_bytes=payload_bytes,
snapshot_frequencies=frequencies,
repetitions=args.repetitions,
seed=args.seed,
history_limits=history_limits,
read_repetitions=args.read_repetitions,
)
git_sha = common.resolve_git_sha()
rows: list[dict] = []
for index, case in enumerate(cases, start=1):
print(
f"[{index}/{len(cases)}] {case.mode} turns={case.turns} payload={case.payload_bytes} freq={case.snapshot_frequency} repetition={case.repetition}",
file=sys.stderr,
)
worker_args = ["--worker-case", json.dumps(asdict(case), separators=(",", ":"))]
if args.profile_dir is not None:
worker_args.extend(["--worker-profile", str(args.profile_dir / _profile_filename(case))])
rows.append(
common.run_child_case(
script=Path(__file__).resolve(),
worker_args=worker_args,
failure_row=lambda error, c=case: _failure_row(c, error),
timeout_seconds=args.timeout_seconds,
git_sha=git_sha,
)
)
validate_cross_mode_rows(rows)
args.output.parent.mkdir(parents=True, exist_ok=True)
with args.output.open("w", encoding="utf-8") as output_file:
for row in rows:
output_file.write(json.dumps(row, ensure_ascii=False, separators=(",", ":")) + "\n")
failures = sum(1 for row in rows if not row.get("success"))
print(f"Wrote {len(rows)} result row(s) to {args.output}; failures={failures}", file=sys.stderr)
return 1 if failures else 0
# ---------------------------------------------------------------------------
# Worker: production-shaped run + read phases (single event loop)
# ---------------------------------------------------------------------------
import asyncio # noqa: E402
import hashlib # noqa: E402
import time # noqa: E402
from collections import defaultdict # noqa: E402
from typing import Any # noqa: E402
from unittest.mock import patch # noqa: E402
import httpx # noqa: E402
from langchain_core.language_models.fake_chat_models import FakeMessagesListChatModel # noqa: E402
from langchain_core.messages import AIMessage, HumanMessage # noqa: E402
from langchain_core.outputs import ChatGeneration, ChatResult # noqa: E402
class _ScriptedBenchModel(FakeMessagesListChatModel):
"""Deterministic model: call k returns a fixed-size AIMessage.
Content and IDs depend only on the call index, so full and delta modes
produce byte-identical message streams. Serves the ainvoke path via
BaseChatModel's default ``_agenerate`` fallback (``run_in_executor`` over
``_generate``); streaming is not supported.
"""
def __init__(self, *, payload_bytes: int) -> None:
super().__init__(responses=[])
self._bench_payload_bytes = payload_bytes
self._bench_call_index = 0
def bind_tools(self, tools: Any, **kwargs: Any) -> Any:
return self
def _generate(self, messages: Any, stop: Any = None, run_manager: Any = None, **kwargs: Any) -> ChatResult:
index = self._bench_call_index
self._bench_call_index += 1
message = AIMessage(id=f"bench-ai-{index:08d}", content="x" * self._bench_payload_bytes)
return ChatResult(generations=[ChatGeneration(message=message)])
class _TimingSaver:
"""Delegating saver wrapper recording write timing.
LangGraph issues ``aput``/``aput_writes`` as concurrent asyncio tasks, so
per-call latencies overlap (queueing on the single sqlite connection is
counted in every concurrent call). Summing them double-counts and can
exceed the turn wall clock. ``intervals`` keeps absolute (start, end) per
call so the run phase can compute the merged busy time the honest,
wall-bounded write cost. ``timings_ms`` remains the cumulative-per-method
ledger used for read-path attribution.
"""
def __init__(self, saver: Any) -> None:
self._saver = saver
self.timings_ms: dict[str, float] = defaultdict(float)
self.intervals: list[tuple[str, float, float]] = []
def __getattr__(self, name: str) -> Any:
return getattr(self._saver, name)
async def _timed(self, label: str, fn: Any, *args: Any, **kwargs: Any) -> Any:
start = time.perf_counter()
try:
return await fn(*args, **kwargs)
finally:
end = time.perf_counter()
self.timings_ms[label] += (end - start) * 1000
self.intervals.append((label, start * 1000, end * 1000))
async def aput(self, *args: Any, **kwargs: Any) -> Any:
return await self._timed("aput", self._saver.aput, *args, **kwargs)
async def aput_writes(self, *args: Any, **kwargs: Any) -> Any:
return await self._timed("aput_writes", self._saver.aput_writes, *args, **kwargs)
async def aget_tuple(self, *args: Any, **kwargs: Any) -> Any:
return await self._timed("aget_tuple", self._saver.aget_tuple, *args, **kwargs)
def alist(self, *args: Any, **kwargs: Any) -> Any:
# alist is an async generator on real savers; wrap it to time the
# full iteration when consumed.
agen = self._saver.alist(*args, **kwargs)
timing = self.timings_ms
async def _wrapped() -> Any:
start_total = time.perf_counter()
try:
async for item in agen:
yield item
finally:
timing["alist"] += (time.perf_counter() - start_total) * 1000
return _wrapped()
def _thread_id(case: ProductionCase) -> str:
return f"checkpoint-production-{case.seed}-{case.repetition}"
async def _run_phase(case: ProductionCase, saver: Any, timing: _TimingSaver, work_dir: Path) -> dict[str, Any]:
"""Seed the thread through the real lead-agent graph; measure per-turn latency."""
from deerflow.agents.lead_agent.agent import make_lead_agent
from deerflow.config.app_config import AppConfig, set_app_config
from deerflow.runtime.checkpoint_mode import inject_checkpoint_mode
app_config = AppConfig.model_validate(
{
# Mirrors tests/test_lead_agent_model_resolution.py::_make_model:
# `use` is a required ModelConfig field, so a provider/api_key
# style entry does not validate.
"models": [
{
"name": "bench-model",
"display_name": "bench-model",
"description": None,
"use": "langchain_openai:ChatOpenAI",
"model": "bench",
"supports_thinking": False,
"supports_vision": False,
}
],
"sandbox": {"use": "deerflow.sandbox.local:LocalSandboxProvider"},
# Title and memory middleware would fire background LLM calls and
# debounce timers; both pollute per-turn latency, so disable them.
"title": {"enabled": False},
"memory": {"enabled": False},
"database": {
"backend": "sqlite",
"sqlite_dir": str(work_dir),
"checkpoint_channel_mode": case.mode,
"checkpoint_delta_snapshot_frequency": case.snapshot_frequency or 1000,
},
}
)
set_app_config(app_config)
model = _ScriptedBenchModel(payload_bytes=case.payload_bytes)
def _fake_create_chat_model(**kwargs: Any) -> Any:
return model
with patch("deerflow.agents.lead_agent.agent.create_chat_model", _fake_create_chat_model):
graph = make_lead_agent({"configurable": {}})
# Attach the timed saver post-construction, mirroring the run worker
# (deerflow.runs.worker assigns ``agent.checkpointer`` the same way), so
# every checkpoint write flows through ``timing``.
graph.checkpointer = timing
config: dict[str, Any] = {"configurable": {"thread_id": _thread_id(case)}}
inject_checkpoint_mode(config, case.mode)
turn_ms: list[float] = []
write_ms: list[float] = []
write_methods = {"aput", "aput_writes"}
for turn in range(case.turns):
intervals_before = len(timing.intervals)
start = time.perf_counter()
await graph.ainvoke(
{"messages": [HumanMessage(id=f"bench-h-{turn:08d}", content="y" * case.payload_bytes)]},
config,
)
turn_ms.append((time.perf_counter() - start) * 1000)
write_ms.append(_merged_busy_ms(iv for iv in timing.intervals[intervals_before:] if iv[0] in write_methods))
return {"graph": graph, "config": config, "turn_ms": turn_ms[WARMUP_TURNS:], "write_ms": write_ms[WARMUP_TURNS:]}
def _merged_busy_ms(intervals: Any) -> float:
"""Union of (start, end) write intervals in ms — wall-bounded write cost.
Concurrent write tasks overlap; merging their intervals yields the time
the event loop was actually inside aput/aput_writes, which cannot exceed
the turn wall clock the way a naive latency sum does.
"""
merged: list[list[float]] = []
for _label, start, end in sorted(intervals, key=lambda iv: iv[1]):
if merged and start <= merged[-1][1]:
if end > merged[-1][1]:
merged[-1][1] = end
else:
merged.append([start, end])
return sum(end - start for start, end in merged)
# ``ThreadHistoryRequest.limit`` is validated with ``le=100``: the production
# route cannot serve a larger page. Configured history limits above the cap
# are measured at the cap and labelled by the *effective* limit so no metric
# field ever claims a request the API would reject.
_HISTORY_ROUTE_MAX_LIMIT = 100
def _wire_messages_digest(messages: Any) -> str:
"""Digest wire-format messages, excluding volatile checkpoint ids/timestamps."""
canonical = [[message.get("type"), message.get("id"), message.get("content") if isinstance(message.get("content"), str) else json.dumps(message.get("content"), sort_keys=True)] for message in messages or []]
payload = json.dumps(canonical, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
return hashlib.sha256(payload).hexdigest()
def _combined_digest(digests: list[str]) -> str:
return hashlib.sha256("".join(digests).encode("utf-8")).hexdigest()
def _make_gateway_app(saver: Any, mode: str, store: Any) -> Any:
"""Minimal authed gateway app, mirroring tests/_router_auth_helpers.py.
Reimplemented inline: benchmark scripts must not import from tests/.
"""
from unittest.mock import MagicMock
from uuid import uuid4
from fastapi import FastAPI, Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from app.gateway.auth.models import User
from app.gateway.authz import AuthContext, Permissions
from app.gateway.routers import threads
from deerflow.persistence.thread_meta.memory import MemoryThreadMetaStore
from deerflow.runtime.user_context import reset_current_user, set_current_user
permissions = [Permissions.THREADS_READ, Permissions.THREADS_WRITE, Permissions.THREADS_DELETE]
class _StubAuthMiddleware(BaseHTTPMiddleware):
async def dispatch(self, request: Request, call_next: Any) -> Response:
user = User(email="bench@example.com", password_hash="x", system_role="user", id=uuid4())
request.state.user = user
request.state.auth = AuthContext(user=user, permissions=list(permissions))
# Mirror the production AuthMiddleware: the user contextvar must be
# set or thread-meta lookups hit user_id=AUTO and raise per request.
token = set_current_user(user)
try:
return await call_next(request)
finally:
reset_current_user(token)
app = FastAPI()
app.add_middleware(_StubAuthMiddleware)
app.state.store = store
app.state.checkpointer = saver
app.state.thread_store = MemoryThreadMetaStore(store)
app.state.checkpoint_channel_mode = mode
app.state.run_event_store = MagicMock()
app.include_router(threads.router)
return app
async def _read_phase(case: ProductionCase, saver: Any, timing: _TimingSaver) -> dict[str, Any]:
"""Measure state/history endpoint latency through the real route stack.
Runs in the same event loop as the run phase (the AsyncSqliteSaver is
bound to its creating loop), driving the real threads router over an
ASGI transport. ``saver`` must be the RAW saver, not the timing
wrapper: ``Pregel.aget_state`` gates delta replay behind
``isinstance(self.checkpointer, BaseCheckpointSaver)``
(langgraph/pregel/main.py), so the duck-typed ``_TimingSaver``
silently materializes delta channels as empty. The wrapper therefore
stays on the run phase only and the row records
``saver_read_attribution: False``; the ``*_saver_ms`` fields remain in
the schema but report 0.0.
"""
from langgraph.store.memory import InMemoryStore
from app.gateway import services as gateway_services
store = InMemoryStore()
app = _make_gateway_app(saver, case.mode, store)
graph_build_ms: list[float] = []
original_resolve = gateway_services.resolve_agent_factory
real_factory = original_resolve(None)
def _timed_factory(config: Any = None) -> Any:
start = time.perf_counter()
graph = real_factory(config)
graph_build_ms.append((time.perf_counter() - start) * 1000)
return graph
# The accessor graph cache is keyed (assistant_id, mode); the factory
# identity check compares the callable, so one stable wrapper instance.
gateway_services.resolve_agent_factory = lambda assistant_id=None: _timed_factory
# The accessor graph is compiled lazily inside the first request and
# rebuilds the lead agent, so the scripted model must be patched in for
# the whole read phase, exactly as in _run_phase.
model = _ScriptedBenchModel(payload_bytes=case.payload_bytes)
def _fake_create_chat_model(**kwargs: Any) -> Any:
return model
thread_id = _thread_id(case)
transport = httpx.ASGITransport(app=app)
result: dict[str, Any] = {}
try:
with patch("deerflow.agents.lead_agent.agent.create_chat_model", _fake_create_chat_model):
async with httpx.AsyncClient(transport=transport, base_url="http://bench") as client:
async def _timed_request(label: str, method: str, url: str, **kwargs: Any) -> tuple[Any, float, float]:
del label # timing only; labels aid debugging failures
saver_before = timing.timings_ms["aget_tuple"] + timing.timings_ms["alist"]
start = time.perf_counter()
response = await client.request(method, url, **kwargs)
elapsed_ms = (time.perf_counter() - start) * 1000
# The read phase drives the RAW saver (see docstring), so
# the timing wrapper never advances here: saver_ms — and
# thus every read-path *_saver_ms field — is structurally 0.0.
saver_ms = timing.timings_ms["aget_tuple"] + timing.timings_ms["alist"] - saver_before
response.raise_for_status()
return response, elapsed_ms, saver_ms
# --- state endpoint: one cold sample after clearing the accessor cache
gateway_services._state_accessor_graph_cache.clear()
graph_build_ms.clear()
response, state_cold_ms, state_cold_saver_ms = await _timed_request("state-cold", "GET", f"/api/threads/{thread_id}/state")
# Cold/warm contract, asserted structurally: the cold request
# MUST build the accessor graph exactly once (otherwise the
# cold sample did not exercise the cold path), and warm
# requests MUST NOT build again (otherwise the cache is not
# hit and warm samples silently measure the cold path).
if len(graph_build_ms) != 1:
raise AssertionError(f"state cold read triggered {len(graph_build_ms)} accessor graph builds; expected exactly 1")
cold_graph_build_ms = graph_build_ms[0]
state_payload = response.json()
state_digest = _wire_messages_digest((state_payload.get("values") or {}).get("messages"))
state_warm: list[float] = []
for _ in range(case.read_repetitions):
_response, elapsed_ms, _saver_ms = await _timed_request("state-warm", "GET", f"/api/threads/{thread_id}/state")
state_warm.append(elapsed_ms)
if len(graph_build_ms) != 1:
raise AssertionError("state warm reads triggered an accessor graph rebuild; cache was not hit")
# --- history endpoint per limit, same cold/warm split
history_metrics: dict[str, Any] = {}
history_digests: list[str] = []
effective_limits = list(dict.fromkeys(min(limit, _HISTORY_ROUTE_MAX_LIMIT) for limit in case.history_limits))
for limit_index, limit in enumerate(effective_limits):
gateway_services._state_accessor_graph_cache.clear()
response, cold_ms, cold_saver_ms = await _timed_request(
f"history-cold-{limit}",
"POST",
f"/api/threads/{thread_id}/history",
json={"limit": limit},
)
expected_builds = 2 + limit_index
if len(graph_build_ms) != expected_builds:
raise AssertionError(f"history cold read (limit={limit}) triggered {len(graph_build_ms)} total accessor graph builds; expected {expected_builds}")
entries = response.json()
digest_source: list[Any] = []
for entry in entries:
digest_source.extend((entry.get("values") or {}).get("messages") or [])
history_digests.append(_wire_messages_digest(digest_source))
warm: list[float] = []
for _ in range(case.read_repetitions):
_r, elapsed_ms, _s = await _timed_request(
f"history-warm-{limit}",
"POST",
f"/api/threads/{thread_id}/history",
json={"limit": limit},
)
warm.append(elapsed_ms)
if len(graph_build_ms) != expected_builds:
raise AssertionError(f"history warm reads (limit={limit}) triggered an accessor graph rebuild; cache was not hit")
history_metrics[f"history_limit_{limit}_cold_ms"] = cold_ms
history_metrics[f"history_limit_{limit}_cold_saver_ms"] = cold_saver_ms
history_metrics[f"history_limit_{limit}_warm_p50_ms"] = common.percentile(warm, 50)
history_metrics[f"history_limit_{limit}_warm_p95_ms"] = common.percentile(warm, 95)
result.update(
{
"state_cold_ms": state_cold_ms,
"state_cold_saver_ms": state_cold_saver_ms,
"state_cold_graph_build_ms": cold_graph_build_ms,
"state_warm_p50_ms": common.percentile(state_warm, 50),
"state_warm_p95_ms": common.percentile(state_warm, 95),
"state_wire_sha256": state_digest,
"history_wire_sha256": _combined_digest(history_digests),
"saver_read_attribution": False,
**history_metrics,
}
)
finally:
gateway_services.resolve_agent_factory = original_resolve
return result
def _storage_file_stats(db_path: Path) -> dict[str, Any]:
return {
"db_bytes": common.file_size(db_path),
"wal_bytes": common.file_size(Path(f"{db_path}-wal")),
"shm_bytes": common.file_size(Path(f"{db_path}-shm")),
}
def _run_case(case: ProductionCase, *, work_dir: Path) -> dict:
"""Run one benchmark case: run phase + read phase in a single event loop."""
row = _base_row(case)
db_path = work_dir / "production-benchmark.sqlite"
async def _async_body() -> dict[str, Any]:
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
async with AsyncSqliteSaver.from_conn_string(str(db_path)) as saver:
await saver.setup()
timing = _TimingSaver(saver)
run_info = await _run_phase(case, saver, timing, work_dir)
# Seed correctness: materialize through the mode-matched accessor.
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor
accessor = CheckpointStateAccessor.bind(run_info["graph"], saver, mode=case.mode)
snapshot = await accessor.aget(run_info["config"])
messages = list(snapshot.values.get("messages", []))
seed_digest = common.canonical_messages_digest(messages)
read_metrics = await _read_phase(case, saver, timing)
storage_metrics = _storage_file_stats(db_path)
metrics = {
"run_turn_p50_ms": common.percentile(run_info["turn_ms"], 50),
"run_turn_p95_ms": common.percentile(run_info["turn_ms"], 95),
"checkpoint_write_p50_ms": common.percentile(run_info["write_ms"], 50),
"checkpoint_write_p95_ms": common.percentile(run_info["write_ms"], 95),
"checkpoint_write_p99_ms": common.percentile(run_info["write_ms"], 99),
"message_count": len(messages),
"seed_content_sha256": seed_digest,
"peak_rss_bytes": common.peak_rss_bytes(),
**storage_metrics,
**read_metrics,
}
return metrics
try:
row.update(asyncio.run(_async_body()))
except Exception as exc:
row["success"] = False
row["error"] = common.safe_error(exc, work_dir=work_dir)
return row
if __name__ == "__main__":
raise SystemExit(main())

View File

@ -0,0 +1,208 @@
#!/usr/bin/env python3
"""Shared infrastructure for the checkpoint benchmark family.
Pure, stateless helpers plus the child-process controller protocol used by
both bench_channels.py (storage microbenchmark) and bench_production.py
(production-shaped benchmark). Scripts in this folder import this module via
a sibling sys.path insert; it is not a package.
"""
from __future__ import annotations
import cProfile
import hashlib
import json
import os
import statistics
import subprocess
import sys
import time
from collections.abc import Callable
from pathlib import Path
from typing import Any, Literal
from langchain_core.messages import AnyMessage
try:
import resource
except ImportError: # pragma: no cover - Windows only
resource = None # type: ignore[assignment]
GIT_SHA_ENV = "DEERFLOW_CHECKPOINT_BENCH_GIT_SHA"
def parse_positive_int_csv(value: str, *, option: str) -> list[int]:
if not value or value.startswith(",") or value.endswith(",") or ",," in value:
raise ValueError(f"{option} must be a comma-separated list of positive integers")
result: list[int] = []
seen: set[int] = set()
duplicates: list[int] = []
try:
parsed = [int(part.strip()) for part in value.split(",")]
except ValueError as exc:
raise ValueError(f"{option} must be a comma-separated list of positive integers") from exc
if any(item <= 0 for item in parsed):
raise ValueError(f"{option} values must be positive integers")
for item in parsed:
if item not in seen:
result.append(item)
seen.add(item)
elif item not in duplicates:
duplicates.append(item)
if duplicates:
print(
f"{option}: ignored duplicate value(s): {', '.join(str(item) for item in duplicates)}; use --repetitions for repeated samples.",
file=sys.stderr,
)
return result
def parse_choice_csv(value: str, *, option: str, choices: tuple[str, ...]) -> list[str]:
if not value or value.startswith(",") or value.endswith(",") or ",," in value:
raise ValueError(f"{option} must contain one or more of: {', '.join(choices)}")
result: list[str] = []
duplicates: list[str] = []
for raw in value.split(","):
item = raw.strip()
if item not in choices:
raise ValueError(f"{option} contains unsupported value {item!r}; expected: {', '.join(choices)}")
if item not in result:
result.append(item)
elif item not in duplicates:
duplicates.append(item)
if duplicates:
print(
f"{option}: ignored duplicate value(s): {', '.join(duplicates)}; use --repetitions for repeated samples.",
file=sys.stderr,
)
return result
def canonical_messages_digest(messages: list[AnyMessage]) -> str:
canonical = [
{
"id": message.id,
"type": message.type,
"content": message.content,
}
for message in messages
]
payload = json.dumps(canonical, ensure_ascii=False, separators=(",", ":"), sort_keys=True).encode("utf-8")
return hashlib.sha256(payload).hexdigest()
def percentile(values: list[float], percentile: float) -> float:
if not values:
return 0.0
ordered = sorted(values)
if len(ordered) == 1:
return ordered[0]
rank = (len(ordered) - 1) * percentile / 100
lower = int(rank)
upper = min(lower + 1, len(ordered) - 1)
fraction = rank - lower
return ordered[lower] * (1 - fraction) + ordered[upper] * fraction
def window_median(values: list[float], window: Literal["first", "middle", "last"]) -> float:
if not values:
return 0.0
width = max(1, len(values) // 10)
if window == "first":
selected = values[:width]
elif window == "last":
selected = values[-width:]
else:
center = len(values) // 2
start = max(0, center - width // 2)
selected = values[start : start + width]
return statistics.median(selected)
def resolve_git_sha() -> str:
try:
return subprocess.run(
["git", "rev-parse", "HEAD"],
# scripts/benchmark/checkpoint/checkpoint_bench_common.py -> repo root
cwd=Path(__file__).resolve().parents[4],
check=True,
capture_output=True,
text=True,
timeout=5,
).stdout.strip()
except (OSError, subprocess.SubprocessError):
return "unknown"
def safe_error(error: BaseException | str, *, work_dir: Path | None = None) -> str:
message = str(error).replace(str(Path.home()), "<home>")
if work_dir is not None:
message = message.replace(str(work_dir), "<work-dir>")
return message[:2000]
def file_size(path: Path) -> int:
try:
return path.stat().st_size
except FileNotFoundError:
return 0
def peak_rss_bytes() -> int | None:
if resource is None:
return None
peak_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
return int(peak_rss) if sys.platform == "darwin" else int(peak_rss * 1024)
def run_profiled(fn: Callable[..., dict[str, Any]], *args: Any, profile_path: Path, **kwargs: Any) -> dict[str, Any]:
"""Run *fn* under cProfile, mark the row, and dump stats."""
profile_path.parent.mkdir(parents=True, exist_ok=True)
profiler = cProfile.Profile()
row = profiler.runcall(fn, *args, **kwargs)
row["profiled"] = True
profiler.dump_stats(profile_path)
return row
def run_child_case(
*,
script: Path,
worker_args: list[str],
failure_row: Callable[[str], dict[str, Any]],
timeout_seconds: float,
git_sha: str,
) -> dict[str, Any]:
"""Run one benchmark case in a fresh child process; return its JSONL row.
The child prints exactly one JSON row on its last stdout line. Any
protocol failure is converted into a failure row via *failure_row*.
"""
command = [sys.executable, str(script), *worker_args]
started = time.perf_counter()
child_env = os.environ.copy()
child_env[GIT_SHA_ENV] = git_sha
try:
completed = subprocess.run(
command,
check=False,
capture_output=True,
text=True,
timeout=timeout_seconds,
env=child_env,
)
except subprocess.TimeoutExpired:
return failure_row(f"child process timed out after {timeout_seconds:g} seconds")
child_process_ms = (time.perf_counter() - started) * 1000
output_lines = [line for line in completed.stdout.splitlines() if line.strip()]
if not output_lines:
return failure_row(f"child process returned {completed.returncode} without a result")
try:
row = json.loads(output_lines[-1])
except json.JSONDecodeError:
return failure_row(f"child process returned {completed.returncode} with malformed JSON")
row["child_process_ms"] = child_process_ms
if completed.returncode != 0 and row.get("success"):
row["success"] = False
row["error"] = f"child process exited with status {completed.returncode}"
return row

View File

@ -8,9 +8,9 @@ values below 1 mean delta used less time or storage.
Examples::
python scripts/benchmark/summarize_checkpoint_channels.py results.jsonl
python scripts/benchmark/summarize_checkpoint_channels.py results/*.jsonl --json
python scripts/benchmark/summarize_checkpoint_channels.py results.jsonl \
python scripts/benchmark/checkpoint/summarize_channels.py results.jsonl
python scripts/benchmark/checkpoint/summarize_channels.py results/*.jsonl --json
python scripts/benchmark/checkpoint/summarize_channels.py results.jsonl \
--metrics write_total_ms,cold_read_ms,logical_checkpoint_bytes --csv
"""

View File

@ -0,0 +1,199 @@
#!/usr/bin/env python3
"""Summarize paired full/delta production-benchmark JSONL results.
Full rows carry snapshot_frequency=None; each delta frequency group pairs
against the full row at the same (turns, payload_bytes, repetition). Ratios
are always delta/full. Decision metrics: snapshot_write_spike (delta
checkpoint_write p99/p50), cache_effect_ms (state cold - warm p50), and
checkpoint_write_share (checkpoint_write p50 / run_turn p50). Unless
--metrics is given, history_limit_*_cold_ms / history_limit_*_warm_p50_ms
fields found in the input rows are appended to the default metric list.
Example::
python scripts/benchmark/checkpoint/summarize_production.py results.jsonl
"""
from __future__ import annotations
import argparse
import csv
import json
import re
import statistics
import sys
from collections import defaultdict
from pathlib import Path
from typing import Any
DEFAULT_METRICS = [
"run_turn_p50_ms",
"run_turn_p95_ms",
"checkpoint_write_p50_ms",
"checkpoint_write_p95_ms",
"state_cold_ms",
"state_warm_p50_ms",
"state_warm_p95_ms",
"db_bytes",
"peak_rss_bytes",
]
_GROUP_FIELDS = ("turns", "payload_bytes", "snapshot_frequency")
_INPUT_INDEX_FIELD = "_benchmark_input_index"
_HISTORY_METRIC_RE = re.compile(r"history_limit_(\d+)_(?:cold_ms|warm_p50_ms)")
def _discover_history_metrics(rows: list[dict[str, Any]]) -> list[str]:
"""Auto-discover per-limit history metrics present in the input rows."""
discovered = {key for row in rows for key in row if _HISTORY_METRIC_RE.fullmatch(key)}
return sorted(discovered, key=lambda key: (int(_HISTORY_METRIC_RE.fullmatch(key).group(1)), key)) # type: ignore[union-attr]
def _load_jsonl(paths: list[Path]) -> list[dict[str, Any]]:
rows: list[dict[str, Any]] = []
errors: list[str] = []
for input_index, path in enumerate(paths):
with path.open("r", encoding="utf-8") as input_file:
for line_number, raw_line in enumerate(input_file, start=1):
line = raw_line.strip()
if not line or line.startswith("#"):
continue
try:
row = json.loads(line)
except json.JSONDecodeError as exc:
errors.append(f"{path}:{line_number}: {exc.msg}")
continue
if not isinstance(row, dict):
errors.append(f"{path}:{line_number}: expected a JSON object")
continue
row[_INPUT_INDEX_FIELD] = input_index
rows.append(row)
if errors:
raise ValueError("Malformed JSONL row(s):\n" + "\n".join(errors))
return rows
def _numeric(row: dict[str, Any], metric: str) -> float | None:
value = row.get(metric)
if isinstance(value, bool) or not isinstance(value, (int, float)):
return None
return float(value)
def _medians(rows: list[dict[str, Any]], metric: str) -> float | None:
values = [v for row in rows if (v := _numeric(row, metric)) is not None]
return statistics.median(values) if values else None
def _summarize(rows: list[dict[str, Any]], *, metrics: list[str]) -> list[dict[str, Any]]:
# Full rows are frequency-less; index them for broadcast pairing.
full_rows: dict[tuple, dict[str, Any]] = {}
delta_groups: dict[tuple, dict[tuple, dict[str, Any]]] = defaultdict(dict)
failed_rows: list[dict[str, Any]] = []
for row in rows:
if row.get("profiled"):
continue
if not row.get("success"):
failed_rows.append(row)
continue
rep_key = (row.get(_INPUT_INDEX_FIELD), row.get("repetition"), row.get("turns"), row.get("payload_bytes"))
if row.get("mode") == "full":
full_rows[rep_key] = row
elif row.get("mode") == "delta":
group_key = tuple(row.get(field) for field in _GROUP_FIELDS)
delta_groups[group_key][rep_key] = row
summaries: list[dict[str, Any]] = []
for group_key in sorted(delta_groups, key=str):
deltas = delta_groups[group_key]
pairs = [(full_rows[k], d) for k, d in deltas.items() if k in full_rows]
if not pairs:
continue
summary: dict[str, Any] = {field: group_key[index] for index, field in enumerate(_GROUP_FIELDS)}
summary["paired_repetitions"] = len(pairs)
# Per-group failures: a failed delta row counts toward the group whose
# (turns, payload_bytes, snapshot_frequency) it matches; a failed full
# row (frequency-less) counts toward every group at the same
# (turns, payload_bytes).
summary["failed_rows"] = sum(1 for row in failed_rows if row.get("turns") == group_key[0] and row.get("payload_bytes") == group_key[1] and (row.get("snapshot_frequency") is None or row.get("snapshot_frequency") == group_key[2]))
full_pair_rows = [pair[0] for pair in pairs]
delta_pair_rows = [pair[1] for pair in pairs]
for metric in metrics:
full_median = _medians(full_pair_rows, metric)
delta_median = _medians(delta_pair_rows, metric)
if full_median is None or delta_median is None:
continue
summary[f"full_{metric}"] = round(full_median, 6)
summary[f"delta_{metric}"] = round(delta_median, 6)
summary[f"ratio_{metric}"] = round(delta_median / full_median, 6) if full_median != 0 else None
for label, group_rows in (("full", full_pair_rows), ("delta", delta_pair_rows)):
p50 = _medians(group_rows, "checkpoint_write_p50_ms")
p99 = _medians(group_rows, "checkpoint_write_p99_ms")
run_p50 = _medians(group_rows, "run_turn_p50_ms")
cold = _medians(group_rows, "state_cold_ms")
warm = _medians(group_rows, "state_warm_p50_ms")
if p50 and p99 is not None:
summary[f"{label}_snapshot_write_spike"] = round(p99 / p50, 6)
if p50 is not None and run_p50:
summary[f"{label}_checkpoint_write_share"] = round(p50 / run_p50, 6)
if cold is not None and warm is not None:
summary[f"{label}_cache_effect_ms"] = round(cold - warm, 6)
summaries.append(summary)
return summaries
def _columns(rows: list[dict[str, Any]]) -> list[str]:
preferred = [*_GROUP_FIELDS, "paired_repetitions", "failed_rows"]
extras = sorted({key for row in rows for key in row if key not in preferred})
return [field for field in preferred if any(field in row for row in rows)] + extras
def _print_table(rows: list[dict[str, Any]]) -> None:
if not rows:
print("(no comparable full/delta pairs)")
return
columns = _columns(rows)
widths = {column: len(column) for column in columns}
for row in rows:
for column in columns:
widths[column] = max(widths[column], len(str(row.get(column, ""))))
print(" ".join(column.rjust(widths[column]) for column in columns))
for row in rows:
print(" ".join(str(row.get(column, "")).rjust(widths[column]) for column in columns))
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description="Summarize production-shaped checkpoint benchmark results")
parser.add_argument("inputs", nargs="+", type=Path)
parser.add_argument("--metrics", default=None, help="comma-separated metric fields; defaults to the built-in list plus auto-discovered history_limit_* metrics")
output = parser.add_mutually_exclusive_group()
output.add_argument("--json", action="store_true")
output.add_argument("--csv", action="store_true")
args = parser.parse_args(argv)
try:
rows = _load_jsonl(args.inputs)
except (OSError, ValueError) as exc:
print(str(exc), file=sys.stderr)
return 1
if args.metrics is None:
metrics = [*DEFAULT_METRICS, *_discover_history_metrics(rows)]
else:
metrics = [metric.strip() for metric in args.metrics.split(",") if metric.strip()]
if not metrics:
parser.error("--metrics must contain at least one field")
summaries = _summarize(rows, metrics=metrics)
if args.json:
json.dump(summaries, sys.stdout, ensure_ascii=False, indent=2)
print()
elif args.csv:
writer = csv.DictWriter(sys.stdout, fieldnames=_columns(summaries), extrasaction="ignore")
writer.writeheader()
writer.writerows(summaries)
else:
_print_table(summaries)
return 0 if summaries else 1
if __name__ == "__main__":
raise SystemExit(main())

View File

@ -6,7 +6,7 @@ workloads, and concurrency levels. Outputs JSONL for aggregation.
Usage::
python scripts/benchmark/bench_sandbox_provider.py \\
python scripts/benchmark/sandbox/bench_provider.py \\
--provider boxlite \\
--scenario warm_same_thread \\
--workload noop \\
@ -14,7 +14,7 @@ Usage::
--concurrency 4 \\
--output results.jsonl
python scripts/benchmark/bench_sandbox_provider.py \\
python scripts/benchmark/sandbox/bench_provider.py \\
--provider boxlite \\
--scenario cold_unique_thread \\
--no-warmpool \\

View File

@ -3,9 +3,9 @@
Usage::
python scripts/benchmark/summarize_bench.py results.jsonl
python scripts/benchmark/summarize_bench.py results/*.jsonl --group provider,scenario,workload
python scripts/benchmark/summarize_bench.py results.jsonl --csv > summary.csv
python scripts/benchmark/sandbox/summarize.py results.jsonl
python scripts/benchmark/sandbox/summarize.py results/*.jsonl --group provider,scenario,workload
python scripts/benchmark/sandbox/summarize.py results.jsonl --csv > summary.csv
"""
from __future__ import annotations
@ -148,7 +148,7 @@ def _print_table(rows: list[dict[str, Any]], fmt: str = "plain") -> None:
def main(argv: list[str] | None = None) -> int:
p = argparse.ArgumentParser(description="Aggregate JSONL benchmark results")
p.add_argument("inputs", nargs="+", help="JSONL file(s) from bench_sandbox_provider.py")
p.add_argument("inputs", nargs="+", help="JSONL file(s) from sandbox/bench_provider.py")
p.add_argument(
"--group",
default="provider,scenario,workload,concurrency",

View File

@ -0,0 +1,56 @@
"""Regression anchor: DingTalk ``receive_file`` must not block the event loop.
``_receive_single_file`` prepares the thread directories, resolves the uploads
dir, scans it for the uniqueness claim, and writes the attachment all blocking
filesystem IO that must run inside ``asyncio.to_thread`` (and sandbox sync must
go through ``acquire_async`` + an offloaded ``update_file``). This anchor drives
the real ``receive_file`` under the strict Blockbuster gate; if any of that
regresses back onto the event loop, Blockbuster raises ``BlockingError``.
The ``Paths`` construction is offloaded only because ``Paths.__init__`` resolves
paths synchronously; the surface under test (``receive_file``'s persist path) is
exercised on the event loop, not bypassed. The download itself is mocked the
network leg is httpx-async and not the subject here.
"""
from __future__ import annotations
import asyncio
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
pytestmark = pytest.mark.asyncio
async def test_receive_file_persist_does_not_block_event_loop(tmp_path, monkeypatch) -> None:
from app.channels.dingtalk import DingTalkChannel
from app.channels.message_bus import MessageBus
from deerflow.config.paths import Paths
paths = await asyncio.to_thread(Paths, str(tmp_path))
monkeypatch.setattr("app.channels.dingtalk.get_paths", lambda: paths)
async def _acquire_async(thread_id, user_id=None):
return "local"
monkeypatch.setattr(
"app.channels.dingtalk.get_sandbox_provider",
lambda: SimpleNamespace(acquire_async=_acquire_async, get=lambda sid: None),
)
channel = DingTalkChannel(MessageBus(), config={})
channel._download_by_code = AsyncMock(return_value=b"DATA")
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="hi",
thread_ts="m",
files=[{"type": "file", "download_code": "dc", "filename": "a.pdf"}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
assert "/uploads/a.pdf" in out.text
assert out.files == []

View File

@ -113,6 +113,20 @@ def _reset_frozen_checkpoint_channel_mode(monkeypatch):
yield
@pytest.fixture(autouse=True)
def _reset_frozen_delta_snapshot_frequency(monkeypatch):
"""Reset the process-global frozen delta snapshot frequency between tests.
Same restart-required semantics as the checkpoint channel mode freeze
above: ``make_lead_agent`` freezes it from the app config, and the suite
builds many agents with different configs in one process.
"""
from deerflow.agents import thread_state
monkeypatch.setattr(thread_state, "_frozen_delta_snapshot_frequency", None)
yield
@pytest.fixture(autouse=True)
def _restore_title_config_singleton():
"""Reset ``_title_config`` to its pristine default after every test.

View File

@ -5,6 +5,17 @@ from starlette.testclient import TestClient
from app.gateway.auth_middleware import AuthMiddleware, _is_public
from app.gateway.csrf_middleware import CSRFMiddleware
from deerflow.config.authorization_config import AuthorizationConfig
@pytest.fixture(autouse=True)
def _default_route_authorization_config(monkeypatch):
"""Keep minimal middleware apps independent of a repository config.yaml."""
monkeypatch.setattr(
"app.gateway.authz._get_route_authorization_config",
lambda: AuthorizationConfig(),
)
# ── _is_public unit tests ─────────────────────────────────────────────────

View File

@ -0,0 +1,355 @@
"""Route-level authorization tests for the Gateway permission decorators."""
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from fastapi import FastAPI, Request
from fastapi.testclient import TestClient
from app.gateway.auth.models import User
from app.gateway.auth_middleware import AuthMiddleware
from app.gateway.authz import (
Permissions,
_authenticate,
_get_cached_route_provider,
require_permission,
resolve_route_permissions,
)
from deerflow.authz.provider import AuthzDecision, AuthzReason
from deerflow.authz.rbac import RbacAuthorizationProvider
from deerflow.config.authorization_config import AuthorizationConfig, AuthorizationProviderConfig
class _RecordingProvider:
name = "recording"
def __init__(
self,
*,
denied: set[str] | None = None,
errors: set[str] | None = None,
) -> None:
self.denied = denied or set()
self.errors = errors or set()
self.requests = []
def authorize(self, request):
raise AssertionError("route authorization must use the async provider API")
async def aauthorize(self, request):
self.requests.append(request)
if request.target in self.errors:
raise RuntimeError(f"provider failed for {request.target}")
allowed = request.target not in self.denied
return AuthzDecision(
allow=allowed,
reasons=[AuthzReason(code="authz.allowed" if allowed else "authz.denied")],
)
def filter_resources(self, principal, resource_type, candidates):
raise AssertionError("route authorization must preserve per-action requests")
def _user(**overrides):
values = {
"id": "user-123",
"system_role": "user",
"oauth_provider": "github",
"oauth_id": "oauth-456",
}
values.update(overrides)
return SimpleNamespace(**values)
def _enable_authorization(monkeypatch, provider, *, fail_closed: bool = True) -> None:
config = AuthorizationConfig(
enabled=True,
fail_closed=fail_closed,
default_role="user",
)
monkeypatch.setattr("app.gateway.authz._get_route_authorization_config", lambda: config)
# Bypass the provider cache so each test gets its own provider instance.
monkeypatch.setattr("app.gateway.authz._get_cached_route_provider", lambda c: provider)
@pytest.mark.asyncio
async def test_route_permissions_disabled_preserves_all_permissions(monkeypatch):
config = AuthorizationConfig(enabled=False)
monkeypatch.setattr("app.gateway.authz._get_route_authorization_config", lambda: config)
# Bypass cache + ensure provider is never resolved when disabled.
cached = AsyncMock(side_effect=AssertionError("disabled authorization must not resolve a provider"))
monkeypatch.setattr("app.gateway.authz._get_cached_route_provider", cached)
permissions = await resolve_route_permissions(_user(), is_internal=False)
assert permissions == [
Permissions.THREADS_READ,
Permissions.THREADS_WRITE,
Permissions.THREADS_DELETE,
Permissions.RUNS_CREATE,
Permissions.RUNS_READ,
Permissions.RUNS_CANCEL,
]
cached.assert_not_called()
@pytest.mark.asyncio
async def test_route_permissions_use_async_provider_and_trusted_principal(monkeypatch):
provider = _RecordingProvider(denied={Permissions.THREADS_DELETE, Permissions.RUNS_CANCEL})
_enable_authorization(monkeypatch, provider)
permissions = await resolve_route_permissions(_user(), is_internal=True)
assert permissions == [
Permissions.THREADS_READ,
Permissions.THREADS_WRITE,
Permissions.RUNS_CREATE,
Permissions.RUNS_READ,
]
assert [(request.resource, request.action, request.target) for request in provider.requests] == [
("route", "read", Permissions.THREADS_READ),
("route", "write", Permissions.THREADS_WRITE),
("route", "delete", Permissions.THREADS_DELETE),
("route", "create", Permissions.RUNS_CREATE),
("route", "read", Permissions.RUNS_READ),
("route", "cancel", Permissions.RUNS_CANCEL),
]
principal = provider.requests[0].principal
assert principal.user_id == "user-123"
assert principal.role == "user"
assert principal.oauth_provider == "github"
assert principal.oauth_id == "oauth-456"
assert principal.is_internal is True
@pytest.mark.asyncio
async def test_route_permissions_fail_closed_denies_only_the_failed_permission(monkeypatch):
provider = _RecordingProvider(errors={Permissions.RUNS_CANCEL})
_enable_authorization(monkeypatch, provider, fail_closed=True)
permissions = await resolve_route_permissions(_user(), is_internal=False)
assert permissions == [
Permissions.THREADS_READ,
Permissions.THREADS_WRITE,
Permissions.THREADS_DELETE,
Permissions.RUNS_CREATE,
Permissions.RUNS_READ,
]
@pytest.mark.asyncio
async def test_route_permissions_fail_open_allows_the_failed_permission(monkeypatch):
provider = _RecordingProvider(errors={Permissions.RUNS_CANCEL})
_enable_authorization(monkeypatch, provider, fail_closed=False)
permissions = await resolve_route_permissions(_user(), is_internal=False)
assert permissions == [
Permissions.THREADS_READ,
Permissions.THREADS_WRITE,
Permissions.THREADS_DELETE,
Permissions.RUNS_CREATE,
Permissions.RUNS_READ,
Permissions.RUNS_CANCEL,
]
@pytest.mark.asyncio
@pytest.mark.parametrize(
("fail_closed", "expected"),
[
(True, []),
(
False,
[
Permissions.THREADS_READ,
Permissions.THREADS_WRITE,
Permissions.THREADS_DELETE,
Permissions.RUNS_CREATE,
Permissions.RUNS_READ,
Permissions.RUNS_CANCEL,
],
),
],
)
async def test_route_permissions_apply_failure_mode_to_provider_resolution(monkeypatch, fail_closed, expected):
config = AuthorizationConfig(
enabled=True,
fail_closed=fail_closed,
default_role="user",
)
monkeypatch.setattr("app.gateway.authz._get_route_authorization_config", lambda: config)
def fail_cached(c):
raise ValueError("invalid provider configuration")
monkeypatch.setattr("app.gateway.authz._get_cached_route_provider", fail_cached)
assert await resolve_route_permissions(_user(), is_internal=False) == expected
@pytest.mark.asyncio
async def test_route_permissions_use_builtin_rbac_route_policy(monkeypatch):
provider = RbacAuthorizationProvider(
roles={
"user": {
"routes": {
"allow": [Permissions.THREADS_READ, Permissions.RUNS_READ],
}
}
}
)
_enable_authorization(monkeypatch, provider)
permissions = await resolve_route_permissions(_user(), is_internal=False)
assert permissions == [Permissions.THREADS_READ, Permissions.RUNS_READ]
@pytest.mark.asyncio
async def test_authenticate_uses_route_permission_resolution(monkeypatch):
user = User(email="route-authz@example.com", password_hash="hash")
permission_resolver = AsyncMock(return_value=[Permissions.THREADS_READ])
monkeypatch.setattr("app.gateway.deps.get_optional_user_from_request", AsyncMock(return_value=user))
monkeypatch.setattr("app.gateway.authz.resolve_route_permissions", permission_resolver)
request = SimpleNamespace(state=SimpleNamespace())
auth_context = await _authenticate(request)
assert auth_context.user is user
assert auth_context.permissions == [Permissions.THREADS_READ]
permission_resolver.assert_awaited_once_with(user, is_internal=False)
def _make_middleware_app() -> FastAPI:
app = FastAPI()
app.add_middleware(AuthMiddleware)
@app.get("/api/threads")
@require_permission("threads", "read")
async def read_threads(request: Request):
return {"ok": True}
@app.delete("/api/threads")
@require_permission("threads", "delete")
async def delete_threads(request: Request):
return {"ok": True}
return app
def test_auth_middleware_stamps_provider_derived_permissions(monkeypatch):
monkeypatch.setenv("DEER_FLOW_AUTH_DISABLED", "1")
permission_resolver = AsyncMock(return_value=[Permissions.THREADS_READ])
monkeypatch.setattr("app.gateway.auth_middleware.resolve_route_permissions", permission_resolver)
with TestClient(_make_middleware_app()) as client:
assert client.get("/api/threads").status_code == 200
assert client.delete("/api/threads").status_code == 403
assert permission_resolver.await_count == 2
for call in permission_resolver.await_args_list:
assert call.kwargs == {"is_internal": False}
def test_auth_middleware_marks_internal_route_principal(monkeypatch):
from app.gateway.internal_auth import create_internal_auth_headers
permission_resolver = AsyncMock(return_value=[Permissions.THREADS_READ])
monkeypatch.setattr("app.gateway.auth_middleware.resolve_route_permissions", permission_resolver)
with TestClient(_make_middleware_app()) as client:
response = client.get("/api/threads", headers=create_internal_auth_headers())
assert response.status_code == 200
permission_resolver.assert_awaited_once()
assert permission_resolver.await_args.kwargs == {"is_internal": True}
# ── Provider cache tests ────────────────────────────────────────────────
class TestRouteProviderCache:
"""Verify the provider cache returns the same instance for unchanged config
and re-resolves when config content changes."""
def test_same_config_returns_same_provider(self):
"""Calling twice with the same config object returns the same instance."""
import app.gateway.authz as authz_module
# Reset cache
authz_module._route_provider_cache.clear()
authz_module._route_provider_config_id = None
authz_module._route_provider_config_sig = None
config = AuthorizationConfig(
enabled=True,
provider=AuthorizationProviderConfig(
use="deerflow.authz.rbac:RbacAuthorizationProvider",
config={"roles": {"user": {"routes": {"allow": "*"}}}},
),
)
p1 = _get_cached_route_provider(config)
p2 = _get_cached_route_provider(config)
assert p1 is not None
assert p2 is p1
def test_changed_config_returns_new_provider(self):
"""A config with different content triggers re-resolution."""
import app.gateway.authz as authz_module
authz_module._route_provider_cache.clear()
authz_module._route_provider_config_id = None
authz_module._route_provider_config_sig = None
config1 = AuthorizationConfig(
enabled=True,
provider=AuthorizationProviderConfig(
use="deerflow.authz.rbac:RbacAuthorizationProvider",
config={"roles": {"user": {"routes": {"allow": "*"}}}},
),
)
config2 = AuthorizationConfig(
enabled=True,
provider=AuthorizationProviderConfig(
use="deerflow.authz.rbac:RbacAuthorizationProvider",
config={"roles": {"user": {"routes": {"allow": []}}}},
),
)
p1 = _get_cached_route_provider(config1)
p2 = _get_cached_route_provider(config2)
assert p1 is not None
assert p2 is not None
assert p1 is not p2
def test_same_content_different_object_reuses_provider(self):
"""Same content in a new object (e.g. hot-reload with no changes) reuses provider."""
import app.gateway.authz as authz_module
authz_module._route_provider_cache.clear()
authz_module._route_provider_config_id = None
authz_module._route_provider_config_sig = None
config1 = AuthorizationConfig(
enabled=True,
provider=AuthorizationProviderConfig(
use="deerflow.authz.rbac:RbacAuthorizationProvider",
config={"roles": {"user": {"routes": {"allow": "*"}}}},
),
)
# Same content, different object
config2 = AuthorizationConfig(
enabled=True,
provider=AuthorizationProviderConfig(
use="deerflow.authz.rbac:RbacAuthorizationProvider",
config={"roles": {"user": {"routes": {"allow": "*"}}}},
),
)
p1 = _get_cached_route_provider(config1)
p2 = _get_cached_route_provider(config2)
assert p1 is p2

View File

@ -9,7 +9,7 @@ import pytest
def _load_module():
path = Path(__file__).resolve().parents[1] / "scripts/benchmark/bench_checkpoint_channels.py"
path = Path(__file__).resolve().parents[1] / "scripts/benchmark/checkpoint/bench_channels.py"
spec = importlib.util.spec_from_file_location("bench_checkpoint_channels", path)
assert spec is not None
module = importlib.util.module_from_spec(spec)

View File

@ -0,0 +1,237 @@
"""Harness tests for the production-shaped checkpoint benchmark."""
from __future__ import annotations
import importlib.util
import sys
from pathlib import Path
import pytest
def _load_module():
path = Path(__file__).resolve().parents[1] / "scripts/benchmark/checkpoint/bench_production.py"
spec = importlib.util.spec_from_file_location("bench_production", path)
assert spec is not None
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
bench = _load_module()
def test_expand_cases_sweeps_frequency_for_delta_only() -> None:
cases = bench.expand_cases(
modes=["full", "delta"],
turn_counts=[10],
payload_bytes=[128],
snapshot_frequencies=[10, 50],
repetitions=1,
)
delta = [c for c in cases if c.mode == "delta"]
full = [c for c in cases if c.mode == "full"]
assert sorted(c.snapshot_frequency for c in delta) == [10, 50]
assert [c.snapshot_frequency for c in full] == [None]
assert all(c.history_limits == (10, 50, 200) for c in cases)
def test_expand_cases_alternates_mode_order() -> None:
cases = bench.expand_cases(
modes=["full", "delta"],
turn_counts=[10, 100],
payload_bytes=[128],
snapshot_frequencies=[10],
repetitions=2,
seed=1,
)
grouped_modes: dict[tuple[int, int], list[str]] = {}
for case in cases:
grouped_modes.setdefault((case.repetition, case.turns), []).append(case.mode)
mode_orders = list(grouped_modes.values())
assert all(set(order) == {"full", "delta"} for order in mode_orders)
assert all(current == list(reversed(previous)) for previous, current in zip(mode_orders, mode_orders[1:], strict=False))
def test_case_rejects_nonpositive_turns() -> None:
with pytest.raises(ValueError, match="turns"):
bench.ProductionCase(
mode="full",
turns=0,
payload_bytes=128,
snapshot_frequency=None,
history_limits=(10,),
read_repetitions=5,
repetition=0,
seed=1,
)
@pytest.mark.parametrize("turns", [1, 2])
def test_case_rejects_turns_consumed_by_warmup(turns: int) -> None:
with pytest.raises(ValueError, match="warm-up"):
bench.ProductionCase(
mode="full",
turns=turns,
payload_bytes=128,
snapshot_frequency=None,
history_limits=(10,),
read_repetitions=5,
repetition=0,
seed=1,
)
def test_case_rejects_frequency_on_full_mode() -> None:
with pytest.raises(ValueError, match="snapshot_frequency"):
bench.ProductionCase(
mode="full",
turns=10,
payload_bytes=128,
snapshot_frequency=10,
history_limits=(10,),
read_repetitions=5,
repetition=0,
seed=1,
)
def test_cross_mode_digest_gate_fails_both_rows_on_mismatch() -> None:
base = {"turns": 10, "payload_bytes": 128, "repetition": 0, "success": True, "seed_content_sha256": "a", "state_wire_sha256": "b", "history_wire_sha256": "c"}
full = {**base, "mode": "full", "snapshot_frequency": None}
delta = {**base, "mode": "delta", "snapshot_frequency": 10}
rows = [dict(full), dict(delta)]
bench.validate_cross_mode_rows(rows)
assert all(r["success"] for r in rows)
delta_bad = {**delta, "seed_content_sha256": "DIFFERENT"}
rows = [dict(full), delta_bad]
bench.validate_cross_mode_rows(rows)
assert not any(r["success"] for r in rows)
assert all(r["error"] == "cross-mode materialized state mismatch" for r in rows)
def test_cross_mode_gate_tolerates_missing_pair() -> None:
row = {"mode": "full", "snapshot_frequency": None, "turns": 10, "payload_bytes": 128, "repetition": 0, "success": True, "seed_content_sha256": "a", "state_wire_sha256": "b", "history_wire_sha256": "c"}
rows = [dict(row)]
bench.validate_cross_mode_rows(rows)
assert rows[0]["success"] is True
def test_scripted_model_is_deterministic_across_instances() -> None:
from langchain_core.messages import HumanMessage
first = bench._ScriptedBenchModel(payload_bytes=16)
second = bench._ScriptedBenchModel(payload_bytes=16)
for _ in range(3):
a = first.invoke([HumanMessage(content="hi", id="h0")])
b = second.invoke([HumanMessage(content="hi", id="h0")])
assert a.id == b.id and a.content == b.content and len(a.content) == 16
def test_timing_saver_records_cumulative_ms() -> None:
import asyncio
from langgraph.checkpoint.base import empty_checkpoint
from langgraph.checkpoint.memory import InMemorySaver
timing = bench._TimingSaver(InMemorySaver())
config = {"configurable": {"thread_id": "t", "checkpoint_ns": "", "checkpoint_id": "c1"}}
async def go() -> None:
await timing.aput(config, empty_checkpoint(), {"source": "input", "step": 0, "writes": {}}, {})
await timing.aput_writes(config, [("ch", "v")], "task")
asyncio.run(go())
assert timing.timings_ms["aput"] >= 0
assert timing.timings_ms["aput_writes"] >= 0
assert timing._saver is not None
def _reset_checkpoint_freezes(monkeypatch: pytest.MonkeyPatch) -> None:
from deerflow.agents import thread_state
from deerflow.config.app_config import reset_app_config
from deerflow.runtime import checkpoint_mode
reset_app_config()
monkeypatch.setattr(checkpoint_mode, "_frozen_checkpoint_channel_mode", None)
monkeypatch.setattr(thread_state, "_frozen_delta_snapshot_frequency", None)
@pytest.mark.parametrize("mode,frequency", [("full", None), ("delta", 10)])
def test_run_case_succeeds_end_to_end_small(tmp_path, monkeypatch, mode, frequency) -> None:
_reset_checkpoint_freezes(monkeypatch)
from app.gateway import services as gateway_services
from deerflow.config.app_config import reset_app_config
gateway_services._state_accessor_graph_cache.clear()
case = bench.ProductionCase(
mode=mode,
turns=3,
payload_bytes=32,
snapshot_frequency=frequency,
history_limits=(2,),
read_repetitions=2,
repetition=0,
seed=1,
)
row = bench._run_case(case, work_dir=tmp_path)
assert row["success"], row.get("error")
assert row["message_count"] > 0
assert row["seed_content_sha256"]
assert row["state_warm_p50_ms"] >= 0
# Wire digests must reflect real messages: a duck-typed checkpointer on
# app.state silently materializes delta channels as empty (Pregel gates
# replay behind isinstance(checkpointer, BaseCheckpointSaver)), which
# otherwise passes every assertion above vacuously.
empty_wire = bench._wire_messages_digest([])
assert row["state_wire_sha256"] != empty_wire
assert row["history_wire_sha256"] != bench._combined_digest([empty_wire])
assert row["wal_bytes"] > 0
assert row["shm_bytes"] > 0
# Write cost is merged busy time over concurrent write tasks: it must not
# exceed the turn wall clock the way a naive latency sum can.
assert row["checkpoint_write_p50_ms"] <= row["run_turn_p50_ms"]
gateway_services._state_accessor_graph_cache.clear()
reset_app_config()
def test_run_case_fails_when_warm_cache_is_not_hit(tmp_path, monkeypatch) -> None:
"""The cold/warm contract is structural: if the accessor cache never holds,
warm samples silently measure the cold path the case row must fail."""
_reset_checkpoint_freezes(monkeypatch)
from app.gateway import services as gateway_services
from deerflow.config.app_config import reset_app_config
gateway_services._state_accessor_graph_cache.clear()
# Bypass the cache entirely: every accessor resolution rebuilds the graph,
# so warm reads trip the contract assertion.
monkeypatch.setattr(
gateway_services,
"_state_accessor_graph",
lambda agent_factory, assistant_id, mode, config: agent_factory(config=config),
)
case = bench.ProductionCase(
mode="full",
turns=3,
payload_bytes=32,
snapshot_frequency=None,
history_limits=(2,),
read_repetitions=2,
repetition=0,
seed=1,
)
row = bench._run_case(case, work_dir=tmp_path)
assert not row["success"]
assert "cache was not hit" in row["error"]
gateway_services._state_accessor_graph_cache.clear()
reset_app_config()
def test_merged_busy_ms_unions_overlapping_intervals() -> None:
intervals = [("aput", 0.0, 10.0), ("aput_writes", 5.0, 12.0), ("aput", 20.0, 25.0)]
assert bench._merged_busy_ms(iter(intervals)) == 17.0
assert bench._merged_busy_ms(iter([])) == 0.0

View File

@ -17,8 +17,8 @@ def _load_module(name: str, relative: str):
return module
bench = _load_module("bench_sandbox_provider", "scripts/benchmark/bench_sandbox_provider.py")
summarize = _load_module("summarize_bench", "scripts/benchmark/summarize_bench.py")
bench = _load_module("bench_provider", "scripts/benchmark/sandbox/bench_provider.py")
summarize = _load_module("summarize", "scripts/benchmark/sandbox/summarize.py")
class _FakeProvider:

View File

@ -1264,7 +1264,7 @@ class TestChannelManager:
outbound_received.append(msg)
if msg.is_final:
key = manager._inbound_dedupe_key(_inbound())
key_present_during_final_publish.append(key in manager._recent_inbound_events)
key_present_during_final_publish.append(key in manager._inbound_dedupe_store._store)
bus.subscribe_outbound(capture_outbound)
@ -1428,7 +1428,8 @@ class TestChannelManager:
else:
CHANNEL_RUN_POLICY[channel_name] = original
def test_github_redelivery_is_deduped_like_other_channels(self, tmp_path):
@pytest.mark.asyncio
async def test_github_redelivery_is_deduped_like_other_channels(self, tmp_path):
"""A redelivered GitHub webhook must dispatch the agent only once.
PR #3584 added inbound dedupe for the IM channels; the GitHub channel
@ -1467,18 +1468,18 @@ class TestChannelManager:
assert ChannelManager._inbound_dedupe_key(_gh("d1")) == ("github", "zhfeng/llm-gateway", "zhfeng/llm-gateway", "d1:alice:reviewer")
# First delivery fires; an identical redelivery of the same GUID is dropped.
assert manager._is_duplicate_inbound(_gh("d1")) is False
assert manager._is_duplicate_inbound(_gh("d1")) is True
assert await manager._is_duplicate_inbound(_gh("d1")) is False
assert await manager._is_duplicate_inbound(_gh("d1")) is True
# A genuinely new delivery still fires.
assert manager._is_duplicate_inbound(_gh("d2")) is False
assert await manager._is_duplicate_inbound(_gh("d2")) is False
# A second agent fanned out from the SAME delivery is not cross-deduped.
assert manager._is_duplicate_inbound(_gh("d1", agent="coder")) is False
assert await manager._is_duplicate_inbound(_gh("d1", agent="coder")) is False
# A second user's SAME-named agent on the SAME delivery is not
# cross-deduped either. A helper still stamping the old 2-part
# (delivery, agent) id could not even express this case — it would
# collide with the very first assertion's "d1"+"reviewer" key and
# silently drop this user's run (willem-bd, PR #4104 review).
assert manager._is_duplicate_inbound(_gh("d1", owner_user_id="bob")) is False
assert await manager._is_duplicate_inbound(_gh("d1", owner_user_id="bob")) is False
def test_dispatch_loop_releases_dedupe_key_when_handling_fails(self, tmp_path):
"""A transient handling failure must not black-hole a provider redelivery (ShenAC #1)."""
@ -8881,6 +8882,7 @@ class TestTelegramStreaming:
bot = SimpleNamespace()
bot.sent = []
bot.edited = []
bot.rich = []
bot.next_message_id = 100
async def send_message(**kwargs):
@ -8896,8 +8898,15 @@ class TestTelegramStreaming:
result.message_id = kwargs["message_id"]
return result
async def do_api_request(endpoint, api_kwargs):
bot.rich.append((endpoint, api_kwargs))
result = {"message_id": bot.next_message_id}
bot.next_message_id += 1
return result
bot.send_message = send_message
bot.edit_message_text = edit_message_text
bot.do_api_request = do_api_request
mock_app.bot = bot
ch._application = mock_app
return ch, bot
@ -9194,6 +9203,104 @@ class TestTelegramStreaming:
_run(go())
def test_final_uses_telegram_rich_markdown_when_enabled(self):
async def go():
ch, bot = self._make_channel_with_bot()
ch.config["rich_messages"] = True
markdown = "# Result\n\n**Bold** and `code`"
await ch.send(OutboundMessage(channel_name="telegram", chat_id="12345", thread_id="t1", text=markdown, is_final=True))
assert bot.rich == [
(
"sendRichMessage",
{"chat_id": 12345, "rich_message": {"markdown": markdown}},
)
]
assert bot.sent == []
_run(go())
def test_final_replaces_plain_stream_with_rich_message(self, monkeypatch):
async def go():
ch, bot = self._make_channel_with_bot()
ch.config["rich_messages"] = True
monkeypatch.setattr("app.channels.telegram._monotonic", lambda: 1000.0)
await ch._send_running_reply("12345", 42)
await ch.send(OutboundMessage(channel_name="telegram", chat_id="12345", thread_id="t1", text="**partial", is_final=False, thread_ts="42"))
await ch.send(OutboundMessage(channel_name="telegram", chat_id="12345", thread_id="t1", text="**final**", is_final=True, thread_ts="42"))
assert bot.rich == [
(
"editMessageText",
{
"chat_id": 12345,
"message_id": 100,
"rich_message": {"markdown": "**final**"},
},
)
]
assert [message["text"] for message in bot.sent] == ["Working on it..."]
_run(go())
def test_rich_message_bad_request_falls_back_to_plain_text_once(self):
from telegram.error import BadRequest
async def go():
ch, bot = self._make_channel_with_bot()
ch.config["rich_messages"] = True
async def reject_rich(endpoint, api_kwargs):
raise BadRequest("Can't parse rich message")
bot.do_api_request = reject_rich
await ch.send(OutboundMessage(channel_name="telegram", chat_id="12345", thread_id="t1", text="**answer**", is_final=True))
assert [message["text"] for message in bot.sent] == ["**answer**"]
_run(go())
def test_rich_message_retryable_failure_falls_back_to_plain_text_once(self):
async def go():
ch, bot = self._make_channel_with_bot()
ch.config["rich_messages"] = True
async def fail_rich(endpoint, api_kwargs):
raise RuntimeError("network failed")
bot.do_api_request = fail_rich
await ch.send(
OutboundMessage(channel_name="telegram", chat_id="12345", thread_id="t1", text="**answer**", is_final=True),
_max_retries=1,
)
assert [message["text"] for message in bot.sent] == ["**answer**"]
_run(go())
def test_stream_rich_message_bad_request_falls_back_to_one_plain_text_edit(self, monkeypatch):
from telegram.error import BadRequest
async def go():
ch, bot = self._make_channel_with_bot()
ch.config["rich_messages"] = True
monkeypatch.setattr("app.channels.telegram._monotonic", lambda: 1000.0)
await ch._send_running_reply("12345", 42)
async def fail_rich(endpoint, api_kwargs):
raise BadRequest("Can't parse rich message")
bot.do_api_request = fail_rich
await ch.send(OutboundMessage(channel_name="telegram", chat_id="12345", thread_id="t1", text="**answer**", is_final=True, thread_ts="42"))
assert [message["text"] for message in bot.sent] == ["Working on it..."]
assert bot.edited == [{"chat_id": 12345, "message_id": 100, "text": "**answer**"}]
_run(go())
def test_final_edit_retries_once_after_rate_limit(self, monkeypatch):
async def go():
ch, bot = self._make_channel_with_bot()

View File

@ -249,6 +249,31 @@ def test_direct_langgraph_request_cannot_select_delta_in_full_process(
assert CHECKPOINT_MODE_METADATA_KEY not in config["metadata"]
def test_make_lead_agent_freezes_delta_snapshot_frequency_from_app_config(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from deerflow.agents import thread_state
from deerflow.agents.lead_agent import agent as lead_agent
from deerflow.config.app_config import AppConfig
app_config = AppConfig.model_validate(
{
"sandbox": {"use": "deerflow.sandbox.local.provider:LocalSandboxProvider"},
"database": {"checkpoint_delta_snapshot_frequency": 7},
}
)
monkeypatch.setattr(lead_agent, "get_app_config", lambda: app_config)
monkeypatch.setattr(
lead_agent,
"_make_lead_agent",
lambda config, *, app_config: object(),
)
lead_agent.make_lead_agent({"configurable": {}})
assert thread_state._frozen_delta_snapshot_frequency == 7
def test_gateway_runtime_app_config_can_supply_its_frozen_internal_mode(
monkeypatch: pytest.MonkeyPatch,
) -> None:

View File

@ -1015,6 +1015,7 @@ class TestClientCheckpointerFallback:
model_mock = MagicMock()
config_mock = MagicMock()
config_mock.models = [model_mock]
config_mock.database.checkpoint_delta_snapshot_frequency = 1000
config_mock.get_model_config.return_value = MagicMock(supports_vision=False)
config_mock.checkpointer = None
@ -1054,6 +1055,7 @@ class TestClientCheckpointerFallback:
model_mock = MagicMock()
config_mock = MagicMock()
config_mock.models = [model_mock]
config_mock.database.checkpoint_delta_snapshot_frequency = 1000
config_mock.get_model_config.return_value = MagicMock(supports_vision=False)
config_mock.checkpointer = None

View File

@ -215,6 +215,396 @@ class TestHumanInputPayload:
assert "options" not in payload
class TestFormPayload:
"""v2 protocol: fields render a structured form card."""
def _fields(self):
return [
{"name": "amount", "label": "Amount", "type": "number", "required": True},
{"name": "category", "label": "Category", "type": "select", "options": ["travel", "meals"], "required": True},
{"name": "receipts", "label": "Receipts", "type": "multi_select", "options": ["A-1", "A-2"]},
{"name": "note", "type": "textarea"},
]
def test_form_payload_with_native_fields(self, middleware):
payload = middleware._build_human_input_payload(
{
"question": "Please provide the expense details.",
"clarification_type": "missing_info",
"fields": self._fields(),
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["version"] == 2
assert payload["input_mode"] == "form"
assert payload["fields"] == [
{"name": "amount", "label": "Amount", "type": "number", "required": True},
{
"name": "category",
"label": "Category",
"type": "select",
"required": True,
"options": [
{"id": "category-option-1", "label": "travel", "value": "travel"},
{"id": "category-option-2", "label": "meals", "value": "meals"},
],
},
{
"name": "receipts",
"label": "Receipts",
"type": "multi_select",
"required": False,
"options": [
{"id": "receipts-option-1", "label": "A-1", "value": "A-1"},
{"id": "receipts-option-2", "label": "A-2", "value": "A-2"},
],
},
{"name": "note", "label": "note", "type": "textarea", "required": False},
]
def test_form_payload_with_json_string_fields(self, middleware):
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": json.dumps([{"name": "amount", "type": "number"}]),
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["input_mode"] == "form"
assert payload["fields"] == [{"name": "amount", "label": "amount", "type": "number", "required": False}]
def test_form_takes_precedence_over_options(self, middleware):
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [{"name": "amount"}],
"options": ["dev", "prod"],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["input_mode"] == "form"
assert "options" not in payload
def test_unknown_field_type_degrades_to_text(self, middleware):
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [{"name": "amount", "type": "slider"}],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["fields"][0]["type"] == "text"
def test_select_without_options_degrades_to_text(self, middleware):
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [{"name": "category", "type": "select"}],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["fields"][0]["type"] == "text"
assert "options" not in payload["fields"][0]
def test_any_invalid_field_entry_degrades_whole_form(self, middleware):
"""Atomic validation: a structurally broken entry must not silently
vanish from an otherwise-rendered form the whole form degrades."""
for bad_entry in ["not-a-dict", {"label": "no name"}, {"name": " "}]:
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [bad_entry, {"name": "amount", "type": "number"}],
"options": ["dev", "prod"],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["version"] == 1
assert payload["input_mode"] == "choice_with_other"
assert "fields" not in payload
def test_duplicate_field_names_degrade_whole_form(self, middleware):
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [
{"name": "amount", "type": "number"},
{"name": "amount", "type": "text"},
],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["version"] == 1
assert payload["input_mode"] == "free_text"
assert "fields" not in payload
def test_reserved_prototype_field_names_degrade_whole_form(self, middleware):
"""`__proto__`/`constructor`/... collide with JS Object.prototype in the
frontend form-value store and must be rejected server-side."""
for reserved in ["__proto__", "constructor", "toString", "hasOwnProperty"]:
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [{"name": reserved, "type": "text"}],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["input_mode"] == "free_text", reserved
assert "fields" not in payload
def test_field_options_are_trimmed_deduped_and_blank_dropped(self, middleware):
"""Backend must never emit blank option labels — the frontend parser
rejects the whole payload on them."""
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [{"name": "category", "type": "select", "options": [" travel ", "", " ", "travel", "meals"]}],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert [option["label"] for option in payload["fields"][0]["options"]] == ["travel", "meals"]
def test_field_count_over_cap_degrades_whole_form(self, middleware):
fields = [{"name": f"field_{i}", "type": "text"} for i in range(17)]
payload = middleware._build_human_input_payload(
{"question": "Details please", "clarification_type": "missing_info", "fields": fields},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["input_mode"] == "free_text"
assert "fields" not in payload
def test_option_count_over_cap_degrades_whole_form(self, middleware):
options = [f"option-{i}" for i in range(25)]
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [{"name": "category", "type": "select", "options": options}],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["input_mode"] == "free_text"
assert "fields" not in payload
def test_overlong_field_text_degrades_whole_form(self, middleware):
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [{"name": "amount", "label": "x" * 201, "type": "number"}],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["input_mode"] == "free_text"
assert "fields" not in payload
def test_top_level_options_are_trimmed_and_blank_dropped(self, middleware):
payload = middleware._build_human_input_payload(
{
"question": "Which env?",
"clarification_type": "approach_choice",
"options": [" dev ", "", "prod", "dev"],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["options"] == [
{"id": "option-1", "label": "dev", "value": "dev"},
{"id": "option-2", "label": "prod", "value": "prod"},
]
def test_unhashable_field_type_degrades_to_text_not_crash(self, middleware):
"""`type: []` / `type: {}` are legal JSON from a model; an unhashable
membership check would raise TypeError and (with return_direct=True)
end the turn with an error instead of any card or fallback."""
for bad_type in [[], {}, ["select"], {"t": "select"}]:
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [{"name": "amount", "type": bad_type}],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["fields"][0]["type"] == "text"
def test_unhashable_clarification_type_does_not_crash(self, middleware):
"""Same crash class as unhashable field types: `clarification_type: []`
must not raise in the icon lookup or payload builder."""
from langgraph.types import Command
request = SimpleNamespace(
tool_call={
"name": "ask_clarification",
"id": "call-clarify-1",
"args": {"question": "q?", "clarification_type": [], "fields": [{"name": "amount", "type": []}]},
},
runtime=None,
)
result = middleware.wrap_tool_call(request, lambda _req: pytest.fail("handler should not be called"))
assert isinstance(result, Command)
message = result.update["messages"][0]
assert message.artifact["human_input"]["input_mode"] == "form"
def test_serialized_fields_over_byte_budget_degrade_whole_form(self, middleware):
"""Per-item caps alone allow a form whose IM text fallback exceeds
channel limits (Slack 40k chars, Feishu ~30KB card); a total serialized
byte budget must bound the whole definition."""
fields = [
{
"name": f"field_{i}",
"label": "" * 190,
"type": "select",
"options": ["" * 190 + str(j) for j in range(20)],
}
for i in range(16)
]
payload = middleware._build_human_input_payload(
{"question": "Details please", "clarification_type": "missing_info", "fields": fields},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["input_mode"] == "free_text"
assert "fields" not in payload
def test_accepted_forms_keep_text_fallback_under_channel_limits(self, middleware):
"""Boundary: any form the byte budget accepts must produce an IM text
fallback deliverable to the strictest supported channel (~30KB)."""
# Grow a form until the budget rejects it; the largest accepted form's
# fallback must stay under the channel bound.
largest_accepted_fallback = None
for count in range(1, 17):
fields = [
{
"name": f"field_{i}",
"label": "" * 150,
"type": "select",
"options": ["" * 100 + str(j) for j in range(10)],
}
for i in range(count)
]
args = {"question": "Q", "clarification_type": "missing_info", "fields": fields}
payload = middleware._build_human_input_payload(args, tool_call_id="c", request_id="clarification:c")
if payload["input_mode"] != "form":
break
largest_accepted_fallback = middleware._format_clarification_message(args)
assert largest_accepted_fallback is not None
assert len(largest_accepted_fallback.encode("utf-8")) < 30_000
def test_required_accepts_integer_serialization(self, middleware):
"""Some providers emit 1/0 for booleans; `required: 1` must not
silently flip a model-intended required field to optional."""
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [
{"name": "amount", "type": "number", "required": 1},
{"name": "note", "type": "text", "required": 0},
],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["fields"][0]["required"] is True
assert payload["fields"][1]["required"] is False
def test_field_placeholder_is_preserved(self, middleware):
payload = middleware._build_human_input_payload(
{
"question": "Details please",
"clarification_type": "missing_info",
"fields": [{"name": "note", "placeholder": "Optional remarks"}],
},
tool_call_id="call-abc",
request_id="clarification:call-abc",
)
assert payload["fields"][0]["placeholder"] == "Optional remarks"
def test_form_fallback_text_lists_fields(self, middleware):
result = middleware._format_clarification_message(
{
"question": "Please provide the expense details.",
"clarification_type": "missing_info",
"fields": self._fields(),
}
)
assert "1. Amount (required)" in result
assert "2. Category (required) — options: travel / meals" in result
assert "3. Receipts — options: A-1 / A-2 (multiple allowed)" in result
assert "4. note" in result
assert "Please reply with a value for each field." in result
class TestClarificationToolSchema:
"""The tool schema must expose the v2 form parameter (request-side only)."""
def test_tool_exposes_fields_argument(self):
from deerflow.tools.builtins.clarification_tool import ask_clarification_tool
schema = ask_clarification_tool.args
assert "fields" in schema
# The reply protocol stays v1 (text/option): no top-level multi_select
# mode — a standalone multi-select question is a one-field form.
assert "multi_select" not in schema
def test_fields_item_schema_is_typed(self):
"""The provider-facing schema must expose the field item shape (typed
via ClarificationFormField), not an opaque object relying on the
docstring alone."""
from langchain_core.utils.function_calling import convert_to_openai_tool
from deerflow.tools.builtins.clarification_tool import ask_clarification_tool
parameters = convert_to_openai_tool(ask_clarification_tool)["function"]["parameters"]
items = parameters["properties"]["fields"]["anyOf"][0]["items"]
assert items["required"] == ["name"]
assert sorted(items["properties"].keys()) == ["label", "name", "options", "placeholder", "required", "type"]
assert items["properties"]["type"]["enum"] == ["text", "textarea", "number", "select", "multi_select", "checkbox", "date"]
class TestClarificationCommandIdempotency:
"""Clarification tool-call retries should not duplicate messages in state."""

View File

@ -52,6 +52,7 @@ def mock_app_config():
config.skills.container_path = "/mnt/skills"
config.tool_search.enabled = False
config.database.checkpoint_channel_mode = "full"
config.database.checkpoint_delta_snapshot_frequency = 1000
config.authorization = AuthorizationConfig(enabled=False)
return config
@ -145,6 +146,23 @@ class TestClientInit:
):
DeerFlowClient()
def test_delta_snapshot_frequency_is_frozen_from_app_config(self, mock_app_config):
from typing import get_type_hints
from langgraph.channels import DeltaChannel
from deerflow.agents import thread_state
mock_app_config.database.checkpoint_channel_mode = "delta"
mock_app_config.database.checkpoint_delta_snapshot_frequency = 7
with patch("deerflow.client.get_app_config", return_value=mock_app_config):
DeerFlowClient()
schema = thread_state.get_thread_state_schema("delta")
hint = get_type_hints(schema, include_extras=True)["messages"]
channel = next(item for item in hint.__metadata__ if isinstance(item, DeltaChannel))
assert channel.snapshot_frequency == 7
# ---------------------------------------------------------------------------
# list_models / list_skills / get_memory

View File

@ -13,6 +13,7 @@ from langgraph.checkpoint.memory import InMemorySaver
from langgraph.graph import StateGraph
from langgraph.graph.message import REMOVE_ALL_MESSAGES, add_messages
from deerflow.agents import thread_state
from deerflow.agents.thread_state import (
DeltaThreadState,
ThreadState,
@ -21,6 +22,7 @@ from deerflow.agents.thread_state import (
merge_message_writes,
normalize_middleware_state_schemas,
)
from deerflow.runtime.checkpoint_mode import CheckpointModeReconfigurationError
def _fold(state: list, writes: list) -> list:
@ -287,6 +289,20 @@ def test_delta_adaptation_replaces_agent_state_message_reducer() -> None:
assert any(isinstance(item, DeltaChannel) for item in hint.__metadata__)
def test_delta_adaptation_cache_varies_by_resolved_snapshot_frequency(monkeypatch) -> None:
monkeypatch.setattr(thread_state, "_frozen_delta_snapshot_frequency", None)
thread_state._adapt_state_schema_for_mode.cache_clear()
default_adapted = adapt_state_schema_for_mode(AgentState, "delta")
thread_state.freeze_delta_snapshot_frequency(7)
configured_adapted = adapt_state_schema_for_mode(AgentState, "delta")
assert configured_adapted is not default_adapted
hint = get_type_hints(configured_adapted, include_extras=True)["messages"]
channel = next(item for item in hint.__metadata__ if isinstance(item, DeltaChannel))
assert channel.snapshot_frequency == 7
def test_agents_package_exports_delta_thread_state() -> None:
from deerflow.agents import DeltaThreadState as ExportedDeltaThreadState
@ -369,3 +385,31 @@ def test_production_message_forms_keep_assigned_ids_across_delta_replay(write: o
assert first[0].id is not None
assert second[0].id == first[0].id
def test_delta_snapshot_frequency_defaults_to_1000(monkeypatch) -> None:
monkeypatch.setattr(thread_state, "_frozen_delta_snapshot_frequency", None)
assert thread_state.resolved_delta_snapshot_frequency() == 1000
def test_freeze_delta_snapshot_frequency_is_process_frozen(monkeypatch) -> None:
monkeypatch.setattr(thread_state, "_frozen_delta_snapshot_frequency", None)
assert thread_state.freeze_delta_snapshot_frequency(50) == 50
assert thread_state.freeze_delta_snapshot_frequency(50) == 50
with pytest.raises(CheckpointModeReconfigurationError):
thread_state.freeze_delta_snapshot_frequency(10)
def test_get_thread_state_schema_honors_frozen_frequency(monkeypatch) -> None:
monkeypatch.setattr(thread_state, "_frozen_delta_snapshot_frequency", None)
thread_state.freeze_delta_snapshot_frequency(7)
hint = get_type_hints(thread_state.get_thread_state_schema("delta"), include_extras=True)["messages"]
channel = next(item for item in hint.__metadata__ if isinstance(item, DeltaChannel))
assert channel.snapshot_frequency == 7
def test_database_config_default_snapshot_frequency() -> None:
from deerflow.config.database_config import DatabaseConfig
assert DatabaseConfig().checkpoint_delta_snapshot_frequency == 1000
assert DatabaseConfig(checkpoint_delta_snapshot_frequency=25).checkpoint_delta_snapshot_frequency == 25

View File

@ -3,6 +3,7 @@
from __future__ import annotations
import asyncio
import logging
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock
@ -22,6 +23,7 @@ from app.channels.dingtalk import (
_normalize_conversation_type,
)
from app.channels.message_bus import InboundMessageType, MessageBus, OutboundMessage
from deerflow.config.paths import VIRTUAL_PATH_PREFIX
def _run(coro):
@ -1691,3 +1693,841 @@ class TestCardMode:
assert channel._card_track_ids == {}
_run(go())
# ---------------------------------------------------------------------------
# Inbound file support
# ---------------------------------------------------------------------------
def _make_picture_message(*, download_code: str = "dc_img", **kwargs):
"""Build a mock DingTalk ``picture`` message carrying an ImageContent."""
msg = _make_chatbot_message(text="", message_type="picture", **kwargs)
msg.image_content = SimpleNamespace(download_code=download_code)
return msg
def _make_file_message(*, download_code: str = "dc_file", file_name: str | None = "report.xlsx", **kwargs):
"""Build a mock DingTalk ``file`` (document) message.
dingtalk_stream does not parse ``file`` messages, so the descriptor lives on
the raw callback payload the handler stashes as ``_df_raw_data``.
"""
msg = _make_chatbot_message(text="", message_type="file", **kwargs)
content: dict = {"downloadCode": download_code}
if file_name is not None:
content["fileName"] = file_name
msg._df_raw_data = {"msgtype": "file", "content": content}
return msg
def _patch_uploads(monkeypatch, uploads_dir, *, sandbox_id="local", sandbox=None):
monkeypatch.setattr(
"app.channels.dingtalk.get_paths",
lambda: SimpleNamespace(
ensure_thread_dirs=lambda thread_id, user_id=None: None,
sandbox_uploads_dir=lambda thread_id, user_id=None: uploads_dir,
),
)
monkeypatch.setattr("app.channels.dingtalk.get_effective_user_id", lambda: "default")
async def _acquire_async(thread_id, user_id=None):
return sandbox_id
monkeypatch.setattr(
"app.channels.dingtalk.get_sandbox_provider",
lambda: SimpleNamespace(acquire_async=_acquire_async, get=lambda sid: sandbox),
)
class TestExtractFiles:
def test_text_message_has_no_files(self):
assert DingTalkChannel._extract_files(_make_chatbot_message(text="hi")) == []
def test_picture_message(self):
files = DingTalkChannel._extract_files(_make_picture_message(download_code="dc_1"))
assert files == [{"type": "image", "download_code": "dc_1", "filename": "image.png"}]
def test_picture_message_missing_download_code_ignored(self):
msg = _make_chatbot_message(text="", message_type="picture")
msg.image_content = SimpleNamespace(download_code=None)
assert DingTalkChannel._extract_files(msg) == []
def test_rich_text_images(self):
msg = _make_chatbot_message(text="", message_type="richText")
msg.get_image_list = lambda: ["dc_a", "dc_b"]
files = DingTalkChannel._extract_files(msg)
assert files == [
{"type": "image", "download_code": "dc_a", "filename": "image_0.png"},
{"type": "image", "download_code": "dc_b", "filename": "image_1.png"},
]
def test_rich_text_get_image_list_failure_is_logged_not_silent(self, caplog):
"""An SDK parse failure must be distinguishable from "no inline images"."""
msg = _make_chatbot_message(text="", message_type="richText")
def _boom():
raise RuntimeError("sdk failure")
msg.get_image_list = _boom
with caplog.at_level(logging.WARNING, logger="app.channels.dingtalk"):
assert DingTalkChannel._extract_files(msg) == []
assert any("failed to read inline images" in r.message for r in caplog.records)
def test_file_message_from_raw_data(self):
files = DingTalkChannel._extract_files(_make_file_message(download_code="dc_doc", file_name="报价.xlsx"))
assert files == [{"type": "file", "download_code": "dc_doc", "filename": "报价.xlsx"}]
def test_file_message_default_filename(self):
files = DingTalkChannel._extract_files(_make_file_message(download_code="dc_doc", file_name=None))
assert files == [{"type": "file", "download_code": "dc_doc", "filename": "file.bin"}]
def test_file_message_without_raw_data_ignored(self):
msg = _make_chatbot_message(text="", message_type="file")
assert DingTalkChannel._extract_files(msg) == []
class TestOnChatbotMessageFiles:
def test_picture_message_publishes_inbound_with_files(self):
async def go():
bus = MessageBus()
bus.publish_inbound = AsyncMock()
channel = DingTalkChannel(bus, config={})
channel._client_id = "test_key"
channel._main_loop = asyncio.get_event_loop()
channel._running = True
channel._send_running_reply = AsyncMock()
channel._on_chatbot_message(_make_picture_message(download_code="dc_pic", sender_staff_id="u1", message_id="m1"))
await asyncio.sleep(0.1)
bus.publish_inbound.assert_awaited_once()
inbound = bus.publish_inbound.await_args.args[0]
assert inbound.msg_type == InboundMessageType.CHAT
assert inbound.files == [{"type": "image", "download_code": "dc_pic", "filename": "image.png"}]
_run(go())
def test_file_message_publishes_inbound_with_files(self):
async def go():
bus = MessageBus()
bus.publish_inbound = AsyncMock()
channel = DingTalkChannel(bus, config={})
channel._client_id = "test_key"
channel._main_loop = asyncio.get_event_loop()
channel._running = True
channel._send_running_reply = AsyncMock()
channel._on_chatbot_message(_make_file_message(download_code="dc_doc", file_name="data.csv", sender_staff_id="u1", message_id="m1"))
await asyncio.sleep(0.1)
bus.publish_inbound.assert_awaited_once()
inbound = bus.publish_inbound.await_args.args[0]
assert inbound.files == [{"type": "file", "download_code": "dc_doc", "filename": "data.csv"}]
_run(go())
def test_message_without_text_or_files_dropped(self):
async def go():
bus = MessageBus()
bus.publish_inbound = AsyncMock()
channel = DingTalkChannel(bus, config={})
channel._client_id = "test_key"
channel._main_loop = asyncio.get_event_loop()
channel._running = True
# picture message whose image carries no download_code -> no files, no text
msg = _make_chatbot_message(text="", message_type="picture")
msg.image_content = SimpleNamespace(download_code=None)
channel._on_chatbot_message(msg)
await asyncio.sleep(0.1)
bus.publish_inbound.assert_not_awaited()
_run(go())
class TestReceiveFile:
def test_no_files_is_noop(self):
async def go():
channel = DingTalkChannel(MessageBus(), config={})
msg = channel._make_inbound(chat_id="c", user_id="u", text="hi", thread_ts="m")
out = await channel.receive_file(msg, "t1")
assert out is msg
assert out.text == "hi"
_run(go())
def test_downloads_persists_and_prepends_path(self, tmp_path, monkeypatch):
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._client_id = "robot"
channel._download_by_code = AsyncMock(return_value=b"XLSXDATA")
uploads = tmp_path / "uploads"
uploads.mkdir()
_patch_uploads(monkeypatch, uploads)
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="please translate",
thread_ts="m",
files=[{"type": "file", "download_code": "dc1", "filename": "quote.xlsx"}],
)
out = await channel.receive_file(msg, "thread-1", user_id="default")
assert out.files == []
assert (uploads / "quote.xlsx").read_bytes() == b"XLSXDATA"
assert f"{VIRTUAL_PATH_PREFIX}/uploads/quote.xlsx" in out.text
assert out.text.endswith("please translate")
channel._download_by_code.assert_awaited_once_with("dc1")
_run(go())
def test_pure_file_message_text_is_only_the_path(self, tmp_path, monkeypatch):
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._download_by_code = AsyncMock(return_value=b"IMG")
uploads = tmp_path / "uploads"
uploads.mkdir()
_patch_uploads(monkeypatch, uploads)
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="",
thread_ts="m",
files=[{"type": "image", "download_code": "dc_i", "filename": "image.png"}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
assert out.text == f"{VIRTUAL_PATH_PREFIX}/uploads/image.png"
_run(go())
def test_download_failure_surfaces_marker(self):
"""A failed download must stay visible to the agent, not vanish silently."""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._download_by_code = AsyncMock(return_value=None)
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="hi",
thread_ts="m",
files=[{"type": "file", "download_code": "dc1", "filename": "x.pdf"}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
assert out.files == []
assert out.text == "[failed to load file: x.pdf]\n\nhi"
_run(go())
def test_token_failure_yields_marker_not_exception(self):
"""An auth failure during download must degrade to a marker, not an exception.
The manager calls ``receive_file`` without a try, so anything escaping
here kills the whole chat turn with no reply at all.
"""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._client_id = "robot"
channel._get_access_token = AsyncMock(side_effect=ValueError("bad token response"))
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="hi",
thread_ts="m",
files=[{"type": "file", "download_code": "dc", "filename": "x.pdf"}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
assert out.text == "[failed to load file: x.pdf]\n\nhi"
assert out.files == []
_run(go())
def test_unexpected_error_becomes_marker(self):
"""Any unforeseen per-attachment error must surface as a marker, not escape."""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._receive_single_file = AsyncMock(side_effect=RuntimeError("boom"))
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="hi",
thread_ts="m",
files=[{"type": "file", "download_code": "dc", "filename": "x.pdf"}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
assert out.text == "[failed to load file: x.pdf]\n\nhi"
assert out.files == []
_run(go())
def test_failure_marker_sanitizes_hostile_filename(self):
"""A hostile filename must not forge extra lines in msg.text or bloat it.
``fileName`` is webhook data: embedding it raw in the marker would let a
newline fake a standalone ``/mnt/user-data/uploads/...`` line, and an
over-long name would balloon the message text.
"""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._download_by_code = AsyncMock(return_value=None)
hostile = "evil\n/mnt/user-data/uploads/fake.pdf\n" + "a" * 300 + ".pdf"
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="hi",
thread_ts="m",
files=[{"type": "file", "download_code": "dc", "filename": hostile}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
lines = out.text.splitlines()
assert len(lines) == 3, out.text
assert lines[0].startswith("[failed to load file: ")
assert len(lines[0]) <= 120
assert lines[1] == ""
assert lines[2] == "hi"
_run(go())
def test_download_failure_without_filename_uses_type_only_marker(self):
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._download_by_code = AsyncMock(return_value=None)
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="",
thread_ts="m",
files=[{"type": "image", "download_code": "dc1", "filename": ""}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
assert out.text == "[failed to load image]"
_run(go())
def test_partial_failure_keeps_successful_path(self, tmp_path, monkeypatch):
"""One bad attachment must not cost the agent the one that did load."""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
uploads = tmp_path / "uploads"
uploads.mkdir()
_patch_uploads(monkeypatch, uploads)
channel._download_by_code = AsyncMock(side_effect=[b"OK", None])
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="two files",
thread_ts="m",
files=[
{"type": "file", "download_code": "good", "filename": "ok.pdf"},
{"type": "file", "download_code": "bad", "filename": "broken.pdf"},
],
)
out = await channel.receive_file(msg, "t1", user_id="default")
assert f"{VIRTUAL_PATH_PREFIX}/uploads/ok.pdf" in out.text
assert "[failed to load file: broken.pdf]" in out.text
assert out.text.endswith("two files")
_run(go())
def test_traversal_filename_is_sanitized(self, tmp_path, monkeypatch):
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._download_by_code = AsyncMock(return_value=b"D")
uploads = tmp_path / "uploads"
uploads.mkdir()
_patch_uploads(monkeypatch, uploads)
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="",
thread_ts="m",
files=[{"type": "file", "download_code": "dc", "filename": "../../etc/passwd"}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
written = list(uploads.iterdir())
assert len(written) == 1
assert written[0].name == "passwd"
assert out.text == f"{VIRTUAL_PATH_PREFIX}/uploads/passwd"
_run(go())
def test_rejected_filename_falls_back_to_generated_name(self, tmp_path, monkeypatch):
"""A filename normalize_filename rejects must land on the generated name.
``".."`` raises inside ``normalize_filename`` (unlike ``../../etc/passwd``,
whose basename ``passwd`` is accepted), so this exercises the except branch.
"""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._download_by_code = AsyncMock(return_value=b"D")
uploads = tmp_path / "uploads"
uploads.mkdir()
_patch_uploads(monkeypatch, uploads)
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="",
thread_ts="m",
files=[{"type": "file", "download_code": "abc123456789", "filename": ".."}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
written = list(uploads.iterdir())
assert len(written) == 1
assert written[0].name == "dingtalk_abc123456789.bin"
assert out.text == f"{VIRTUAL_PATH_PREFIX}/uploads/dingtalk_abc123456789.bin"
_run(go())
def test_fallback_name_sanitizes_download_code(self, tmp_path, monkeypatch):
"""download_code is webhook data: it must not inject path separators."""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._download_by_code = AsyncMock(return_value=b"D")
uploads = tmp_path / "uploads"
uploads.mkdir()
_patch_uploads(monkeypatch, uploads)
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="",
thread_ts="m",
files=[{"type": "image", "download_code": "../../evil", "filename": ".."}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
written = list(uploads.iterdir())
assert len(written) == 1
assert "/" not in written[0].name
assert written[0].name == "dingtalk_evil.png"
assert out.text == f"{VIRTUAL_PATH_PREFIX}/uploads/dingtalk_evil.png"
_run(go())
def test_second_picture_does_not_overwrite_first(self, tmp_path, monkeypatch):
"""Two picture messages both generate "image.png" and must not collide.
The path handed to the agent must still hold the bytes it referred to.
"""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
uploads = tmp_path / "uploads"
uploads.mkdir()
_patch_uploads(monkeypatch, uploads)
channel._download_by_code = AsyncMock(return_value=b"FIRST")
first = channel._make_inbound(
chat_id="c",
user_id="u",
text="",
thread_ts="m1",
files=[{"type": "image", "download_code": "dc1", "filename": "image.png"}],
)
out_first = await channel.receive_file(first, "t1", user_id="default")
channel._download_by_code = AsyncMock(return_value=b"SECOND")
second = channel._make_inbound(
chat_id="c",
user_id="u",
text="",
thread_ts="m2",
files=[{"type": "image", "download_code": "dc2", "filename": "image.png"}],
)
out_second = await channel.receive_file(second, "t1", user_id="default")
assert out_first.text != out_second.text
first_name = out_first.text.rsplit("/", 1)[-1]
second_name = out_second.text.rsplit("/", 1)[-1]
# Each advertised path must still resolve to its own bytes.
assert (uploads / first_name).read_bytes() == b"FIRST"
assert (uploads / second_name).read_bytes() == b"SECOND"
_run(go())
def test_multiple_images_in_one_message_get_distinct_paths(self, tmp_path, monkeypatch):
async def go():
channel = DingTalkChannel(MessageBus(), config={})
uploads = tmp_path / "uploads"
uploads.mkdir()
_patch_uploads(monkeypatch, uploads)
channel._download_by_code = AsyncMock(side_effect=[b"A", b"B"])
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="",
thread_ts="m",
files=[
{"type": "image", "download_code": "dc1", "filename": "image_0.png"},
{"type": "image", "download_code": "dc2", "filename": "image_0.png"},
],
)
out = await channel.receive_file(msg, "t1", user_id="default")
paths = out.text.split("\n")
assert len(paths) == 2
assert paths[0] != paths[1]
assert len(list(uploads.iterdir())) == 2
assert {p.read_bytes() for p in uploads.iterdir()} == {b"A", b"B"}
_run(go())
def test_duplicate_document_names_do_not_collide(self, tmp_path, monkeypatch):
"""The same real filename sent twice must not clobber the earlier upload."""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
uploads = tmp_path / "uploads"
uploads.mkdir()
_patch_uploads(monkeypatch, uploads)
channel._download_by_code = AsyncMock(return_value=b"V1")
m1 = channel._make_inbound(chat_id="c", user_id="u", text="", thread_ts="m1", files=[{"type": "file", "download_code": "d1", "filename": "quote.xlsx"}])
await channel.receive_file(m1, "t1", user_id="default")
channel._download_by_code = AsyncMock(return_value=b"V2")
m2 = channel._make_inbound(chat_id="c", user_id="u", text="", thread_ts="m2", files=[{"type": "file", "download_code": "d2", "filename": "quote.xlsx"}])
await channel.receive_file(m2, "t1", user_id="default")
assert (uploads / "quote.xlsx").read_bytes() == b"V1"
assert len(list(uploads.iterdir())) == 2
_run(go())
def test_write_does_not_follow_planted_symlink(self, tmp_path, monkeypatch):
"""A symlink planted at the destination must not be written through.
Upload dirs can be mounted into local sandboxes, so a sandbox process can
leave a symlink at a future upload name; following it would let a
gateway-privileged write land outside the bucket.
"""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
uploads = tmp_path / "uploads"
uploads.mkdir()
outside = tmp_path / "outside.txt"
(uploads / "image.png").symlink_to(outside)
_patch_uploads(monkeypatch, uploads)
channel._download_by_code = AsyncMock(return_value=b"PWNED")
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="",
thread_ts="m",
files=[{"type": "image", "download_code": "dc", "filename": "image.png"}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
assert not outside.exists()
assert "[failed to load image: image.png]" in out.text
_run(go())
def test_missing_sandbox_after_acquire_yields_marker(self, tmp_path, monkeypatch):
"""A non-local sandbox that cannot be resolved must not yield a path.
The agent's sandbox cannot see the file in that case, so handing it the
virtual path would point at nothing readable; mirror Feishu and surface
a failed-load marker instead.
"""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._download_by_code = AsyncMock(return_value=b"BYTES")
uploads = tmp_path / "uploads"
uploads.mkdir()
_patch_uploads(monkeypatch, uploads, sandbox_id="aio:box1", sandbox=None)
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="hi",
thread_ts="m",
files=[{"type": "file", "download_code": "dc", "filename": "a.pdf"}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
assert out.text == "[failed to load file: a.pdf]\n\nhi"
_run(go())
def test_update_file_failure_yields_marker(self, tmp_path, monkeypatch):
"""A failed non-local sandbox sync must not yield a path either.
Same failure mode as the missing-sandbox case: the bytes never reached
the agent's sandbox, so the virtual path would read as nothing. Mirrors
Feishu, whose sync except-branch returns the failure marker.
"""
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._download_by_code = AsyncMock(return_value=b"BYTES")
uploads = tmp_path / "uploads"
uploads.mkdir()
broken_sandbox = MagicMock()
broken_sandbox.update_file.side_effect = RuntimeError("sandbox transport down")
_patch_uploads(monkeypatch, uploads, sandbox_id="aio:box1", sandbox=broken_sandbox)
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="hi",
thread_ts="m",
files=[{"type": "file", "download_code": "dc", "filename": "a.pdf"}],
)
out = await channel.receive_file(msg, "t1", user_id="default")
assert out.text == "[failed to load file: a.pdf]\n\nhi"
_run(go())
def test_non_local_sandbox_is_synced(self, tmp_path, monkeypatch):
async def go():
channel = DingTalkChannel(MessageBus(), config={})
channel._download_by_code = AsyncMock(return_value=b"BYTES")
uploads = tmp_path / "uploads"
uploads.mkdir()
fake_sandbox = MagicMock()
_patch_uploads(monkeypatch, uploads, sandbox_id="aio:box1", sandbox=fake_sandbox)
msg = channel._make_inbound(
chat_id="c",
user_id="u",
text="",
thread_ts="m",
files=[{"type": "file", "download_code": "dc", "filename": "a.pdf"}],
)
await channel.receive_file(msg, "t1", user_id="default")
fake_sandbox.update_file.assert_called_once_with(f"{VIRTUAL_PATH_PREFIX}/uploads/a.pdf", b"BYTES")
_run(go())
class TestDownloadByCode:
def test_two_step_download(self):
async def go():
from unittest.mock import patch
channel = DingTalkChannel(MessageBus(), config={})
channel._client_id = "robot_x"
channel._get_access_token = AsyncMock(return_value="tok")
class PostResponse:
status_code = 200
@staticmethod
def json():
return {"downloadUrl": "https://dl.dingtalk/xyz"}
captured: dict = {}
class FakeStream:
async def __aenter__(self):
return self
async def __aexit__(self, *a):
pass
def raise_for_status(self):
pass
async def aiter_bytes(self):
yield b"FILE"
yield b"BYTES"
class FakeClient:
async def __aenter__(self):
return self
async def __aexit__(self, *a):
pass
async def post(self, url, **kwargs):
captured["post_url"] = url
captured["json"] = kwargs.get("json")
return PostResponse()
def stream(self, method, url, **kwargs):
captured["get_url"] = url
return FakeStream()
with patch("app.channels.dingtalk.httpx.AsyncClient", return_value=FakeClient()):
data = await channel._download_by_code("dl_code")
assert data == b"FILEBYTES"
assert captured["post_url"].endswith("/v1.0/robot/messageFiles/download")
assert captured["json"] == {"downloadCode": "dl_code", "robotCode": "robot_x"}
assert captured["get_url"] == "https://dl.dingtalk/xyz"
_run(go())
def test_oversized_download_is_dropped(self, monkeypatch):
"""Inbound bytes are buffered in memory; a file over the cap must be refused.
Outbound uploads already enforce a size cap without an inbound one, a
single large chat attachment balloons gateway memory.
"""
async def go():
from unittest.mock import patch
channel = DingTalkChannel(MessageBus(), config={})
channel._client_id = "robot_x"
channel._get_access_token = AsyncMock(return_value="tok")
monkeypatch.setattr("app.channels.dingtalk._MAX_INBOUND_FILE_SIZE_BYTES", 10)
class PostResponse:
status_code = 200
@staticmethod
def json():
return {"downloadUrl": "https://dl.dingtalk/big"}
class FakeStream:
async def __aenter__(self):
return self
async def __aexit__(self, *a):
pass
def raise_for_status(self):
pass
async def aiter_bytes(self):
for _ in range(4):
yield b"x" * 8
class FakeClient:
async def __aenter__(self):
return self
async def __aexit__(self, *a):
pass
async def post(self, url, **kwargs):
return PostResponse()
def stream(self, method, url, **kwargs):
return FakeStream()
with patch("app.channels.dingtalk.httpx.AsyncClient", return_value=FakeClient()):
assert await channel._download_by_code("dl") is None
_run(go())
def test_post_non_200_returns_none(self):
async def go():
from unittest.mock import patch
channel = DingTalkChannel(MessageBus(), config={})
channel._client_id = "robot_x"
channel._get_access_token = AsyncMock(return_value="tok")
class PostResponse:
status_code = 403
text = "forbidden"
class FakeClient:
async def __aenter__(self):
return self
async def __aexit__(self, *a):
pass
async def post(self, url, **kwargs):
return PostResponse()
with patch("app.channels.dingtalk.httpx.AsyncClient", return_value=FakeClient()):
assert await channel._download_by_code("dl") is None
_run(go())
def test_missing_download_url_returns_none(self):
async def go():
from unittest.mock import patch
channel = DingTalkChannel(MessageBus(), config={})
channel._client_id = "robot_x"
channel._get_access_token = AsyncMock(return_value="tok")
class PostResponse:
status_code = 200
@staticmethod
def json():
return {}
class FakeClient:
async def __aenter__(self):
return self
async def __aexit__(self, *a):
pass
async def post(self, url, **kwargs):
return PostResponse()
with patch("app.channels.dingtalk.httpx.AsyncClient", return_value=FakeClient()):
assert await channel._download_by_code("dl") is None
_run(go())
class TestHandlerStashesRawData:
def test_process_stashes_raw_callback_payload(self):
pytest.importorskip("dingtalk_stream")
async def go():
bus = MessageBus()
channel = DingTalkChannel(bus, config={})
captured: dict = {}
channel._on_chatbot_message = lambda m: captured.update(msg=m)
handler = _DingTalkMessageHandler(channel)
cb = MagicMock()
cb.data = {
"msgtype": "file",
"content": {"downloadCode": "dc_doc", "fileName": "a.xlsx"},
"senderStaffId": "u1",
"conversationType": "1",
"msgId": "m1",
}
await handler.process(cb)
msg = captured["msg"]
assert getattr(msg, "_df_raw_data", None) == cb.data
# _extract_files can now read the document descriptor from the stash.
assert DingTalkChannel._extract_files(msg) == [{"type": "file", "download_code": "dc_doc", "filename": "a.xlsx"}]
_run(go())

View File

@ -175,6 +175,65 @@ class FakeSandboxClass:
return self.list_return
class FakeOwnershipStore:
"""Shared lease state with per-provider identities for reconciliation tests."""
supports_cross_process = True
def __init__(
self,
leases: dict[str, tuple[str, str]],
*,
owner_id: str,
lock: threading.Lock | None = None,
) -> None:
self._leases = leases
self._lock = lock or threading.Lock()
self.owner_id = owner_id
def take(self, sandbox_id: str) -> bool:
with self._lock:
current = self._leases.get(sandbox_id)
if current is not None and current[1] == "del":
return False
self._leases[sandbox_id] = (self.owner_id, "own")
return True
def claim(self, sandbox_id: str, *, for_destroy: bool = False) -> bool:
with self._lock:
current = self._leases.get(sandbox_id)
if current is not None and current[0] != self.owner_id:
return False
if current is not None and current[1] == "del" and not for_destroy:
return False
self._leases[sandbox_id] = (self.owner_id, "del" if for_destroy else "own")
return True
def renew(self, sandbox_id: str):
from deerflow.community.aio_sandbox.ownership import RenewOutcome
with self._lock:
current = self._leases.get(sandbox_id)
if current is None:
return RenewOutcome.LAPSED
if current == (self.owner_id, "own"):
return RenewOutcome.RENEWED
return RenewOutcome.LOST
def release(self, sandbox_id: str) -> None:
with self._lock:
if self._leases.get(sandbox_id, (None,))[0] == self.owner_id:
self._leases.pop(sandbox_id, None)
def owner(self, sandbox_id: str) -> str | None:
with self._lock:
current = self._leases.get(sandbox_id)
return current[0] if current is not None else None
def close(self) -> None:
return None
def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_policy: str = "wait", acquire_timeout: int = 30, burst_limit: int = 0) -> Any:
"""Build a ``E2BSandboxProvider`` instance bypassing ``__init__``."""
mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider")
@ -192,6 +251,15 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_poli
provider._transitioning_slots = 0
provider._capacity_cond = threading.Condition(provider._lock)
provider._shutdown_called = False
provider._owner_id = "owner-a"
provider._ownership = FakeOwnershipStore({}, owner_id=provider._owner_id)
provider._ownership_config = SimpleNamespace(renewal_interval_seconds=60.0)
provider._owned_sandbox_ids = set()
provider._acquire_inflight = set()
provider._orphan_first_seen = {}
provider._maintenance_stop = threading.Event()
provider._lease_thread = None
provider._reconcile_thread = None
provider._config = {
"api_key": "test-key",
"template": "code-interpreter-v1",
@ -204,6 +272,12 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_poli
"burst_limit": burst_limit,
"mounts": [],
"environment": {},
"reconciliation_interval_seconds": 60.0,
"reconciliation_grace_seconds": 30.0,
"reconciliation_orphan_ttl_seconds": 3600.0,
"reconciliation_max_pages": 10,
"reconciliation_max_items": 100,
"reconciliation_max_seconds": 10.0,
}
return provider
@ -412,6 +486,84 @@ def test_kill_and_close_swallows_kill_exceptions():
client.kill = explode
p._kill_and_close(sb)
assert "fake-sb-1" not in p._owned_sandbox_ids
assert p._ownership.owner("fake-sb-1") is None
def test_kill_and_close_skips_peer_owned_sandbox():
p = _make_provider()
client = FakeClient(sandbox_id="peer-owned")
sandbox = _make_sandbox(client, sandbox_id="peer-owned")
p._ownership = FakeOwnershipStore(
{"peer-owned": ("owner-peer", "own")},
owner_id=p._owner_id,
)
p._kill_and_close(sandbox)
assert client.killed is False
assert client.closed is True
assert p._ownership.owner("peer-owned") == "owner-peer"
def test_startup_reconciliation_runs_in_background_without_blocking_caller(monkeypatch):
p = _make_provider()
entered = threading.Event()
release = threading.Event()
def reconcile():
entered.set()
assert release.wait(timeout=2)
monkeypatch.setattr(p, "_reconcile_remote_sandboxes", reconcile)
started = time.monotonic()
p._start_maintenance_threads()
elapsed = time.monotonic() - started
assert entered.wait(timeout=1)
assert elapsed < 0.5
release.set()
p._maintenance_stop.set()
for thread in (p._lease_thread, p._reconcile_thread):
assert thread is not None
thread.join(timeout=2)
def test_refresh_owned_leases_reclaims_lapsed_lease():
p = _make_provider()
client = FakeClient(sandbox_id="sb-lapsed")
sandbox = _make_sandbox(client, sandbox_id="sb-lapsed")
p._sandboxes["sb-lapsed"] = sandbox
p._thread_sandboxes[("u1", "t1")] = "sb-lapsed"
p._owned_sandbox_ids.add("sb-lapsed")
p._refresh_owned_leases()
assert p._ownership.owner("sb-lapsed") == p._owner_id
assert p.get("sb-lapsed") is sandbox
assert client.closed is False
def test_refresh_owned_leases_forgets_peer_owned_sandbox():
p = _make_provider()
client = FakeClient(sandbox_id="sb-lost")
sandbox = _make_sandbox(client, sandbox_id="sb-lost")
p._sandboxes["sb-lost"] = sandbox
p._thread_sandboxes[("u1", "t1")] = "sb-lost"
p._owned_sandbox_ids.add("sb-lost")
p._ownership = FakeOwnershipStore(
{"sb-lost": ("owner-peer", "own")},
owner_id=p._owner_id,
)
p._refresh_owned_leases()
assert p.get("sb-lost") is None
assert ("u1", "t1") not in p._thread_sandboxes
assert "sb-lost" not in p._owned_sandbox_ids
assert client.closed is True
def test_reuse_in_process_sandbox_returns_cached_id_on_healthy_reuse():
@ -614,6 +766,239 @@ def test_discover_remote_sandbox_skips_dead_candidate(monkeypatch):
assert client.closed is True
def test_discover_remote_sandbox_tries_later_candidate_when_first_is_dead(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
fake_cls.list_return = [
_info("sb-a-dead", "u1", "t1"),
_info("sb-b-live", "u1", "t1"),
]
dead = FakeClient(
sandbox_id="sb-a-dead",
commands=FakeCommandsAPI([FakeCommandsAPI.GONE]),
)
live = FakeClient(sandbox_id="sb-b-live")
fake_cls.connect_factory = lambda sid, **_kw: dead if sid == "sb-a-dead" else live
assert p._discover_remote_sandbox("t1", user_id="u1") == "sb-b-live"
assert dead.closed is True
assert p._thread_sandboxes[("u1", "t1")] == "sb-b-live"
assert "sb-b-live" in p._owned_sandbox_ids
def test_reconcile_defers_duplicate_with_live_peer_lease(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
fake_cls.list_return = [
_info("sb-canonical", "u1", "t1"),
_info("sb-duplicate", "u1", "t1"),
]
clients = {
"sb-canonical": FakeClient(sandbox_id="sb-canonical"),
"sb-duplicate": FakeClient(sandbox_id="sb-duplicate"),
}
fake_cls.connect_factory = lambda sid, **_kw: clients[sid]
leases = {"sb-duplicate": ("owner-peer", "own")}
p._ownership = FakeOwnershipStore(leases, owner_id=p._owner_id)
p._config["reconciliation_grace_seconds"] = 0.0
stats = p._reconcile_remote_sandboxes(now=time.monotonic())
assert stats.adopted == 1
assert stats.deferred == 1
assert stats.killed == 0
assert clients["sb-duplicate"].killed is False
def test_reconcile_kills_unowned_duplicate_after_grace(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
fake_cls.list_return = [
_info("sb-canonical", "u1", "t1"),
_info("sb-duplicate", "u1", "t1"),
]
clients = {
"sb-canonical": FakeClient(sandbox_id="sb-canonical"),
"sb-duplicate": FakeClient(sandbox_id="sb-duplicate"),
}
fake_cls.connect_factory = lambda sid, **_kw: clients[sid]
p._config["reconciliation_grace_seconds"] = 5.0
first = p._reconcile_remote_sandboxes(now=100.0)
second = p._reconcile_remote_sandboxes(now=106.0)
assert first.deferred == 1
assert first.killed == 0
assert second.killed == 1
assert clients["sb-duplicate"].killed is True
def test_competing_reconcilers_only_one_kills_duplicate(monkeypatch):
shared_leases: dict[str, tuple[str, str]] = {}
shared_lock = threading.Lock()
providers = [_make_provider(), _make_provider()]
providers[0]._owner_id = "owner-a"
providers[1]._owner_id = "owner-b"
clients = {
"sb-canonical": FakeClient(sandbox_id="sb-canonical"),
"sb-duplicate": FakeClient(sandbox_id="sb-duplicate"),
}
kill_count = 0
def kill_duplicate() -> None:
nonlocal kill_count
kill_count += 1
clients["sb-duplicate"].killed = True
clients["sb-duplicate"].commands._responses.append(FakeCommandsAPI.GONE)
clients["sb-duplicate"].kill = kill_duplicate
for provider in providers:
fake_cls = _install_fake_sdk(monkeypatch, provider)
fake_cls.list_return = [
_info("sb-canonical", "u1", "t1"),
_info("sb-duplicate", "u1", "t1"),
]
fake_cls.connect_factory = lambda sid, **_kw: clients[sid]
provider._ownership = FakeOwnershipStore(
shared_leases,
owner_id=provider._owner_id,
lock=shared_lock,
)
provider._config["reconciliation_grace_seconds"] = 0.0
results = [provider._reconcile_remote_sandboxes(now=100.0) for provider in providers]
assert sum(result.killed for result in results) == 1
assert kill_count == 1
def test_reconcile_honors_item_budget(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
fake_cls.list_return = [_info(f"sb-{index}", "u1", f"t-{index}") for index in range(5)]
p._config["reconciliation_max_items"] = 2
stats = p._reconcile_remote_sandboxes(now=100.0)
assert stats.discovered == 2
assert len(fake_cls.connect_calls) == 2
def test_reconcile_honors_page_budget(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
paginator = _FakePaginator(
[
[_info("sb-first", "u1", "t1")],
[_info("sb-never-read", "u2", "t2")],
]
)
fake_cls.list_return = paginator
p._config["reconciliation_max_pages"] = 1
stats = p._reconcile_remote_sandboxes(now=100.0)
assert stats.discovered == 1
assert stats.budget_exhausted is True
assert paginator.calls == 1
assert [call[0] for call in fake_cls.connect_calls] == ["sb-first"]
def test_reconcile_honors_wall_clock_budget(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
fake_cls.list_return = [_info("sb-never-probed", "u1", "t1")]
p._config["reconciliation_max_seconds"] = 0.5
provider_mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider")
ticks = iter([0.0, 0.0, 1.0, 1.0])
monkeypatch.setattr(provider_mod.time, "monotonic", lambda: next(ticks))
stats = p._reconcile_remote_sandboxes(now=100.0)
assert stats.budget_exhausted is True
assert fake_cls.connect_calls == []
def test_reconcile_adopts_canonical_after_restart_loses_local_state(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
fake_cls.list_return = [_info("sb-existing", "u1", "t1")]
stats = p._reconcile_remote_sandboxes(now=100.0)
assert stats.adopted == 1
assert p._thread_sandboxes[("u1", "t1")] == "sb-existing"
assert "sb-existing" in p._owned_sandbox_ids
def test_reconcile_bootstrap_failure_clears_inflight_after_peer_take(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
sandbox_id = "sb-racy-bootstrap"
fake_cls.list_return = [_info(sandbox_id, "u1", "t1")]
shared_leases: dict[str, tuple[str, str]] = {}
shared_lock = threading.Lock()
p._ownership = FakeOwnershipStore(
shared_leases,
owner_id=p._owner_id,
lock=shared_lock,
)
peer_ownership = FakeOwnershipStore(
shared_leases,
owner_id="owner-peer",
lock=shared_lock,
)
def peer_takes_before_bootstrap_fails(_command: str) -> SimpleNamespace:
assert peer_ownership.take(sandbox_id) is True
return SimpleNamespace(stdout="", stderr="permission denied", exit_code=1)
client = FakeClient(
sandbox_id=sandbox_id,
commands=FakeCommandsAPI(
[
SimpleNamespace(stdout="ok", stderr="", exit_code=0),
peer_takes_before_bootstrap_fails,
]
),
)
fake_cls.connect_factory = lambda _sid, **_kw: client
stats = p._reconcile_remote_sandboxes(now=100.0)
assert stats.adopted == 0
assert stats.deferred == 1
assert sandbox_id not in p._acquire_inflight
assert sandbox_id not in p._unowned_remote_ops_in_progress
assert p._reserved_slots == 0
assert p._ownership.owner(sandbox_id) == "owner-peer"
assert client.killed is False
assert client.closed is True
def test_reconcile_kills_metadata_orphan_only_after_ttl(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
fake_cls.list_return = [
SimpleNamespace(
sandbox_id="sb-orphan",
metadata={"deer_flow_provider": "e2b_sandbox_provider"},
)
]
client = FakeClient(sandbox_id="sb-orphan")
fake_cls.connect_factory = lambda _sid, **_kw: client
p._config["reconciliation_orphan_ttl_seconds"] = 5.0
first = p._reconcile_remote_sandboxes(now=100.0)
second = p._reconcile_remote_sandboxes(now=106.0)
assert first.deferred == 1
assert first.killed == 0
assert second.killed == 1
assert client.killed is True
def test_discover_remote_sandbox_discards_candidate_when_bootstrap_fails(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
@ -635,6 +1020,65 @@ def test_discover_remote_sandbox_discards_candidate_when_bootstrap_fails(monkeyp
assert ("u1", "t1") not in p._thread_sandboxes
def test_discovery_claims_ownership_before_bootstrap_cleanup(monkeypatch):
events: list[str] = []
class RecordingOwnershipStore(FakeOwnershipStore):
def take(self, sandbox_id: str) -> bool:
events.append("take")
return super().take(sandbox_id)
def claim(self, sandbox_id: str, *, for_destroy: bool = False) -> bool:
events.append("claim-destroy" if for_destroy else "claim")
return super().claim(sandbox_id, for_destroy=for_destroy)
def release(self, sandbox_id: str) -> None:
events.append("release")
super().release(sandbox_id)
p = _make_provider()
p._ownership = RecordingOwnershipStore(
{"sb-broken": ("owner-peer", "own")},
owner_id=p._owner_id,
)
fake_cls = _install_fake_sdk(monkeypatch, p)
fake_cls.list_return = [_info("sb-broken", "u1", "t1")]
client = FakeClient(
sandbox_id="sb-broken",
commands=FakeCommandsAPI(
[
SimpleNamespace(stdout="ok", stderr="", exit_code=0),
SimpleNamespace(stdout="", stderr="permission denied", exit_code=1),
]
),
)
client.kill = lambda: events.append("kill")
fake_cls.connect_factory = lambda _sid, **_kw: client
assert p._discover_remote_sandbox("t1", user_id="u1") is None
assert events == ["take", "claim-destroy", "kill", "release"]
def test_bootstrap_failure_does_not_kill_without_destroy_lease(monkeypatch):
p = _make_provider()
client = FakeClient(
sandbox_id="sb-peer",
commands=FakeCommandsAPI([SimpleNamespace(stdout="", stderr="permission denied", exit_code=1)]),
)
p._ownership = FakeOwnershipStore(
{"sb-peer": ("owner-peer", "own")},
owner_id=p._owner_id,
)
error, remote_destroyed = p._bootstrap_or_discard(client, "sb-peer")
assert error is not None
assert remote_destroyed is False
assert client.killed is False
assert client.closed is True
assert p._ownership.owner("sb-peer") == "owner-peer"
def test_kill_client_returns_exception_without_raising():
p = _make_provider()
client = FakeClient()
@ -728,6 +1172,48 @@ def test_evict_oldest_warm_keeps_slot_when_kill_lookup_raises(monkeypatch):
assert client.closed is True
assert p._eviction_tombstones == {"sb-warm"}
assert p._transitioning_slots == 1
assert "sb-warm" not in p._warm_pool
assert "sb-warm" not in p._owned_sandbox_ids
assert p._ownership.owner("sb-warm") is None
def test_evict_oldest_warm_defers_peer_owned_sandbox(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
p._warm_pool["sb-peer"] = ("seed", 12345.0)
p._ownership = FakeOwnershipStore(
{"sb-peer": ("owner-peer", "own")},
owner_id=p._owner_id,
)
assert p._evict_oldest_warm() == "sb-peer"
assert fake_cls.connect_calls == []
assert "sb-peer" not in p._warm_pool
assert p._eviction_tombstones == set()
assert p._evictions_in_progress == set()
assert p._transitioning_slots == 0
assert p._ownership.owner("sb-peer") == "owner-peer"
def test_evict_oldest_warm_releases_destroy_lease_after_kill_failure(monkeypatch):
p = _make_provider()
fake_cls = _install_fake_sdk(monkeypatch, p)
client = FakeClient(sandbox_id="sb-warm")
client.kill = MagicMock(side_effect=RuntimeError("transient kill failure"))
fake_cls.connect_factory = lambda _sid, **_kw: client
p._warm_pool["sb-warm"] = ("seed", 12345.0)
p._owned_sandbox_ids.add("sb-warm")
p._ownership.take("sb-warm")
assert p._evict_oldest_warm() is None
assert client.closed is True
assert "sb-warm" not in p._warm_pool
assert p._eviction_tombstones == {"sb-warm"}
assert p._transitioning_slots == 1
assert "sb-warm" not in p._owned_sandbox_ids
assert p._ownership.owner("sb-warm") is None
def test_evict_oldest_warm_uses_kill_helper_and_closes_client(monkeypatch):
@ -937,6 +1423,23 @@ def test_release_skips_warm_pool_when_sync_reveals_dead_vm(monkeypatch, tmp_path
assert client.killed is True
def test_shutdown_only_kills_sandboxes_owned_by_current_instance(monkeypatch):
p = _make_provider()
owned_client = FakeClient(sandbox_id="sb-owned")
peer_client = FakeClient(sandbox_id="sb-peer")
p._sandboxes = {
"sb-owned": _make_sandbox(owned_client),
"sb-peer": _make_sandbox(peer_client),
}
p._owned_sandbox_ids = {"sb-owned"}
p.shutdown()
assert owned_client.killed is True
assert peer_client.killed is False
assert peer_client.closed is True
def _setup_paths(monkeypatch, tmp_path):
paths_mod = importlib.import_module("deerflow.config.paths")
monkeypatch.setattr(paths_mod, "get_paths", lambda: Paths(base_dir=tmp_path), raising=False)

View File

@ -11,6 +11,7 @@ from __future__ import annotations
import asyncio
import logging
import threading
from contextlib import asynccontextmanager
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch
@ -116,7 +117,7 @@ def test_lifespan_sweeps_upload_staging_files_on_startup():
stop_channel_service.assert_awaited_once()
async def _run_lifespan_with_memory_flush(*, enabled: bool, flush_return: bool) -> MagicMock:
async def _run_lifespan_with_memory_flush(*, enabled: bool, flush_return: bool | Exception) -> MagicMock:
"""Drive lifespan with a spied memory manager.shutdown_flush.
Returns the manager mock so the caller can assert the shutdown flush was
@ -146,7 +147,10 @@ async def _run_lifespan_with_memory_flush(*, enabled: bool, flush_return: bool)
return fake_service
manager = MagicMock()
manager.shutdown_flush.return_value = flush_return
if isinstance(flush_return, Exception):
manager.shutdown_flush.side_effect = flush_return
else:
manager.shutdown_flush.return_value = flush_return
with (
patch("app.gateway.app.get_app_config", return_value=startup_config),
@ -191,6 +195,13 @@ def test_lifespan_skips_memory_flush_when_disabled() -> None:
manager.shutdown_flush.assert_not_called()
def test_lifespan_closes_memory_manager_when_flush_raises() -> None:
"""Derived retrieval resources are released even when queue drain fails."""
manager = asyncio.run(_run_lifespan_with_memory_flush(enabled=True, flush_return=RuntimeError("flush failed")))
manager.shutdown_flush.assert_called_once_with(5.0)
manager.close.assert_called_once_with()
# ── startup warm-up log accuracy ────────────────────────────────────────────
@ -256,3 +267,113 @@ def test_lifespan_warns_when_warm_returns_false(caplog) -> None:
manager = asyncio.run(_run_lifespan_with_warm_return(False))
manager.warm.assert_called_once_with()
assert any(r.levelno == logging.WARNING and "warm-up failed" in r.message for r in caplog.records)
async def _run_lifespan_with_slow_retrieval_warm() -> float:
from app.gateway.app import lifespan
app = FastAPI()
startup_config = SimpleNamespace(
log_level="INFO",
memory=SimpleNamespace(
token_counting="char",
enabled=True,
shutdown_flush_timeout_seconds=5.0,
),
)
fake_service = MagicMock()
fake_service.get_status.return_value = {}
release_rebuild = threading.Event()
manager = MagicMock()
manager.warm_retrieval.side_effect = lambda: release_rebuild.wait(5.0) or True
manager.warm.return_value = True
manager.shutdown_flush.return_value = True
async def fake_start(_startup_config):
return fake_service
with (
patch("app.gateway.app.get_app_config", return_value=startup_config),
patch("app.gateway.app.get_gateway_config", return_value=MagicMock(host="x", port=0)),
patch("app.gateway.app.langgraph_runtime", _noop_langgraph_runtime),
patch("app.gateway.app.auth.close_oidc_service", AsyncMock()),
patch("app.channels.service.start_channel_service", side_effect=fake_start),
patch("app.channels.service.stop_channel_service", AsyncMock()),
patch("deerflow.agents.memory.get_memory_manager", return_value=manager),
):
context = lifespan(app)
loop = asyncio.get_running_loop()
started_at = loop.time()
try:
await asyncio.wait_for(context.__aenter__(), timeout=1.0)
startup_elapsed = loop.time() - started_at
finally:
release_rebuild.set()
await context.__aexit__(None, None, None)
return startup_elapsed
def test_lifespan_does_not_wait_for_retrieval_rebuild_before_serving() -> None:
assert asyncio.run(_run_lifespan_with_slow_retrieval_warm()) < 1.0
async def _run_shutdown_with_blocked_retrieval_warm() -> tuple[float, MagicMock]:
from app.gateway.app import lifespan
app = FastAPI()
startup_config = SimpleNamespace(
log_level="INFO",
memory=SimpleNamespace(
token_counting="char",
enabled=True,
shutdown_flush_timeout_seconds=5.0,
),
)
fake_service = MagicMock()
fake_service.get_status.return_value = {}
rebuild_started = threading.Event()
release_rebuild = threading.Event()
manager = MagicMock()
def block_rebuild() -> bool:
rebuild_started.set()
release_rebuild.wait(5.0)
return True
manager.warm_retrieval.side_effect = block_rebuild
manager.warm.return_value = True
manager.shutdown_flush.return_value = True
async def fake_start(_startup_config, **_kwargs):
return fake_service
with (
patch("app.gateway.app.get_app_config", return_value=startup_config),
patch("app.gateway.app.get_gateway_config", return_value=MagicMock(host="x", port=0)),
patch("app.gateway.app.langgraph_runtime", _noop_langgraph_runtime),
patch("app.gateway.app._RETRIEVAL_WARM_SHUTDOWN_TIMEOUT_SECONDS", 0.01),
patch("app.gateway.app.auth.close_oidc_service", AsyncMock()),
patch("app.channels.service.start_channel_service", side_effect=fake_start),
patch("app.channels.service.stop_channel_service", AsyncMock()),
patch("deerflow.agents.memory.get_memory_manager", return_value=manager),
):
context = lifespan(app)
await context.__aenter__()
assert await asyncio.to_thread(rebuild_started.wait, 1.0)
loop = asyncio.get_running_loop()
started_at = loop.time()
try:
await context.__aexit__(None, None, None)
finally:
release_rebuild.set()
shutdown_elapsed = loop.time() - started_at
return shutdown_elapsed, manager
def test_lifespan_preserves_flush_budget_when_retrieval_warm_is_still_running() -> None:
shutdown_elapsed, manager = asyncio.run(_run_shutdown_with_blocked_retrieval_warm())
assert shutdown_elapsed < 1.0
manager.shutdown_flush.assert_called_once_with(5.0)
manager.close.assert_not_called()

View File

@ -1171,8 +1171,8 @@ async def test_dedupe_identity_distinguishes_same_agent_name_across_users(base_d
from app.channels.store import ChannelStore
manager = ChannelManager(bus=MessageBus(), store=ChannelStore(path=base_dir / "dedupe-store.json"))
assert manager._is_duplicate_inbound(by_owner["alice"]) is False
assert manager._is_duplicate_inbound(by_owner["bob"]) is False
assert await manager._is_duplicate_inbound(by_owner["alice"]) is False
assert await manager._is_duplicate_inbound(by_owner["bob"]) is False
@pytest.mark.asyncio
@ -1208,12 +1208,12 @@ async def test_missing_delivery_header_leaves_dedupe_open(base_dir: Path) -> Non
await fanout_event(bus, "pull_request", "", payload)
(first,) = await _drain(bus)
assert first.metadata["message_id"] is None
assert manager._is_duplicate_inbound(first) is False
assert await manager._is_duplicate_inbound(first) is False
await fanout_event(bus, "pull_request", "", payload)
(second,) = await _drain(bus)
assert second.metadata["message_id"] is None
assert manager._is_duplicate_inbound(second) is False
assert await manager._is_duplicate_inbound(second) is False
# ---------------------------------------------------------------------------

View File

@ -6,6 +6,7 @@ import pytest
from deerflow.runtime import RunManager, RunStatus
from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.runs.manager import EditReplayVisibility
from deerflow.runtime.runs.store.memory import MemoryRunStore
@ -288,3 +289,121 @@ async def test_run_manager_batch_history_methods_allow_explicit_unscoped_access(
assert sources == {"source-alice", "source-bob"}
assert set(records) == {"regen-alice", "regen-bob"}
@pytest.mark.anyio
async def test_memory_run_store_lists_edit_replay_runs_unbounded_and_owner_scoped():
store = MemoryRunStore()
for index in range(105):
await store.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success")
await store.put(
"edit-success",
thread_id="t1",
user_id="alice",
status="success",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"},
)
await store.put(
"edit-error",
thread_id="t1",
user_id="alice",
status="error",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-error"},
)
await store.put(
"edit-bob",
thread_id="t1",
user_id="bob",
status="success",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-bob"},
)
rows = await store.list_edit_regenerate_runs("t1", user_id="alice")
assert [row["run_id"] for row in rows] == ["edit-success", "edit-error"]
@pytest.mark.anyio
async def test_run_repository_lists_edit_replay_runs_unbounded_and_owner_scoped(tmp_path):
from deerflow.persistence.engine import close_engine, get_session_factory, init_engine
from deerflow.persistence.run import RunRepository
await init_engine("sqlite", url=f"sqlite+aiosqlite:///{tmp_path / 'runs.db'}", sqlite_dir=str(tmp_path))
try:
repo = RunRepository(get_session_factory())
for index in range(105):
await repo.put(f"normal-{index}", thread_id="t1", user_id="alice", status="success")
await repo.put(
"edit-success",
thread_id="t1",
user_id="alice",
status="success",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"},
)
await repo.put(
"edit-error",
thread_id="t1",
user_id="alice",
status="error",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-error"},
)
await repo.put(
"edit-bob",
thread_id="t1",
user_id="bob",
status="success",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-bob"},
)
rows = await repo.list_edit_regenerate_runs("t1", user_id="alice")
assert [row["run_id"] for row in rows] == ["edit-success", "edit-error"]
finally:
await close_engine()
@pytest.mark.anyio
async def test_run_manager_computes_edit_replay_visibility_by_latest_attempt():
manager = RunManager()
running = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-running"},
)
running.status = RunStatus.running
success = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-success"},
)
success.status = RunStatus.success
failed = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-failed"},
)
failed.status = RunStatus.error
older_success = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retried"},
)
older_success.status = RunStatus.success
newer_failed = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retried"},
)
newer_failed.status = RunStatus.interrupted
older_failed_then_success = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retry-success"},
)
older_failed_then_success.status = RunStatus.error
newer_success_after_failure = await manager.create(
"t1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-retry-success"},
)
newer_success_after_failure.status = RunStatus.success
visibility = await manager.list_edit_replay_visibility("t1", user_id=None)
assert visibility == EditReplayVisibility(
hidden_source_run_ids={"source-running", "source-success", "source-retry-success"},
hidden_attempt_run_ids={failed.run_id, newer_failed.run_id, older_failed_then_success.run_id},
)

View File

@ -22,3 +22,10 @@ def test_read_human_input_response_preserves_non_empty_value():
assert response is not None
assert response["value"] == " staging "
def test_read_human_input_response_rejects_unknown_kind():
raw = _text_response("staging")
raw["response_kind"] = "unknown"
assert read_human_input_response({"human_input_response": raw}) is None

View File

@ -0,0 +1,252 @@
"""Unit tests for the inbound dedupe store (issue #4120).
Covers config resolution (memory / postgres / auto), the multi-worker
misconfiguration WARNING, and the Postgres store's atomic insert / fail-open
behavior with a mocked session factory.
"""
from __future__ import annotations
from unittest.mock import AsyncMock, MagicMock
import pytest
from app.channels.dedupe_store import (
MemoryInboundDedupeStore,
PostgresInboundDedupeStore,
make_inbound_dedupe_store,
)
class _FakeDedupe:
def __init__(self, backend: str) -> None:
self.backend = backend # plain string, mimics DedupeStorageBackend.value
class _FakeDb:
def __init__(self, backend: str) -> None:
self.backend = backend
class _FakeApp:
def __init__(self, dedupe: str, db: str) -> None:
self.dedupe_storage = _FakeDedupe(dedupe)
self.database = _FakeDb(db)
def test_factory_default_is_memory():
assert isinstance(make_inbound_dedupe_store(None), MemoryInboundDedupeStore)
def test_factory_explicit_memory_is_memory():
app = _FakeApp("memory", "sqlite")
assert isinstance(make_inbound_dedupe_store(app), MemoryInboundDedupeStore)
def test_factory_auto_with_sqlite_resolves_memory():
app = _FakeApp("auto", "sqlite")
assert isinstance(make_inbound_dedupe_store(app), MemoryInboundDedupeStore)
def test_factory_auto_with_postgres_and_multi_worker_resolves_postgres_store(monkeypatch):
# 'auto' shares Postgres state whenever the DB is Postgres, even with a single
# worker: the common K8s case is N pods x 1 worker sharing one DB, and gating on
# the in-pod worker count would silently disable cross-pod dedupe there. See the
# sibling ..._regardless_of_workers test for the documented behavior.
monkeypatch.setenv("GATEWAY_WORKERS", "2")
app = _FakeApp("auto", "postgres")
assert isinstance(make_inbound_dedupe_store(app), PostgresInboundDedupeStore)
def test_factory_auto_with_postgres_resolves_postgres_store_regardless_of_workers():
# 'auto' shares Postgres dedupe state whenever the DB is Postgres, including
# the common Kubernetes case of many replicas each running a single worker
# (GATEWAY_WORKERS=1) that still share one Postgres DB. Cross-pod dedupe must
# not depend on the in-pod worker count (issue #4120).
app = _FakeApp("auto", "postgres")
assert isinstance(make_inbound_dedupe_store(app), PostgresInboundDedupeStore)
def test_factory_explicit_postgres_resolves_postgres_store():
app = _FakeApp("postgres", "postgres")
assert isinstance(make_inbound_dedupe_store(app), PostgresInboundDedupeStore)
def test_factory_explicit_postgres_without_postgres_db_falls_back_to_memory_and_warns(monkeypatch, caplog):
import logging
# Explicit 'postgres' but the application DB is not Postgres: must warn and
# fall back to the in-process store rather than silently disabling dedupe.
app = _FakeApp("postgres", "sqlite")
with caplog.at_level(logging.WARNING):
store = make_inbound_dedupe_store(app)
assert isinstance(store, MemoryInboundDedupeStore)
assert any("dedupe_storage=postgres requires database.backend='postgres'" in r.message for r in caplog.records)
def test_factory_warns_when_memory_explicit_under_multi_worker(monkeypatch, caplog):
import logging
monkeypatch.setenv("GATEWAY_WORKERS", "2")
app = _FakeApp("memory", "sqlite")
with caplog.at_level(logging.WARNING):
store = make_inbound_dedupe_store(app)
assert isinstance(store, MemoryInboundDedupeStore)
assert any("dedupe_storage=memory with GATEWAY_WORKERS>1" in r.message for r in caplog.records)
def test_factory_warns_when_auto_resolves_memory_under_multi_worker(monkeypatch, caplog):
import logging
monkeypatch.setenv("GATEWAY_WORKERS", "2")
app = _FakeApp("auto", "sqlite")
with caplog.at_level(logging.WARNING):
store = make_inbound_dedupe_store(app)
assert isinstance(store, MemoryInboundDedupeStore)
assert any("Multi-worker deployment detected but dedupe_storage=auto" in r.message for r in caplog.records)
# ---------------------------------------------------------------------------
# PostgresInboundDedupeStore (unit, mocked session factory)
# ---------------------------------------------------------------------------
def _fake_session_factory(*, insert_rowcount: int = 1, execute_raises: BaseException | None = None):
"""Build a fake async session factory whose single session returns the
given upsert result (fetchone -> a row when admitted, None on live conflict)
or raises for ``execute``.
``try_record`` issues up to two executes: (0) the conditional upsert
``INSERT ... ON CONFLICT DO UPDATE ... RETURNING`` (admitted -> row returned)
and, only when admitted, (1) the cross-table lazy cleanup ``DELETE``. On a
live conflict only the upsert runs (no row returned -> no cleanup).
"""
insert_result = MagicMock()
insert_result.rowcount = insert_rowcount
# A truthy fetchone means the key was admitted (new delivery or an expired row
# re-admitted); None means a live conflict (duplicate). Mirrors the RETURNING
# channel contract.
insert_result.fetchone.return_value = ("c",)
session = AsyncMock()
if execute_raises is not None:
session.execute = AsyncMock(side_effect=execute_raises)
else:
session.execute = AsyncMock(
side_effect=[
insert_result, # conditional upsert result (fetchone -> row when admitted)
MagicMock(), # lazy cleanup DELETE result (only when admitted)
]
)
begin_cm = AsyncMock()
begin_cm.__aenter__ = AsyncMock(return_value=None)
begin_cm.__aexit__ = AsyncMock(return_value=False)
session.begin = MagicMock(return_value=begin_cm)
factory_cm = AsyncMock()
factory_cm.__aenter__ = AsyncMock(return_value=session)
factory_cm.__aexit__ = AsyncMock(return_value=False)
factory = MagicMock()
factory.return_value = factory_cm
return factory, session
@pytest.mark.asyncio
async def test_postgres_try_record_new_then_conflict():
factory, session = _fake_session_factory(insert_rowcount=1)
store = PostgresInboundDedupeStore(session_factory=factory)
key = ("github", "repo", "repo", "d1:uA:agentX")
# First delivery is inserted -> proceed (not a duplicate).
assert await store.try_record(key) is False
# Redelivery (same key) on a still-live row: ON CONFLICT -> DO UPDATE WHERE
# fails -> no row returned -> duplicate -> drop. Only the upsert runs.
session.execute.side_effect = [MagicMock(fetchone=MagicMock(return_value=None))]
assert await store.try_record(key) is True
@pytest.mark.asyncio
async def test_postgres_try_record_uses_atomic_on_conflict_and_lazy_cleanup():
factory, session = _fake_session_factory()
store = PostgresInboundDedupeStore(session_factory=factory)
await store.try_record(("slack", "T1", "C1", "123.456"))
# call 0: conditional upsert (INSERT ... ON CONFLICT DO UPDATE ... WHERE TTL);
# call 1: lazy cross-table cleanup (only on admit).
upsert_sql = session.execute.call_args_list[0].args[0].text
cleanup_sql = session.execute.call_args_list[1].args[0].text
assert "ON CONFLICT" in upsert_sql
assert "DO UPDATE SET first_seen = now()" in upsert_sql
assert "first_seen < now()" in upsert_sql # TTL reclamation condition
assert "RETURNING" in upsert_sql
assert "make_interval" in cleanup_sql
@pytest.mark.asyncio
async def test_postgres_try_record_reclaims_expired_unreleased_row():
# Reproduces the P1 defect: a key's row has passed the TTL but was never
# released (e.g. a crashed run). With the conditional upsert, the conflict
# fires DO UPDATE (WHERE first_seen < TTL is true), refreshes first_seen and
# RETURNS the row, so the redelivery is re-admitted (proceed, not dropped).
factory, session = _fake_session_factory()
store = PostgresInboundDedupeStore(session_factory=factory)
key = ("slack", "T1", "C1", "123.456")
# The upsert result returns a row, meaning the expired row was re-admitted
# (proceed, not a duplicate).
assert await store.try_record(key) is False
upsert_sql = session.execute.call_args_list[0].args[0].text
# The reclaim is part of the atomic upsert: a conditional DO UPDATE gated on
# the TTL, not a separate PK-scoped DELETE.
assert "ON CONFLICT" in upsert_sql
assert "DO UPDATE SET first_seen = now()" in upsert_sql
assert "first_seen < now() - make_interval" in upsert_sql
# Exactly two executes: the upsert (admitted/refreshed=True) + lazy cleanup.
assert len(session.execute.call_args_list) == 2
@pytest.mark.asyncio
async def test_postgres_try_record_alive_row_still_deduped_without_cleanup():
# A still-live (not expired) row for this key must remain a duplicate, and
# because no row is returned the lazy cleanup is skipped (only 1 execute).
factory, session = _fake_session_factory()
store = PostgresInboundDedupeStore(session_factory=factory)
key = ("slack", "T1", "C1", "123.456")
session.execute.side_effect = [
MagicMock(fetchone=MagicMock(return_value=None)), # upsert: live conflict, no row
]
assert await store.try_record(key) is True
assert len(session.execute.call_args_list) == 1
@pytest.mark.asyncio
async def test_postgres_try_record_fail_open_on_exception():
factory, _ = _fake_session_factory(execute_raises=RuntimeError("db down"))
store = PostgresInboundDedupeStore(session_factory=factory)
# Storage error must NOT drop the webhook: fail open = proceed (return False).
assert await store.try_record(("discord", "G1", "C1", "111")) is False
@pytest.mark.asyncio
async def test_postgres_release_deletes_key():
factory, session = _fake_session_factory()
store = PostgresInboundDedupeStore(session_factory=factory)
await store.release(("telegram", "chat1", "chat1", "55"))
sql = session.execute.call_args_list[0].args[0].text
assert "DELETE FROM webhook_deliveries WHERE channel = " in sql
assert "AND message_id = " in sql
@pytest.mark.asyncio
async def test_postgres_try_record_fail_open_when_no_session_factory(monkeypatch):
# When Postgres is selected but no session factory is available, the store
# must fail open (allow the message) rather than crash startup/handling.
import deerflow.persistence.engine as engine_mod
monkeypatch.setattr(engine_mod, "get_session_factory", lambda: None)
store = PostgresInboundDedupeStore() # no injected factory -> resolved lazily
# No DB available must NOT drop the message: fail open = proceed (return False).
assert await store.try_record(("discord", "G1", "C1", "111")) is False
def test_factory_resolves_postgres_store_when_db_is_postgres():
app = _FakeApp("postgres", "postgres")
store = make_inbound_dedupe_store(app)
assert isinstance(store, PostgresInboundDedupeStore)

View File

@ -273,6 +273,9 @@ def test_make_lead_agent_empty_skills_passed_correctly(monkeypatch):
supports_thinking = False
mock_app_config = MagicMock()
# make_lead_agent freezes the delta snapshot frequency from the app config;
# a bare MagicMock attribute cannot survive the freeze's positivity check.
mock_app_config.database.checkpoint_delta_snapshot_frequency = 1000
mock_app_config.get_model_config.return_value = MockModelConfig()
monkeypatch.setattr(lead_agent_module, "get_app_config", lambda: mock_app_config)
@ -315,6 +318,9 @@ def test_make_lead_agent_custom_skill_allowlist_does_not_activate_tool_policy(mo
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: [NamedTool("task"), NamedTool("bash"), NamedTool("read_file"), NamedTool("web_search")])
mock_app_config = MagicMock()
# make_lead_agent freezes the delta snapshot frequency from the app config;
# a bare MagicMock attribute cannot survive the freeze's positivity check.
mock_app_config.database.checkpoint_delta_snapshot_frequency = 1000
mock_app_config.get_model_config.return_value = SimpleNamespace(supports_thinking=False, supports_vision=False)
mock_app_config.tool_search.enabled = True
mock_app_config.skills.container_path = "/mnt/skills"
@ -357,6 +363,9 @@ def test_make_lead_agent_all_legacy_skills_preserve_all_tools(monkeypatch):
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: [NamedTool("bash"), NamedTool("read_file")])
mock_app_config = MagicMock()
# make_lead_agent freezes the delta snapshot frequency from the app config;
# a bare MagicMock attribute cannot survive the freeze's positivity check.
mock_app_config.database.checkpoint_delta_snapshot_frequency = 1000
mock_app_config.get_model_config.return_value = SimpleNamespace(supports_thinking=False, supports_vision=False)
monkeypatch.setattr(lead_agent_module, "get_app_config", lambda: mock_app_config)
@ -410,6 +419,9 @@ def test_make_lead_agent_passive_empty_skill_policy_preserves_mcp_and_other_tool
)
mock_app_config = MagicMock()
# make_lead_agent freezes the delta snapshot frequency from the app config;
# a bare MagicMock attribute cannot survive the freeze's positivity check.
mock_app_config.database.checkpoint_delta_snapshot_frequency = 1000
mock_app_config.get_model_config.return_value = SimpleNamespace(supports_thinking=False, supports_vision=False)
mock_app_config.tool_search.enabled = True
mock_app_config.tool_search.auto_promote_top_k = 3
@ -456,6 +468,9 @@ def test_default_lead_agent_does_not_apply_installed_skill_allowlists(monkeypatc
)
mock_app_config = MagicMock()
# make_lead_agent freezes the delta snapshot frequency from the app config;
# a bare MagicMock attribute cannot survive the freeze's positivity check.
mock_app_config.database.checkpoint_delta_snapshot_frequency = 1000
mock_app_config.get_model_config.return_value = SimpleNamespace(supports_thinking=False, supports_vision=False)
mock_app_config.tool_search.enabled = True
mock_app_config.skills.container_path = "/mnt/skills"
@ -485,6 +500,9 @@ def test_make_lead_agent_fails_closed_when_skill_policy_load_fails(monkeypatch):
monkeypatch.setattr(lead_agent_module, "load_agent_config", lambda x: AgentConfig(name="test", skills=["restricted"]))
mock_app_config = MagicMock()
# make_lead_agent freezes the delta snapshot frequency from the app config;
# a bare MagicMock attribute cannot survive the freeze's positivity check.
mock_app_config.database.checkpoint_delta_snapshot_frequency = 1000
mock_app_config.get_model_config.return_value = SimpleNamespace(supports_thinking=False, supports_vision=False)
def fail_storage(*args, **kwargs):
@ -530,6 +548,9 @@ def test_make_lead_agent_drops_update_agent_on_github_channel(monkeypatch):
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: [NamedTool("bash"), NamedTool("read_file")])
mock_app_config = MagicMock()
# make_lead_agent freezes the delta snapshot frequency from the app config;
# a bare MagicMock attribute cannot survive the freeze's positivity check.
mock_app_config.database.checkpoint_delta_snapshot_frequency = 1000
mock_app_config.get_model_config.return_value = SimpleNamespace(supports_thinking=False, supports_vision=False)
monkeypatch.setattr(lead_agent_module, "get_app_config", lambda: mock_app_config)
@ -566,6 +587,9 @@ def test_make_lead_agent_keeps_update_agent_on_non_webhook_channels(monkeypatch)
monkeypatch.setattr("deerflow.tools.get_available_tools", lambda **kwargs: [NamedTool("bash")])
mock_app_config = MagicMock()
# make_lead_agent freezes the delta snapshot frequency from the app config;
# a bare MagicMock attribute cannot survive the freeze's positivity check.
mock_app_config.database.checkpoint_delta_snapshot_frequency = 1000
mock_app_config.get_model_config.return_value = SimpleNamespace(supports_thinking=False, supports_vision=False)
monkeypatch.setattr(lead_agent_module, "get_app_config", lambda: mock_app_config)

View File

@ -0,0 +1,289 @@
"""Regression coverage for DeerMem's #4279 RetrievalPort integration."""
from __future__ import annotations
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from unittest.mock import MagicMock
import pytest
from deerflow.agents.memory.backends.deermem.deer_mem import DeerMem
from deerflow.agents.memory.backends.deermem.deermem.config import DeerMemConfig
from deerflow.agents.memory.backends.deermem.deermem.core.retrieval import (
FTS5Retrieval,
FTS5RetrievalAdapter,
_is_advanced_query,
_jieba_available,
create_fts5_retrieval,
)
from deerflow.agents.memory.backends.deermem.deermem.core.storage import FileMemoryStorage
def _fact(fact_id: str, content: str, *, category: str = "context") -> dict:
return {
"id": fact_id,
"content": content,
"category": category,
"confidence": 0.8,
"createdAt": "2026-07-21T00:00:00Z",
"source": {"type": "test", "threadId": None},
}
def test_natural_language_hyphens_are_not_treated_as_fts5_syntax() -> None:
assert not _is_advanced_query("real-time co-pilot node -js +python")
@pytest.mark.skipif(not _jieba_available, reason="install the optional memory-zh extra")
def test_chinese_subphrase_search_uses_jieba_tokenization(tmp_path: Path) -> None:
adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
scope = {"userId": "alice", "agentName": "agent-a"}
try:
adapter.upsert(_fact("zh", "中文检索支持验证"), scope=scope, path="")
results = adapter.search("检索", scopes=[scope], top_k=5, mode="fts5", filters=None)
assert [item["fact"]["id"] for item in results] == ["zh"]
finally:
adapter.close()
def test_score_time_decay_and_advanced_phrase_query(tmp_path: Path) -> None:
engine = FTS5Retrieval(tmp_path / "facts.sqlite3")
try:
assert engine._compute_final_score(1.0, 0.5, "2026-07-21T00:00:00Z") > engine._compute_final_score(1.0, 0.5, "2020-01-01T00:00:00Z")
engine.index_fact("phrase", "alpha beta", scope_user='"alice"', scope_agent='"agent-a"')
assert engine.search('"alpha beta"', scope_user='"alice"', scope_agent='"agent-a"')
finally:
engine.close()
def test_warm_retrieval_rebuilds_the_complete_index(tmp_path: Path) -> None:
manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
_, fact_id = manager.create_fact("warm retrieval fact", user_id="alice")
assert manager.warm_retrieval()
assert any(fact["id"] == fact_id for fact in manager.search("warm", user_id="alice"))
def test_warm_retrieval_accepts_partial_fact_failures(tmp_path: Path) -> None:
manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
manager._storage.rebuild_index = lambda scopes=None: {"supported": True, "indexed": 2, "failed": 1} # type: ignore[method-assign]
assert manager.warm_retrieval()
assert manager._retrieval_fully_warmed
def test_warm_retrieval_retries_after_fatal_rebuild_failure(tmp_path: Path) -> None:
manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
manager._storage.rebuild_index = lambda scopes=None: {"supported": True, "indexed": 0, "failed": 1, "fatal": True} # type: ignore[method-assign]
assert not manager.warm_retrieval()
assert not manager._retrieval_fully_warmed
def test_lazy_warm_rebuilds_each_requested_scope(tmp_path: Path) -> None:
manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
calls: list[list[dict[str, str | None]]] = []
manager._storage.rebuild_index = lambda scopes=None: calls.append(scopes or []) or {"supported": True, "failed": 0} # type: ignore[method-assign]
scopes = [{"userId": "alice", "agentName": "a"}, {"userId": "bob", "agentName": "b"}]
manager._ensure_retrieval_scopes(scopes)
assert calls == [[scopes[0]], [scopes[1]]]
def test_lazy_warm_does_not_retry_partial_fact_failures(tmp_path: Path) -> None:
manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
rebuild = MagicMock(return_value={"supported": True, "indexed": 2, "failed": 1})
manager._storage.rebuild_index = rebuild # type: ignore[method-assign]
scopes = [{"userId": "alice", "agentName": "a"}]
manager._ensure_retrieval_scopes(scopes)
manager._ensure_retrieval_scopes(scopes)
rebuild.assert_called_once_with(scopes)
def test_deermem_close_releases_retrieval_connection(tmp_path: Path) -> None:
manager = DeerMem(backend_config={"storage_path": str(tmp_path), "token_counting": "char"})
close_storage = MagicMock()
manager._storage.close = close_storage # type: ignore[attr-defined]
manager.close()
close_storage.assert_called_once_with()
def test_file_storage_close_releases_retrieval_connection(tmp_path: Path) -> None:
retrieval = MagicMock()
storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path)), retrieval=retrieval)
storage.close()
retrieval.close.assert_called_once_with()
def test_factory_recreates_corrupt_persistent_index(tmp_path: Path) -> None:
index_dir = tmp_path / ".retrieval"
index_dir.mkdir()
db_path = index_dir / "memory-fts5.sqlite3"
db_path.write_bytes(b"not a sqlite database")
adapter = create_fts5_retrieval(DeerMemConfig(storage_path=str(tmp_path)))
assert adapter is not None
try:
adapter.upsert(
_fact("recovered", "recreated derived index"),
scope={"userId": "alice", "agentName": "agent-a"},
path="",
)
assert adapter.search(
"recreated",
scopes=[{"userId": "alice", "agentName": "agent-a"}],
top_k=5,
mode="fts5",
filters=None,
)
finally:
adapter.close()
def test_adapter_isolates_scopes_even_when_fact_ids_repeat(tmp_path: Path) -> None:
adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
try:
adapter.upsert(_fact("same", "Alice private alpha"), scope={"userId": "alice", "agentName": "__default__"}, path="")
adapter.upsert(_fact("same", "Bob private beta"), scope={"userId": "bob", "agentName": "__default__"}, path="")
alice = adapter.search("alpha", scopes=[{"userId": "alice", "agentName": "__default__"}], top_k=5, mode="hybrid", filters=None)
bob = adapter.search("alpha", scopes=[{"userId": "bob", "agentName": "__default__"}], top_k=5, mode="hybrid", filters=None)
assert [item["fact"]["content"] for item in alice] == ["Alice private alpha"]
assert bob == []
finally:
adapter.close()
def test_adapter_persists_across_restart_and_supports_category_filter(tmp_path: Path) -> None:
db_path = tmp_path / "facts.sqlite3"
scope = {"userId": "alice", "agentName": "agent-a"}
first = FTS5RetrievalAdapter(db_path)
first.upsert(_fact("one", "Python deployment preference", category="preference"), scope=scope, path="")
first.upsert(_fact("two", "Python deployment context", category="context"), scope=scope, path="")
first.close()
second = FTS5RetrievalAdapter(db_path)
try:
results = second.search("Python deployment", scopes=[scope], top_k=5, mode="hybrid", filters={"category": "preference"})
assert [item["fact"]["id"] for item in results] == ["one"]
assert results[0]["matchType"] == "fts5"
finally:
second.close()
def test_file_storage_incremental_notifications_update_and_remove_index(tmp_path: Path) -> None:
adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path)), retrieval=adapter)
scope = {"userId": "alice", "agentName": "agent-a"}
try:
storage.upsert_fact(_fact("one", "initial searchable value"), user_id="alice", agent_name="agent-a", expected_manifest_revision=0)
assert storage.search_facts("initial", scopes=[scope])[0]["fact"]["id"] == "one"
updated = storage.get_fact("one", user_id="alice", agent_name="agent-a")
assert updated is not None
updated["content"] = "replacement searchable value"
storage.upsert_fact(
updated,
user_id="alice",
agent_name="agent-a",
expected_manifest_revision=1,
expected_fact_revision=updated["revision"],
)
assert storage.search_facts("replacement", scopes=[scope])[0]["fact"]["id"] == "one"
assert storage.search_facts("initial", scopes=[scope]) == []
storage.delete_fact("one", user_id="alice", agent_name="agent-a")
assert storage.search_facts("replacement", scopes=[scope]) == []
finally:
adapter.close()
def test_full_rebuild_removes_stale_rows_after_markdown_delete(tmp_path: Path) -> None:
adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path)), retrieval=adapter)
scope = {"userId": "alice", "agentName": "agent-a"}
try:
storage.upsert_fact(_fact("one", "stale markdown value"), user_id="alice", agent_name="agent-a", expected_manifest_revision=0)
memory_path = storage._get_memory_file_path("agent-a", user_id="alice")
fact_path = next(memory_path.parent.glob("agents/agent-a/facts/**/*.md"))
fact_path.unlink()
result = storage.rebuild_index()
assert result == {"supported": True, "indexed": 0, "failed": 0}
assert storage.search_facts("stale", scopes=[scope]) == []
finally:
adapter.close()
def test_reload_rebuilds_index_after_out_of_band_markdown_edit(tmp_path: Path) -> None:
adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path)), retrieval=adapter)
scope = {"userId": "alice", "agentName": "agent-a"}
try:
storage.upsert_fact(_fact("one", "original markdown value"), user_id="alice", agent_name="agent-a", expected_manifest_revision=0)
memory_path = storage._get_memory_file_path("agent-a", user_id="alice")
fact_path = next(memory_path.parent.glob("agents/agent-a/facts/**/*.md"))
raw = fact_path.read_text(encoding="utf-8").replace("original markdown value", "edited markdown value")
fact_path.write_text(raw, encoding="utf-8")
storage.reload("agent-a", user_id="alice")
assert storage.search_facts("edited", scopes=[scope])[0]["fact"]["id"] == "one"
assert storage.search_facts("original", scopes=[scope]) == []
finally:
adapter.close()
def test_failed_notification_marks_scope_dirty_and_rebuilds_on_search(tmp_path: Path) -> None:
adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
storage = FileMemoryStorage(DeerMemConfig(storage_path=str(tmp_path)), retrieval=adapter)
original_upsert = adapter.upsert
failed_once = True
def fail_once(fact, *, scope, path):
nonlocal failed_once
if failed_once:
failed_once = False
raise RuntimeError("simulated index outage")
original_upsert(fact, scope=scope, path=path)
adapter.upsert = fail_once # type: ignore[method-assign]
try:
storage.upsert_fact(_fact("one", "recoverable indexed value"), user_id="alice", agent_name="agent-a", expected_manifest_revision=0)
results = storage.search_facts("recoverable", scopes=[{"userId": "alice", "agentName": "agent-a"}])
assert results[0]["fact"]["id"] == "one"
finally:
adapter.close()
def test_concurrent_upserts_are_searchable(tmp_path: Path) -> None:
adapter = FTS5RetrievalAdapter(tmp_path / "facts.sqlite3")
scope = {"userId": "alice", "agentName": "agent-a"}
try:
def write(index: int) -> None:
adapter.upsert(_fact(f"fact-{index}", f"concurrent memory item {index}"), scope=scope, path="")
with ThreadPoolExecutor(max_workers=8) as pool:
list(pool.map(write, range(20)))
results = adapter.search("concurrent memory", scopes=[scope], top_k=20, mode="hybrid", filters=None)
assert {item["fact"]["id"] for item in results} == {f"fact-{index}" for index in range(20)}
finally:
adapter.close()
def test_deermem_create_and_restart_use_retrieval_adapter(tmp_path: Path) -> None:
config = {"storage_path": str(tmp_path), "token_counting": "char"}
manager = DeerMem(backend_config=config)
_, fact_id = manager.create_fact("BM25 retrieval survives restart", category="knowledge", user_id="alice")
assert fact_id is not None
assert any(fact["id"] == fact_id for fact in manager.search("retrieval BM25", user_id="alice"))
restarted = DeerMem(backend_config=config)
assert any(fact["id"] == fact_id for fact in restarted.search("restart retrieval", user_id="alice"))

View File

@ -157,7 +157,7 @@ async def test_migration_dedupes_duplicate_active_rows_before_unique_index(tmp_p
with sqlite3.connect(db_path) as raw:
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
# Bootstrap upgrades through the later revisions after 0004.
assert version_row[0] == "0009_feedback_tags"
assert version_row[0] == "0010_feedback_tags"
# Sanity: the invariant the index enforces is now true — at most one
# active row per thread.

View File

@ -169,7 +169,7 @@ async def test_migration_supersedes_duplicate_active_runs_before_unique_index(tm
with sqlite3.connect(db_path) as raw:
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
assert version_row[0] == "0009_feedback_tags"
assert version_row[0] == "0010_feedback_tags"
# Sanity: the invariant the index enforces now holds — at most one
# active row per task_id.

View File

@ -0,0 +1,53 @@
"""Migration ``0009_webhook_dedupe`` regression test (issue #4120).
Verifies the migration creates ``webhook_deliveries`` with the composite
primary key (channel, workspace_id, chat_id, message_id) and no legacy
``dedupe_key`` column, and that re-running it is idempotent.
"""
from __future__ import annotations
from pathlib import Path
import pytest
import sqlalchemy as sa
from sqlalchemy.ext.asyncio import create_async_engine
import deerflow.persistence.models # noqa: F401 -- registers ORM models
from deerflow.persistence.base import Base
from deerflow.persistence.bootstrap import bootstrap_schema
from deerflow.persistence.engine import close_engine, init_engine
pytestmark = pytest.mark.asyncio
async def test_migration_0009_creates_composite_pk_table_and_is_idempotent(tmp_path: Path) -> None:
db_path = tmp_path / "deer.db"
url = f"sqlite+aiosqlite:///{db_path}"
engine = create_async_engine(url)
try:
# Seed all baseline tables, then drop ONLY webhook_deliveries so the
# 0009 upgrade actually exercises its create_table path (not the
# idempotent early-return). Stamp at 0004 so bootstrap upgrades to head.
sync = sa.create_engine(f"sqlite:///{db_path}")
Base.metadata.create_all(sync)
with sync.begin() as conn:
conn.execute(sa.text("DROP TABLE IF EXISTS webhook_deliveries"))
conn.execute(sa.text("CREATE TABLE IF NOT EXISTS alembic_version (version_num VARCHAR(32) NOT NULL)"))
conn.execute(sa.text("DELETE FROM alembic_version"))
conn.execute(sa.text("INSERT INTO alembic_version (version_num) VALUES ('0004_run_ownership')"))
await init_engine("sqlite", url=url, sqlite_dir=str(tmp_path))
# Runs upgrade head -> executes 0009.create_table.
await bootstrap_schema(engine, backend="sqlite")
async with engine.connect() as conn:
cols = {row["name"] for row in await conn.run_sync(lambda c: sa.inspect(c).get_columns("webhook_deliveries"))}
# Composite PK columns only; the old single-column ``dedupe_key`` must
# NOT exist (it is illegal in Postgres TEXT and caused schema drift).
assert cols == {"channel", "workspace_id", "chat_id", "message_id", "first_seen"}
# Idempotent: re-running bootstrap at head must not raise (table exists).
await bootstrap_schema(engine, backend="sqlite")
finally:
await close_engine()

View File

@ -0,0 +1,170 @@
"""Integration test: two ChannelManagers share one Postgres dedupe table.
Mirrors ``test_multi_worker_run_ownership.py``'s "two managers sharing a DB"
pattern. Skipped unless a real Postgres backend is configured via
``DEDUPE_TEST_POSTGRES_URL`` (the dev container runs sqlite), so it only runs
where cross-pod dedupe actually matters CI with Postgres, or a local run with
a throwaway Postgres. See issue #4120.
"""
from __future__ import annotations
import os
import uuid
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import pytest
import pytest_asyncio
from app.channels.dedupe_store import INBOUND_DEDUPE_TTL_SECONDS, PostgresInboundDedupeStore
POSTGRES_URL = os.environ.get("DEDUPE_TEST_POSTGRES_URL")
# libpq/psycopg-only query parameters that asyncpg's ``connect()`` rejects. CI
# hands over ``TEST_POSTGRES_URI`` as ``postgresql://...?sslmode=disable``; left
# in place these raise ``TypeError: connect() got an unexpected keyword argument
# 'sslmode'`` once the URL is routed to the asyncpg driver, so strip them here.
_LIBPQ_ONLY_QUERY_KEYS = {"sslmode", "channel_binding"}
def _asyncpg_url(url: str | None) -> str | None:
"""Normalize a sync ``postgresql://`` URL to the async driver SQLAlchemy needs.
``init_engine`` builds an *async* engine, so a bare ``postgresql://`` would fall
back to the synchronous psycopg2 driver and fail. Accept either form so the test
runs whether CI hands over ``TEST_POSTGRES_URI`` (``postgresql://...``) or a
caller passes ``postgresql+asyncpg://...`` directly. libpq-only query params
(e.g. ``sslmode``) are dropped because asyncpg does not understand them.
"""
if not url:
return url
if url.startswith("postgresql://"):
url = "postgresql+asyncpg://" + url[len("postgresql://") :]
parts = urlsplit(url)
if parts.query:
kept = [(k, v) for k, v in parse_qsl(parts.query, keep_blank_values=True) if k not in _LIBPQ_ONLY_QUERY_KEYS]
url = urlunsplit(parts._replace(query=urlencode(kept)))
return url
pytestmark = pytest.mark.skipif(
not POSTGRES_URL,
reason="requires DEDUPE_TEST_POSTGRES_URL (real Postgres for cross-pod dedupe)",
)
@pytest_asyncio.fixture(autouse=True)
async def _postgres_engine():
"""Initialize the Postgres engine within the test's event loop.
Doing this inside the loop (not at import time) avoids the "connection
attached to a different loop" error that arises when the engine is built in
a loop that pytest-asyncio later closes.
"""
from deerflow.persistence.engine import close_engine, init_engine
await init_engine("postgres", url=_asyncpg_url(POSTGRES_URL))
try:
yield
finally:
await close_engine()
@pytest.mark.asyncio
async def test_two_stores_share_dedupe_state_across_pods():
from deerflow.persistence.base import Base
from deerflow.persistence.engine import get_engine, get_session_factory
from deerflow.persistence.webhook_delivery.model import WebhookDeliveryRow
engine = get_engine()
sf = get_session_factory()
# Ensure the table exists regardless of whether the alembic migration ran.
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all, tables=[WebhookDeliveryRow.__table__], checkfirst=True)
store_a = PostgresInboundDedupeStore(session_factory=sf)
store_b = PostgresInboundDedupeStore(session_factory=sf)
unique = uuid.uuid4().hex
key = ("github", "repo", "repo", f"d-{unique}:uA:agentX")
try:
# First pod records the delivery and proceeds.
assert await store_a.try_record(key) is False
# A redelivery landing on the second pod hits the same table -> duplicate.
assert await store_b.try_record(key) is True
finally:
await store_a.release(key)
@pytest.mark.asyncio
async def test_manager_injects_shared_store_and_dedupes_cross_pod():
"""End-to-end: two ChannelManagers wired to the *same* Postgres store drop a
redelivery routed to the second manager, exactly as the production dispatch
loop would (``if await self._is_duplicate_inbound(msg): continue``)."""
from app.channels.manager import ChannelManager
from app.channels.message_bus import InboundMessage, MessageBus
from app.channels.store import ChannelStore
from deerflow.persistence.engine import get_session_factory
sf = get_session_factory()
shared_store = PostgresInboundDedupeStore(session_factory=sf)
manager_a = ChannelManager(bus=MessageBus(), store=ChannelStore(), inbound_dedupe_store=shared_store)
manager_b = ChannelManager(bus=MessageBus(), store=ChannelStore(), inbound_dedupe_store=shared_store)
msg = InboundMessage(
channel_name="github",
chat_id="repo",
user_id="alice",
text="@bot review",
topic_id="7:agent",
workspace_id="repo",
metadata={"message_id": "d-crosspod:agent"},
)
try:
# First manager records the delivery and lets it through.
assert await manager_a._is_duplicate_inbound(msg) is False
# Same redelivery on the second manager hits the shared table -> dropped.
assert await manager_b._is_duplicate_inbound(msg) is True
finally:
await manager_a._release_inbound_dedupe_key(msg)
@pytest.mark.asyncio
async def test_expired_unreleased_row_is_reclaimed_on_next_redelivery():
"""P1 end-to-end: a row past the TTL that was never released must be
reclaimed so the next redelivery is re-admitted (not dropped forever)."""
from sqlalchemy import text
from deerflow.persistence.base import Base
from deerflow.persistence.engine import get_engine, get_session_factory
from deerflow.persistence.webhook_delivery.model import WebhookDeliveryRow
engine = get_engine()
sf = get_session_factory()
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all, tables=[WebhookDeliveryRow.__table__], checkfirst=True)
store = PostgresInboundDedupeStore(session_factory=sf)
unique = uuid.uuid4().hex
channel, workspace_id, chat_id, message_id = ("github", "repo", "repo", f"d-{unique}:expired")
key = (channel, workspace_id, chat_id, message_id)
try:
# Seed an already-expired, unreleased row for this key.
async with sf() as session:
async with session.begin():
await session.execute(
text("INSERT INTO webhook_deliveries (channel, workspace_id, chat_id, message_id, first_seen) VALUES (:c, :w, :ch, :m, now() - make_interval(secs => :age))"),
{"c": channel, "w": workspace_id, "ch": chat_id, "m": message_id, "age": INBOUND_DEDUPE_TTL_SECONDS + 1},
)
# The expired row is reclaimed -> redelivery re-admitted (not a duplicate).
assert await store.try_record(key) is False
# Its first_seen is refreshed to ~now (well within the TTL).
async with sf() as session:
row = (
await session.execute(
text("SELECT (now() - first_seen) < make_interval(secs => :ttl) AS fresh FROM webhook_deliveries WHERE channel = :c AND workspace_id = :w AND chat_id = :ch AND message_id = :m"),
{"c": channel, "w": workspace_id, "ch": chat_id, "m": message_id, "ttl": INBOUND_DEDUPE_TTL_SECONDS},
)
).fetchone()
assert row is not None and row[0] is True
finally:
await store.release(key)

View File

@ -48,7 +48,7 @@ from deerflow.persistence.migrations._helpers import _normalize_default
asyncio_test = pytest.mark.asyncio
HEAD = "0009_feedback_tags"
HEAD = "0010_feedback_tags"
BASELINE = "0001_baseline"

View File

@ -28,7 +28,7 @@ from deerflow.persistence.bootstrap import bootstrap_schema
pytestmark = pytest.mark.asyncio
HEAD = "0009_feedback_tags"
HEAD = "0010_feedback_tags"
def _url(tmp_path: Path) -> str:

View File

@ -76,7 +76,7 @@ async def test_legacy_database_recovers_token_usage_column(tmp_path: Path) -> No
cols = {row[1] for row in raw.execute("PRAGMA table_info(runs)").fetchall()}
assert "token_usage_by_model" in cols
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
assert version_row[0] == "0009_feedback_tags"
assert version_row[0] == "0010_feedback_tags"
# And the read path that originally 500'd must now succeed.
sf = get_session_factory()
@ -116,6 +116,6 @@ async def test_legacy_database_with_manual_alter_still_bootstraps(tmp_path: Path
# No duplicate column -- list, not set, to catch dupes.
assert cols.count("token_usage_by_model") == 1
version_row = raw.execute("SELECT version_num FROM alembic_version").fetchone()
assert version_row[0] == "0009_feedback_tags"
assert version_row[0] == "0010_feedback_tags"
finally:
await close_engine()

View File

@ -116,6 +116,28 @@ def test_run_request_keeps_supported_modes_and_compatibility_defaults() -> None:
assert body.if_not_exists == "create"
@pytest.mark.parametrize(
"artifact_delivery",
[
{"required": True, "tool": "present_files"},
{},
None,
"present_files",
[],
],
)
def test_run_request_rejects_removed_artifact_delivery_override(
client: TestClient,
artifact_delivery: Any,
) -> None:
response = client.post(
"/api/runs/stream",
json={"artifact_delivery": artifact_delivery},
)
assert response.status_code == 422
def _sdk_default_payload(method: str) -> dict[str, Any]:
"""Capture the body a stock ``langgraph_sdk`` client sends for a default run."""
from langgraph_sdk.client import get_client

View File

@ -13,7 +13,7 @@ from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.runs.manager import RunManager
from deerflow.runtime.runs.schemas import RunStatus
from deerflow.runtime.runs.store.memory import MemoryRunStore
from deerflow.runtime.runs.worker import RunContext, run_agent
from deerflow.runtime.runs.worker import RunContext, _delivery_content_with_outputs, run_agent
def _make_bridge():
@ -25,6 +25,28 @@ async def _delivery_events(store: MemoryRunEventStore, thread_id: str, run_id: s
return [e for e in events if e["event_type"] == "run.delivery"]
def test_delivery_verification_treats_presented_directory_as_covering_produced_files():
content = {
"presented": 1,
"paths": ["/mnt/user-data/outputs/site"],
"by_tool": {"present_files": ["/mnt/user-data/outputs/site"]},
}
delivery = _delivery_content_with_outputs(
content,
[
"/mnt/user-data/outputs/site/index.html",
"/mnt/user-data/outputs/site/assets/style.css",
],
)
assert delivery["matched_paths"] == [
"/mnt/user-data/outputs/site/index.html",
"/mnt/user-data/outputs/site/assets/style.css",
]
assert delivery["satisfied"] is True
@pytest.mark.anyio
async def test_delivery_event_records_present_files_paths_on_success():
run_manager = RunManager()
@ -95,6 +117,201 @@ async def test_delivery_event_presented_zero_without_artifact_production():
assert fetched.status == RunStatus.success
@pytest.mark.anyio
async def test_changed_outputs_succeed_when_a_produced_output_is_presented(monkeypatch):
run_manager = RunManager()
record = await run_manager.create("thread-1")
store = MemoryRunEventStore()
monkeypatch.setattr(
"deerflow.runtime.runs.worker._produced_output_paths",
AsyncMock(return_value=["/mnt/user-data/outputs/report.md"]),
)
class DummyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
journal = config["context"]["__run_journal"]
ai = AIMessage(content="", tool_calls=[{"id": "call_1", "name": "present_files", "args": {}}])
journal._remember_current_run_tool_calls(ai, caller="lead_agent")
journal.on_tool_end(
Command(
update={
"artifacts": ["/mnt/user-data/outputs/report.md"],
"messages": [ToolMessage("Successfully presented files", tool_call_id="call_1")],
}
),
run_id=uuid4(),
)
yield {"messages": []}
await run_agent(
_make_bridge(),
run_manager,
record,
ctx=RunContext(checkpointer=None, event_store=store),
agent_factory=lambda *, config: DummyAgent(),
graph_input={},
config={},
)
delivery = await _delivery_events(store, "thread-1", record.run_id)
assert delivery[0]["content"] == {
"presented": 1,
"paths": ["/mnt/user-data/outputs/report.md"],
"by_tool": {"present_files": ["/mnt/user-data/outputs/report.md"]},
"verification": {
"source": "outputs_changed",
"requirement": "present_files_matches_produced_output",
},
"produced_paths": ["/mnt/user-data/outputs/report.md"],
"presented_paths": ["/mnt/user-data/outputs/report.md"],
"matched_paths": ["/mnt/user-data/outputs/report.md"],
"stage": "presented",
"satisfied": True,
}
assert record.status == RunStatus.success
@pytest.mark.anyio
async def test_changed_outputs_fail_closed_when_not_presented(monkeypatch):
run_manager = RunManager()
record = await run_manager.create("thread-1")
store = MemoryRunEventStore()
monkeypatch.setattr(
"deerflow.runtime.runs.worker._produced_output_paths",
AsyncMock(return_value=["/mnt/user-data/outputs/report.md"]),
)
class ProseOnlyAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
yield {"messages": [AIMessage(content="SESSION SUMMARY")]}
await run_agent(
_make_bridge(),
run_manager,
record,
ctx=RunContext(checkpointer=None, event_store=store),
agent_factory=lambda *, config: ProseOnlyAgent(),
graph_input={},
config={},
)
delivery = await _delivery_events(store, "thread-1", record.run_id)
assert delivery[0]["content"] == {
"presented": 0,
"paths": [],
"by_tool": {},
"verification": {
"source": "outputs_changed",
"requirement": "present_files_matches_produced_output",
},
"produced_paths": ["/mnt/user-data/outputs/report.md"],
"presented_paths": [],
"matched_paths": [],
"stage": "not_started",
"satisfied": False,
}
assert record.status == RunStatus.error
assert record.error == "Artifact delivery incomplete: no produced output artifact was presented"
assert record.stop_reason is None
@pytest.mark.anyio
async def test_changed_outputs_succeed_when_one_of_multiple_outputs_is_presented(monkeypatch):
run_manager = RunManager()
record = await run_manager.create("thread-1")
store = MemoryRunEventStore()
monkeypatch.setattr(
"deerflow.runtime.runs.worker._produced_output_paths",
AsyncMock(
return_value=[
"/mnt/user-data/outputs/report.md",
"/mnt/user-data/outputs/appendix.md",
]
),
)
class PartiallyPresentingAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
journal = config["context"]["__run_journal"]
journal._remember_current_run_tool_calls(
AIMessage(content="", tool_calls=[{"id": "call_1", "name": "present_files", "args": {}}]),
caller="lead_agent",
)
journal.on_tool_end(
Command(
update={
"artifacts": ["/mnt/user-data/outputs/report.md"],
"messages": [ToolMessage("Successfully presented files", tool_call_id="call_1")],
}
),
run_id=uuid4(),
)
yield {"messages": []}
await run_agent(
_make_bridge(),
run_manager,
record,
ctx=RunContext(checkpointer=None, event_store=store),
agent_factory=lambda *, config: PartiallyPresentingAgent(),
graph_input={},
config={},
)
delivery = (await _delivery_events(store, "thread-1", record.run_id))[0]["content"]
assert delivery["stage"] == "presented"
assert delivery["presented_paths"] == ["/mnt/user-data/outputs/report.md"]
assert delivery["matched_paths"] == ["/mnt/user-data/outputs/report.md"]
assert delivery["satisfied"] is True
assert record.status == RunStatus.success
@pytest.mark.anyio
async def test_changed_outputs_fail_when_present_files_only_presents_an_unrelated_file(monkeypatch):
run_manager = RunManager()
record = await run_manager.create("thread-1")
store = MemoryRunEventStore()
monkeypatch.setattr(
"deerflow.runtime.runs.worker._produced_output_paths",
AsyncMock(return_value=["/mnt/user-data/outputs/report.md"]),
)
class UnrelatedPresentingAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
journal = config["context"]["__run_journal"]
journal._remember_current_run_tool_calls(
AIMessage(content="", tool_calls=[{"id": "call_1", "name": "present_files", "args": {}}]),
caller="lead_agent",
)
journal.on_tool_end(
Command(
update={
"artifacts": ["/mnt/user-data/outputs/old-report.md"],
"messages": [ToolMessage("Successfully presented files", tool_call_id="call_1")],
}
),
run_id=uuid4(),
)
yield {"messages": []}
await run_agent(
_make_bridge(),
run_manager,
record,
ctx=RunContext(checkpointer=None, event_store=store),
agent_factory=lambda *, config: UnrelatedPresentingAgent(),
graph_input={},
config={},
)
delivery = (await _delivery_events(store, "thread-1", record.run_id))[0]["content"]
assert delivery["stage"] == "mismatched"
assert delivery["presented_paths"] == ["/mnt/user-data/outputs/old-report.md"]
assert delivery["matched_paths"] == []
assert delivery["satisfied"] is False
assert record.status == RunStatus.error
@pytest.mark.anyio
async def test_fenced_worker_leaves_delivery_receipt_to_peer_recovery():
"""A stale worker must not finalize the singleton delivery receipt."""
@ -337,6 +554,53 @@ async def test_delivery_write_failure_preserves_real_durable_terminal_status():
assert (await run_store.get(record.run_id))["status"] == "success"
@pytest.mark.anyio
async def test_produced_artifact_delivery_fails_closed_when_receipt_cannot_be_persisted(monkeypatch):
class FailingReceiptStore(MemoryRunEventStore):
async def put_if_absent(self, **kwargs):
raise RuntimeError("event store unavailable")
run_store = MemoryRunStore()
run_manager = RunManager(store=run_store)
record = await run_manager.create("thread-1")
monkeypatch.setattr(
"deerflow.runtime.runs.worker._produced_output_paths",
AsyncMock(return_value=["/mnt/user-data/outputs/report.md"]),
)
class PresentingAgent:
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
journal = config["context"]["__run_journal"]
journal._remember_current_run_tool_calls(
AIMessage(content="", tool_calls=[{"id": "call_1", "name": "present_files", "args": {}}]),
caller="lead_agent",
)
journal.on_tool_end(
Command(
update={
"artifacts": ["/mnt/user-data/outputs/report.md"],
"messages": [ToolMessage("Successfully presented files", tool_call_id="call_1")],
}
),
run_id=uuid4(),
)
yield {"messages": []}
await run_agent(
_make_bridge(),
run_manager,
record,
ctx=RunContext(checkpointer=None, event_store=FailingReceiptStore()),
agent_factory=lambda *, config: PresentingAgent(),
graph_input={},
config={},
)
assert record.status == RunStatus.error
assert record.error == "Artifact delivery verification failed: terminal delivery receipt could not be persisted"
assert (await run_store.get(record.run_id))["status"] == "error"
@pytest.mark.anyio
async def test_delivery_event_emitted_when_checkpoint_preflight_fails(monkeypatch):
run_manager = RunManager()

View File

@ -41,6 +41,7 @@ def anyio_backend():
class _DeltaChannelState(TypedDict):
messages: Annotated[list[AnyMessage], DeltaChannel(merge_message_writes, snapshot_frequency=1000)]
title: str | None
class _FullChannelState(TypedDict):
@ -226,12 +227,18 @@ async def test_run_agent_streams_from_the_linearized_delta_resume(tmp_path):
branch_writer = CheckpointStateAccessor.bind(mutation_graph, checkpointer, mode="delta")
replay_base_config = await branch_writer.aupdate(
_run_config("worker-branch"),
{"messages": Overwrite(list(source_pre_turn.values["messages"]))},
{
"messages": Overwrite(list(source_pre_turn.values["messages"])),
"title": "Original title",
},
as_node="branch",
)
await branch_writer.aupdate(
replay_base_config,
{"messages": Overwrite(list(source_head.values["messages"]))},
{
"messages": Overwrite(list(source_head.values["messages"])),
"title": "User renamed title",
},
as_node="branch",
)
@ -247,6 +254,10 @@ async def test_run_agent_streams_from_the_linearized_delta_resume(tmp_path):
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
thread_store = SimpleNamespace(
update_display_name=AsyncMock(),
update_status=AsyncMock(),
)
created_graphs: list[Any] = []
def agent_factory(*, config):
@ -260,9 +271,12 @@ async def test_run_agent_streams_from_the_linearized_delta_resume(tmp_path):
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=checkpointer, checkpoint_channel_mode="delta"),
ctx=RunContext(checkpointer=checkpointer, thread_store=thread_store, checkpoint_channel_mode="delta"),
agent_factory=agent_factory,
graph_input={"messages": [HumanMessage(content="q2", id="h2")]},
graph_input={
"messages": [HumanMessage(content="q2", id="h2")],
"title": "User renamed title",
},
config=_run_config("worker-branch", base.config["configurable"]["checkpoint_id"]),
stream_modes=["values"],
),
@ -273,6 +287,8 @@ async def test_run_agent_streams_from_the_linearized_delta_resume(tmp_path):
final_accessor = CheckpointStateAccessor.bind(created_graphs[-1], checkpointer, mode="delta")
final = await final_accessor.aget(_run_config("worker-branch"))
assert _ids(final) == ["h1", "a1", "h2", "a2-new"]
assert final.values["title"] == "User renamed title"
thread_store.update_display_name.assert_awaited_once_with("worker-branch", "User renamed title")
async def test_run_agent_serializes_resume_preparation_with_checkpoint_writes(monkeypatch):

View File

@ -814,6 +814,150 @@ async def test_run_agent_marks_llm_error_fallback_as_error_status():
bridge.publish_end.assert_awaited_once_with(record.run_id)
@pytest.mark.anyio
async def test_run_agent_rolls_back_failed_edit_replay_and_publishes_restored_values():
run_manager = RunManager()
record = await run_manager.create(
"thread-1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-run"},
)
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
before_messages = [HumanMessage(id="h1", content="original"), AIMessage(id="a1", content="answer")]
class DummyCheckpointer:
async def aget_tuple(self, _config):
return SimpleNamespace(
config={
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": "checkpoint-1",
}
},
checkpoint={"id": "checkpoint-1", "channel_values": {}},
metadata={"source": "loop"},
pending_writes=[],
)
class DummyAgent:
async def aget_state(self, _config):
return SimpleNamespace(
values={"messages": before_messages},
config={
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": "checkpoint-1",
}
},
parent_config=None,
metadata={},
next=(),
tasks=(),
created_at=None,
)
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
raise RuntimeError("edit replay failed")
if False:
yield # pragma: no cover - keep this an async generator
error_status_finalizing_states: list[bool] = []
original_set_status = run_manager.set_status
async def _set_status(run_id, status, **kwargs):
if run_id == record.run_id and status == RunStatus.error:
error_status_finalizing_states.append(record.finalizing)
return await original_set_status(run_id, status, **kwargs)
run_manager.set_status = _set_status # type: ignore[method-assign]
with patch(
"deerflow.runtime.runs.worker._rollback_to_pre_run_checkpoint",
new_callable=AsyncMock,
) as rollback:
rollback.return_value = True
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=DummyCheckpointer()),
agent_factory=lambda *, config: DummyAgent(),
graph_input={},
config={},
)
fetched = await run_manager.get(record.run_id)
assert fetched is not None
assert fetched.status == RunStatus.error
assert error_status_finalizing_states == [True]
rollback.assert_awaited_once()
publish_events = [call_args.args[1] for call_args in bridge.publish.await_args_list]
assert "error" in publish_events
assert "values" in publish_events
assert publish_events.index("values") > publish_events.index("error")
bridge.publish_end.assert_awaited_once_with(record.run_id)
@pytest.mark.anyio
async def test_failed_edit_replay_does_not_publish_restored_values_when_snapshot_capture_failed():
run_manager = RunManager()
record = await run_manager.create(
"thread-1",
metadata={"replay_kind": "edit", "regenerate_from_run_id": "source-run"},
)
bridge = SimpleNamespace(
publish=AsyncMock(),
publish_end=AsyncMock(),
cleanup=AsyncMock(),
)
class DummyCheckpointer:
async def aget_tuple(self, _config):
return SimpleNamespace(
config={
"configurable": {
"thread_id": "thread-1",
"checkpoint_ns": "",
"checkpoint_id": "checkpoint-1",
}
},
checkpoint={"id": "checkpoint-1", "channel_values": {}},
metadata={"source": "loop"},
pending_writes=[],
)
class DummyAgent:
async def aget_state(self, _config):
raise RuntimeError("snapshot capture failed")
async def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False):
raise RuntimeError("edit replay failed")
if False:
yield # pragma: no cover - keep this an async generator
await run_agent(
bridge,
run_manager,
record,
ctx=RunContext(checkpointer=DummyCheckpointer()),
agent_factory=lambda *, config: DummyAgent(),
graph_input={},
config={},
)
fetched = await run_manager.get(record.run_id)
assert fetched is not None
assert fetched.status == RunStatus.error
publish_events = [call_args.args[1] for call_args in bridge.publish.await_args_list]
assert "error" in publish_events
assert "values" not in publish_events
bridge.publish_end.assert_awaited_once_with(record.run_id)
@pytest.mark.anyio
async def test_run_agent_defaults_root_run_name_from_assistant_id():
run_manager = RunManager()

View File

@ -9,7 +9,7 @@ import pytest
def _load_module():
path = Path(__file__).resolve().parents[1] / "scripts/benchmark/summarize_checkpoint_channels.py"
path = Path(__file__).resolve().parents[1] / "scripts/benchmark/checkpoint/summarize_channels.py"
spec = importlib.util.spec_from_file_location("summarize_checkpoint_channels", path)
assert spec is not None
module = importlib.util.module_from_spec(spec)

View File

@ -0,0 +1,128 @@
"""Tests for the production benchmark summarizer."""
from __future__ import annotations
import importlib.util
import json
import sys
from pathlib import Path
def _load_module():
path = Path(__file__).resolve().parents[1] / "scripts/benchmark/checkpoint/summarize_production.py"
spec = importlib.util.spec_from_file_location("summarize_production", path)
assert spec is not None
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
summ = _load_module()
def _row(mode, rep, *, freq=None, turns=100, payload=128, run_p50=100.0, write_p50=10.0, write_p99=20.0, state_cold=50.0, state_warm=5.0, success=True):
return {
"mode": mode,
"repetition": rep,
"snapshot_frequency": freq,
"turns": turns,
"payload_bytes": payload,
"success": success,
"profiled": False,
"run_turn_p50_ms": run_p50,
"checkpoint_write_p50_ms": write_p50,
"checkpoint_write_p99_ms": write_p99,
"state_cold_ms": state_cold,
"state_warm_p50_ms": state_warm,
}
def test_pairs_full_row_against_every_delta_frequency(tmp_path) -> None:
path = tmp_path / "rows.jsonl"
rows = [
_row("full", 0),
_row("delta", 0, freq=10, run_p50=80.0),
_row("delta", 0, freq=100, run_p50=90.0),
]
path.write_text("\n".join(json.dumps(r) for r in rows) + "\n")
summaries = summ._summarize(summ._load_jsonl([path]), metrics=["run_turn_p50_ms"])
assert len(summaries) == 2
by_freq = {s["snapshot_frequency"]: s for s in summaries}
assert by_freq[10]["ratio_run_turn_p50_ms"] == 0.8
assert by_freq[100]["ratio_run_turn_p50_ms"] == 0.9
def test_decision_metrics_emitted(tmp_path) -> None:
path = tmp_path / "rows.jsonl"
rows = [_row("full", 0), _row("delta", 0, freq=10, write_p50=10.0, write_p99=50.0)]
path.write_text("\n".join(json.dumps(r) for r in rows) + "\n")
summaries = summ._summarize(summ._load_jsonl([path]), metrics=["run_turn_p50_ms"])
assert summaries[0]["delta_snapshot_write_spike"] == 5.0
assert summaries[0]["delta_cache_effect_ms"] == 45.0
assert summaries[0]["full_cache_effect_ms"] == 45.0
def test_failed_rows_counted_per_group(tmp_path) -> None:
path = tmp_path / "rows.jsonl"
rows = [
_row("full", 0),
_row("delta", 0, freq=10),
_row("delta", 0, freq=100),
_row("delta", 1, freq=10, success=False),
_row("delta", 2, freq=10, success=False),
_row("delta", 1, freq=100, success=False),
# Different turns: matches no group below.
_row("delta", 0, freq=10, turns=200, success=False),
# A failed full row counts toward every group at the same (turns, payload_bytes).
_row("full", 1, success=False),
]
path.write_text("\n".join(json.dumps(r) for r in rows) + "\n")
summaries = summ._summarize(summ._load_jsonl([path]), metrics=["run_turn_p50_ms"])
by_freq = {s["snapshot_frequency"]: s for s in summaries}
assert by_freq[10]["failed_rows"] == 3
assert by_freq[100]["failed_rows"] == 2
def test_history_limit_metrics_auto_discovered(tmp_path, capsys) -> None:
path = tmp_path / "rows.jsonl"
full = _row("full", 0)
delta = _row("delta", 0, freq=10)
for row, cold, warm in ((full, 30.0, 12.0), (delta, 60.0, 24.0)):
row["history_limit_64_cold_ms"] = cold
row["history_limit_64_warm_p50_ms"] = warm
path.write_text("\n".join(json.dumps(r) for r in (full, delta)) + "\n")
assert summ.main([str(path), "--json"]) == 0
out = json.loads(capsys.readouterr().out)
assert out[0]["ratio_history_limit_64_cold_ms"] == 2.0
assert out[0]["delta_history_limit_64_warm_p50_ms"] == 24.0
def test_explicit_metrics_override_skips_discovery(tmp_path, capsys) -> None:
path = tmp_path / "rows.jsonl"
full = _row("full", 0)
delta = _row("delta", 0, freq=10)
for row in (full, delta):
row["history_limit_64_cold_ms"] = 30.0
path.write_text("\n".join(json.dumps(r) for r in (full, delta)) + "\n")
assert summ.main([str(path), "--json", "--metrics", "run_turn_p50_ms"]) == 0
out = json.loads(capsys.readouterr().out)
assert "delta_history_limit_64_cold_ms" not in out[0]
def test_checkpoint_write_share_emitted_and_zero_guarded(tmp_path) -> None:
path = tmp_path / "rows.jsonl"
rows = [
_row("full", 0, run_p50=100.0, write_p50=10.0),
_row("delta", 0, freq=10, run_p50=80.0, write_p50=20.0),
_row("full", 0, turns=200, run_p50=0.0, write_p50=10.0),
_row("delta", 0, freq=10, turns=200, run_p50=0.0, write_p50=20.0),
]
path.write_text("\n".join(json.dumps(r) for r in rows) + "\n")
summaries = summ._summarize(summ._load_jsonl([path]), metrics=["run_turn_p50_ms"])
by_turns = {s["turns"]: s for s in summaries}
assert by_turns[100]["full_checkpoint_write_share"] == 0.1
assert by_turns[100]["delta_checkpoint_write_share"] == 0.25
assert "full_checkpoint_write_share" not in by_turns[200]
assert "delta_checkpoint_write_share" not in by_turns[200]

View File

@ -0,0 +1,933 @@
"""Unit tests for the Tenki community sandbox provider.
These run in CI without ``tenki-sandbox`` installed: they cover the lazy-import
error path, provider lifecycle, path-safety guards, the native ``fs`` file
round-trip, warm-pool mechanics, and scope resolution none of which need a live
sandbox. A single opt-in integration test (``test_integration_real_sandbox``)
exercises a real Tenki microVM end to end when ``TENKI_API_KEY`` is set.
"""
from __future__ import annotations
import errno
import os
import shlex
import sys
import threading
import time
import types
import pytest
from deerflow.community.tenki.provider import _BOOTSTRAP_TIMEOUT, TenkiSandboxProvider, _import_client
from deerflow.community.tenki.sandbox import TenkiSandbox
# ── Fake Tenki SDK ────────────────────────────────────────────────────
class _FakeResult:
def __init__(self, exit_code: int = 0, stdout: bytes = b"", stderr: bytes = b"") -> None:
self.exit_code = exit_code
self.stdout = stdout
self.stderr = stderr
@property
def stdout_text(self) -> str:
return self.stdout.decode(errors="replace")
@property
def stderr_text(self) -> str:
return self.stderr.decode(errors="replace")
class _FakeSessionTerminatedError(RuntimeError):
"""Stand-in whose class name matches a Tenki terminal error."""
# Rename so ``type(e).__name__`` matches the adapter's terminal-name set.
_FakeSessionTerminatedError.__name__ = "SessionTerminatedError"
class _FakeMissingFileError(Exception):
"""Stand-in for ``tenki_sandbox.errors.FileNotFoundError`` (matched by name)."""
_FakeMissingFileError.__name__ = "FileNotFoundError"
class _FakeFileInfo:
def __init__(self, path: str, size: int) -> None:
self.path = path
self.size = size
self.is_dir = False
class _FakeFS:
"""In-memory stand-in for Tenki's native ``sandbox.fs`` API."""
_READ_FRAME = 64 # small, so multi-frame reads are exercised
def __init__(self, owner: _FakeSandbox) -> None:
self._owner = owner
def _guard(self) -> None:
if self._owner.fs_error is not None:
raise self._owner.fs_error
def mkdir(self, path: str, *, recursive: bool = True, mode: int = 0o755) -> None:
self._guard()
self._owner.dirs.add(path)
def read_bytes(self, path: str) -> bytes:
self._guard()
if path not in self._owner.files:
raise _FakeMissingFileError(f"no such file or directory: {path}")
return self._owner.files[path]
def read_text(self, path: str, *, encoding: str = "utf-8") -> str:
return self.read_bytes(path).decode(encoding, errors="replace")
def read_stream(self, path: str, *, offset: int = 0, length: int = 0, chunk_bytes: int = 0):
# A generator, like the real SDK: errors surface while iterating.
data = self.read_bytes(path)
for i in range(0, len(data), self._READ_FRAME):
yield data[i : i + self._READ_FRAME]
def write_stream(self, path: str, chunks, *, mode: int = 0o644, truncate: bool = True, sync: bool = False) -> None:
self._guard()
frames = list(chunks)
self._owner.write_frame_counts.append(len(frames))
self._owner.files[path] = b"".join(frames)
def stat(self, path: str) -> _FakeFileInfo:
self._guard()
if path not in self._owner.files:
raise _FakeMissingFileError(f"no such file or directory: {path}")
return _FakeFileInfo(path, len(self._owner.files[path]))
class _FakeSandbox:
"""A fake tenki ``Sandbox``: a native ``fs`` API plus the handful of shell
commands the adapter still emits for search, backed by one in-memory
filesystem so file and search round-trips are exercised for real.
"""
def __init__(
self,
*,
exec_error: Exception | None = None,
close_error: Exception | None = None,
fs_error: Exception | None = None,
wait_ready_error: Exception | None = None,
) -> None:
self.id = "remote-session-id"
self.files: dict[str, bytes] = {}
self.dirs: set[str] = set()
self.write_frame_counts: list[int] = []
self.exec_calls: list[dict] = []
self.exec_error = exec_error
self.close_error = close_error
self.fs_error = fs_error
self.wait_ready_error = wait_ready_error
self.wait_ready_calls = 0
self.closed = False
self.fs = _FakeFS(self)
def wait_ready(self, timeout: float = 180) -> None:
self.wait_ready_calls += 1
if self.wait_ready_error is not None:
raise self.wait_ready_error
def exec(self, *argv: str, cwd=None, env=None, timeout=None):
self.exec_calls.append({"argv": argv, "cwd": cwd, "env": env, "timeout": timeout})
if self.exec_error is not None:
raise self.exec_error
if argv[:2] == ("sh", "-lc"):
return self._run_script(argv[2])
return _FakeResult()
def _run_script(self, script: str) -> _FakeResult:
if script == "echo ok":
return _FakeResult(stdout=b"ok\n")
if "BOOTSTRAP_OK" in script: # provider create-time bootstrap script
return _FakeResult(stdout=b"BOOTSTRAP_OK\n")
if script.startswith("find "):
root = shlex.split(script)[1]
hits = [p for p in self.files if p == root or p.startswith(f"{root.rstrip('/')}/")]
return _FakeResult(stdout=("\n".join(hits) + "\n").encode() if hits else b"")
if script.startswith("grep "):
# grep <flags> -e <pattern> <root> 2>/dev/null | head -N
tokens = shlex.split(script)
needle, root = tokens[tokens.index("-e") + 1], tokens[tokens.index("-e") + 2]
lines = []
for path, blob in sorted(self.files.items()):
if not path.startswith(root):
continue
for n, text in enumerate(blob.decode(errors="replace").splitlines(), start=1):
if needle in text:
lines.append(f"{path}:{n}:{text}")
return _FakeResult(stdout=("\n".join(lines) + "\n").encode() if lines else b"")
return _FakeResult()
def close(self):
self.closed = True
if self.close_error is not None:
raise self.close_error
class _FakeProject:
def __init__(self, id: str, name: str) -> None:
self.id = id
self.name = name
class _FakeWorkspace:
def __init__(self, id: str, name: str, projects: list[_FakeProject]) -> None:
self.id = id
self.name = name
self.projects = projects
class _FakeIdentity:
def __init__(self, workspaces: list[_FakeWorkspace]) -> None:
self.workspaces = workspaces
class _FakeClient:
def __init__(self, *, workspaces=None, sandbox_factory=None, **kwargs) -> None:
self.create_count = 0
self.create_kwargs: list[dict] = []
self._sandbox_factory = sandbox_factory or (lambda: _FakeSandbox())
self.last_sandbox: _FakeSandbox | None = None
self._by_id: dict[str, _FakeSandbox] = {}
self._workspaces = workspaces if workspaces is not None else [_FakeWorkspace("ws1", "Workspace", [_FakeProject("proj1", "Project")])]
def who_am_i(self):
return _FakeIdentity(self._workspaces)
def create(self, **kwargs):
self.create_count += 1
self.create_kwargs.append(kwargs)
sandbox = self._sandbox_factory()
sandbox.id = f"sb{self.create_count}"
self._by_id[sandbox.id] = sandbox
self.last_sandbox = sandbox
return sandbox
def get(self, sandbox_id):
return self._by_id[sandbox_id]
# ── Config stub ───────────────────────────────────────────────────────
def _stub_config(sandbox_attrs=None):
attrs = sandbox_attrs or {}
return types.SimpleNamespace(sandbox=types.SimpleNamespace(**attrs))
def _install(monkeypatch, *, client=None, config_attrs=None):
"""Construct a provider with get_app_config + _import_client stubbed."""
monkeypatch.setattr(
"deerflow.community.tenki.provider.get_app_config",
lambda: _stub_config(config_attrs),
)
if client is not None:
monkeypatch.setattr(
"deerflow.community.tenki.provider._import_client",
lambda: lambda **kw: client,
)
provider = TenkiSandboxProvider()
return provider
def _no_tenki(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setitem(sys.modules, "tenki_sandbox", None)
# ── Lazy import ────────────────────────────────────────────────────────
def test_import_client_missing_raises_actionable(monkeypatch: pytest.MonkeyPatch) -> None:
_no_tenki(monkeypatch)
with pytest.raises(ImportError, match=r"deerflow-harness\[tenki\]"):
_import_client()
def test_acquire_without_tenki_raises_and_shuts_down_cleanly(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr("deerflow.community.tenki.provider.get_app_config", lambda: _stub_config())
_no_tenki(monkeypatch)
provider = TenkiSandboxProvider()
try:
with pytest.raises(ImportError, match=r"deerflow-harness\[tenki\]"):
provider.acquire("thread-1", user_id="u")
finally:
provider.shutdown()
provider.shutdown() # idempotent
# ── Adapter: guards, env, output, terminal detection ──────────────────
def test_guard_traversal() -> None:
assert TenkiSandbox._guard_traversal("/mnt/user-data/workspace/a.txt") == "/mnt/user-data/workspace/a.txt"
assert TenkiSandbox._guard_traversal("relative/ok.txt") == "relative/ok.txt"
with pytest.raises(PermissionError):
TenkiSandbox._guard_traversal("/mnt/user-data/../etc/passwd")
with pytest.raises(ValueError):
TenkiSandbox._guard_traversal("")
def test_resolve_path_remaps_virtual_prefix_to_home() -> None:
box = TenkiSandbox("sb", _FakeSandbox(), home_dir="/home/tenki")
assert box._resolve_path("/mnt/user-data") == "/home/tenki"
assert box._resolve_path("/mnt/user-data/workspace/a.txt") == "/home/tenki/workspace/a.txt"
# Non-virtual absolute paths pass through unchanged.
assert box._resolve_path("/etc/hostname") == "/etc/hostname"
with pytest.raises(PermissionError):
box._resolve_path("/mnt/user-data/../etc/passwd")
def test_download_file_guards_reject_before_touching_sandbox() -> None:
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake)
with pytest.raises(PermissionError):
box.download_file("/etc/passwd") # outside the virtual prefix
with pytest.raises(PermissionError):
box.download_file("/mnt/user-data/../etc/passwd") # traversal
assert fake.exec_calls == [] # guards raised before any exec
def test_execute_command_rejects_invalid_env_key() -> None:
box = TenkiSandbox("sb", _FakeSandbox())
with pytest.raises(ValueError, match=r"POSIX"):
box.execute_command("echo hi", env={"BAD KEY": "x"})
def test_execute_command_formats_stdout_and_forwards_env_timeout() -> None:
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake, default_env={"BASE": "1"})
out = box.execute_command("echo ok", env={"EXTRA": "2"}, timeout=5)
assert out == "ok\n" # combined output is returned verbatim (matches boxlite/e2b)
call = fake.exec_calls[-1]
assert call["argv"] == ("sh", "-lc", "echo ok")
assert call["env"] == {"BASE": "1", "EXTRA": "2"}
assert call["timeout"] == 5
assert call["cwd"] is None # no forced cwd; runs in sandbox default dir
def test_execute_command_returns_error_as_text() -> None:
box = TenkiSandbox("sb", _FakeSandbox(exec_error=RuntimeError("boom")))
assert box.execute_command("echo hi") == "Error: boom"
def test_execute_command_closed_returns_error() -> None:
box = TenkiSandbox("sb", _FakeSandbox())
box.close()
assert box.execute_command("echo hi") == "Error: sandbox has been closed"
def test_terminal_failure_triggers_callback() -> None:
invalidated: list[tuple[str, str]] = []
box = TenkiSandbox(
"sb",
_FakeSandbox(exec_error=_FakeSessionTerminatedError("session gone")),
on_terminal_failure=lambda sid, reason: invalidated.append((sid, reason)),
)
out = box.execute_command("echo hi")
assert out == "Error: session gone"
assert invalidated == [("sb", "session gone")]
def test_regular_error_does_not_trigger_terminal_callback() -> None:
invalidated: list[tuple[str, str]] = []
box = TenkiSandbox(
"sb",
_FakeSandbox(exec_error=RuntimeError("user command failed")),
on_terminal_failure=lambda sid, reason: invalidated.append((sid, reason)),
)
box.execute_command("echo hi")
assert invalidated == []
def test_transient_transport_error_is_not_retried_and_not_evicted() -> None:
# exec is not idempotent, so a transient transport error must NOT be retried
# (re-running could double a command's side effects or duplicate a base64
# write chunk). It surfaces as text, and — not being a terminal session
# error — does not evict the sandbox.
invalidated: list[tuple[str, str]] = []
fake = _FakeSandbox(exec_error=RuntimeError("UNAVAILABLE: Socket closed"))
box = TenkiSandbox("sb", fake, on_terminal_failure=lambda sid, reason: invalidated.append((sid, reason)))
out = box.execute_command("echo ok")
assert out.startswith("Error:")
assert len(fake.exec_calls) == 1 # executed exactly once — no retry
assert invalidated == [] # transient, not terminal → sandbox not evicted
def test_connection_error_is_terminal() -> None:
invalidated: list[tuple[str, str]] = []
box = TenkiSandbox(
"sb",
_FakeSandbox(exec_error=ConnectionError("reset")),
on_terminal_failure=lambda sid, reason: invalidated.append((sid, reason)),
)
box.execute_command("echo hi")
assert invalidated == [("sb", "reset")]
def test_close_is_idempotent() -> None:
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake)
box.close()
box.close() # idempotent
assert box.is_closed is True
assert fake.closed is True
def test_close_failure_leaves_sandbox_retryable() -> None:
# Marking the adapter closed before a failed termination would strand a
# running (billed) microVM with no handle left to retry it.
fake = _FakeSandbox(close_error=RuntimeError("terminate rejected"))
box = TenkiSandbox("sb", fake)
with pytest.raises(RuntimeError, match="terminate rejected"):
box.close()
assert box.is_closed is False # still ours to terminate
fake.close_error = None
box.close()
assert box.is_closed is True
def test_close_treats_terminal_error_as_already_gone() -> None:
fake = _FakeSandbox(close_error=_FakeSessionTerminatedError("session gone"))
box = TenkiSandbox("sb", fake)
box.close() # nothing left to terminate → not an error
assert box.is_closed is True
# ── Adapter: native fs file round-trip ────────────────────────────────
def test_file_round_trip_text() -> None:
box = TenkiSandbox("sb", _FakeSandbox())
box.write_file("/mnt/user-data/workspace/note.txt", "hello world")
assert box.read_file("/mnt/user-data/workspace/note.txt") == "hello world"
def test_write_creates_parent_directory() -> None:
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake)
box.write_file("/mnt/user-data/workspace/nested/note.txt", "hi")
assert "/home/tenki/workspace/nested" in fake.dirs
def test_file_round_trip_large_binary_streams_in_frames() -> None:
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake)
data = bytes(range(256)) * 8192 # 2 MB → more than one 1 MiB upload frame
box.update_file("/mnt/user-data/outputs/blob.bin", data)
assert fake.write_frame_counts[-1] > 1
assert box.download_file("/mnt/user-data/outputs/blob.bin") == data
def test_append_accumulates() -> None:
box = TenkiSandbox("sb", _FakeSandbox())
box.write_file("/mnt/user-data/workspace/log.txt", "a")
box.write_file("/mnt/user-data/workspace/log.txt", "b", append=True)
assert box.read_file("/mnt/user-data/workspace/log.txt") == "ab"
def test_append_to_missing_file_creates_it() -> None:
box = TenkiSandbox("sb", _FakeSandbox())
box.write_file("/mnt/user-data/workspace/new.txt", "first", append=True)
assert box.read_file("/mnt/user-data/workspace/new.txt") == "first"
def test_read_missing_file_returns_error() -> None:
box = TenkiSandbox("sb", _FakeSandbox())
assert box.read_file("/mnt/user-data/workspace/nope.txt").startswith("Error:")
def test_download_missing_file_raises_oserror() -> None:
box = TenkiSandbox("sb", _FakeSandbox())
with pytest.raises(OSError):
box.download_file("/mnt/user-data/outputs/nope.bin")
def test_download_cap_counts_bytes_actually_received(monkeypatch) -> None:
# The cap must not rely on a separate size probe: a file that grows between
# the probe and the read would slip past it.
monkeypatch.setattr("deerflow.community.tenki.sandbox._MAX_DOWNLOAD_SIZE", 100)
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake)
fake.files["/home/tenki/outputs/big.bin"] = b"x" * 500
with pytest.raises(OSError) as excinfo:
box.download_file("/mnt/user-data/outputs/big.bin")
assert excinfo.value.errno == errno.EFBIG
def test_download_size_cap_does_not_evict_sandbox(monkeypatch) -> None:
# Hitting the size cap is a client-side limit, not a dead session — the
# sandbox must stay live.
monkeypatch.setattr("deerflow.community.tenki.sandbox._MAX_DOWNLOAD_SIZE", 100)
invalidated: list[tuple[str, str]] = []
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake, on_terminal_failure=lambda sid, reason: invalidated.append((sid, reason)))
fake.files["/home/tenki/outputs/big.bin"] = b"x" * 500
with pytest.raises(OSError):
box.download_file("/mnt/user-data/outputs/big.bin")
assert invalidated == []
def test_download_terminal_transport_error_evicts_sandbox() -> None:
# A ConnectionError mid-stream is an OSError subclass and terminal; the old
# `except OSError: raise` re-raised it before _note_failure, so a session
# that died mid-download never got evicted.
invalidated: list[tuple[str, str]] = []
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake, on_terminal_failure=lambda sid, reason: invalidated.append((sid, reason)))
box.write_file("/mnt/user-data/outputs/f.bin", "payload")
fake.fs_error = ConnectionError("connection reset mid-stream")
with pytest.raises(ConnectionError):
box.download_file("/mnt/user-data/outputs/f.bin")
assert invalidated == [("sb", "connection reset mid-stream")]
def test_fs_terminal_failure_evicts_sandbox() -> None:
invalidated: list[tuple[str, str]] = []
fake = _FakeSandbox(fs_error=_FakeSessionTerminatedError("session gone"))
box = TenkiSandbox("sb", fake, on_terminal_failure=lambda sid, reason: invalidated.append((sid, reason)))
with pytest.raises(Exception, match="session gone"):
box.write_file("/mnt/user-data/workspace/a.txt", "x")
assert invalidated == [("sb", "session gone")]
# ── Adapter: returned paths stay caller-facing ────────────────────────
def test_search_results_use_virtual_paths() -> None:
box = TenkiSandbox("sb", _FakeSandbox())
box.write_file("/mnt/user-data/workspace/pkg/mod.py", "def foo():\n return 42\n")
assert box.list_dir("/mnt/user-data/workspace") == ["/mnt/user-data/workspace/pkg/mod.py"]
found, truncated = box.glob("/mnt/user-data/workspace", "**/*.py")
assert found == ["/mnt/user-data/workspace/pkg/mod.py"] and not truncated
matches, _ = box.grep("/mnt/user-data/workspace", "return")
assert [m.path for m in matches] == ["/mnt/user-data/workspace/pkg/mod.py"]
assert matches[0].line_number == 2
# Results feed straight back into the file APIs.
assert box.read_file(found[0]).startswith("def foo()")
def test_grep_glob_keeps_its_directory_prefix() -> None:
box = TenkiSandbox("sb", _FakeSandbox())
box.write_file("/mnt/user-data/workspace/src/a.js", "const x = 1;\n")
box.write_file("/mnt/user-data/workspace/vendor/b.js", "const x = 2;\n")
matches, _ = box.grep("/mnt/user-data/workspace", "const x", glob="src/*.js")
assert [m.path for m in matches] == ["/mnt/user-data/workspace/src/a.js"]
def test_grep_passes_capital_h_so_single_file_matches_parse() -> None:
# Without -H, `grep -r` on a path that resolves to a single file prints
# "line:text" and the file:line:text unpack drops every match. The flag must
# always be present regardless of the other options.
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake)
box.write_file("/mnt/user-data/workspace/a.txt", "needle here\n")
box.grep("/mnt/user-data/workspace", "needle")
grep_scripts = [c["argv"][2] for c in fake.exec_calls if c["argv"][:2] == ("sh", "-lc") and c["argv"][2].startswith("grep ")]
assert grep_scripts and "-H" in shlex.split(grep_scripts[0])
def _grep_script(fake: _FakeSandbox) -> list[str]:
scripts = [c["argv"][2] for c in fake.exec_calls if c["argv"][:2] == ("sh", "-lc") and c["argv"][2].startswith("grep ")]
assert scripts, "no grep command was issued"
return shlex.split(scripts[0])
def test_grep_literal_uses_fixed_string_flag() -> None:
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake)
box.write_file("/mnt/user-data/workspace/a.txt", "a.b\n")
box.grep("/mnt/user-data/workspace", "a.b", literal=True)
tokens = _grep_script(fake)
assert "-F" in tokens and "-E" not in tokens # -F matches the pattern literally
assert "-i" in tokens # case-insensitive is still the default
def test_grep_case_sensitive_omits_ignore_case_flag() -> None:
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake)
box.write_file("/mnt/user-data/workspace/a.txt", "Needle\n")
box.grep("/mnt/user-data/workspace", "Needle", case_sensitive=True)
tokens = _grep_script(fake)
assert "-i" not in tokens # case-sensitive → no -i
assert "-E" in tokens
def test_glob_include_dirs_adds_directory_type_to_find() -> None:
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake)
box.glob("/mnt/user-data/workspace", "*", include_dirs=True)
find_scripts = [c["argv"][2] for c in fake.exec_calls if c["argv"][:2] == ("sh", "-lc") and c["argv"][2].startswith("find ")]
assert find_scripts and "-type d" in find_scripts[-1] # dirs requested, not just files
def test_list_dir_forwards_max_depth() -> None:
fake = _FakeSandbox()
box = TenkiSandbox("sb", fake)
box.list_dir("/mnt/user-data/workspace", max_depth=4)
find_scripts = [c["argv"][2] for c in fake.exec_calls if c["argv"][:2] == ("sh", "-lc") and c["argv"][2].startswith("find ")]
assert find_scripts and "-maxdepth 4" in find_scripts[-1]
def test_paths_outside_the_virtual_prefix_are_reported_as_is() -> None:
box = TenkiSandbox("sb", _FakeSandbox())
assert box._virtual_path("/etc/hostname") == "/etc/hostname"
assert box._virtual_path("/home/tenki") == "/mnt/user-data"
# ── Provider: id derivation ───────────────────────────────────────────
def test_sandbox_id_deterministic(monkeypatch):
provider = _install(monkeypatch)
assert provider._sandbox_id("t1", "u1") == provider._sandbox_id("t1", "u1")
assert len(provider._sandbox_id("t1", "u1")) == 16 # 64-bit, like community/e2b_sandbox
def test_sandbox_id_distinct_users_and_threads(monkeypatch):
provider = _install(monkeypatch)
assert provider._sandbox_id("t1", "u1") != provider._sandbox_id("t1", "u2")
assert provider._sandbox_id("t1", "u1") != provider._sandbox_id("t2", "u1")
def test_idle_timeout_zero_disables_reaper(monkeypatch):
provider = _install(monkeypatch, config_attrs={"idle_timeout": 0})
assert provider._config["idle_timeout"] == 0
assert provider._idle_checker_thread is None
provider.shutdown()
def test_load_config_rejects_invalid_environment_key(monkeypatch):
# The config `environment` is merged into every command; a bad key must
# fail fast at load time, like the per-call env in execute_command, not
# surface as a confusing SDK error at create/exec time.
monkeypatch.setattr("deerflow.community.tenki.provider.get_app_config", lambda: _stub_config({"environment": {"bad-key": "1"}}))
with pytest.raises(ValueError, match=r"POSIX"):
TenkiSandboxProvider()
def test_load_config_accepts_valid_environment_key(monkeypatch):
provider = _install(monkeypatch, config_attrs={"environment": {"PYTHONUNBUFFERED": "1"}})
assert provider._config["environment"] == {"PYTHONUNBUFFERED": "1"}
provider.shutdown()
# ── Provider: acquire / create / warm pool ────────────────────────────
def test_create_passes_prefixed_name_and_scope(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
sid = provider.acquire("thread-1", user_id="u1")
assert sid in provider._sandboxes
kwargs = client.create_kwargs[0]
assert kwargs["name"].startswith("deer-flow-tenki-")
assert kwargs["project_id"] == "proj1"
assert kwargs["workspace_id"] == "ws1"
provider.shutdown()
def test_create_waits_client_side_and_configures_lifetime(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client, config_attrs={"max_duration": 7200})
provider.acquire("thread-1", user_id="u1")
kwargs = client.create_kwargs[0]
# wait=False keeps the handle on our side of create() so a readiness failure
# is still terminable (see the next test).
assert kwargs["wait"] is False
assert client.last_sandbox.wait_ready_calls == 1
# Without an explicit lifetime, Tenki reaps the sandbox out from under a
# long-lived thread at its default (~30 min).
assert kwargs["max_duration"] == 7200
assert kwargs["sticky"] is False
provider.shutdown()
def test_create_default_lifetime_is_explicit(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
provider.acquire("thread-1", user_id="u1")
assert client.create_kwargs[0]["max_duration"] == pytest.approx(4 * 60 * 60)
provider.shutdown()
def test_create_terminates_the_microvm_when_readiness_fails(monkeypatch):
client = _FakeClient(sandbox_factory=lambda: _FakeSandbox(wait_ready_error=RuntimeError("never became ready")))
provider = _install(monkeypatch, client=client)
with pytest.raises(RuntimeError, match="never became ready"):
provider.acquire("thread-1", user_id="u1")
# The session exists remotely even though create() failed — it must not leak.
assert client.last_sandbox.closed is True
assert provider._sandboxes == {}
provider.shutdown()
def test_create_survives_bootstrap_failure(monkeypatch, caplog):
# Bootstrap is best-effort: the file APIs still work via the home remap, so a
# non-zero bootstrap must warn, not fail the acquire.
class _BootstrapFailSandbox(_FakeSandbox):
def _run_script(self, script: str) -> _FakeResult:
if "BOOTSTRAP_OK" in script:
return _FakeResult(exit_code=1, stderr=b"permission denied")
return super()._run_script(script)
client = _FakeClient(sandbox_factory=lambda: _BootstrapFailSandbox())
provider = _install(monkeypatch, client=client)
with caplog.at_level("WARNING"):
sid = provider.acquire("thread-1", user_id="u1")
box = provider.get(sid)
assert box is not None and not box.is_closed # usable despite bootstrap failure
assert any("bootstrap" in r.message.lower() for r in caplog.records)
provider.shutdown()
def test_bootstrap_is_non_interactive_and_time_bounded(monkeypatch):
# The bootstrap runs under the per-scope acquire lock, so it must not hang:
# `sudo -n` fails fast instead of blocking on a password prompt, and the exec
# carries a timeout so any other stall drops to the warning path.
client = _FakeClient()
provider = _install(monkeypatch, client=client)
provider.acquire("thread-1", user_id="u1")
bootstrap = next(c for c in client.last_sandbox.exec_calls if "BOOTSTRAP_OK" in (c["argv"][2] if len(c["argv"]) > 2 else ""))
assert "sudo -n ln -sfn" in bootstrap["argv"][2]
# Pin the actual bounded value, not merely "some timeout": a regression to
# timeout=0 (no timeout in some SDKs) would slip past an `is not None` check.
assert bootstrap["timeout"] == _BOOTSTRAP_TIMEOUT
provider.shutdown()
def test_release_parks_in_warm_pool(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
sid = provider.acquire("thread-1", user_id="u1")
provider.release(sid)
assert sid not in provider._sandboxes
assert sid in provider._warm_pool
sandbox, _ = provider._warm_pool[sid]
assert not sandbox.is_closed # microVM not terminated
provider.shutdown()
def test_acquire_reclaims_from_warm_pool(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
sid1 = provider.acquire("thread-1", user_id="u1")
provider.release(sid1)
sid2 = provider.acquire("thread-1", user_id="u1")
assert sid1 == sid2
assert client.create_count == 1 # reused, not recreated
assert sid2 in provider._sandboxes
provider.shutdown()
def test_different_threads_dont_reclaim_each_other(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
sid_a = provider.acquire("thread-a", user_id="u1")
provider.release(sid_a)
sid_b = provider.acquire("thread-b", user_id="u1")
assert sid_b != sid_a
assert sid_a in provider._warm_pool
assert sid_b in provider._sandboxes
provider.shutdown()
def test_warm_pool_reclaim_failed_health_check_creates_new(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
sid1 = provider.acquire("thread-1", user_id="u1")
provider.release(sid1)
# Kill the warm sandbox so the health check fails.
sandbox, _ = provider._warm_pool[sid1]
sandbox.close()
sid2 = provider.acquire("thread-1", user_id="u1")
assert sid2 == sid1 # same deterministic id
assert client.create_count == 2 # a fresh sandbox was created
replacement = provider.get(sid2)
assert replacement is not None and not replacement.is_closed
provider.shutdown()
def test_terminal_failure_evicts_active_sandbox(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
sid = provider.acquire("thread-1", user_id="u1")
box = provider.get(sid)
assert box is not None
# Only fail command-time, not the create-time mkdir bootstrap.
client.last_sandbox.exec_error = _FakeSessionTerminatedError("gone")
box.execute_command("echo hi") # terminal failure → invalidate
assert provider.get(sid) is None
provider.shutdown()
def test_concurrent_same_thread_acquire_creates_one_sandbox(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
original_create = provider._create_sandbox
create_started = threading.Event()
def slow_create(sandbox_id: str):
create_started.set()
time.sleep(0.1)
return original_create(sandbox_id)
provider._create_sandbox = slow_create # type: ignore[method-assign]
results: list[str] = []
def worker():
results.append(provider.acquire("thread-1", user_id="u1"))
first = threading.Thread(target=worker)
second = threading.Thread(target=worker)
first.start()
assert create_started.wait(timeout=2)
second.start()
first.join(timeout=2)
second.join(timeout=2)
assert len(results) == 2
assert results[0] == results[1]
assert client.create_count == 1
provider.shutdown()
def test_release_during_shutdown_closes_instead_of_reparking(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
sid = provider.acquire("thread-1", user_id="u1")
box = provider._sandboxes[sid]
with provider._lock:
provider._shutdown_called = True
provider.release(sid)
assert sid not in provider._sandboxes
assert sid not in provider._warm_pool
assert box.is_closed
def test_reset_parks_running_for_later_cleanup(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
sid_active = provider.acquire("thread-active", user_id="u1")
sid_warm = provider.acquire("thread-warm", user_id="u1")
provider.release(sid_warm)
active_box = provider._sandboxes[sid_active]
provider.reset()
assert provider._sandboxes == {}
assert provider._warm_pool[sid_active][0] is active_box
assert provider._thread_sandboxes == {}
assert not active_box.is_closed
provider.shutdown()
assert active_box.is_closed
def test_idle_reaper_destroys_expired_warm(monkeypatch):
client = _FakeClient()
monkeypatch.setattr(TenkiSandboxProvider, "IDLE_CHECK_INTERVAL", 0.1)
provider = _install(monkeypatch, client=client)
sid = provider.acquire("thread-1", user_id="u1")
provider.release(sid)
warm_box = provider._warm_pool[sid][0]
provider._warm_pool[sid] = (warm_box, time.time() - 9999)
time.sleep(0.3)
assert sid not in provider._warm_pool
assert warm_box.is_closed
provider.shutdown()
def test_replica_enforcement_evicts_oldest_warm(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client, config_attrs={"replicas": 2})
sid_a = provider.acquire("thread-a", user_id="u1")
provider.release(sid_a)
sid_b = provider.acquire("thread-b", user_id="u1")
provider.release(sid_b)
box_a = provider._warm_pool[sid_a][0]
provider._warm_pool[sid_a] = (box_a, time.time() - 100)
provider._warm_pool[sid_b] = (provider._warm_pool[sid_b][0], time.time())
provider.acquire("thread-c", user_id="u1")
assert sid_a not in provider._warm_pool
assert box_a.is_closed
assert sid_b in provider._warm_pool
provider.shutdown()
def test_shutdown_destroys_all_and_stops_reaper(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
sid_active = provider.acquire("thread-1", user_id="u1")
sid_warm = provider.acquire("thread-2", user_id="u1")
provider.release(sid_warm)
box_active = provider._sandboxes[sid_active]
box_warm = provider._warm_pool[sid_warm][0]
checker = provider._idle_checker_thread
provider.shutdown()
assert provider._idle_checker_stop.is_set()
assert checker is not None and not checker.is_alive()
assert box_active.is_closed and box_warm.is_closed
assert provider._sandboxes == {} and provider._warm_pool == {}
# ── Provider: scope resolution ─────────────────────────────────────────
def test_scope_auto_resolves_single(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client)
provider.acquire("thread-1", user_id="u1")
assert client.create_kwargs[0]["project_id"] == "proj1"
assert client.create_kwargs[0]["workspace_id"] == "ws1"
provider.shutdown()
def test_explicit_project_id_skips_lookup(monkeypatch):
client = _FakeClient()
provider = _install(monkeypatch, client=client, config_attrs={"project_id": "explicit"})
provider.acquire("thread-1", user_id="u1")
assert client.create_kwargs[0]["project_id"] == "explicit"
provider.shutdown()
def test_ambiguous_project_raises(monkeypatch):
client = _FakeClient(workspaces=[_FakeWorkspace("ws1", "W", [_FakeProject("p1", "A"), _FakeProject("p2", "B")])])
provider = _install(monkeypatch, client=client)
with pytest.raises(ValueError, match="project_id"):
provider.acquire("thread-1", user_id="u1")
provider.shutdown()
# ── Live integration (opt-in) ──────────────────────────────────────────
@pytest.mark.skipif(
not os.getenv("TENKI_API_KEY") and not os.getenv("TENKI_AUTH_TOKEN"),
reason="requires a real Tenki API key (TENKI_API_KEY); network integration test",
)
def test_integration_real_sandbox(monkeypatch):
monkeypatch.setattr("deerflow.community.tenki.provider.get_app_config", lambda: _stub_config())
provider = TenkiSandboxProvider()
try:
sid = provider.acquire("it-thread", user_id="it-user")
box = provider.get(sid)
assert box is not None
assert "42" in box.execute_command("python3 -c 'print(6 * 7)'")
box.write_file("/mnt/user-data/workspace/it.txt", "tenki-e2e")
assert box.read_file("/mnt/user-data/workspace/it.txt") == "tenki-e2e"
finally:
provider.shutdown()

View File

@ -16,6 +16,7 @@ from fastapi.testclient import TestClient
from app.gateway.routers import thread_runs
from deerflow.domain.feedback import Feedback
from deerflow.runtime.runs.manager import EditReplayVisibility
def _make_app(messages, feedback_grouped):
@ -33,6 +34,8 @@ def _make_app(messages, feedback_grouped):
# list_thread_messages also calls run_manager.list_by_thread to inject
# turn durations; stub it to return no runs so that path stays inert.
run_manager = MagicMock()
run_manager.list_successful_regenerate_sources = AsyncMock(return_value=set())
run_manager.list_edit_replay_visibility = AsyncMock(return_value=EditReplayVisibility())
run_manager.list_by_thread = AsyncMock(return_value=[])
app.state.run_manager = run_manager

View File

@ -16,18 +16,28 @@ from deerflow.domain.feedback import Feedback
from deerflow.runtime import RunRecord
from deerflow.runtime.events.store.memory import MemoryRunEventStore
from deerflow.runtime.journal import build_branch_history_seed_events
from deerflow.runtime.runs.manager import EditReplayVisibility
def _make_app(event_store: MemoryRunEventStore, *, superseded: set[str] | None = None, records=None, feedback=None):
def _make_app(
event_store: MemoryRunEventStore,
*,
superseded: set[str] | None = None,
edit_visibility: EditReplayVisibility | None = None,
records=None,
feedback=None,
):
app = make_authed_test_app()
app.include_router(thread_runs.router)
app.state.run_event_store = event_store
run_manager = AsyncMock()
run_manager.list_successful_regenerate_sources.return_value = superseded or set()
run_manager.list_edit_replay_visibility.return_value = edit_visibility or EditReplayVisibility()
run_manager.get_many_by_thread.return_value = records or {}
app.state.run_manager = run_manager
feedback_service = AsyncMock()
feedback_service.latest_for_runs.return_value = feedback or {}
feedback_service.latest_per_run_in_thread.return_value = feedback or {}
app.state.feedback_service = feedback_service
return app
@ -188,6 +198,76 @@ def test_thread_page_filters_all_successfully_superseded_runs_before_filling():
assert body["next_before_seq"] is None
def test_thread_page_applies_edit_replay_visibility_before_filling():
store = MemoryRunEventStore()
async def seed():
await _put_message(store, "source-success", "ai", "source-success-answer")
await _put_message(store, "edit-success", "ai", "edit-success-answer")
await _put_message(store, "source-failed", "ai", "source-failed-answer")
await _put_message(store, "edit-failed", "ai", "edit-failed-answer")
asyncio.run(seed())
app = _make_app(
store,
edit_visibility=EditReplayVisibility(
hidden_source_run_ids={"source-success"},
hidden_attempt_run_ids={"edit-failed"},
),
)
with TestClient(app) as client:
response = client.get("/api/threads/thread-1/messages/page?limit=4")
body = response.json()
assert [row["run_id"] for row in body["data"]] == ["edit-success", "source-failed"]
assert body["has_more"] is False
assert body["next_before_seq"] is None
def test_legacy_thread_messages_apply_regenerate_and_edit_replay_visibility():
store = MemoryRunEventStore()
async def seed():
await _put_message(store, "regen-source", "ai", "regen-source-answer")
await _put_message(store, "source-success", "ai", "source-success-answer")
await _put_message(store, "edit-success", "ai", "edit-success-answer")
await _put_message(store, "source-failed", "ai", "source-failed-answer")
await _put_message(store, "edit-failed", "ai", "edit-failed-answer")
asyncio.run(seed())
app = _make_app(
store,
superseded={"regen-source"},
edit_visibility=EditReplayVisibility(
hidden_source_run_ids={"source-success"},
hidden_attempt_run_ids={"edit-failed"},
),
)
with TestClient(app) as client:
response = client.get("/api/threads/thread-1/messages?limit=5")
assert [row["run_id"] for row in response.json()] == ["edit-success", "source-failed"]
def test_legacy_thread_messages_applies_limit_after_visibility_filtering(monkeypatch):
monkeypatch.setattr(thread_runs, "THREAD_MESSAGE_LEGACY_SCAN_BATCH", 2)
store = MemoryRunEventStore()
async def seed():
await _put_message(store, "visible-old-1", "ai", "visible-old-1-answer")
await _put_message(store, "visible-old-2", "ai", "visible-old-2-answer")
await _put_message(store, "hidden-new-1", "ai", "hidden-new-1-answer")
await _put_message(store, "hidden-new-2", "ai", "hidden-new-2-answer")
asyncio.run(seed())
app = _make_app(store, superseded={"hidden-new-1", "hidden-new-2"})
with TestClient(app) as client:
response = client.get("/api/threads/thread-1/messages?limit=2")
assert response.status_code == 200
assert [row["run_id"] for row in response.json()] == ["visible-old-1", "visible-old-2"]
def test_thread_page_keeps_earlier_branch_turns_when_the_last_one_is_regenerated():
"""Regenerating a branch's inherited answer must not delete the turns before
it (#4458).
@ -254,7 +334,7 @@ def test_thread_page_logs_when_scan_cursor_does_not_advance(caplog):
assert "Thread message scan cursor did not advance" in caplog.text
assert "thread_id=thread-1" in caplog.text
assert "scan_before=10" in caplog.text
assert "next_scan_before=10" in caplog.text
assert "next_cursor=10" in caplog.text
assert "row_count=1" in caplog.text

View File

@ -1,5 +1,6 @@
"""Tests for ThreadMetaRepository (SQLAlchemy-backed)."""
import asyncio
import logging
import pytest
@ -160,6 +161,20 @@ class TestThreadMetaRepository:
assert record["metadata"] == {"a": 1, THREAD_PINNED_METADATA_KEY: True}
assert record["updated_at"] == original
@pytest.mark.anyio
async def test_concurrent_metadata_updates_preserve_disjoint_keys(self, repo):
for index in range(10):
thread_id = f"concurrent-{index}"
await repo.create(thread_id, metadata={"base": index}, user_id=None)
await asyncio.gather(
repo.update_metadata(thread_id, {"left": index}, user_id=None),
repo.update_metadata(thread_id, {"right": index}, user_id=None),
)
record = await repo.get(thread_id, user_id=None)
assert record["metadata"] == {"base": index, "left": index, "right": index}
@pytest.mark.anyio
async def test_search_orders_pinned_threads_before_newer_unpinned_threads(self, repo):
await repo.create("older-pinned", metadata={THREAD_PINNED_METADATA_KEY: True})

View File

@ -7,7 +7,7 @@ from unittest.mock import AsyncMock, patch
import pytest
from fastapi import HTTPException
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
from langgraph.checkpoint.base import empty_checkpoint, uuid6
from langgraph.checkpoint.memory import InMemorySaver
@ -15,7 +15,16 @@ from deerflow.runtime import RunStatus
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
def _checkpoint(checkpoint_id: str, messages: list[object], *, metadata: dict | None = None):
def _checkpoint(
checkpoint_id: str,
messages: list[object],
*,
metadata: dict | None = None,
goal: dict | None = None,
):
channel_values = {"messages": messages}
if goal is not None:
channel_values["goal"] = goal
return SimpleNamespace(
config={
"configurable": {
@ -25,7 +34,7 @@ def _checkpoint(checkpoint_id: str, messages: list[object], *, metadata: dict |
"checkpoint_map": None,
}
},
checkpoint={"channel_values": {"messages": messages}},
checkpoint={"channel_values": channel_values},
metadata=metadata or {},
)
@ -162,6 +171,9 @@ class FakeRunManager:
async def list_by_thread(self, thread_id, *, user_id=None, limit=100):
return self.records[:limit]
async def get(self, run_id, *, user_id=None):
return next((record for record in self.records if record.run_id == run_id), None)
def _request(checkpointer, event_store, *, run_manager=None, user_id="user-1"):
from app.gateway.auth_disabled import AUTH_SOURCE_SESSION
@ -390,12 +402,46 @@ def test_prepare_regenerate_payload_returns_clean_input_and_base_checkpoint():
"regenerate_from_run_id": "run-old",
"regenerate_checkpoint_id": "ckpt-base",
}
assert "title" not in response.input
regenerated_human = response.input["messages"][0]
assert regenerated_human["id"] == "human-1"
assert regenerated_human["content"] == [{"type": "text", "text": "/data-analysis analyze data.csv"}]
assert regenerated_human["additional_kwargs"] == {"files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}]}
def test_prepare_regenerate_payload_preserves_latest_thread_title():
from app.gateway.routers import thread_runs
human = HumanMessage(id="human-1", content="question")
ai = AIMessage(id="ai-1", content="answer v1")
base = _checkpoint("ckpt-base", [])
latest = _checkpoint("ckpt-ai", [human, ai])
latest.checkpoint["channel_values"]["title"] = "User renamed title"
checkpointer = FakeCheckpointer([latest, base])
event_store = FakeEventStore(
[
{
"run_id": "run-old",
"event_type": "llm.ai.response",
"category": "message",
"content": {"id": "ai-1", "type": "ai", "content": "answer v1"},
"metadata": {"caller": "lead_agent"},
}
]
)
response = asyncio.run(
thread_runs._prepare_regenerate_payload(
"thread-1",
"ai-1",
_request(checkpointer, event_store),
)
)
assert response.checkpoint["checkpoint_id"] == "ckpt-base"
assert response.input["title"] == "User renamed title"
def test_prepare_regenerate_payload_does_not_mutate_legacy_single_checkpoint_branch():
from app.gateway.routers.thread_runs import _prepare_regenerate_payload
@ -500,6 +546,338 @@ def test_prepare_regenerate_payload_rejects_legacy_branch_when_source_checkpoint
assert len(branch_history) == 1
def test_prepare_edit_regenerate_payload_returns_new_human_and_edit_metadata():
from app.gateway.routers import thread_runs
human = HumanMessage(
id="human-1",
content="<uploaded_files>injected</uploaded_files>\n\noriginal question",
name="researcher",
additional_kwargs={
ORIGINAL_USER_CONTENT_KEY: "original question",
"files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}],
"referenced_message_contexts": [{"message_id": "ai-prev", "quote": "quoted"}],
"hide_from_ui": False,
"run_id": "old-run",
"timestamp": "2026-07-22T00:00:00Z",
"middleware_private": "do-not-copy",
},
)
ai = AIMessage(id="ai-1", content="answer v1")
base = _checkpoint("ckpt-base", [])
after_human = _checkpoint("ckpt-human", [human])
latest = _checkpoint("ckpt-ai", [human, ai])
checkpointer = FakeCheckpointer([latest, after_human, base])
event_store = FakeEventStore(
[
{
"run_id": "run-old",
"event_type": "llm.ai.response",
"category": "message",
"content": {"id": "ai-1", "type": "ai", "content": "answer v1"},
"metadata": {"caller": "lead_agent"},
}
]
)
run_manager = FakeRunManager(
[
SimpleNamespace(
run_id="run-old",
status=RunStatus.success,
metadata={},
last_ai_message="answer v1",
)
]
)
response = asyncio.run(
thread_runs._prepare_edit_regenerate_payload(
"thread-1",
"human-1",
" updated question\nwith details ",
_request(checkpointer, event_store, run_manager=run_manager),
)
)
assert response.checkpoint == {
"checkpoint_ns": "",
"checkpoint_id": "ckpt-base",
"checkpoint_map": None,
}
assert response.target_run_id == "run-old"
assert response.replacement_human_message_id != "human-1"
assert response.source_message_ids == ["human-1", "ai-1"]
assert response.metadata == {
"replay_kind": "edit",
"regenerate_from_message_id": "ai-1",
"regenerate_from_run_id": "run-old",
"regenerate_checkpoint_id": "ckpt-base",
"edit_from_message_id": "human-1",
"edit_message_id": response.replacement_human_message_id,
"edit_version_group_id": "human-1",
}
replacement = response.input["messages"][0]
assert replacement == {
"type": "human",
"id": response.replacement_human_message_id,
"name": "researcher",
"content": [{"type": "text", "text": "updated question\nwith details"}],
"additional_kwargs": {
"files": [{"filename": "data.csv", "path": "/mnt/user-data/uploads/data.csv"}],
"referenced_message_contexts": [{"message_id": "ai-prev", "quote": "quoted"}],
},
}
@pytest.mark.parametrize(
("replacement_text", "detail"),
[
(" \n\t", "Edited message cannot be empty"),
(" original question ", "Edited message is unchanged"),
],
)
def test_prepare_edit_regenerate_payload_rejects_empty_or_unchanged_text(replacement_text: str, detail: str):
from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload
human = HumanMessage(id="human-1", content="original question")
ai = AIMessage(id="ai-1", content="answer")
checkpointer = FakeCheckpointer(
[
_checkpoint("ckpt-ai", [human, ai]),
_checkpoint("ckpt-human", [human]),
_checkpoint("ckpt-base", []),
]
)
event_store = FakeEventStore(
[
{
"run_id": "run-old",
"event_type": "llm.ai.response",
"category": "message",
"content": {"id": "ai-1", "type": "ai", "content": "answer"},
"metadata": {"caller": "lead_agent"},
}
]
)
run_manager = FakeRunManager([SimpleNamespace(run_id="run-old", status=RunStatus.success, metadata={}, last_ai_message="answer")])
with pytest.raises(HTTPException) as exc:
asyncio.run(
_prepare_edit_regenerate_payload(
"thread-1",
"human-1",
replacement_text,
_request(checkpointer, event_store, run_manager=run_manager),
)
)
assert exc.value.status_code == 409
assert exc.value.detail == detail
def test_prepare_edit_regenerate_payload_requires_latest_human_turn():
from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload
old_human = HumanMessage(id="human-old", content="old question")
old_ai = AIMessage(id="ai-old", content="old answer")
latest_human = HumanMessage(id="human-latest", content="latest question")
latest_ai = AIMessage(id="ai-latest", content="latest answer")
checkpointer = FakeCheckpointer(
[
_checkpoint("ckpt-latest", [old_human, old_ai, latest_human, latest_ai]),
_checkpoint("ckpt-base", []),
]
)
with pytest.raises(HTTPException) as exc:
asyncio.run(
_prepare_edit_regenerate_payload(
"thread-1",
"human-old",
"edited old question",
_request(checkpointer, FakeEventStore([])),
)
)
assert exc.value.status_code == 409
assert exc.value.detail == "Only the latest completed user turn can be edited"
def test_prepare_edit_regenerate_payload_allows_answered_historical_clarification():
from app.gateway.routers import thread_runs
old_human = HumanMessage(id="human-old", content="ambiguous request")
clarification = ToolMessage(
id="tool-clarify",
tool_call_id="call-clarify",
content="need input",
artifact={
"human_input": {
"version": 1,
"kind": "human_input_request",
"request_id": "clarify-1",
"prompt": "Which format?",
}
},
)
clarification_answer = HumanMessage(
id="human-clarify-answer",
content="Use Markdown",
additional_kwargs={
"hide_from_ui": True,
"human_input_response": {
"version": 1,
"kind": "human_input_response",
"request_id": "clarify-1",
"source": "user",
"value": "Use Markdown",
},
},
)
latest_human = HumanMessage(id="human-latest", content="latest question")
latest_ai = AIMessage(id="ai-latest", content="latest answer")
historical_messages = [old_human, clarification, clarification_answer]
checkpointer = FakeCheckpointer(
[
_checkpoint("ckpt-ai", [*historical_messages, latest_human, latest_ai]),
_checkpoint("ckpt-human", [*historical_messages, latest_human]),
_checkpoint("ckpt-base", historical_messages),
]
)
event_store = FakeEventStore(
[
{
"run_id": "run-latest",
"event_type": "llm.ai.response",
"category": "message",
"content": {"id": "ai-latest", "type": "ai", "content": "latest answer"},
"metadata": {"caller": "lead_agent"},
}
]
)
run_manager = FakeRunManager(
[
SimpleNamespace(
run_id="run-latest",
status=RunStatus.success,
metadata={},
last_ai_message="latest answer",
)
]
)
response = asyncio.run(
thread_runs._prepare_edit_regenerate_payload(
"thread-1",
"human-latest",
"updated latest question",
_request(checkpointer, event_store, run_manager=run_manager),
)
)
assert response.checkpoint["checkpoint_id"] == "ckpt-base"
assert response.target_run_id == "run-latest"
assert response.source_message_ids == ["human-latest", "ai-latest"]
def test_prepare_edit_regenerate_payload_rejects_open_clarification_turn():
from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload
human = HumanMessage(id="human-1", content="question")
tool = ToolMessage(
id="tool-1",
tool_call_id="call-1",
content="need input",
artifact={"human_input": {"request_id": "clarify-1", "status": "pending"}},
)
checkpointer = FakeCheckpointer(
[
_checkpoint("ckpt-tool", [human, tool]),
_checkpoint("ckpt-human", [human]),
_checkpoint("ckpt-base", []),
]
)
with pytest.raises(HTTPException) as exc:
asyncio.run(
_prepare_edit_regenerate_payload(
"thread-1",
"human-1",
"updated question",
_request(checkpointer, FakeEventStore([])),
)
)
assert exc.value.status_code == 409
assert exc.value.detail == "Only completed assistant text turns can be edited"
def test_prepare_edit_regenerate_payload_rejects_active_goal():
from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload
human = HumanMessage(id="human-1", content="question")
ai = AIMessage(id="ai-1", content="answer")
checkpointer = FakeCheckpointer(
[
_checkpoint("ckpt-ai", [human, ai], goal={"status": "active", "objective": "finish"}),
_checkpoint("ckpt-human", [human]),
_checkpoint("ckpt-base", []),
]
)
with pytest.raises(HTTPException) as exc:
asyncio.run(
_prepare_edit_regenerate_payload(
"thread-1",
"human-1",
"updated question",
_request(checkpointer, FakeEventStore([])),
)
)
assert exc.value.status_code == 409
assert exc.value.detail == "Cannot edit while a goal is active"
def test_prepare_edit_regenerate_payload_requires_successful_source_run():
from app.gateway.routers.thread_runs import _prepare_edit_regenerate_payload
human = HumanMessage(id="human-1", content="question")
ai = AIMessage(id="ai-1", content="answer")
checkpointer = FakeCheckpointer(
[
_checkpoint("ckpt-ai", [human, ai]),
_checkpoint("ckpt-human", [human]),
_checkpoint("ckpt-base", []),
]
)
event_store = FakeEventStore(
[
{
"run_id": "run-old",
"event_type": "llm.ai.response",
"category": "message",
"content": {"id": "ai-1", "type": "ai", "content": "answer"},
"metadata": {"caller": "lead_agent"},
}
]
)
run_manager = FakeRunManager([SimpleNamespace(run_id="run-old", status=RunStatus.error, metadata={}, last_ai_message="answer")])
with pytest.raises(HTTPException) as exc:
asyncio.run(
_prepare_edit_regenerate_payload(
"thread-1",
"human-1",
"updated question",
_request(checkpointer, event_store, run_manager=run_manager),
)
)
assert exc.value.status_code == 409
assert exc.value.detail == "Only successful assistant runs can be edited and rerun"
def test_prepare_regenerate_uses_materialized_history_when_raw_messages_are_omitted():
from app.gateway.routers.thread_runs import _prepare_regenerate_payload

View File

@ -7,6 +7,7 @@ from _router_auth_helpers import make_authed_test_app
from fastapi.testclient import TestClient
from app.gateway.routers import thread_runs
from deerflow.runtime.runs.manager import EditReplayVisibility
def _make_app():
@ -20,6 +21,8 @@ def _make_app():
app.state.run_event_store = event_store
run_manager = MagicMock()
run_manager.list_successful_regenerate_sources = AsyncMock(return_value=set())
run_manager.list_edit_replay_visibility = AsyncMock(return_value=EditReplayVisibility())
run_manager.list_by_thread = AsyncMock(return_value=[])
app.state.run_manager = run_manager
return app
@ -61,6 +64,10 @@ def test_read_endpoints_reject_non_positive_seq_cursors(path: str, cursor: str,
def test_read_endpoints_accept_positive_limits_and_hit_store():
app = _make_app()
app.state.run_event_store.list_messages.return_value = [
{"seq": 1, "run_id": "run-1", "event_type": "llm.human.input", "content": {"type": "human", "id": "h1"}},
{"seq": 2, "run_id": "run-1", "event_type": "llm.human.input", "content": {"type": "human", "id": "h2"}},
]
with TestClient(app) as client:
thread_messages = client.get("/api/threads/thread-1/messages", params={"limit": 1})
run_messages = client.get("/api/threads/thread-1/runs/run-1/messages", params={"limit": 1})
@ -69,7 +76,8 @@ def test_read_endpoints_accept_positive_limits_and_hit_store():
assert thread_messages.status_code == 200
assert run_messages.status_code == 200
assert run_events.status_code == 200
app.state.run_event_store.list_messages.assert_awaited_once_with("thread-1", limit=1, before_seq=None, after_seq=None)
assert len(thread_messages.json()) == 1
app.state.run_event_store.list_messages.assert_awaited_once_with("thread-1", limit=thread_runs.THREAD_MESSAGE_LEGACY_SCAN_BATCH, before_seq=None, user_id=None)
app.state.run_event_store.list_messages_by_run.assert_awaited_once_with(
"thread-1",
"run-1",

View File

@ -147,6 +147,7 @@ def _mock_app_config():
model.model_dump.return_value = {"name": "test-model", "use": "langchain_openai:ChatOpenAI"}
config = MagicMock()
config.models = [model]
config.database.checkpoint_delta_snapshot_frequency = 1000
return config

Some files were not shown because too many files have changed in this diff Show More