mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-28 08:56:13 +00:00
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:
commit
fd74553f91
6
.github/workflows/backend-unit-tests.yml
vendored
6
.github/workflows/backend-unit-tests.yml
vendored
@ -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
|
||||
|
||||
18
README.md
18
README.md
@ -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.
|
||||
|
||||
|
||||
@ -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; 1–300). 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; 1–300). 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/`)
|
||||
|
||||
|
||||
276
backend/app/channels/dedupe_store.py
Normal file
276
backend/app/channels/dedupe_store.py
Normal 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()
|
||||
@ -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"
|
||||
|
||||
@ -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 --------------------------------------------------
|
||||
|
||||
|
||||
@ -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,
|
||||
)
|
||||
|
||||
|
||||
@ -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}
|
||||
|
||||
@ -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")
|
||||
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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]:
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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.
|
||||
|
||||
@ -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)
|
||||
|
||||
|
||||
@ -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,
|
||||
*,
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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
|
||||
@ -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
|
||||
|
||||
@ -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]]:
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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),
|
||||
)
|
||||
|
||||
@ -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}")
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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)
|
||||
|
||||
97
backend/packages/harness/deerflow/community/tenki/README.md
Normal file
97
backend/packages/harness/deerflow/community/tenki/README.md
Normal 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.
|
||||
@ -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",
|
||||
]
|
||||
450
backend/packages/harness/deerflow/community/tenki/provider.py
Normal file
450
backend/packages/harness/deerflow/community/tenki/provider.py
Normal 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")
|
||||
464
backend/packages/harness/deerflow/community/tenki/sandbox.py
Normal file
464
backend/packages/harness/deerflow/community/tenki/sandbox.py
Normal 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]
|
||||
@ -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``
|
||||
|
||||
@ -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."),
|
||||
|
||||
@ -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."
|
||||
),
|
||||
),
|
||||
)
|
||||
@ -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."
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
|
||||
@ -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(
|
||||
|
||||
@ -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")
|
||||
@ -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
|
||||
|
||||
@ -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",
|
||||
]
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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``.
|
||||
"""
|
||||
@ -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"),
|
||||
)
|
||||
@ -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:
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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)}
|
||||
|
||||
@ -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]:
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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,
|
||||
|
||||
@ -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"]
|
||||
|
||||
@ -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 = [
|
||||
|
||||
@ -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:
|
||||
716
backend/scripts/benchmark/checkpoint/bench_production.py
Executable file
716
backend/scripts/benchmark/checkpoint/bench_production.py
Executable 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())
|
||||
208
backend/scripts/benchmark/checkpoint/checkpoint_bench_common.py
Executable file
208
backend/scripts/benchmark/checkpoint/checkpoint_bench_common.py
Executable 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
|
||||
@ -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
|
||||
"""
|
||||
|
||||
199
backend/scripts/benchmark/checkpoint/summarize_production.py
Executable file
199
backend/scripts/benchmark/checkpoint/summarize_production.py
Executable 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())
|
||||
@ -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 \\
|
||||
@ -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",
|
||||
56
backend/tests/blocking_io/test_dingtalk_receive_file.py
Normal file
56
backend/tests/blocking_io/test_dingtalk_receive_file.py
Normal 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 == []
|
||||
@ -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.
|
||||
|
||||
@ -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 ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
355
backend/tests/test_authorization_route_permissions.py
Normal file
355
backend/tests/test_authorization_route_permissions.py
Normal 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
|
||||
@ -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)
|
||||
|
||||
237
backend/tests/test_bench_checkpoint_production.py
Normal file
237
backend/tests/test_bench_checkpoint_production.py
Normal 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
|
||||
@ -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:
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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."""
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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())
|
||||
|
||||
@ -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)
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@ -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},
|
||||
)
|
||||
|
||||
@ -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
|
||||
|
||||
252
backend/tests/test_inbound_dedupe.py
Normal file
252
backend/tests/test_inbound_dedupe.py
Normal 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)
|
||||
@ -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)
|
||||
|
||||
|
||||
289
backend/tests/test_memory_retrieval_adapter.py
Normal file
289
backend/tests/test_memory_retrieval_adapter.py
Normal 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"))
|
||||
@ -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.
|
||||
|
||||
@ -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.
|
||||
|
||||
53
backend/tests/test_migration_0009_webhook_dedupe.py
Normal file
53
backend/tests/test_migration_0009_webhook_dedupe.py
Normal 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()
|
||||
170
backend/tests/test_multi_pod_inbound_dedupe.py
Normal file
170
backend/tests/test_multi_pod_inbound_dedupe.py
Normal 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)
|
||||
@ -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"
|
||||
|
||||
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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)
|
||||
|
||||
128
backend/tests/test_summarize_checkpoint_production.py
Normal file
128
backend/tests/test_summarize_checkpoint_production.py
Normal 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]
|
||||
933
backend/tests/test_tenki_provider.py
Normal file
933
backend/tests/test_tenki_provider.py
Normal 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()
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
|
||||
@ -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})
|
||||
|
||||
@ -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
|
||||
|
||||
|
||||
@ -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",
|
||||
|
||||
@ -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
Loading…
x
Reference in New Issue
Block a user