mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-08-04 11:58:36 +00:00
Merge branch 'main' into rayhpeng/schedule-hexagonal-domain
Resolves six modify/delete conflicts by keeping the deletions: main touched the pre-hexagonal scheduler while this branch removes it. Both of those commits are carried onto the new path rather than dropped: - #4607 (once-schedule UTC normalization) was reproduced against the hexagonal domain and fixed there in its own commit -- `next_after` had the same offset bug the old `schedules.py` did. - #4589 (unified thread-id validation) is applied to the new router: the two request models and the thread-scoped list route now take `ThreadId` instead of `str`. Response models keep plain `str`, since route-addressable legacy ids stay readable by design. `test_thread_id_route_contract.py` swept routers by last path segment, which cannot import one that lives in its own package, so it collected nothing for the schedule slice; it now records the full dotted path and overrides `get_schedule_service` for the same reason it already overrides `get_config` -- dependency solving precedes path-param validation, so an unconfigured service 503s before the 422 under test. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
commit
6333da9849
@ -13,6 +13,14 @@ INFOQUEST_API_KEY=your-infoquest-api-key
|
||||
# Leave unset when using the unified nginx endpoint, e.g. http://localhost:2026.
|
||||
# GATEWAY_CORS_ORIGINS=http://localhost:3000,http://127.0.0.1:3000
|
||||
|
||||
# Host interface the Docker stack publishes its entry port on. Defaults to
|
||||
# 127.0.0.1 (loopback only), matching the local-trusted-environment deployment
|
||||
# model documented in README.md -- the agent can execute commands.
|
||||
# Set 0.0.0.0 only when the host is protected by your own TLS/auth front door
|
||||
# or firewall, and complete first-run setup before it becomes reachable.
|
||||
# BIND_HOST=0.0.0.0
|
||||
# PORT=2026
|
||||
|
||||
# Optional:
|
||||
# FIRECRAWL_API_KEY=your-firecrawl-api-key
|
||||
# VOLCENGINE_API_KEY=your-volcengine-api-key
|
||||
|
||||
9
.github/workflows/backend-unit-tests.yml
vendored
9
.github/workflows/backend-unit-tests.yml
vendored
@ -61,6 +61,15 @@ jobs:
|
||||
--health-interval 10s
|
||||
--health-timeout 5s
|
||||
--health-retries 5
|
||||
redis:
|
||||
image: redis:7-alpine
|
||||
ports:
|
||||
- 6379:6379
|
||||
options: >-
|
||||
--health-cmd "redis-cli ping"
|
||||
--health-interval 10s
|
||||
--health-timeout 5s
|
||||
--health-retries 5
|
||||
|
||||
env:
|
||||
TEST_POSTGRES_URI: postgresql://deerflow:deerflow@localhost:5432/deerflow_test?sslmode=disable
|
||||
|
||||
12
AGENTS.md
12
AGENTS.md
@ -35,6 +35,18 @@ Nginx is the single public entry: it serves the frontend and proxies `/api/langg
|
||||
to the Gateway's LangGraph runtime, rewriting it to Gateway's native `/api/*` routes; all
|
||||
other `/api/*` go straight to the Gateway REST routers. See
|
||||
[backend/AGENTS.md](backend/AGENTS.md) for the runtime and router detail.
|
||||
It compresses HTML and configured textual assets, while deliberately leaving SSE,
|
||||
fonts, images, audio, and video uncompressed at the proxy layer.
|
||||
|
||||
Both compose files publish that entry as `"${BIND_HOST:-127.0.0.1}:${PORT:-2026}:2026"`
|
||||
— **loopback by default**, matching the README's documented deployment model. A bare
|
||||
`"${PORT}:2026"` binds `0.0.0.0`, which does not.
|
||||
Nginx itself listens `default_server` on IPv4+IPv6 and the
|
||||
Gateway binds `0.0.0.0:8001` inside the container on purpose — both are container-
|
||||
internal; the published nginx port is the entire external surface, and the Gateway's
|
||||
`8001` is deliberately not published. Any new published port needs an explicit bind
|
||||
address; `backend/tests/test_compose_default_bind_host.py` pins this for every service
|
||||
in both compose files.
|
||||
|
||||
## Repository Map
|
||||
|
||||
|
||||
47
CHANGELOG.md
47
CHANGELOG.md
@ -12,6 +12,18 @@ This section accumulates work toward the **2.1.0** milestone
|
||||
|
||||
### ⚠ Breaking changes
|
||||
|
||||
- **skills:** Sandboxes now reserve `/mnt/skills` for managed enabled-only
|
||||
projections. `DEER_FLOW_HOST_SKILLS_PATH` and `SKILLS_HOST_PATH` are no longer
|
||||
used; Docker/AIO and hostPath deployments derive projection paths from
|
||||
`DEER_FLOW_HOST_BASE_DIR`. E2B operator mounts targeting `/mnt/skills` or any
|
||||
child path are skipped with a warning so they cannot shadow the managed
|
||||
projection; move extra E2B content to a different container path. User
|
||||
projections re-read global enable state from disk so toggles propagate across
|
||||
Gateway workers on the next sandbox acquire. Existing E2B sandboxes retain
|
||||
their creation-time snapshot until they are recreated. PVC-backed provisioner
|
||||
deployments still mount the operator-supplied PVC snapshot directly, so
|
||||
disabled-skill filesystem isolation does not apply in PVC mode until dynamic
|
||||
PVC materialization is implemented. ([#4178])
|
||||
- **sandbox:** E2B now enforces `sandbox.replicas` as a process-local capacity
|
||||
limit. The default `wait` policy waits for `acquire_timeout`, then fails the
|
||||
agent turn. DeerFlow does not retry the turn automatically. Use `burst` with
|
||||
@ -100,6 +112,8 @@ This section accumulates work toward the **2.1.0** milestone
|
||||
- **runtime:** Dual-mode checkpoint storage with LangGraph `DeltaChannel` cuts
|
||||
thread storage from O(N²) to near-linear for long research/coding runs.
|
||||
([#4292])
|
||||
- **runtime:** Delta-mode checkpoint history cache (memory/redis) with O(1)
|
||||
incremental composition, configured via `database.checkpoint_cache`.
|
||||
- **agent:** Config-declared lead-agent middlewares let deployments add custom
|
||||
`AgentMiddleware` classes without patching the runtime chain. ([#3964])
|
||||
- **agents:** Per-agent model and generation settings (`temperature`,
|
||||
@ -210,6 +224,17 @@ This section accumulates work toward the **2.1.0** milestone
|
||||
|
||||
### Changed
|
||||
|
||||
- **frontend performance:** Keep the public root and localized docs static;
|
||||
lazy-load closed workspace panels and editor/highlighter dependencies;
|
||||
incrementally derive streamed message state; bound streaming Markdown work;
|
||||
virtualize long message and chat lists; pause offscreen decorative effects;
|
||||
and enforce representative route JS/CSS budgets.
|
||||
- **browser:** Negotiate binary Browser Live JPEG frames, retain the legacy
|
||||
JSON/base64 protocol for older clients, coalesce presentation to the latest
|
||||
frame per refresh, and revoke replaced object URLs.
|
||||
- **artifacts:** Stream regular text artifacts with HTTP byte-range support and
|
||||
limit the initial Web UI preview to 1 MiB until the user explicitly loads the
|
||||
complete file.
|
||||
- **sandbox:** The Helm chart now defaults per-sandbox Services to `ClusterIP`
|
||||
instead of `NodePort`, so the code-execution sandbox is reachable only inside
|
||||
the cluster via Service DNS (`http://sandbox-<id>-svc.<ns>.svc.cluster.local`)
|
||||
@ -245,6 +270,26 @@ This section accumulates work toward the **2.1.0** milestone
|
||||
|
||||
### Fixed
|
||||
|
||||
- **artifacts:** Keep explicit full-file loading scoped to the source thread, so a same-path artifact in another conversation keeps its 1 MiB preview.
|
||||
- **sandbox:** `SandboxAuditMiddleware` no longer blocks ordinary command
|
||||
substitution that only captures output. The rule now judges *position* instead
|
||||
of matching any `$(`: `x=$(curl url)`, `echo $(curl url)`, an argument, and a
|
||||
`for` word list all run normally, while a substitution in command position
|
||||
(`$(curl url)`, after a `|`/`&&`/`;`, behind leading assignments or an
|
||||
`env`/`nohup`/`time` style wrapper, or as an `eval`/`source` argument) still
|
||||
blocks because it executes fetched content. An interpreter's code-string flag
|
||||
(`bash -c`, `python -c`, `perl -e`, `node -p`, `php -r`, and the `<<<`
|
||||
here-string) is treated as an execution context wherever it appears, so
|
||||
`bash -c "$(curl url)"` blocks; `source <(curl url)` and the backtick spelling
|
||||
of `eval`/`source` now block too, neither of which was detected before. An
|
||||
unquoted newline separates statements like `;`, so `echo hi` followed by a
|
||||
new line starting `$(curl url)` blocks as well, while heredoc bodies are
|
||||
consumed as data — writing a file whose content happens to start a line with
|
||||
`$(curl url)` is not a command.
|
||||
Variable expansions whose name merely starts with a risky executable
|
||||
(`$shell`, `$bashrc`, `$python_version`) and lookalike binaries
|
||||
(`shellcheck`, `shasum`) are no longer false positives.
|
||||
([#4611])
|
||||
- **mcp:** Isolate Settings > Tools enable/disable updates to one MCP server, so
|
||||
an unrelated disallowed stdio command no longer blocks every switch; allow
|
||||
disabling a disallowed target while still rejecting its re-enable, preserve
|
||||
@ -1264,6 +1309,7 @@ with **180 merged pull requests** since the first 2.0 milestone tag.
|
||||
[#4170]: https://github.com/bytedance/deer-flow/pull/4170
|
||||
[#4171]: https://github.com/bytedance/deer-flow/pull/4171
|
||||
[#4174]: https://github.com/bytedance/deer-flow/pull/4174
|
||||
[#4178]: https://github.com/bytedance/deer-flow/pull/4178
|
||||
[#4181]: https://github.com/bytedance/deer-flow/pull/4181
|
||||
[#4187]: https://github.com/bytedance/deer-flow/pull/4187
|
||||
[#4188]: https://github.com/bytedance/deer-flow/pull/4188
|
||||
@ -1359,3 +1405,4 @@ with **180 merged pull requests** since the first 2.0 milestone tag.
|
||||
[#4469]: https://github.com/bytedance/deer-flow/pull/4469
|
||||
[#4471]: https://github.com/bytedance/deer-flow/pull/4471
|
||||
[#4516]: https://github.com/bytedance/deer-flow/pull/4516
|
||||
[#4611]: https://github.com/bytedance/deer-flow/issues/4611
|
||||
|
||||
79
README.md
79
README.md
@ -264,6 +264,14 @@ The checkpoint storage settings `database.checkpoint_channel_mode` and
|
||||
both are frozen when the process first builds an agent (including through
|
||||
`DeerFlowClient`) and require a process restart to change safely.
|
||||
|
||||
The optional `database.checkpoint_cache` section (delta channel mode only)
|
||||
caches materialized checkpoint histories: `type` is `memory` (default) or
|
||||
`redis`, and `max_entries: 0` disables the cache. The `redis` backend is
|
||||
Gateway/async-only; the sync TUI/embedded path supports `memory` only. The
|
||||
cache is performance-only — results are identical with it disabled — so it is
|
||||
never frozen and workers sharing one checkpoint database may safely run
|
||||
different cache settings.
|
||||
|
||||
> [!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.
|
||||
|
||||
@ -401,6 +409,7 @@ See the [Sandbox Configuration Guide](backend/docs/CONFIGURATION.md#sandbox) to
|
||||
DeerFlow supports configurable MCP servers and skills to extend its capabilities.
|
||||
For HTTP/SSE MCP servers, OAuth token flows are supported (`client_credentials`, `refresh_token`).
|
||||
For stdio MCP servers, per-tool call timeouts can be configured with `tool_call_timeout`.
|
||||
MCP tool names are prefixed with `<server_name>_` by default to prevent collisions across servers. If a server already namespaces its own tools, set `tool_name_prefix: false` on that server in `extensions_config.json` to keep the original names. Disable the prefix only when the resulting names remain unique across all enabled servers.
|
||||
Settings > Tools updates one MCP server at a time: an invalid stdio command on one server no longer blocks toggling another, while enabling that invalid server remains protected by the command allowlist and surfaces the backend validation message in the UI.
|
||||
Targeted updates accept both DeerFlow's `type` field and the MCP-spec `transport` field for SSE/HTTP servers.
|
||||
Runtime MCP and skill updates replace `extensions_config.json` atomically, so an interrupted write cannot leave the shared configuration truncated or partially written.
|
||||
@ -718,6 +727,8 @@ An enabled skill's `allowed-tools` policy applies only after that skill is expli
|
||||
|
||||
When you install `.skill` archives through the Gateway, DeerFlow accepts standard optional frontmatter metadata such as `version`, `author`, and `compatibility` instead of rejecting otherwise valid external skills.
|
||||
|
||||
Disabling a skill also removes it from the sandbox filesystem view, so shell commands and structured file tools follow the same enabled state. Local, Docker/AIO, hostPath provisioner, and newly created E2B sandboxes source `/mnt/skills` from enabled-only projections that update when public, custom, legacy, or managed integration skills are toggled, edited, created, deleted, or installed. Managed integration packages remain shared, while their projected filesystem visibility follows each user's enabled state. Multi-worker Gateways re-read on-disk enable state while rebuilding user projections, so a toggle handled by one worker is honored by another worker's next sandbox acquire. Existing E2B sandboxes retain their creation-time snapshot until they are recreated. PVC-backed provisioner skills keep their configured PVC snapshot/layout for now; dynamic PVC materialization is tracked separately.
|
||||
|
||||
Managed integrations install shared read-only skill packs without mixing them
|
||||
into custom skills. The Lark/Feishu CLI integration is available under
|
||||
`Settings → Integrations → Lark / Feishu CLI`; an administrator installs or
|
||||
@ -889,6 +900,8 @@ The Web UI shows the active goal above the composer. The same command is availab
|
||||
|
||||
Use `/compact` in the Web UI composer to summarize older context for the current thread. DeerFlow keeps the full chat visible, but future model calls use the compacted summary plus recent messages. The command is ignored when there is not enough history to compact, and it is blocked while the thread has a run in flight, including when that run is owned by another Gateway worker. If a multi-worker reservation loses its lease, DeerFlow cancels the checkpoint writer before the replacing run proceeds and returns a retryable conflict after cleanup. Thread-title edits are serialized through the same state-write boundary and show a conflict without closing the rename dialog when a run is active.
|
||||
|
||||
The chat header also shows a context-window gauge when the selected model has a positive `context_window` configured. It estimates the latest materialized checkpoint's message tokens and keeps the previous same-thread percentage visible while data refetches, independently of the cumulative token-usage setting.
|
||||
|
||||
### Sub-Agents
|
||||
|
||||
Sub-agents are an optimization, not the default response to a complex request.
|
||||
@ -907,7 +920,17 @@ Use `burst` with `burst_limit` to permit bounded extra VMs. The `wait` and
|
||||
`reject` policies use only `replicas`. The `reject` policy can remove one warm
|
||||
VM before it returns an error.
|
||||
|
||||
`replicas` limits one Gateway process. It does not limit all Gateway processes.
|
||||
With in-memory ownership, `replicas` limits one Gateway process. With Redis
|
||||
ownership, E2B shares one capacity Hash between workers using the same
|
||||
`sandbox.ownership.key_prefix`; `replicas` (plus a configured burst) is then a
|
||||
deployment-wide hard limit. Use one unique prefix and the same effective limit
|
||||
per deployment. To change the limit, stop its Gateways, delete the capacity
|
||||
Hash, and restart; mismatched workers fail closed.
|
||||
|
||||
The Hash counts remote VMs and in-flight creates, repairs interrupted creates
|
||||
from E2B metadata, grace-protects stale inventory omissions, and blocks new
|
||||
creates while Redis or initial inventory is unavailable. Run Redis with persistence, non-evicting memory, and HA.
|
||||
|
||||
E2B acquisition uses a bounded executor. Waiting acquisitions do not use the
|
||||
default asyncio executor.
|
||||
|
||||
@ -929,7 +952,13 @@ Image bytes loaded for a vision-model call are transient: DeerFlow removes the h
|
||||
|
||||
After each run, DeerFlow records a workspace change summary for the run-owned `workspace` and `outputs` directories. The Web UI shows a compact "files changed" badge on the assistant turn; opening it reveals created, modified, and deleted files with text diffs when safe to display. Uploads are excluded because they are user inputs, not agent-generated changes. Large, binary, or sensitive-looking files are shown as metadata only.
|
||||
|
||||
Files presented through `present_files` remain part of the thread's artifact state, and the Web UI restores the artifact panel and selected document after a page refresh. The currently selected formal artifact is refreshed once when the run finishes so edits become visible without a manual reload.
|
||||
Files presented through `present_files` remain part of the thread's artifact state, and the Web UI restores the artifact panel and selected document after a page refresh. The currently selected formal artifact is refreshed once when the run finishes so edits become visible without a manual reload. Existing UTF-8 text artifacts under `/mnt/user-data/outputs` can also be edited and explicitly saved from the panel on Unix and Windows while the thread is idle; saves use content revisions to prevent overwriting agent changes.
|
||||
|
||||
Text artifacts are streamed with HTTP byte-range support. The Web UI initially
|
||||
loads at most 1 MiB, shows the preview size when a file is larger, and waits for
|
||||
an explicit **Load full file** action before fetching the remainder or mounting
|
||||
the full code editor. Active HTML, XHTML, and SVG artifacts remain forced
|
||||
downloads at the Gateway boundary.
|
||||
|
||||
With `AioSandboxProvider`, shell execution runs inside isolated containers. With `LocalSandboxProvider`, file tools still map to per-thread directories on the host, but host `bash` is disabled by default because it is not a secure isolation boundary. Re-enable host bash only for fully trusted local workflows. Host bash commands have a wall-clock timeout, and long-lived processes should be started in the background with output redirected to a workspace log.
|
||||
|
||||
@ -965,6 +994,11 @@ uv run playwright install chromium
|
||||
|
||||
Then uncomment the `group: browser` tool entries in `config.yaml` (`browser_navigate`, `browser_snapshot`, `browser_click`, `browser_type`, `browser_get_text`, `browser_back`, `browser_screenshot`, `browser_close`). `make dev` / Docker startup detects an enabled `browser_navigate` tool and preserves the `browser` extra on dependency syncs. The Gateway fails startup if browser control is configured but Playwright is missing, and `/api/features` hides the Browser UI unless the backend can actually serve it. Keep `headless: true` and `allow_private_addresses: false` for anything but local, trusted debugging. Attaching to an existing Chrome with `cdp_url` cannot enforce DeerFlow's subresource/redirect SSRF guard and therefore fails closed unless `allow_unguarded_cdp: true` explicitly acknowledges that risk; use it only with a trusted local browser. Browser sessions are process-local; keep `GATEWAY_WORKERS=1` while this tool group is enabled because ordinary uvicorn worker dispatch does not provide thread affinity.
|
||||
|
||||
The workspace Browser Live client negotiates binary JPEG WebSocket frames,
|
||||
keeps only the newest pending frame per display refresh, and revokes replaced
|
||||
object URLs. Gateway control messages remain JSON, and clients that do not
|
||||
request the binary capability retain the legacy JSON/base64 frame protocol.
|
||||
|
||||
### Context Engineering
|
||||
|
||||
**Isolated Sub-Agent Context**: Each sub-agent runs in its own isolated context. This means that the sub-agent will not be able to see the context of the main agent or other sub-agents. This is important to ensure that the sub-agent is able to focus on the task at hand and not be distracted by the context of the main agent or other sub-agents.
|
||||
@ -1000,6 +1034,8 @@ requires an explicit local-development opt-in. See the
|
||||
|
||||
Memory updates now skip duplicate fact entries at apply time, so repeated preferences and context do not accumulate endlessly across sessions.
|
||||
|
||||
In the default DeerMem `middleware` mode, automatic extraction now classifies every proposed fact by scope, durability, and authority before a deterministic write gate accepts it. Only durable, descriptive user-level facts are stored; current-thread or project constraints and one-time action permissions stay in conversation state. User-global summaries require both user scope and descriptive authority, contradiction removals are scope-gated, and a replacement-dependent removal is applied only when its replacement actually survives validation and storage. These classification labels are extraction-only metadata, add no extra LLM call, and are not written into the fact files. The explicit CRUD tools in `memory.mode: tool` remain a separate, model-directed path. Deployments that override the bundled DeerMem prompts via `memory.backend_config.prompts_dir` must add the new classification fields to their custom templates (the `memory_update` fact/summary/removal formats and the `consolidation` consolidated-fact schema): the write gate fails closed, so an un-migrated template stops every extraction-driven fact, summary, and removal write, surfacing only through the `rejected_by_scope_gate` metrics and the high-rejection-rate warning.
|
||||
|
||||
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.
|
||||
|
||||
Memory injection follows the configured operation mode. In `middleware` mode, DeerMem injects the user-global summaries and the selected agent's facts. In `tool` mode, the automatic `<memory>` block contains only the global `user` and `history` summaries; agent facts are retrieved explicitly through `memory_search`, avoiding duplicate automatic and tool-returned fact context. Setting `memory.injection_enabled: false` still disables the entire block in either mode.
|
||||
@ -1032,6 +1068,18 @@ DeerFlow is model-agnostic — it works with any LLM that implements the OpenAI-
|
||||
|
||||
DeerFlow can be used as an embedded Python library without running the full HTTP services. The `DeerFlowClient` provides direct in-process access to all agent and Gateway capabilities, returning the same response schemas as the HTTP Gateway API. The HTTP Gateway also exposes `DELETE /api/threads/{thread_id}` to remove DeerFlow-managed local thread data after the LangGraph thread itself has been deleted:
|
||||
|
||||
Thread IDs may be supplied by callers and do not have to be UUIDs. Explicit
|
||||
IDs must contain 1–64 ASCII letters, digits, hyphens, or underscores
|
||||
(`^[A-Za-z0-9_-]{1,64}$`). DeerFlow generates a UUID only when `thread_id` is
|
||||
omitted or `None`; an explicitly supplied empty string is invalid.
|
||||
Existing route-addressable threads created under older, looser rules remain
|
||||
readable and deletable, but cannot start new runs or create new filesystem or
|
||||
sandbox state. Legacy deletion skips local path cleanup when the ID is not
|
||||
safe under the canonical contract. For canonical legacy threads whose
|
||||
conversation exists only in LangGraph checkpoints, DeerFlow seeds an empty
|
||||
run-event feed from the checkpoint before the first new run so
|
||||
`/messages/page` keeps both the old and new turns.
|
||||
|
||||
```python
|
||||
from deerflow.client import DeerFlowClient
|
||||
|
||||
@ -1122,6 +1170,27 @@ DeerFlow has key high-privilege capabilities including **system command executio
|
||||
- **Unauthorized illegal invocation**: Agent functionality could be discovered by unauthorized third parties or malicious internet scanners, triggering bulk unauthorized requests that execute high-risk operations such as system commands and file read/write, potentially causing serious security consequences.
|
||||
- **Compliance and legal risks**: If the agent is illegally invoked to conduct cyberattacks, data theft, or other illegal activities, it may result in legal liability and compliance risks.
|
||||
|
||||
### Gateway Admin Is Equivalent to Code Execution
|
||||
|
||||
An admin can register stdio MCP servers, which run commands inside the Gateway
|
||||
container. The API restricts them to an allowlist (`npx`, `uvx` by default,
|
||||
extended via `DEER_FLOW_MCP_STDIO_COMMAND_ALLOWLIST`) and rejects arguments and
|
||||
environment variables that would evaluate arbitrary code. That is defense in
|
||||
depth, not a boundary: these launchers exist to fetch and run remote packages,
|
||||
so **treat Gateway admin as equivalent to code execution on the host** and grant
|
||||
it accordingly.
|
||||
|
||||
### Deployment Defaults
|
||||
|
||||
The Docker stack publishes its entry port on `127.0.0.1` only, matching the
|
||||
local-trusted-environment model described above. To reach it from another
|
||||
machine, set `BIND_HOST` in `.env` (e.g. `BIND_HOST=0.0.0.0`) — and only after
|
||||
putting the security measures below in place.
|
||||
|
||||
**Complete first-run setup before the host becomes reachable.** A fresh
|
||||
instance has no accounts yet, so create the admin account through `/setup`
|
||||
immediately after starting any deployment that is not loopback-only.
|
||||
|
||||
### Security Recommendations
|
||||
|
||||
**Note: We strongly recommend deploying DeerFlow in a local trusted network environment.** If you need cross-device or cross-network deployment, you must implement strict security measures, such as:
|
||||
@ -1151,6 +1220,12 @@ The JSON includes compact review records with `priority`, `location`,
|
||||
`blocking_call`, `event_loop_exposure`, `reason`, and `code`.
|
||||
Gateway artifact serving now forces active web content types (`text/html`, `application/xhtml+xml`, `image/svg+xml`) to download as attachments instead of inline rendering, reducing XSS risk for generated artifacts.
|
||||
|
||||
Frontend route asset budgets can be checked with `cd frontend && pnpm
|
||||
perf:check`. The command measures `/login` from a normal production build, then
|
||||
performs a production static-demo build for the fixture-backed workspace routes.
|
||||
It measures the unique JavaScript and CSS referenced by representative routes
|
||||
and writes the detailed result to `.next/performance-results.json`.
|
||||
|
||||
## License
|
||||
|
||||
This project is open source and available under the [MIT License](./LICENSE).
|
||||
|
||||
@ -651,6 +651,8 @@ DeerFlow 不只是“会说它能做”,它是真的有一台自己的“电
|
||||
|
||||
跨 session 使用时,DeerFlow 会逐步积累关于你的持久 memory,包括你的个人偏好、知识背景,以及长期沉淀下来的工作习惯。你用得越多,它越了解你的写作风格、技术栈和重复出现的工作流。memory 保存在本地,控制权也始终在你手里。
|
||||
|
||||
默认 DeerMem `middleware` 模式会先判断候选信息的作用域、持久性和授权属性,再由确定性写入门决定是否保存。只有稳定、描述性的用户级事实能进入长期 memory;当前对话或项目的约束、一次性操作授权仍留在对话状态中。用户全局 summary 必须同时具有用户级作用域和描述性授权属性,基于矛盾的删除也会经过作用域保护;如果删除依赖一条替代事实,只有替代事实真正通过校验并保留下来后才执行删除。这些分类字段只用于本次抽取,不写入 fact 文件,也不增加 LLM 调用次数。`memory.mode: tool` 的显式 CRUD 仍是独立的模型直写路径。如果通过 `memory.backend_config.prompts_dir` 覆盖了内置抽取模板,必须同步在自定义模板中加入新的分类字段(`memory_update` 的 fact/summary/removal 格式与 `consolidation` 的合并 fact 结构):写入门是 fail closed 的,未迁移的旧模板会导致所有抽取驱动的 fact、summary 与删除写入停止,只能通过 `rejected_by_scope_gate` 指标和高拒绝率告警发现。
|
||||
|
||||
## 推荐模型
|
||||
|
||||
DeerFlow 对模型没有强绑定,只要实现了 OpenAI 兼容 API 的 LLM,理论上都可以接入。不过在下面这些能力上表现更强的模型,通常会更适合 DeerFlow:
|
||||
|
||||
@ -267,6 +267,11 @@ Blocking-IO runtime gate (`tests/blocking_io/`):
|
||||
`test_uploads_router.py` (locks Gateway upload/list/delete endpoints
|
||||
offloading upload directory creation, staged writes, chmod/cleanup,
|
||||
directory scans/deletes, and remote sandbox sync off the event loop);
|
||||
`test_feishu_receive_file.py` (locks Feishu attachment path preparation and
|
||||
persistence plus remote sandbox acquisition/sync off the event loop, and
|
||||
skips redundant sandbox sync when thread data is already mounted);
|
||||
`test_channel_outbound_files.py` (locks Feishu, Telegram, and WeCom outbound
|
||||
attachment open/read/hash work off the event loop);
|
||||
`test_openviking_memory_backend.py` (locks the OpenViking backend's async
|
||||
add/context/search entrypoints offloading synchronous HTTP and watermark
|
||||
filesystem IO); and
|
||||
@ -309,6 +314,13 @@ Agentic browser sessions are process-local. The Gateway startup safety gate reje
|
||||
uvicorn worker dispatch does not provide thread affinity for browser tools, REST
|
||||
navigation, and the Live WebSocket.
|
||||
|
||||
Browser Live screenshots remain JPEG bytes inside the harness and the Gateway's
|
||||
bounded, drop-oldest frame queue. WebSocket clients that request
|
||||
`frame_format=binary` receive binary messages; control metadata remains JSON.
|
||||
The legacy no-parameter protocol still base64-encodes frames into JSON at the
|
||||
Gateway boundary for backward compatibility. Unknown `frame_format` values
|
||||
receive a JSON error and close code 1008.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Harness / App Split
|
||||
@ -390,7 +402,7 @@ Lead-agent middlewares are assembled in strict order across three functions: the
|
||||
7. **DanglingToolCallMiddleware** - Injects placeholder ToolMessages for AIMessage tool_calls that lack responses (e.g., user interruption), preserving raw provider tool-call payloads in `additional_kwargs["tool_calls"]`; malformed tool-call names and arguments are sanitized in the model-bound request so strict OpenAI-compatible providers do not reject the next request
|
||||
8. **LLMErrorHandlingMiddleware** - Normalizes provider/model invocation failures into recoverable assistant-facing errors before later stages run
|
||||
9. **Authorization / GuardrailMiddleware** - Up to two independent pre-tool-call gates run here. When `authorization.enabled`, the `AuthorizationProvider` instance already used for Layer 1 capability filtering is wrapped by `GuardrailAuthorizationAdapter` and reused for Layer 2 execution checks. A generated `tool_search` bypasses the adapter's second provider call only when the current build has a concrete deferred setup; its catalog was already filtered by Layer 1, and an ordinary same-named tool without that deferred setup receives no exemption. When `guardrails.enabled`, the explicitly configured `GuardrailProvider` is appended after authorization and still evaluates every call, including `tool_search`. Authorization therefore runs outermost and can deny before an external guardrail call; both use the existing middleware's fail-closed, audit, sync/async, and error-`ToolMessage` behavior. See the authorization RFC and [docs/GUARDRAILS.md](docs/GUARDRAILS.md).
|
||||
10. **SandboxAuditMiddleware** - Audits sandboxed shell/file operations for security logging before tool execution
|
||||
10. **SandboxAuditMiddleware** - Audits sandboxed shell/file operations for security logging before tool execution. Command classification is **defense-in-depth and audit, not a security boundary** — the sandbox itself is the isolation boundary. Command substitution is judged by *position*, not by the presence of `$(`: a substitution in **command position** (`$(curl url)`, `` `curl url` ``, the word after a `|`/`&&`/`;`, or any `eval`/`source` argument) executes fetched or interpreted content and is blocked, while **value position** (`x=$(curl url)`, `echo $(curl url)`, an argument, a `for` word list) only captures output and passes (#4611). `_HIGH_RISK_COMMAND_POSITION_PATTERNS` is therefore matched anchored against each split sub-command, never against the whole compound string, and `_split_compound_command(split_pipes=True)` supplies those sub-commands; rules that span a pipe (`| sh`, `base64 -d | ...`) still rely on `_classify_command`'s whole-command Pass 1. `_COMMAND_POSITION_PREFIX` extends the anchor over leading variable assignments and exec wrappers (`FOO=1 $(curl url)`, `env`/`command`/`builtin`/`exec`/`nohup`/`time`/`sudo`/`doas`), which are still command position; its assignment branch requires whitespace before the substitution, which is exactly what keeps `x=$(curl url)` in value position. Two execution contexts are deliberately **position-blind** and matched against the whole command in Pass 1, because they execute what they receive wherever they appear (including as an argument to something else, e.g. `xargs sh -c "$(curl url)"`): an `eval`/`source` argument, and an interpreter's **code-string flag** — `-c` (shells, `python`), `-e` (`perl`/`ruby`/`node`), `-p` (`perl`/`node`), `-r` (`php`) — plus the here-string (`<<<`) that reaches the same place through stdin. All three substitution spellings (`$(cmd`, `<(cmd`, `` `cmd ``) share one `_RISKY_SUBSTITUTION` opener so a rule cannot cover one spelling and miss another. An unquoted newline splits like `;`, because it separates statements the same way: leaving it joined let `echo hi\n$(curl url)` evade the anchored rules that its `;` spelling triggers. A heredoc body is data rather than statements, so `_split_compound_command` records headers (`<<EOF`, `<<-EOF`, `<<'EOF'`) and consumes their bodies verbatim at the newline that starts them — otherwise a body line beginning with `$(curl url)` would be promoted to a command position the shell never creates. Two things that look like headers must not open one, or a body that never terminates swallows every following statement: `<<<` is a here-string (both a lookahead and a lookbehind are needed, or the trailing `<<` of `<<< "text"` reads as a heredoc with delimiter `text`), and a `<<` inside `$(( ... ))` / `(( ... ))` is a bit shift, so arithmetic depth is tracked alongside the quote flags. That is a heuristic, not shell parsing: it exists only to avoid manufacturing command positions *and* to avoid destroying real ones. An unterminated body consumes the rest of the string; an unclosed `((` only disables heredoc detection, so newlines keep splitting and the failure direction stays towards seeing more command positions rather than fewer. Known, deliberate gaps: process substitution outside `eval`/`source` (`. <(curl u)`) is not detected — closing it would require real shell parsing, which is out of scope for this layer. Two-step forms (`x=$(curl u); eval "$x"`) are inherent rather than incidental: any rule that allows output capture allows the first statement, and connecting it to the later `eval` needs dataflow analysis, not pattern matching. There is currently no config gate: the middleware is appended unconditionally in `_build_runtime_middlewares`, so it applies to both the lead agent and subagents.
|
||||
11. **ReadBeforeWriteMiddleware** - *(optional, if `read_before_write.enabled`, default on)* Outermost write gate (issue #3857): `read_file` stamps a content hash onto its ToolMessage; `write_file` (append/overwrite-existing) and `str_replace` are blocked unless the newest mark for that path matches the file's current hash. Sits outside ToolProgressMiddleware and ToolErrorHandlingMiddleware so a blocked write returns immediately without consuming a ToolProgress slot. Blocked results call `normalize_tool_result` directly to stamp `deerflow_tool_meta` (`recoverable_by_model=True`) before returning, keeping the result well-formed for any outer consumer. Marks live on messages, so summarization dropping the read result invalidates the gate automatically; writes never refresh marks, forcing a re-read between consecutive edits. Gate check + tool execution are serialized per (thread, path) so same-turn parallel writes cannot reuse one stale mark; on sandboxes whose `read_file` reports failures as `"Error: ..."` strings instead of raising (AIO/E2B), uninspectable targets fail open (creation proceeds, no mark stamped)
|
||||
12. **ToolProgressMiddleware** - *(optional, if `tool_progress.enabled`)* State-machine-based stagnation guard (RFC #3177). Outer wrapper around ToolErrorHandlingMiddleware so its `wrap_tool_call` receives results already stamped with `deerflow_tool_meta`. Tracks per-(thread, tool) consecutive "no-new-info" calls across three error categories: (a) `recoverable_by_model=True` (no_results, not_found, permission, Jaccard-duplicate success): ACTIVE → WARNED (terminal — hint re-injected on each subsequent problem); (b) `recoverable_by_model=False, action≠stop` (rate_limited, transient): ACTIVE → WARNED → BLOCKED after `warn_escalation_count` more problems; (c) `recoverable_by_model=False, action=stop` (auth, config, internal): immediately BLOCKED on first occurrence. **Division of labor with LoopDetectionMiddleware:** ToolProgressMiddleware is a result-quality guard — fires after tool execution and blocks specific tools that stop producing new information; LoopDetectionMiddleware is a call-pattern guard — fires after the model responds and hard-stops the whole turn when the model repeatedly issues identical tool_calls. Both can inject HumanMessage hints in the same model call without conflict; neither reads the other's internal state.
|
||||
13. **ToolErrorHandlingMiddleware** - Receives `AppConfig`, converts tool exceptions into error `ToolMessage`s so the run can continue instead of aborting, stamps every result with `deerflow_tool_meta` (status / error_type / recoverable_by_model / recommended_next_action / source) via `tool_result_meta.normalize_tool_result`, stamps structured metadata for task exception wrappers, and stamps skill-read metadata for downstream durable-context capture. Task tool result text is generated from the same status/result/error inputs as the structured metadata so callers do not hand-write a second protocol string.
|
||||
@ -495,15 +507,28 @@ Localhost persistence deliberately reads the direct request `Host` and ignores `
|
||||
| **Memory** (`/api/memory`) | `GET /` - memory data; `POST /reload` - force reload; `GET /config` - config; `GET /status` - config + data |
|
||||
| **Uploads** (`/api/threads/{id}/uploads`) | `POST /` - upload files (auto-converts PDF/PPT/Excel/Word); `GET /list` - list; `DELETE /{filename}` - delete |
|
||||
| **Threads** (`/api/threads/{id}`) | `DELETE /` - remove DeerFlow-managed local thread data after LangGraph thread deletion; `POST /branches` - create a new main-thread branch from a completed assistant turn checkpoint and, when an addressable pre-user replay checkpoint exists, materialize it into the branch namespace so the inherited response remains regeneratable. Workspace files are not checkpointed, so the branch only best-effort copies the current workspace when branching from the **latest** turn (`workspace_clone_mode="current_thread_best_effort"`); branching from an older/historical turn skips the copy (`workspace_clone_mode="skipped_historical_turn"`) so the branch never inherits files that only exist in a later timeline. Thread-scoped runtime channels (`sandbox`, `thread_data`) are not copied onto the branch: the parent's `sandbox_id` binds path mappings and the release lifecycle to the parent's workspace, so the branch lazily acquires its own sandbox instead. Branch creation also seeds the new thread's run-event feed from the branch checkpoint's visible messages (`history_seed_mode` in the response): the thread feed reads run_events, not checkpoints, so without the seed the inherited history disappears from the UI after the branch's first run (#4380). Seeded rows are grouped into one synthetic run per inherited turn (`branch-seed-{thread_id}-{n}`, a new turn opening at every persisted human message, including an allowlisted hidden `ask_clarification` reply) because `run_id` is a turn identity to the feed's consumers, not a provenance tag: regenerating an inherited answer supersedes that row's whole `run_id` in `GET /messages/page`, so one shared id for the entire seed deleted the complete inherited history on a branch's first regenerate (#4458); `GET /goal`, `PUT /goal`, `DELETE /goal` - read, set, and clear the active thread goal; `POST /compact` - manually summarize older active context into `summary_text` and retain the recent message window, blocked while a run is in flight; unexpected failures are logged server-side and return a generic 500 detail |
|
||||
| **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 |
|
||||
| **Artifacts** (`/api/threads/{id}/artifacts`) | `GET /{path}` - stream regular text and binary artifacts with `FileResponse`, including byte-`Range` 206/416 behavior used by bounded text previews and media seeking; 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. `PUT /{path}` atomically replaces an existing UTF-8 text file under `/mnt/user-data/outputs` when its expected SHA-256 still matches; active runs conflict, and non-mounted sandbox providers receive the same update explicitly. Atomic replacement applies the existing POSIX permission handling when descriptor-based APIs are available and otherwise keeps the platform-native temporary-file permissions (Windows). |
|
||||
| **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 completed or interrupted 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; it carries the current thread title the same way, but only when the replay base already has one — an untitled base belongs to a thread the title middleware has not named yet, so pinning the current title there would keep a name generated from the prompt the edit just replaced; `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 |
|
||||
| **Thread Runs** (`/api/threads/{id}/runs`) | `POST /` - create background run; `POST /stream` - create + SSE stream; `POST /wait` - create + block. Before the first journaled run, an empty run-event message feed is seeded from an existing checkpoint head so legacy checkpoint-only history receives earlier thread-global sequence numbers and remains visible after the new run; a thread with no checkpoint or an already-populated feed skips this compatibility path. `POST /regenerate/prepare` - prepare clean input + checkpoint metadata for regenerating the latest completed or interrupted 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; it carries the current thread title the same way, but only when the replay base already has one — an untitled base belongs to a thread the title middleware has not named yet, so pinning the current title there would keep a name generated from the prompt the edit just replaced; `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 plus an optional `context_usage` percentage. Context usage approximately counts messages from the latest materialized thread state through `build_thread_checkpoint_state_accessor`, so full and delta checkpoint modes expose the same input. The percentage uses the latest run's model and its configured `context_window`. |
|
||||
| **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. |
|
||||
| **GitHub Event-Driven Agents** | Custom agents can declare a `github:` block in their `config.yaml` to bind to repos and event triggers. Webhook fan-out publishes one `InboundMessage` per matching binding to the channel bus; `GitHubChannel` routes those messages through `ChannelManager`. The response `dispatch` summarizes matched/fired/skipped agents. |
|
||||
|
||||
Thread identifiers use the shared `deerflow.utils.thread_id` contract
|
||||
`^[A-Za-z0-9_-]{1,64}$`. Caller-provided opaque IDs remain supported; UUIDs
|
||||
are generated only for `None`, while explicit empty strings fail validation.
|
||||
Gateway creation and state-producing request boundaries, embedded-client
|
||||
entry points, filesystem/upload/event-store consumers, scheduled launches,
|
||||
and the standalone Provisioner enforce the same contract before persistence
|
||||
or workspace initialization. Route-addressable legacy IDs remain accepted by
|
||||
pure reads and cleanup/control endpoints. Deleting a noncanonical legacy ID
|
||||
best-effort removes its metadata and checkpoints but deliberately skips local
|
||||
filesystem cleanup, so the raw value is never interpolated into a host path;
|
||||
new runs, workspace/sandbox operations, and other state-producing mutations
|
||||
remain blocked.
|
||||
|
||||
**Workspace change review**: `packages/harness/deerflow/workspace_changes/`
|
||||
captures a pre-run and post-run snapshot of the thread-owned `workspace` and
|
||||
`outputs` directories. `runtime/runs/worker.py` performs the filesystem scan via
|
||||
@ -570,7 +595,7 @@ JSONL event stores when `GATEWAY_WORKERS > 1`.
|
||||
- Store-only hydrated runs are readable history. In multi-worker mode with heartbeat enabled, cancel on a store-only run records `runs.cancel_action` / `cancel_requested_at` while the owner's lease is live; the first action wins even if a retry later lands on the owner. `RunStore.request_cancel()` and owner completion through `finalize_if_not_cancelled()` are competing active-row CAS operations, so an accepted cancel cannot be overwritten by a later success. `RunStore.renew_lease()` renews and observes the request atomically in the SQL implementation. The owner then executes the normal process-local interrupt/rollback and terminal stream path without transferring the lease. An expired owner is still taken over and marked `error`. `wait=true` and cancel-then-stream use the shared bridge to observe owner finalization; a non-standard process-local bridge returns accepted 202 instead of subscribing to an unreachable stream. In single-worker mode (heartbeat off), store-only runs still return 409.
|
||||
- A local worker's `RunRecord.lease_expires_at` is the last durably confirmed ownership deadline. `_renew_leases()` bounds each renewal attempt by that deadline: transient store exceptions remain retryable while it is valid, but an exception or blocked call that reaches expiry sets the process-local `ownership_lost` fence, raises `abort_event`, and cancels the run task. Successful renewals collect durable cancellation actions; after all local renewals have been attempted, heartbeat only signals the corresponding process-local tasks, leaving status writes and rollback cleanup to the worker finalization path. Fenced workers do not perform subsequent journal/delivery-receipt, progress/completion/status, checkpoint/thread-metadata, or `on_run_completed` writes; the peer recovery path owns the terminal receipt. `RunStore.update_run_completion()` also refuses to replace a different terminal status, closing the peer-takeover/late-finalization race. `grace_seconds` delays peer reclamation for clock skew but is not extra execution time for an owner that can no longer confirm its lease. Already-committed remote tool side effects remain outside this local cancellation boundary.
|
||||
- Startup/orphan reconciliation must claim stale active rows with `RunStore.claim_for_takeover()`, not a plain `update_status()`. The final claim re-checks `status` and lease expiry atomically, so a heartbeat renewal between the candidate scan and the recovery write keeps the run active.
|
||||
- Run admission and independent checkpoint writes are first-class thread operations. `runs.operation_kind` distinguishes user-visible `run` rows from internal `checkpoint_write` reservations, while every active kind shares the existing durable active-thread uniqueness constraint. New operation kinds must go through `RunStore.create_thread_operation_atomic()` and `RunManager.reserve_thread_operation()` rather than adding another lock or metadata marker. Live and lease-less reservations are non-interruptible; an expired leased reservation can be reclaimed immediately by interrupt/rollback admission without waiting for orphan reconciliation. Lease-less rows stay fail-closed because the store cannot distinguish a stale row from a live checkpoint writer in another heartbeat-disabled worker; a rare failed delete therefore requires startup reconciliation, and heartbeat-disabled multi-worker deployment remains unsupported. Reservation bodies are attached to their caller task so loss detected by lease renewal cancels the checkpoint writer before it can continue after takeover; the context manager translates that lease-loss cancellation to `ConflictError` after cleanup so Gateway mutation routes return a retryable 409 instead of dropping the HTTP request. The cleanup scope begins immediately after durable admission, including the await that attaches the caller task, so cancellation cannot strand a locally renewed pending reservation. A failed renewal is revalidated under the manager lock before cancellation; if the reservation completed and unregistered while the store update was in flight, its request task must not be cancelled after the checkpoint write. Reservations are excluded from run history/reporting and from run-only helpers such as `list_by_thread()` and `has_inflight()`, release uses the captured owner rather than ambient user context, and local cleanup still runs when the best-effort store delete fails. `RunStore.create_run_atomic()` remains a deprecated compatibility shim for external stores that only admit normal runs; new stores must implement `create_thread_operation_atomic()` to support internal operation kinds.
|
||||
- Run admission and independent writes are first-class thread operations. `runs.operation_kind` distinguishes user-visible `run` rows from internal `checkpoint_write` and `artifact_write` reservations, while every active kind shares the existing durable active-thread uniqueness constraint. New operation kinds must go through `RunStore.create_thread_operation_atomic()` and `RunManager.reserve_thread_operation()` rather than adding another lock or metadata marker. Live and lease-less reservations are non-interruptible; an expired leased reservation can be reclaimed immediately by interrupt/rollback admission without waiting for orphan reconciliation. Lease-less rows stay fail-closed because the store cannot distinguish a stale row from a live writer in another heartbeat-disabled worker; a rare failed delete therefore requires startup reconciliation, and heartbeat-disabled multi-worker deployment remains unsupported. Reservation bodies are attached to their caller task so loss detected by lease renewal cancels the writer before it can continue after takeover; the context manager translates that lease-loss cancellation to `ConflictError` after cleanup so Gateway mutation routes return a retryable 409 instead of dropping the HTTP request. The cleanup scope begins immediately after durable admission, including the await that attaches the caller task, so cancellation cannot strand a locally renewed pending reservation. A failed renewal is revalidated under the manager lock before cancellation; if the reservation completed and unregistered while the store update was in flight, its request task must not be cancelled after the write. Reservations are excluded from run history/reporting and from run-only helpers such as `list_by_thread()` and `has_inflight()`, release uses the captured owner rather than ambient user context, and local cleanup still runs when the best-effort store delete fails. `RunStore.create_run_atomic()` remains a deprecated compatibility shim for external stores that only admit normal runs; new stores must implement `create_thread_operation_atomic()` to support internal operation kinds.
|
||||
- Gateway checkpoint mutations outside run execution must use `services.reserve_checkpoint_write()`, which composes the process-local thread lock with the durable `checkpoint_write` reservation. Manual compaction, `POST /threads/{id}/state`, and both goal mutation routes (`PUT` / `DELETE /threads/{id}/goal`, including creation of a missing goal checkpoint) use this boundary, so an existing run blocks the write and the reservation blocks new reject/interrupt/rollback runs across workers.
|
||||
- `POST /wait` (both thread-scoped and `/api/runs/wait`) drains the stream bridge via `wait_for_run_completion()` instead of bare `await record.task`, so it honours the run's `on_disconnect` setting and cancels the background run on real client disconnect rather than returning a stale checkpoint (issue #3265).
|
||||
- Memory and Redis `StreamBridge` implementations retain only `stream_bridge.queue_maxsize` data events. A syntactically valid `Last-Event-ID` older than the retained watermark, or a live subscriber that falls behind it, yields `StreamGap` before any partial replay. `sse_consumer` maps that control item to an id-less SSE `gap` payload (`stream_replay_gap`) and intentionally leaves the run active; internal `/wait` consumers resume from its latest retained ID because they only need terminal completion. Redis checks bounds plus the non-blocking read in one transaction, using blocking `XREAD` only as a wake-up before repeating the atomic snapshot. For a no-cursor subscriber that established a wait on an empty stream, the first wake response remains provisional until that next snapshot verifies its tail is still retained; this closes the pre-first-delivery trimming window without changing malformed-cursor live tailing. The correctness tradeoff is one three-command snapshot pipeline per poll plus the blocking wake round trip while idle. Malformed cursor behavior remains backend-specific. Memory treats a syntactically numeric cursor below its watermark conservatively as a gap even when the evicted timestamp can no longer be verified; unknown ids at or above the watermark retain the legacy replay-from-earliest policy.
|
||||
@ -611,16 +636,24 @@ that cannot tell sibling branches apart.
|
||||
**Provider Pattern**: `SandboxProvider` with `acquire`, `acquire_async`, `get`, `release` lifecycle. Async agent/tool paths call async sandbox lifecycle hooks so Docker sandbox creation, discovery, cross-process locking, readiness polling, and release stay off the event loop.
|
||||
**Environment policy** (`sandbox/env_policy.py`): `execute_command` no longer inherits the full `os.environ`. `build_sandbox_env()` scrubs secret-looking names (`*KEY*`/`*SECRET*`/`*TOKEN*`/`*PASS*`/`*CREDENTIAL*`) from the inherited environment before layering injected request secrets on top, so platform credentials (e.g. `OPENAI_API_KEY`) never leak into skill subprocesses. Benign vars (`PATH`, `HOME`, `LANG`, `VIRTUAL_ENV`, ...) are preserved.
|
||||
**Implementations**:
|
||||
- `LocalSandboxProvider` - Local filesystem execution. `acquire(thread_id)` returns a per-thread `LocalSandbox` (id `local:{thread_id}`) whose `path_mappings` resolve `/mnt/user-data/{workspace,uploads,outputs}` and `/mnt/acp-workspace` to that thread's host directories, so the public `Sandbox` API honours the `/mnt/user-data` contract uniformly with AIO. `acquire()` / `acquire(None)` keeps the legacy generic singleton (id `local`) for callers without a thread context. Per-thread sandboxes are held in an LRU cache (default 256 entries) guarded by a `threading.Lock`. Legacy global-custom mounts are gated by the same user-scoped skill discovery rule used for prompt/list visibility; providers must not infer visibility from raw directory presence alone.
|
||||
- `AioSandboxProvider` (`packages/harness/deerflow/community/`) - Docker-based isolation. Active-cache and warm-pool entries are checked with the backend during acquire/reuse; definitively dead containers are dropped from all in-process maps so the thread can discover or create a fresh sandbox instead of reusing a stale client. Backend health-check failures are treated as unknown, not dead; local discovery likewise treats an unverifiable container as not adoptable and falls through to create rather than failing acquire. `get()` remains an in-memory lookup for event-loop-safe tool paths — it never touches the ownership store (that would be blocking IO on the event loop); ownership is published on acquire/reclaim and refreshed off the event loop by the dedicated renewal thread (`_renew_owned_leases`). `uses_thread_data_mounts` defaults to backend detection (`LocalContainerBackend=True`, remote/provisioner backends=False), while the optional `sandbox.thread_data_mounts` boolean takes precedence for deployments that guarantee the Gateway and sandbox share the same thread user-data directories. Setting it `true` skips upload-time sandbox acquire/sync; a false positive leaves uploads unavailable to the sandbox. Legacy global-custom mounts follow the same shared visibility helper as local and remote providers. Readiness probes and `agent_sandbox` clients classify loopback/private IPs, single-label cluster hosts, and Docker/Podman internal hostnames as direct control-plane destinations and set `trust_env=False`; external FQDNs and public IPs retain environment proxy support.
|
||||
- `LocalSandboxProvider` - Local filesystem execution. `acquire(thread_id)` returns a per-thread `LocalSandbox` (id `local:{thread_id}`) whose `path_mappings` resolve `/mnt/user-data/{workspace,uploads,outputs}` and `/mnt/acp-workspace` to that thread's host directories, so the public `Sandbox` API honours the `/mnt/user-data` contract uniformly with AIO. `acquire()` / `acquire(None)` keeps the legacy generic singleton (id `local`) for callers without a thread context. Per-thread sandboxes are held in an LRU cache (default 256 entries) guarded by a `threading.Lock`. Public, custom, legacy, and managed integration skill mappings point at stable enabled-only projection roots rather than raw skill directories.
|
||||
- `AioSandboxProvider` (`packages/harness/deerflow/community/`) - Docker-based isolation. Active-cache and warm-pool entries are checked with the backend during acquire/reuse; definitively dead containers are dropped from all in-process maps so the thread can discover or create a fresh sandbox instead of reusing a stale client. Backend health-check failures are treated as unknown, not dead; local discovery likewise treats an unverifiable container as not adoptable and falls through to create rather than failing acquire. `get()` remains an in-memory lookup for event-loop-safe tool paths — it never touches the ownership store (that would be blocking IO on the event loop); ownership is published on acquire/reclaim and refreshed off the event loop by the dedicated renewal thread (`_renew_owned_leases`). `uses_thread_data_mounts` defaults to backend detection (`LocalContainerBackend=True`, remote/provisioner backends=False), while the optional `sandbox.thread_data_mounts` boolean takes precedence for deployments that guarantee the Gateway and sandbox share the same thread user-data directories. Setting it `true` skips upload-time sandbox acquire/sync; a false positive leaves uploads unavailable to the sandbox. Local-container and hostPath-provisioner mounts use the same stable skill projection roots; PVC-backed skills remain governed by the operator-supplied PVC layout until PVC materialization is implemented. Readiness probes and `agent_sandbox` clients classify loopback/private IPs, single-label cluster hosts, and Docker/Podman internal hostnames as direct control-plane destinations and set `trust_env=False`; external FQDNs and public IPs retain environment proxy support.
|
||||
- `E2BSandboxProvider` (`packages/harness/deerflow/community/e2b_sandbox/`) provides E2B remote isolation.
|
||||
New sandboxes receive a one-shot upload from the enabled-only public, custom,
|
||||
legacy, and managed integration projections. Existing E2B VMs keep their
|
||||
creation-time snapshot because E2B has no shared host mount.
|
||||
Acquire and release share a per-user and thread lock. The provider lock does
|
||||
not cover remote IO. `burst_limit` adds capacity only for the `burst` policy.
|
||||
The `wait` policy fails the turn after `acquire_timeout`. The runtime does not
|
||||
retry the turn automatically. E2B acquisition uses a bounded executor.
|
||||
Waiting calls do not consume the default asyncio executor. The `reject`
|
||||
policy can evict one warm VM before it returns an error. `replicas` limits
|
||||
one Gateway process. It does not provide multi-process capacity control.
|
||||
policy can evict one warm VM before it returns an error. With memory
|
||||
ownership, `replicas` limits one Gateway process. Redis ownership shares one
|
||||
`<ownership.key_prefix>:e2b-capacity` Hash, making the limit (plus a bounded
|
||||
burst) deployment-wide. Lua atomically manages VM and in-flight-create
|
||||
entries; missing or unavailable state fails closed. E2B reservation metadata
|
||||
repairs interrupted creates. Inventory replacement is revision-CAS guarded,
|
||||
and incomplete inventories never remove entries; complete omissions get a grace period.
|
||||
Uncertain cleanup keeps a tombstone slot. Shutdown tracks owned remote
|
||||
operation IDs. Discovery can find a VM from another Gateway. Shutdown closes
|
||||
an unowned discovery client without destroying its VM. Release ends its
|
||||
@ -659,7 +692,7 @@ that cannot tell sibling branches apart.
|
||||
- **Teardown join budget covers refresh plus release.** Redis bounds each ownership operation at five seconds, and context exit can catch the heartbeat in one final refresh before its `finally` performs the final release. `_TEARDOWN_JOIN_TIMEOUT_SECONDS` is therefore 12 seconds — greater than both sequential operation bounds — so a normal pair of socket timeouts does not emit the deferred-release warning; a still-running heartbeat continues to own the release safely.
|
||||
- **An absent lease means the same thing on both paths, and reconciliation must say so too.** The `LAPSED` rule above only covers an owner renewing its *own* lease; on its own it does not make state loss safe, because reconciliation reads the same absent key as "orphan, adopt". After a Redis flush (restart without persistence, or eviction under `maxmemory`) every owner is alive and merely pre-renewal-tick, so whichever instance reconciles first would adopt every live container, each real owner's next renewal would report `LOST`, and it would drop a sandbox mid-turn for the adopter to idle-destroy — #4206 through the back door. `_adoptable_after_grace` closes it: an untracked container must be seen unowned (`owner()`, a read-only peek — the atomic `claim()` is still what actually gates adoption) across a full lease TTL before it can be adopted, tracked per container in `_unowned_since`. That rebuilds the delay the flush erased — a live owner republishes within one renewal interval, shorter than the TTL by construction (`ttl_multiplier >= 2`) — while a genuinely crashed owner never republishes, so its containers are still adopted one grace later rather than leaking. A republished lease **resets** the grace; a pausing-only timer would still expire over a live owner's lease. The grace is skipped when `supports_cross_process` is `False`: no peer can hold a lease such a store would show us, so single-instance deployments keep instant orphan cleanup, and a grace could not help a multi-worker gateway on `memory` anyway (peers are invisible to each other's leases with or without it).
|
||||
- **The `memory` store is single-instance only** and says so via `supports_cross_process = False`; the provider logs a warning at startup when the configured store cannot see peers. A multi-worker gateway on `memory` has no cross-process coordination at all — same contract as `stream_bridge`'s memory backend. This is why the redis inference matters: it reads `app_config.stream_bridge` **and** the env var, in the same order the bridge's own resolver does, so any deployment already pointing the bridge at Redis (i.e. every multi-instance one) gets a redis ownership store without extra config.
|
||||
- `get()` stays a pure in-memory lookup and must never call the store (that is blocking filesystem/network IO on the event loop); anchored by `tests/blocking_io/test_aio_sandbox_get.py`, which injects a deliberately-blocking probe store so the anchor keeps its teeth regardless of the configured backend. Tests: `tests/test_sandbox_ownership_store.py` (store contract, defined once for **both** backends — but the redis tier is `@pytest.mark.integration` + opt-in via `DEER_FLOW_TEST_REDIS_URL` and self-skips, and **CI provisions no redis**, so the merge gate runs the memory tier only and the Lua scripts never execute there; drift between the backends is caught only when the suite runs against a live redis. There is no fake-redis tier because the redis exclusion lives in Lua a fake would not execute) and `tests/test_sandbox_orphan_reconciliation.py` (provider behaviour, two providers sharing one store).
|
||||
- `get()` stays a pure in-memory lookup and must never call the store (that is blocking filesystem/network IO on the event loop); anchored by `tests/blocking_io/test_aio_sandbox_get.py`, which injects a deliberately-blocking probe store so the anchor keeps its teeth regardless of the configured backend. Tests: `tests/test_sandbox_ownership_store.py` (store contract, defined once for **both** backends — the redis tier is `@pytest.mark.integration`, uses `DEER_FLOW_TEST_REDIS_URL` when set, and otherwise self-skips without a reachable Redis. Backend CI provisions Redis, so the merge gate executes the real Lua tier; there is no fake-redis tier because a fake would not execute the Lua exclusions) and `tests/test_sandbox_orphan_reconciliation.py` (provider behaviour, two providers sharing one store).
|
||||
- `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()`.
|
||||
@ -670,7 +703,7 @@ that cannot tell sibling branches apart.
|
||||
|
||||
**Virtual Path System**:
|
||||
- Agent sees: `/mnt/user-data/{workspace,uploads,outputs}`, `/mnt/skills`
|
||||
- Physical: `backend/.deer-flow/users/{user_id}/threads/{thread_id}/user-data/...`, `deer-flow/skills/`
|
||||
- Physical: `backend/.deer-flow/users/{user_id}/threads/{thread_id}/user-data/...`; raw skills stay under `deer-flow/skills/` and managed integration storage, while sandboxes read `backend/.deer-flow/skills_view/public/` and `backend/.deer-flow/users/{user_id}/skills_view/{custom,legacy,integrations}/`
|
||||
- Translation: `LocalSandboxProvider` builds per-thread `PathMapping`s for the user-data prefixes at acquire time; `tools.py` keeps `replace_virtual_path()` / `replace_virtual_paths_in_command()` as a defense-in-depth layer (and for path validation). AIO has the directories volume-mounted at the same virtual paths inside its container, so both implementations accept `/mnt/user-data/...` natively.
|
||||
- Detection: `is_local_sandbox()` accepts both `sandbox_id == "local"` (legacy / no-thread) and `sandbox_id.startswith("local:")` (per-thread)
|
||||
|
||||
@ -744,6 +777,7 @@ E2B output sync records remote file versions and actual host file metadata in a
|
||||
- **Lazy initialization**: Tools loaded on first use via `get_cached_mcp_tools()`
|
||||
- **Cache invalidation**: Detects extensions-config changes by comparing the resolved config path and a `(mtime, size, sha256)` content signature against the values recorded at initialization, not a strict mtime `>` comparison. This catches same-second edits, mtime that stays put or moves backward (`git checkout`, `cp -p` / backup restore, `tar` / `rsync`, object-store / network mounts), and a switch to a different config file with an equal-or-older mtime. The signature helper (`config/file_signature.py::get_config_signature`) is shared with `config/app_config.py::get_app_config()` for the sibling runtime-editable config file, rather than each maintaining its own copy. `ExtensionsConfig.resolve_config_path()` raises `FileNotFoundError` for an explicit `config_path`/`DEER_FLOW_EXTENSIONS_CONFIG_PATH` that points at a missing file — an operator-asserted path going missing is a real misconfiguration, so this is intentionally loud for callers that load the config for actual use (e.g. `from_file()` via `get_mcp_tools()`); only the fallback search mode returns `None`. The MCP cache's own path resolution (`mcp/cache.py::_resolve_config_path`) is narrower: it catches that specific `FileNotFoundError` locally and treats it the same as "unconfigured", so this staleness check degrades to "not stale" instead of propagating an exception when a previously-valid explicit/env-var config disappears mid-run
|
||||
- **Transports**: stdio (command-based), SSE, HTTP
|
||||
- **Per-server tool-name prefixing**: `mcpServers.<server>.tool_name_prefix` defaults to `true`, preserving the collision-safe `<server_name>_` prefix. Servers whose tools already carry a stable namespace may set it to `false`; discovery then calls `langchain_mcp_adapters.tools.load_mcp_tools` with that server's flag. Source routing and stdio session-pool wrapping are based on the producing server and transport, never on whether the visible tool name starts with the server prefix.
|
||||
- **OAuth (HTTP/SSE)**: Supports token endpoint flows (`client_credentials`, `refresh_token`) with automatic token refresh + Authorization header injection
|
||||
- **Routing hints**: `extensions_config.json -> mcpServers.<server>.routing` and
|
||||
`tools.<original_tool_name>.routing` are soft preference metadata. The effective
|
||||
@ -756,6 +790,13 @@ E2B output sync records remote file versions and actual host file metadata in a
|
||||
- **Stdio file outputs**: Persistent stdio sessions are scoped by `user_id:thread_id`. For stdio transports only, DeerFlow pins the subprocess default `cwd` to the thread workspace and `TMPDIR`/`TMP`/`TEMP` to `workspace/.mcp/tmp/`, unless the operator explicitly configured `cwd` or temp env values. SSE/HTTP transports skip this filesystem prep entirely.
|
||||
- **Stdio path translation**: MCP-returned local file references are not copied. If a `ResourceLink` or conservative free-text path resolves to an existing file inside the thread's mounted user-data tree, it is translated deterministically to `/mnt/user-data/...`; paths outside that tree remain unchanged.
|
||||
- **Runtime updates**: Gateway API saves to extensions_config.json; the Gateway-embedded runtime detects changes via the resolved-path + content-signature check above, so multi-worker / stale-mtime deployments still pick up an added/removed MCP server without a restart (`PUT /api/mcp/config` keeps whole-payload validation, while `PATCH /api/mcp/config` changes only one server's `enabled` field, normalizes the same `type`/MCP-spec `transport` alias as the runtime config model, and validates the target only when enabling it; either endpoint's reset clears the cache only in its own worker). MCP, skill, and embedded-client writers share `atomic_write_extensions_config()`, which writes and fsyncs a same-directory temporary file before `os.replace()` and preserves an existing file's mode and symlink target; failed serialization or replacement leaves the prior config intact and cleans up the temporary file.
|
||||
- **Stdio launch policy at the HTTP boundary** (`routers/mcp.py::_validate_mcp_update_request`, shared by `PUT` and the enable branch of `PATCH`): a config file may express anything, but the API is untrusted input, so an API-registered stdio server must (a) name a bare executable from the allowlist — `_DEFAULT_MCP_STDIO_COMMAND_ALLOWLIST` = `{npx, uvx}`, extended by `DEER_FLOW_MCP_STDIO_COMMAND_ALLOWLIST`, with path separators, whitespace, and shell metacharacters rejected in `command`; (b) carry no `args` flag in `_ARBITRARY_EXEC_ARGS`; and (c) set no `env` name in `_CODE_INJECTING_ENV_VARS`. Checks (b) and (c) exist because the command check alone names a binary without constraining what that binary runs. The `env` denylist applies to **every** allowlisted command, and both denylists match `--flag=value` as well as `--flag value`. The `args` denylist's **scope depends on the command**, because where a launcher stops parsing its own flags is what decides whether a token is an exec flag at all:
|
||||
|
||||
- For a **package launcher** in `_PACKAGE_LAUNCHERS` (`{npx, uvx}`) only the launcher's own **option region** is screened. `npx`/`uvx` stop parsing their flags at the package name and hand every later token to the spawned server's argv, where `-c` is routinely "config" and `-e` "env" — screening those rejected ordinary third-party servers while covering nothing. A bare `--` ends the region too: only the *first* token after it is the package name. Finding that boundary needs each launcher's option **arity**, since a value is not a positional — `npx -p <pkg> -c '<command>'` **runs** the command (`-p` is `npm exec`'s `--package`, so `<pkg>` is its value and npm keeps parsing), so ending the region at the first non-flag token would walk straight past it. `_NPX_BOOLEAN_ARGS` is generated from `@npmcli/config`'s definitions (npm 10.9.4) minus the `-p` exec override; `_UVX_VALUE_ARGS` comes from `uvx --help` (uv 0.11.1). Regenerate these against a newer launcher rather than hand-editing. The unknown-option default is deliberately **opposite** per launcher, following the exec set rather than symmetry: npx owns real exec flags (`-c`/`--call`), so an unknown option consumes a value and keeps the region open (npm errors on options it does not define, so this cannot reject a working invocation); uvx owns no string-eval flag at all, so its screen is a tripwire, an unknown option consumes nothing, and uv's large boolean surface cannot over-block. uvx's exec set also drops the short spellings, because `-c` is uv's `--constraints` and `-p` its `--python`.
|
||||
- Every **other** command is screened whole, with two extra rules, because it is an interpreter rather than a package runner: `-p` counts as an exec flag there (node's `--print`), and single-dash short-option clusters are decomposed letter by letter so `node -pe` cannot pass a check that only splits on `=`.
|
||||
|
||||
Verdicts are pinned against the real launchers: for npx, every argument vector the validator rejects is one `npx` actually executes, and every vector it allows is one `npx` passes through to the server. `env` screening covers names that execute code **unconditionally** at process startup, e.g. `PYTHONPATH`/`PYTHONHOME`, which run a caller-controlled `sitecustomize.py` at interpreter startup under plain `uvx`. Caller-controlled **search paths** are a weaker, conditional class and are an accepted residual: `LD_LIBRARY_PATH`/`DYLD_LIBRARY_PATH` (conditional on the process loading a shadowable library, and legitimately set by native-dependency servers) and `NODE_PATH` (searched *after* the local `node_modules` chain, so it cannot shadow an installed dependency, and ignored entirely by ESM `import` — it can only supply a CJS module that would otherwise fail to resolve). Do not move a search path into the set: it would make the "unconditional" rule untrue, which is how a defense-in-depth list starts being mistaken for a boundary. Remote transports skip all three — they spawn nothing.
|
||||
**This is defense in depth, not a trust boundary.** `npx`/`uvx` exist to fetch and execute remote packages, so an admin can still point one at a package they published; the boundary is admin authentication plus network reachability. Do not add a check here on the assumption that it makes MCP registration safe for untrusted admins — it does not, and the fix for that is not a bigger denylist.
|
||||
|
||||
### Skills System (`packages/harness/deerflow/skills/`)
|
||||
|
||||
@ -764,6 +805,7 @@ E2B output sync records remote file versions and actual host file metadata in a
|
||||
- **Loading**: `load_skills()` recursively scans public, per-user custom, global integration, and legacy custom locations for `SKILL.md`, parses metadata, and reads enabled state from extensions_config.json plus per-user skill state for non-public categories; that directory is a package boundary, so no nested `SKILL.md` is registered as a runtime skill. SkillScan has a deliberately narrower packaging rule: known eval fixtures are permitted as support data, while other nested `SKILL.md` files are reported as package defects. It parses runtime metadata and reads enabled state from extensions_config.json.
|
||||
- **External reload**: `POST /api/skills/reload` is an admin-only, process-local invalidation hook for trusted MinIO/NFS/CSI writes. `SkillStorage` instances do not cache a catalog — `load_skills()` scans on every call — so the route clears all `(app_config, user_id)` entries and the rendered prompt-section LRU, then waits up to the shared refresh timeout for the existing off-loop single-flight refresh. Each invalidation receives a generation-bound result handle; a successful scan atomically replaces the global enabled-skills cache, while a loader-level failure propagates to the HTTP waiter and preserves the last-known-good global cache. Per-user/config scans capture the refresh version and cannot repopulate shared caches if invalidation occurs while they are loading. A timed-out HTTP wait fails generically while the daemon refresh worker continues. Subsequent runs rescan after a successful reload; active runs keep their existing snapshot. Each Uvicorn worker/Kubernetes Pod must be targeted separately. Direct mount writes bypass install/edit validation, SkillScan, and history, so mounted roots are an operator-controlled trust boundary.
|
||||
- **Tool policy**: Agent `allowed-tools` declarations apply dynamically only to slash-activated skills and skills captured in `ThreadState.skill_context` through configured `read_file` loads; passive enabled skills and custom-agent/subagent skill allowlists remain discoverable without clamping the baseline toolset. Subagents render only skill discovery metadata at startup and reuse the same adjacent `SkillActivationMiddleware` + `SkillToolPolicyMiddleware` pair as the lead; their configured `skills` field limits discovery and activation instead of eagerly loading bodies or unioning policies. Slash policy is dominant for its run, preventing subsequently read skills from widening explicit authority; autonomous captured skills use the existing union only when no slash source exists. `tool_search` and `describe_skill` stay available as framework discovery infrastructure, while every discovered or promoted business tool still requires active-policy permission for schema visibility and execution; `task` likewise requires an explicit declaration. Each active model call intentionally reloads the full live registry so enable/disable changes, frontmatter edits, and custom/public name-shadow winners take effect without a stale TTL or unsafe direct-path cache; all tool calls produced by that model step reuse the resulting source-and-path-signed decision. Registry failures and all-invalid active sets fail closed, while stale individual paths are skipped when another valid skill remains. This is best-effort behavioral scoping, not a hard security boundary: alternate loading paths are not captured and bounded autonomous context may evict entries.
|
||||
- **Sandbox projection**: `skills/projection.py` materializes enabled-only trees at `{base_dir}/skills_view/public` and `{base_dir}/users/{user_id}/skills_view/{custom,legacy,integrations}`. It hardlinks files when possible and falls back to copies across filesystems. Storage writes, archive installs, deletes, and toggles rebuild under a cross-process lock; Gateway boot ensures only the shared public view, while each user view is repaired lazily on first sandbox acquire. Managed integration packages are global, but their projected category is per-user because enabled state is isolated. Rebuilds stage a complete tree and reconcile it with per-file atomic replacement, so unrelated enabled skills remain continuously visible; disable/delete paths remove only the affected package before mutating to preserve fail-closed behavior. User projection rebuilds re-read global enable state from disk instead of the process singleton, so a toggle handled by another Gateway worker is reflected on the next acquire. Gateway public-skill toggles take the public projection lock before the shared `extensions_config_write_lock`, re-read an existing config from disk, persist the full model shape, and rebuild before responding; keep this as one worker-owned critical section so MCP writes cannot interleave and request cancellation cannot release either lock while the worker still runs. The shared public steady-state signature check runs without the global projection lock; stale/error paths take the lock and re-check before rebuilding or clearing. User-scope checks remain serialized per user. Category root inodes remain stable so live bind mounts observe content changes without sandbox recreation. Projection failures clear the affected view before raising.
|
||||
- **Injection (legacy / default)**: Enabled skills are listed in the agent system prompt with full metadata and container paths (`<available_skills>` block). Controlled by `skills.deferred_discovery: false` (default).
|
||||
- **Deferred discovery** (`skills.deferred_discovery: true`): Skills are listed by name only in a compact `<skill_index>` block, keeping the system prompt prefix-cache friendly. The agent calls the `describe_skill` tool at runtime to fetch full metadata for skills it wants to use, then loads the SKILL.md via `read_file`. Two new modules support this path:
|
||||
- `skills/catalog.py` — `SkillCatalog` (immutable, searchable; query forms: `select:a,b`, `+prefix`, free-text regex); `select:` returns all requested skills without a result cap; other modes cap at `MAX_RESULTS=5`.
|
||||
@ -929,6 +971,7 @@ The cached value is reused for both the blocking (`runs.wait`) and streaming (`_
|
||||
runtime.
|
||||
- `memory.mode: tool` skips `MemoryMiddleware` and registers `memory_search`, `memory_add`, `memory_update`, and `memory_delete` on the agent. The model decides when to search, add, update, or delete facts; this is opt-in/experimental and should not be described as better than middleware mode without eval evidence.
|
||||
- Both modes share `FileMemoryStorage`, per-user/per-agent isolation, manual CRUD primitives, and the updater backend. Injection is mode-aware: middleware mode injects global `user`/`history` summaries plus the selected agent's facts, while tool mode injects only the global summaries and leaves every agent fact behind `memory_search` to avoid duplicating automatically injected and retrieval-returned context. `memory.injection_enabled: false` suppresses the complete block in either mode.
|
||||
- Middleware extraction classifies proposed facts with extraction-only `scope`/`durability`/`authority` labels. `_apply_updates` accepts only `user` + `durable` + `descriptive` new/consolidated facts, accepts only wholly user-scoped summary prose with `authority=descriptive`, and rejects missing labels per item without aborting unrelated updates. Contradiction removals use object entries with `id`, `scope`, `reason`, and optional zero-based `replacementFactIndex`; task/project removals fail closed, and a paired removal runs only when the referenced replacement survives the scope/confidence gates, deduplication, and max-fact trim under another fact ID. The labels are not persisted, so no storage migration is required. Staleness removals retain their independent candidate/cap guardrails, while tool-mode CRUD remains outside this extraction gate. Custom `memory.backend_config.prompts_dir` templates (including per-agent overrides) must carry the same classification fields; an un-migrated template makes the fail-closed gate reject every extraction-driven write, observable only through `rejected_by_scope_gate` and the >60% fact-rejection warning.
|
||||
- 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.
|
||||
- `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.
|
||||
@ -1050,7 +1093,8 @@ Checkpointer storage runs in one of two channel modes, selected by `checkpoint_c
|
||||
- `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` / `DELTA_MESSAGES_FIELD` (`DeltaChannel` at the configured `snapshot_frequency`, default 10), 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`
|
||||
- `runtime/checkpoint_cache/` + `runtime/checkpointer/cached_saver.py` — delta-mode checkpoint history cache; checkpoint state reads MUST go through `CheckpointStateAccessor`, and the checkpointer may be a `CachedHistorySaver` wrapper — never rely on concrete saver types
|
||||
- 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`, `tests/test_cached_history_saver.py` + `tests/test_cached_history_saver_integration.py` (history cache)
|
||||
|
||||
**Checkpoint channel benchmark**: `scripts/benchmark/checkpoint/bench_channels.py`
|
||||
runs paired `full`/`delta` message-only StateGraphs in a fresh child process per
|
||||
|
||||
@ -91,6 +91,8 @@ Async task delegation with concurrent execution:
|
||||
LLM-powered persistent context retention across conversations:
|
||||
|
||||
- **Automatic extraction**: Analyzes conversations for user context, facts, and preferences
|
||||
- **Scope-safe writes**: Middleware extraction stores only durable, descriptive user-level facts; global summaries also require descriptive authority, while contradiction removals and consolidated facts fail closed when scope metadata is missing or task/project-local
|
||||
- **Atomic replacements**: A contradiction removal linked to a replacement runs only after the replacement survives scope/confidence gates, deduplication, and fact-limit trimming
|
||||
- **Structured storage**: User context (work, personal, top-of-mind), history, and confidence-scored facts
|
||||
- **Debounced updates**: Batches updates to minimize LLM calls (configurable wait time)
|
||||
- **System prompt injection**: Top facts + context injected into agent prompts
|
||||
|
||||
@ -325,15 +325,23 @@ class FeishuChannel(Channel):
|
||||
logger.exception("[Feishu] failed to upload/send file: %s", attachment.filename)
|
||||
return False
|
||||
|
||||
async def _upload_image(self, path) -> str:
|
||||
"""Upload an image to Feishu and return the image_key."""
|
||||
def _upload_image_sync(self, path):
|
||||
with open(str(path), "rb") as f:
|
||||
request = self._CreateImageRequest.builder().request_body(self._CreateImageRequestBody.builder().image_type("message").image(f).build()).build()
|
||||
response = await asyncio.to_thread(self._api_client.im.v1.image.create, request)
|
||||
return self._api_client.im.v1.image.create(request)
|
||||
|
||||
async def _upload_image(self, path) -> str:
|
||||
"""Upload an image to Feishu and return the image_key."""
|
||||
response = await asyncio.to_thread(self._upload_image_sync, path)
|
||||
if not response.success():
|
||||
raise RuntimeError(f"Feishu image upload failed: code={response.code}, msg={response.msg}")
|
||||
return response.data.image_key
|
||||
|
||||
def _upload_file_sync(self, path, filename: str, file_type: str):
|
||||
with open(str(path), "rb") as f:
|
||||
request = self._CreateFileRequest.builder().request_body(self._CreateFileRequestBody.builder().file_type(file_type).file_name(filename).file(f).build()).build()
|
||||
return self._api_client.im.v1.file.create(request)
|
||||
|
||||
async def _upload_file(self, path, filename: str) -> str:
|
||||
"""Upload a file to Feishu and return the file_key."""
|
||||
suffix = path.suffix.lower() if hasattr(path, "suffix") else ""
|
||||
@ -348,9 +356,7 @@ class FeishuChannel(Channel):
|
||||
else:
|
||||
file_type = "stream"
|
||||
|
||||
with open(str(path), "rb") as f:
|
||||
request = self._CreateFileRequest.builder().request_body(self._CreateFileRequestBody.builder().file_type(file_type).file_name(filename).file(f).build()).build()
|
||||
response = await asyncio.to_thread(self._api_client.im.v1.file.create, request)
|
||||
response = await asyncio.to_thread(self._upload_file_sync, path, filename, file_type)
|
||||
if not response.success():
|
||||
raise RuntimeError(f"Feishu file upload failed: code={response.code}, msg={response.msg}")
|
||||
return response.data.file_key
|
||||
@ -424,10 +430,8 @@ class FeishuChannel(Channel):
|
||||
logger.warning("[Feishu] empty resource content: resource_key=%s, type=%s", file_key, type)
|
||||
return f"Failed to obtain the [{type}]"
|
||||
|
||||
paths = get_paths()
|
||||
effective_user_id = user_id or get_effective_user_id()
|
||||
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()
|
||||
paths = await asyncio.to_thread(get_paths)
|
||||
|
||||
ext = "png" if type == "image" else "bin"
|
||||
raw_filename = getattr(response, "file_name", "") or f"feishu_{file_key[-12:]}.{ext}"
|
||||
@ -439,30 +443,33 @@ class FeishuChannel(Channel):
|
||||
filename = f"{name_part}.{ext}"
|
||||
else:
|
||||
filename = re.sub(r"[./\\]", "_", raw_filename)
|
||||
resolved_target = uploads_dir / filename
|
||||
|
||||
def down_load():
|
||||
# use thread_lock to avoid filename conflicts when writing
|
||||
def _persist():
|
||||
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()
|
||||
resolved_target = uploads_dir / filename
|
||||
# Use thread_lock to avoid filename conflicts when writing.
|
||||
with self._thread_lock:
|
||||
resolved_target.write_bytes(content)
|
||||
return resolved_target
|
||||
|
||||
try:
|
||||
await asyncio.to_thread(down_load)
|
||||
resolved_target = await asyncio.to_thread(_persist)
|
||||
except Exception:
|
||||
logger.exception("[Feishu] failed to persist downloaded resource: %s, type=%s", resolved_target, type)
|
||||
logger.exception("[Feishu] failed to persist downloaded resource: %s, type=%s", filename, type)
|
||||
return f"Failed to obtain the [{type}]"
|
||||
|
||||
virtual_path = f"{VIRTUAL_PATH_PREFIX}/uploads/{resolved_target.name}"
|
||||
|
||||
try:
|
||||
sandbox_provider = get_sandbox_provider()
|
||||
sandbox_id = sandbox_provider.acquire(thread_id, user_id=effective_user_id)
|
||||
if sandbox_id != "local":
|
||||
sandbox_provider = await asyncio.to_thread(get_sandbox_provider)
|
||||
if not getattr(sandbox_provider, "uses_thread_data_mounts", False):
|
||||
sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=effective_user_id)
|
||||
sandbox = sandbox_provider.get(sandbox_id)
|
||||
if sandbox is None:
|
||||
logger.warning("[Feishu] sandbox not found for thread_id=%s", thread_id)
|
||||
return f"Failed to obtain the [{type}]"
|
||||
sandbox.update_file(virtual_path, content)
|
||||
await asyncio.to_thread(sandbox.update_file, virtual_path, content)
|
||||
except Exception:
|
||||
logger.exception("[Feishu] failed to sync resource into non-local sandbox: %s", virtual_path)
|
||||
return f"Failed to obtain the [{type}]"
|
||||
|
||||
@ -42,6 +42,12 @@ MAX_TRACKED_STREAM_MESSAGES = 256
|
||||
_monotonic = time.monotonic
|
||||
|
||||
|
||||
def _load_telegram_input_file(path, filename: str):
|
||||
from telegram import InputFile
|
||||
|
||||
return InputFile(path.read_bytes(), filename=filename)
|
||||
|
||||
|
||||
class TelegramChannel(Channel):
|
||||
"""Telegram bot channel using long-polling.
|
||||
|
||||
@ -385,21 +391,17 @@ class TelegramChannel(Channel):
|
||||
reply_to = self._last_bot_message.get(msg.chat_id)
|
||||
|
||||
try:
|
||||
input_file = await asyncio.to_thread(_load_telegram_input_file, attachment.actual_path, attachment.filename)
|
||||
if attachment.is_image and attachment.size <= 10 * 1024 * 1024:
|
||||
with open(attachment.actual_path, "rb") as f:
|
||||
kwargs: dict[str, Any] = {"chat_id": chat_id, "photo": f}
|
||||
if reply_to:
|
||||
kwargs["reply_to_message_id"] = reply_to
|
||||
sent = await bot.send_photo(**kwargs)
|
||||
kwargs: dict[str, Any] = {"chat_id": chat_id, "photo": input_file}
|
||||
if reply_to:
|
||||
kwargs["reply_to_message_id"] = reply_to
|
||||
sent = await bot.send_photo(**kwargs)
|
||||
else:
|
||||
from telegram import InputFile
|
||||
|
||||
with open(attachment.actual_path, "rb") as f:
|
||||
input_file = InputFile(f, filename=attachment.filename)
|
||||
kwargs = {"chat_id": chat_id, "document": input_file}
|
||||
if reply_to:
|
||||
kwargs["reply_to_message_id"] = reply_to
|
||||
sent = await bot.send_document(**kwargs)
|
||||
kwargs = {"chat_id": chat_id, "document": input_file}
|
||||
if reply_to:
|
||||
kwargs["reply_to_message_id"] = reply_to
|
||||
sent = await bot.send_document(**kwargs)
|
||||
|
||||
self._last_bot_message[msg.chat_id] = sent.message_id
|
||||
logger.info("[Telegram] file sent: %s to chat=%s", attachment.filename, msg.chat_id)
|
||||
|
||||
@ -21,6 +21,18 @@ from app.channels.message_bus import (
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _file_md5(path: str) -> str:
|
||||
md5_hasher = hashlib.md5()
|
||||
with open(path, "rb") as file_obj:
|
||||
for chunk in iter(lambda: file_obj.read(1024 * 1024), b""):
|
||||
md5_hasher.update(chunk)
|
||||
return md5_hasher.hexdigest()
|
||||
|
||||
|
||||
def _open_binary(path: str):
|
||||
return open(path, "rb")
|
||||
|
||||
|
||||
class WeComChannel(Channel):
|
||||
def __init__(self, bus: MessageBus, config: dict[str, Any]) -> None:
|
||||
super().__init__(name="wecom", bus=bus, config=config)
|
||||
@ -428,11 +440,7 @@ class WeComChannel(Channel):
|
||||
logger.warning("[WeCom] invalid total_chunks=%d for %s", total_chunks, filename)
|
||||
return None
|
||||
|
||||
md5_hasher = hashlib.md5()
|
||||
with open(path, "rb") as f:
|
||||
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
||||
md5_hasher.update(chunk)
|
||||
md5 = md5_hasher.hexdigest()
|
||||
md5 = await asyncio.to_thread(_file_md5, path)
|
||||
|
||||
init_req_id = generate_req_id("aibot_upload_media_init")
|
||||
init_body = {
|
||||
@ -448,9 +456,10 @@ class WeComChannel(Channel):
|
||||
logger.warning("[WeCom] upload init returned no upload_id: %s", init_ack)
|
||||
return None
|
||||
|
||||
with open(path, "rb") as f:
|
||||
file_obj = await asyncio.to_thread(_open_binary, path)
|
||||
try:
|
||||
for idx in range(total_chunks):
|
||||
data = f.read(chunk_size)
|
||||
data = await asyncio.to_thread(file_obj.read, chunk_size)
|
||||
if not data:
|
||||
break
|
||||
chunk_req_id = generate_req_id("aibot_upload_media_chunk")
|
||||
@ -460,6 +469,8 @@ class WeComChannel(Channel):
|
||||
"base64_data": base64.b64encode(data).decode("utf-8"),
|
||||
}
|
||||
await self._send_ws_upload_command(chunk_req_id, chunk_body, "aibot_upload_media_chunk")
|
||||
finally:
|
||||
await asyncio.to_thread(file_obj.close)
|
||||
|
||||
finish_req_id = generate_req_id("aibot_upload_media_finish")
|
||||
finish_ack = await self._send_ws_upload_command(finish_req_id, {"upload_id": upload_id}, "aibot_upload_media_finish")
|
||||
|
||||
@ -210,6 +210,12 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
config = get_gateway_config()
|
||||
logger.info(f"Starting API Gateway on {config.host}:{config.port}")
|
||||
|
||||
from deerflow.skills.projection import ensure_public_skill_projection
|
||||
|
||||
public_projection_ready = await asyncio.to_thread(ensure_public_skill_projection, app_config=startup_config)
|
||||
if public_projection_ready:
|
||||
logger.info("Ensured the public skill projection; user projections repair lazily on sandbox acquire")
|
||||
|
||||
# Agent observability (Monocle). Off by default; enabled with
|
||||
# MONOCLE_TRACING. Initialized here at startup — not at import time — so a
|
||||
# plain `import deerflow.agents` never installs a process-global tracer.
|
||||
|
||||
@ -40,7 +40,7 @@ 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.provider import AuthorizationProvider, AuthzDecision, AuthzRequest, Principal
|
||||
from deerflow.authz.runtime import resolve_authorization_provider
|
||||
from deerflow.config.authorization_config import AuthorizationConfig
|
||||
|
||||
@ -263,6 +263,61 @@ async def resolve_route_permissions(user: User, *, is_internal: bool) -> list[st
|
||||
return [p for p in results if p is not None]
|
||||
|
||||
|
||||
class _AuthorizationUnavailable(Exception):
|
||||
"""Raised internally when the provider cannot be resolved for a route check.
|
||||
|
||||
Carries the ``fail_closed`` flag so the caller can decide between deny-all
|
||||
and legacy allow-all without re-reading config.
|
||||
"""
|
||||
|
||||
def __init__(self, *, fail_closed: bool) -> None:
|
||||
self.fail_closed = fail_closed
|
||||
|
||||
|
||||
def resolve_model_authorization(user: User, *, is_internal: bool) -> tuple[AuthorizationProvider | None, Principal | None]:
|
||||
"""Return ``(provider, principal)`` for model-route authorization.
|
||||
|
||||
When authorization is disabled, returns ``(None, None)`` so callers can
|
||||
short-circuit to legacy behavior (all models visible). When enabled,
|
||||
resolves the cached provider and builds a Principal identical to
|
||||
``resolve_route_permissions`` (including the ``INTERNAL_SYSTEM_ROLE``
|
||||
→ ``None`` pop so internal callers fall under ``default_role``).
|
||||
|
||||
Raises ``_AuthorizationUnavailable`` (carrying the ``fail_closed`` flag)
|
||||
when the provider cannot be resolved; callers translate that into the
|
||||
appropriate deny response (empty list / 403).
|
||||
"""
|
||||
config = _get_route_authorization_config()
|
||||
if config.enabled is not True:
|
||||
return None, None
|
||||
|
||||
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 model routes", exc_info=True)
|
||||
raise _AuthorizationUnavailable(fail_closed=config.fail_closed)
|
||||
|
||||
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,
|
||||
)
|
||||
return provider, principal
|
||||
|
||||
|
||||
async def _authenticate(request: Request) -> AuthContext:
|
||||
"""Authenticate request and return AuthContext.
|
||||
|
||||
|
||||
86
backend/app/gateway/context_usage.py
Normal file
86
backend/app/gateway/context_usage.py
Normal file
@ -0,0 +1,86 @@
|
||||
"""Compute the current message-context usage for a thread."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
|
||||
from app.gateway.deps import get_config
|
||||
from app.gateway.services import build_thread_checkpoint_state_accessor
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _count_messages_approximately(messages: list[Any]) -> int:
|
||||
"""Count checkpoint messages with LangChain's network-free heuristic."""
|
||||
if not messages:
|
||||
return 0
|
||||
from langchain_core.messages.utils import count_tokens_approximately
|
||||
|
||||
return int(count_tokens_approximately(messages))
|
||||
|
||||
|
||||
async def _load_checkpoint_messages(accessor: Any, config: dict[str, Any]) -> list[Any]:
|
||||
"""Read materialized messages so full and delta checkpoints behave alike."""
|
||||
snapshot = await accessor.aget(config)
|
||||
values = getattr(snapshot, "values", None) or {}
|
||||
if not isinstance(values, dict):
|
||||
return []
|
||||
return list(values.get("messages") or [])
|
||||
|
||||
|
||||
async def _resolve_thread_model_name(run_store: Any, thread_id: str, app_config: Any) -> str | None:
|
||||
"""Prefer the latest run's model, then fall back to the first configured model."""
|
||||
try:
|
||||
runs = await run_store.list_by_thread(thread_id, limit=1)
|
||||
except Exception:
|
||||
runs = []
|
||||
if runs:
|
||||
latest = runs[0]
|
||||
name = latest.get("model_name") if isinstance(latest, dict) else getattr(latest, "model_name", None)
|
||||
if isinstance(name, str) and name:
|
||||
return name
|
||||
models = getattr(app_config, "models", None) or []
|
||||
return models[0].name if models else None
|
||||
|
||||
|
||||
def build_context_usage_payload(*, token_count: int, max_context_tokens: int | None) -> dict[str, Any]:
|
||||
"""Build the stable API payload for a message count and model capacity."""
|
||||
percentage: float | None = None
|
||||
if max_context_tokens and max_context_tokens > 0:
|
||||
percentage = round(token_count / max_context_tokens * 100, 1)
|
||||
return {
|
||||
"token_count": token_count,
|
||||
"max_context_tokens": max_context_tokens,
|
||||
"percentage": percentage,
|
||||
}
|
||||
|
||||
|
||||
async def build_context_usage(request: Request, thread_id: str, run_store: Any) -> dict[str, Any] | None:
|
||||
"""Return approximate usage for the latest materialized thread checkpoint."""
|
||||
try:
|
||||
app_config = get_config()
|
||||
except HTTPException:
|
||||
return None
|
||||
|
||||
try:
|
||||
accessor, checkpoint_config = await build_thread_checkpoint_state_accessor(request, thread_id=thread_id)
|
||||
messages = await _load_checkpoint_messages(accessor, checkpoint_config)
|
||||
except Exception:
|
||||
logger.warning("Failed to load checkpoint for context usage on thread %s", thread_id, exc_info=True)
|
||||
return None
|
||||
|
||||
try:
|
||||
token_count = await asyncio.to_thread(_count_messages_approximately, messages)
|
||||
except Exception:
|
||||
logger.warning("Failed to count context messages for thread %s", thread_id, exc_info=True)
|
||||
return None
|
||||
|
||||
model_name = await _resolve_thread_model_name(run_store, thread_id, app_config)
|
||||
model_config = app_config.get_model_config(model_name) if model_name else None
|
||||
configured_window = getattr(model_config, "context_window", None) if model_config is not None else None
|
||||
max_context_tokens = int(configured_window) if configured_window else None
|
||||
return build_context_usage_payload(token_count=token_count, max_context_tokens=max_context_tokens)
|
||||
@ -1,17 +1,29 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import logging
|
||||
import mimetypes
|
||||
import os
|
||||
import stat
|
||||
import tempfile
|
||||
import zipfile
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from urllib.parse import quote
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
from fastapi.responses import FileResponse, PlainTextResponse, Response
|
||||
from fastapi.responses import FileResponse, Response
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.gateway.authz import require_permission
|
||||
from app.gateway.deps import get_run_manager
|
||||
from app.gateway.internal_auth import get_trusted_internal_owner_user_id
|
||||
from app.gateway.path_utils import resolve_thread_virtual_path
|
||||
from deerflow.config.paths import make_safe_user_id
|
||||
from deerflow.runtime import ConflictError, ThreadOperationKind
|
||||
from deerflow.runtime.user_context import get_effective_user_id
|
||||
from deerflow.sandbox.sandbox_provider import get_sandbox_provider
|
||||
from deerflow.utils.thread_id import ThreadId
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -25,6 +37,114 @@ ACTIVE_CONTENT_MIME_TYPES = {
|
||||
|
||||
MAX_SKILL_ARCHIVE_MEMBER_BYTES = 16 * 1024 * 1024
|
||||
_SKILL_ARCHIVE_READ_CHUNK_SIZE = 64 * 1024
|
||||
MAX_EDITABLE_ARTIFACT_BYTES = 2 * 1024 * 1024
|
||||
_EDITABLE_OUTPUTS_PREFIX = "mnt/user-data/outputs/"
|
||||
_ARTIFACT_EDIT_TEMP_PREFIX = ".artifact-edit-"
|
||||
|
||||
|
||||
class ArtifactUpdateRequest(BaseModel):
|
||||
content: str
|
||||
expected_sha256: str = Field(pattern=r"^[0-9a-f]{64}$")
|
||||
|
||||
|
||||
class ArtifactUpdateResponse(BaseModel):
|
||||
path: str
|
||||
sha256: str
|
||||
size: int
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def reserve_artifact_write(request: Request, thread_id: str, *, user_id: str) -> AsyncIterator[None]:
|
||||
"""Serialize an artifact edit against runs and other thread mutations."""
|
||||
run_manager = get_run_manager(request)
|
||||
async with run_manager.reserve_thread_operation(
|
||||
thread_id,
|
||||
kind=ThreadOperationKind.artifact_write,
|
||||
user_id=user_id,
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
def _normalize_editable_artifact_path(path: str) -> str:
|
||||
stripped = path.lstrip("/")
|
||||
if not stripped.startswith(_EDITABLE_OUTPUTS_PREFIX):
|
||||
raise HTTPException(status_code=400, detail="Only files in /mnt/user-data/outputs can be edited")
|
||||
if ".skill/" in stripped or stripped.endswith(".skill"):
|
||||
raise HTTPException(status_code=415, detail="Skill archives cannot be edited in the artifacts panel")
|
||||
return f"/{stripped}"
|
||||
|
||||
|
||||
def _load_editable_artifact(actual_path: Path, path: str, expected_sha256: str) -> tuple[bytes, os.stat_result]:
|
||||
try:
|
||||
file_stat = os.lstat(actual_path)
|
||||
except FileNotFoundError:
|
||||
raise HTTPException(status_code=404, detail=f"Artifact not found: {path}") from None
|
||||
if stat.S_ISLNK(file_stat.st_mode):
|
||||
raise HTTPException(status_code=415, detail="Symlinked artifacts cannot be edited")
|
||||
if not stat.S_ISREG(file_stat.st_mode):
|
||||
raise HTTPException(status_code=400, detail=f"Path is not a file: {path}")
|
||||
if file_stat.st_size > MAX_EDITABLE_ARTIFACT_BYTES:
|
||||
raise HTTPException(status_code=413, detail="Artifact is too large to edit")
|
||||
|
||||
current = actual_path.read_bytes()
|
||||
if len(current) > MAX_EDITABLE_ARTIFACT_BYTES:
|
||||
raise HTTPException(status_code=413, detail="Artifact is too large to edit")
|
||||
if b"\x00" in current:
|
||||
raise HTTPException(status_code=415, detail="Binary artifacts cannot be edited")
|
||||
try:
|
||||
current.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
raise HTTPException(status_code=415, detail="Only UTF-8 text artifacts can be edited") from None
|
||||
|
||||
current_sha256 = hashlib.sha256(current).hexdigest()
|
||||
if current_sha256 != expected_sha256:
|
||||
raise HTTPException(status_code=412, detail="Artifact changed since it was opened")
|
||||
return current, file_stat
|
||||
|
||||
|
||||
def _encode_artifact_update(content: str) -> bytes:
|
||||
encoded = content.encode("utf-8")
|
||||
if len(encoded) > MAX_EDITABLE_ARTIFACT_BYTES:
|
||||
raise HTTPException(status_code=413, detail="Artifact is too large to edit")
|
||||
if b"\x00" in encoded:
|
||||
raise HTTPException(status_code=415, detail="Binary content cannot be saved as an artifact")
|
||||
return encoded
|
||||
|
||||
|
||||
def _replace_artifact_atomically(actual_path: Path, content: bytes, file_stat: os.stat_result) -> None:
|
||||
temp_fd, temp_path_str = tempfile.mkstemp(prefix=_ARTIFACT_EDIT_TEMP_PREFIX, dir=actual_path.parent)
|
||||
temp_path = Path(temp_path_str)
|
||||
try:
|
||||
# Preserve ownership where possible and keep replacement permissions
|
||||
# scoped to the owner/group. The shared outputs directory allows a
|
||||
# mounted sandbox to reach the file without making it world-writable.
|
||||
if hasattr(os, "fchown"):
|
||||
try:
|
||||
os.fchown(temp_fd, file_stat.st_uid, file_stat.st_gid)
|
||||
except OSError:
|
||||
logger.debug("Could not preserve artifact ownership: %s", actual_path, exc_info=True)
|
||||
# Windows has no fchmod and uses ACLs rather than POSIX mode bits.
|
||||
# Keep the mkstemp permissions there; retain the existing POSIX
|
||||
# behavior on platforms that expose descriptor-based chmod.
|
||||
if hasattr(os, "fchmod"):
|
||||
os.fchmod(temp_fd, stat.S_IMODE(file_stat.st_mode) | 0o660)
|
||||
with os.fdopen(temp_fd, "wb") as handle:
|
||||
temp_fd = -1
|
||||
handle.write(content)
|
||||
handle.flush()
|
||||
os.fsync(handle.fileno())
|
||||
os.replace(temp_path, actual_path)
|
||||
finally:
|
||||
if temp_fd >= 0:
|
||||
os.close(temp_fd)
|
||||
try:
|
||||
temp_path.unlink()
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
|
||||
|
||||
def _sync_artifact_to_sandbox(sandbox, virtual_path: str, content: bytes) -> None:
|
||||
sandbox.update_file(virtual_path, content)
|
||||
|
||||
|
||||
def _build_content_disposition(disposition_type: str, filename: str) -> str:
|
||||
@ -39,6 +159,51 @@ def _build_attachment_headers(filename: str, extra_headers: dict[str, str] | Non
|
||||
return headers
|
||||
|
||||
|
||||
def _slice_byte_range(content: bytes, range_header: str | None) -> tuple[bytes, int, dict[str, str]]:
|
||||
"""Apply one RFC 9110 byte range to an in-memory archive member."""
|
||||
size = len(content)
|
||||
headers = {"Accept-Ranges": "bytes"}
|
||||
if range_header is None:
|
||||
return content, 200, headers
|
||||
|
||||
def unsatisfied() -> HTTPException:
|
||||
return HTTPException(
|
||||
status_code=416,
|
||||
detail="Requested range is not satisfiable",
|
||||
headers={"Accept-Ranges": "bytes", "Content-Range": f"bytes */{size}"},
|
||||
)
|
||||
|
||||
if not range_header.startswith("bytes=") or "," in range_header:
|
||||
raise unsatisfied()
|
||||
range_spec = range_header.removeprefix("bytes=")
|
||||
if "-" not in range_spec:
|
||||
raise unsatisfied()
|
||||
start_text, end_text = range_spec.split("-", 1)
|
||||
try:
|
||||
if start_text:
|
||||
start = int(start_text)
|
||||
end = size - 1 if not end_text else min(int(end_text), size - 1)
|
||||
else:
|
||||
suffix_length = int(end_text)
|
||||
if suffix_length <= 0:
|
||||
raise unsatisfied()
|
||||
start = max(size - suffix_length, 0)
|
||||
end = size - 1
|
||||
except ValueError as exc:
|
||||
raise unsatisfied() from exc
|
||||
|
||||
if size == 0 or start < 0 or start >= size or end < start:
|
||||
raise unsatisfied()
|
||||
ranged_content = content[start : end + 1]
|
||||
headers.update(
|
||||
{
|
||||
"Content-Range": f"bytes {start}-{end}/{size}",
|
||||
"Content-Length": str(len(ranged_content)),
|
||||
}
|
||||
)
|
||||
return ranged_content, 206, headers
|
||||
|
||||
|
||||
def is_text_file_by_content(path: Path, sample_size: int = 8192) -> bool:
|
||||
"""Check if file is text by examining content for null bytes."""
|
||||
try:
|
||||
@ -119,18 +284,14 @@ def _load_skill_archive_member(actual_skill_path: Path, skill_file_path: str, in
|
||||
return content, mime_type
|
||||
|
||||
|
||||
def _read_artifact_payload(actual_path: Path, path: str, download: bool) -> tuple[str, str | None, bytes | str | None]:
|
||||
def _read_artifact_payload(actual_path: Path, path: str, download: bool) -> tuple[str, str | None]:
|
||||
"""Worker-thread body for the regular branch of ``get_artifact``.
|
||||
|
||||
Stat probes, MIME sniffing (``mimetypes`` lazily stats the system MIME
|
||||
database on first use), and text reads are blocking filesystem IO. Returns
|
||||
a ``(kind, mime_type, payload)`` plan the handler turns into a response on
|
||||
the loop: ``("file", mime, None)`` (attachment / forced-download active
|
||||
content, streamed by ``FileResponse``), ``("inline_file", mime, None)``
|
||||
(inline binary preview — also streamed by ``FileResponse`` so the client
|
||||
can issue byte-``Range`` requests, e.g. to seek within audio/video
|
||||
artifacts instead of always replaying from byte 0), or ``("text", mime,
|
||||
str)``. Behavior/error codes match the previous inline logic.
|
||||
Stat probes and MIME sniffing (``mimetypes`` lazily stats the system MIME
|
||||
database on first use) are blocking filesystem IO. Returns a
|
||||
``(kind, mime_type)`` plan the handler turns into a streamed
|
||||
``FileResponse``. Inline text and binary previews both use FileResponse so
|
||||
clients can request a bounded byte range instead of buffering a whole file.
|
||||
"""
|
||||
if not actual_path.exists():
|
||||
raise HTTPException(status_code=404, detail=f"Artifact not found: {path}")
|
||||
@ -139,15 +300,12 @@ def _read_artifact_payload(actual_path: Path, path: str, download: bool) -> tupl
|
||||
mime_type, _ = mimetypes.guess_type(actual_path)
|
||||
# Active content / explicit download is streamed by FileResponse — no read here.
|
||||
if download or mime_type in ACTIVE_CONTENT_MIME_TYPES:
|
||||
return ("file", mime_type, None)
|
||||
return ("file", mime_type)
|
||||
if mime_type and mime_type.startswith("text/"):
|
||||
return ("text", mime_type, actual_path.read_text(encoding="utf-8"))
|
||||
return ("inline_file", mime_type)
|
||||
if is_text_file_by_content(actual_path):
|
||||
return ("text", mime_type, actual_path.read_text(encoding="utf-8"))
|
||||
# Binary inline preview (images, audio, video, PDFs, ...): stream via
|
||||
# FileResponse instead of buffering the whole file in memory, so it also
|
||||
# gets FileResponse's built-in byte-Range handling (see get_artifact).
|
||||
return ("inline_file", mime_type, None)
|
||||
return ("inline_file", mime_type or "text/plain")
|
||||
return ("inline_file", mime_type)
|
||||
|
||||
|
||||
@router.get(
|
||||
@ -156,7 +314,7 @@ def _read_artifact_payload(actual_path: Path, path: str, download: bool) -> tupl
|
||||
description="Retrieve an artifact file generated by the AI agent. Text and binary files can be viewed inline, while active web content is always downloaded.",
|
||||
)
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
async def get_artifact(thread_id: str, path: str, request: Request, download: bool = False) -> Response:
|
||||
async def get_artifact(thread_id: ThreadId, path: str, request: Request, download: bool = False) -> Response:
|
||||
"""Get an artifact file by its path.
|
||||
|
||||
The endpoint automatically detects file types and returns appropriate content types.
|
||||
@ -218,23 +376,37 @@ async def get_artifact(thread_id: str, path: str, request: Request, download: bo
|
||||
if download or mime_type in ACTIVE_CONTENT_MIME_TYPES:
|
||||
return Response(content=content, media_type=mime_type or "application/octet-stream", headers=_build_attachment_headers(download_name, cache_headers))
|
||||
|
||||
# Archive members are already bounded during extraction. Preserve byte
|
||||
# semantics here so the frontend can request only its preview budget,
|
||||
# including a final partial UTF-8 sequence.
|
||||
request_headers = request.headers if request is not None else {}
|
||||
range_header = None if request_headers.get("if-range") else request_headers.get("range")
|
||||
ranged_content, status_code, range_headers = _slice_byte_range(content, range_header)
|
||||
inline_headers = {**cache_headers, **range_headers}
|
||||
|
||||
if mime_type and mime_type.startswith("text/"):
|
||||
return PlainTextResponse(content=content.decode("utf-8"), media_type=mime_type, headers=cache_headers)
|
||||
return Response(content=ranged_content, status_code=status_code, media_type=mime_type, headers=inline_headers)
|
||||
|
||||
# Default to plain text for unknown types that look like text
|
||||
try:
|
||||
return PlainTextResponse(content=content.decode("utf-8"), media_type="text/plain", headers=cache_headers)
|
||||
content.decode("utf-8")
|
||||
return Response(content=ranged_content, status_code=status_code, media_type="text/plain", headers=inline_headers)
|
||||
except UnicodeDecodeError:
|
||||
return Response(content=content, media_type=mime_type or "application/octet-stream", headers=cache_headers)
|
||||
return Response(
|
||||
content=ranged_content,
|
||||
status_code=status_code,
|
||||
media_type=mime_type or "application/octet-stream",
|
||||
headers=inline_headers,
|
||||
)
|
||||
|
||||
actual_path = await asyncio.to_thread(resolve_thread_virtual_path, thread_id, path, user_id=owner_user_id)
|
||||
|
||||
logger.info(f"Resolving artifact path: thread_id={thread_id}, requested_path={path}, actual_path={actual_path}")
|
||||
|
||||
# Offload path stat + MIME sniff + file reads (all blocking filesystem IO).
|
||||
# Active content and explicit downloads are streamed by FileResponse, so the
|
||||
# worker only reports the kind; inline text/binary payloads are read in-thread.
|
||||
kind, mime_type, payload = await asyncio.to_thread(_read_artifact_payload, actual_path, path, download)
|
||||
# Offload path stat + MIME sniff (blocking filesystem IO). Every regular
|
||||
# artifact response is streamed by FileResponse; the worker only reports
|
||||
# disposition and media type.
|
||||
kind, mime_type = await asyncio.to_thread(_read_artifact_payload, actual_path, path, download)
|
||||
|
||||
if kind == "file":
|
||||
# Always force download for active content types to prevent script
|
||||
@ -242,19 +414,89 @@ async def get_artifact(thread_id: str, path: str, request: Request, download: bo
|
||||
return FileResponse(path=actual_path, filename=actual_path.name, media_type=mime_type, headers=_build_attachment_headers(actual_path.name))
|
||||
|
||||
if kind == "inline_file":
|
||||
# FileResponse (unlike a fully-buffered Response) honors byte-Range
|
||||
# requests. Browsers issue these when seeking an <audio>/<video>
|
||||
# element backed by a remote URL; serving the same bytes through a
|
||||
# plain Response ignores Range headers and always replays from byte
|
||||
# 0, which is why dragging an audio/video artifact's progress bar
|
||||
# reset playback to the start instead of jumping to the new position.
|
||||
# FileResponse honors byte-Range requests for large text previews and
|
||||
# media seeking without buffering the full artifact in the Gateway.
|
||||
return FileResponse(
|
||||
path=actual_path,
|
||||
media_type=mime_type,
|
||||
headers={"Content-Disposition": _build_content_disposition("inline", actual_path.name)},
|
||||
)
|
||||
|
||||
if kind == "text":
|
||||
return PlainTextResponse(content=payload, media_type=mime_type)
|
||||
|
||||
raise AssertionError(f"Unhandled artifact response kind: {kind!r}")
|
||||
|
||||
|
||||
@router.put(
|
||||
"/threads/{thread_id}/artifacts/{path:path}",
|
||||
response_model=ArtifactUpdateResponse,
|
||||
summary="Update Artifact File",
|
||||
description="Replace an existing UTF-8 text artifact after verifying that its content has not changed.",
|
||||
)
|
||||
@require_permission("threads", "write", owner_check=True, require_existing=True)
|
||||
async def update_artifact(
|
||||
thread_id: ThreadId,
|
||||
path: str,
|
||||
body: ArtifactUpdateRequest,
|
||||
request: Request,
|
||||
) -> ArtifactUpdateResponse:
|
||||
"""Update an existing text artifact while the thread has no active run."""
|
||||
virtual_path = _normalize_editable_artifact_path(path)
|
||||
raw_owner_user_id = get_trusted_internal_owner_user_id(request)
|
||||
effective_user_id = make_safe_user_id(raw_owner_user_id) if raw_owner_user_id else get_effective_user_id()
|
||||
|
||||
sandbox_provider = None
|
||||
sandbox_id: str | None = None
|
||||
sandbox = None
|
||||
try:
|
||||
async with reserve_artifact_write(request, thread_id, user_id=effective_user_id):
|
||||
actual_path = await asyncio.to_thread(
|
||||
resolve_thread_virtual_path,
|
||||
thread_id,
|
||||
virtual_path,
|
||||
user_id=effective_user_id,
|
||||
)
|
||||
current, file_stat = await asyncio.to_thread(
|
||||
_load_editable_artifact,
|
||||
actual_path,
|
||||
virtual_path,
|
||||
body.expected_sha256,
|
||||
)
|
||||
updated = _encode_artifact_update(body.content)
|
||||
|
||||
sandbox_provider = get_sandbox_provider()
|
||||
if not bool(getattr(sandbox_provider, "uses_thread_data_mounts", False)):
|
||||
sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=effective_user_id)
|
||||
sandbox = sandbox_provider.get(sandbox_id)
|
||||
if sandbox is None:
|
||||
raise RuntimeError("Failed to acquire sandbox for artifact update")
|
||||
|
||||
try:
|
||||
if sandbox is not None:
|
||||
await asyncio.to_thread(_sync_artifact_to_sandbox, sandbox, virtual_path, updated)
|
||||
await asyncio.to_thread(_replace_artifact_atomically, actual_path, updated, file_stat)
|
||||
except Exception:
|
||||
if sandbox is not None:
|
||||
try:
|
||||
await asyncio.to_thread(_sync_artifact_to_sandbox, sandbox, virtual_path, current)
|
||||
except Exception:
|
||||
logger.exception("Failed to roll back remote artifact after artifact update failure: %s", virtual_path)
|
||||
raise
|
||||
except ConflictError:
|
||||
raise HTTPException(status_code=409, detail="Thread has a run in flight. Save after the run finishes.") from None
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception:
|
||||
logger.exception("Failed to update artifact %s for thread %s", path, thread_id)
|
||||
raise HTTPException(status_code=500, detail="Failed to update artifact") from None
|
||||
finally:
|
||||
if sandbox_id is not None and sandbox_provider is not None:
|
||||
try:
|
||||
await asyncio.to_thread(sandbox_provider.release, sandbox_id)
|
||||
except Exception:
|
||||
logger.warning("Failed to release sandbox after artifact update: %s", sandbox_id, exc_info=True)
|
||||
|
||||
content_sha256 = hashlib.sha256(updated).hexdigest()
|
||||
return ArtifactUpdateResponse(
|
||||
path=virtual_path,
|
||||
sha256=content_sha256,
|
||||
size=len(updated),
|
||||
)
|
||||
|
||||
@ -1,4 +1,5 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import contextlib
|
||||
import json
|
||||
import logging
|
||||
@ -10,6 +11,7 @@ from app.gateway.authz import require_permission
|
||||
from app.gateway.browser_capability import browser_capability
|
||||
from deerflow.config.paths import get_paths
|
||||
from deerflow.runtime.user_context import get_effective_user_id, reset_current_user, set_current_user
|
||||
from deerflow.utils.thread_id import ThreadId
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -76,7 +78,7 @@ async def _browser_thread_owned_by(thread_store, thread_id: str, user_id: str) -
|
||||
description="Steer the thread's live browser session to a URL from the UI and capture a screenshot.",
|
||||
)
|
||||
@require_permission("threads", "write", owner_check=True, require_existing=True)
|
||||
async def navigate_browser(thread_id: str, body: BrowserNavigateRequest, request: Request) -> BrowserNavigateResponse:
|
||||
async def navigate_browser(thread_id: ThreadId, body: BrowserNavigateRequest, request: Request) -> BrowserNavigateResponse:
|
||||
user_id = str(request.state.auth.user.id)
|
||||
thread_store = getattr(request.app.state, "thread_store", None)
|
||||
if thread_store is None or not await _browser_thread_owned_by(thread_store, thread_id, user_id):
|
||||
@ -175,13 +177,40 @@ def _ws_origin_allowed(websocket: WebSocket) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
async def _send_browser_frame(websocket: WebSocket, data: bytes, *, binary: bool) -> None:
|
||||
if binary:
|
||||
await websocket.send_bytes(data)
|
||||
return
|
||||
payload = {"type": "frame", "data": base64.b64encode(data).decode("ascii")}
|
||||
await websocket.send_text(json.dumps(payload))
|
||||
|
||||
|
||||
async def _negotiate_browser_frame_format(websocket: WebSocket) -> bool | None:
|
||||
"""Accept the socket and resolve the optional frame transport capability."""
|
||||
requested_format = websocket.query_params.get("frame_format")
|
||||
await websocket.accept()
|
||||
if requested_format not in {None, "binary"}:
|
||||
await websocket.send_text(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "error",
|
||||
"message": f"Unsupported frame_format: {requested_format}",
|
||||
},
|
||||
),
|
||||
)
|
||||
await websocket.close(code=1008)
|
||||
return None
|
||||
return requested_format == "binary"
|
||||
|
||||
|
||||
@router.websocket("/threads/{thread_id}/browser/stream")
|
||||
async def browser_stream(websocket: WebSocket, thread_id: str) -> None:
|
||||
async def browser_stream(websocket: WebSocket, thread_id: ThreadId) -> None:
|
||||
"""Bidirectional live browser stream.
|
||||
|
||||
Server → client: JSON ``{"type":"frame","data":"<base64 jpeg>"}`` frames
|
||||
captured via CDP screencast. Client → server: input events (click, move,
|
||||
down, up, wheel, key, text, navigate) that drive the live page.
|
||||
Server → client: binary JPEG frames when ``frame_format=binary`` is
|
||||
requested; legacy clients retain JSON base64 frames. Status and navigation
|
||||
metadata remain JSON. Client → server: input events (click, move, down, up,
|
||||
wheel, key, text, navigate) that drive the live page.
|
||||
"""
|
||||
user = await _authenticate_ws(websocket)
|
||||
if user is None:
|
||||
@ -223,11 +252,13 @@ async def browser_stream(websocket: WebSocket, thread_id: str) -> None:
|
||||
await websocket.close(code=4501)
|
||||
return
|
||||
|
||||
await websocket.accept()
|
||||
use_binary_frames = await _negotiate_browser_frame_format(websocket)
|
||||
if use_binary_frames is None:
|
||||
return
|
||||
|
||||
token = set_current_user(user)
|
||||
loop = asyncio.get_running_loop()
|
||||
frame_queue: asyncio.Queue[str] = asyncio.Queue(maxsize=4)
|
||||
frame_queue: asyncio.Queue[bytes] = asyncio.Queue(maxsize=4)
|
||||
send_lock = asyncio.Lock()
|
||||
input_event = asyncio.Event()
|
||||
input_queue: asyncio.Queue[dict] = asyncio.Queue(maxsize=64)
|
||||
@ -238,7 +269,7 @@ async def browser_stream(websocket: WebSocket, thread_id: str) -> None:
|
||||
async with send_lock:
|
||||
await websocket.send_text(json.dumps(payload))
|
||||
|
||||
def _on_frame(data: str) -> None:
|
||||
def _on_frame(data: bytes) -> None:
|
||||
# Invoked on the private Playwright loop; hop to this loop and drop the
|
||||
# oldest frame when the client can't keep up (screencast is lossy).
|
||||
def _enqueue() -> None:
|
||||
@ -294,7 +325,8 @@ async def browser_stream(websocket: WebSocket, thread_id: str) -> None:
|
||||
async def _pump_frames() -> None:
|
||||
while True:
|
||||
data = await frame_queue.get()
|
||||
await _send_payload({"type": "frame", "data": data})
|
||||
async with send_lock:
|
||||
await _send_browser_frame(websocket, data, binary=use_binary_frames)
|
||||
|
||||
async def _send_url() -> None:
|
||||
# Report the page's real URL so the client's address bar reflects the
|
||||
|
||||
@ -14,6 +14,7 @@ from pydantic import BaseModel, Field
|
||||
|
||||
from app.gateway.authz import require_permission
|
||||
from app.gateway.deps import get_current_user, get_feedback_repo, get_run_store
|
||||
from deerflow.utils.thread_id import ThreadId
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/api/threads", tags=["feedback"])
|
||||
@ -61,7 +62,7 @@ class FeedbackStatsResponse(BaseModel):
|
||||
@router.put("/{thread_id}/runs/{run_id}/feedback", response_model=FeedbackResponse)
|
||||
@require_permission("threads", "write", owner_check=True, require_existing=True)
|
||||
async def upsert_feedback(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
run_id: str,
|
||||
body: FeedbackUpsertRequest,
|
||||
request: Request,
|
||||
@ -92,7 +93,7 @@ async def upsert_feedback(
|
||||
@router.delete("/{thread_id}/runs/{run_id}/feedback")
|
||||
@require_permission("threads", "delete", owner_check=True, require_existing=True)
|
||||
async def delete_run_feedback(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
run_id: str,
|
||||
request: Request,
|
||||
) -> dict[str, bool]:
|
||||
@ -112,7 +113,7 @@ async def delete_run_feedback(
|
||||
@router.post("/{thread_id}/runs/{run_id}/feedback", response_model=FeedbackResponse)
|
||||
@require_permission("threads", "write", owner_check=True, require_existing=True)
|
||||
async def create_feedback(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
run_id: str,
|
||||
body: FeedbackCreateRequest,
|
||||
request: Request,
|
||||
@ -145,7 +146,7 @@ async def create_feedback(
|
||||
@router.get("/{thread_id}/runs/{run_id}/feedback", response_model=list[FeedbackResponse])
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
async def list_feedback(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
run_id: str,
|
||||
request: Request,
|
||||
) -> list[dict[str, Any]]:
|
||||
@ -157,7 +158,7 @@ async def list_feedback(
|
||||
@router.get("/{thread_id}/runs/{run_id}/feedback/stats", response_model=FeedbackStatsResponse)
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
async def feedback_stats(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
run_id: str,
|
||||
request: Request,
|
||||
) -> dict[str, Any]:
|
||||
@ -169,7 +170,7 @@ async def feedback_stats(
|
||||
@router.delete("/{thread_id}/runs/{run_id}/feedback/{feedback_id}")
|
||||
@require_permission("threads", "delete", owner_check=True, require_existing=True)
|
||||
async def delete_feedback(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
run_id: str,
|
||||
feedback_id: str,
|
||||
request: Request,
|
||||
|
||||
@ -4,7 +4,7 @@ import logging
|
||||
import os
|
||||
import re
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
from typing import Any, Literal, NamedTuple
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request, status
|
||||
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
||||
@ -32,6 +32,317 @@ _MCP_STDIO_COMMAND_ALLOWLIST_ENV = "DEER_FLOW_MCP_STDIO_COMMAND_ALLOWLIST"
|
||||
_DEFAULT_MCP_STDIO_COMMAND_ALLOWLIST = frozenset({"npx", "uvx"})
|
||||
_SHELL_METACHARS = frozenset(";|&`$<>\n\r")
|
||||
|
||||
# Flags that turn an allowlisted launcher into an arbitrary code evaluator.
|
||||
# Validating only the command name leaves the allowlist naming a binary
|
||||
# without constraining what that binary runs, so these are screened too.
|
||||
# The spellings below mean "evaluate this string" across every launcher an
|
||||
# operator would plausibly allowlist (npx/uvx `--call`, python/sh `-c`,
|
||||
# node/perl/ruby `-e`/`--eval`, node `--print`), plus npx's pass-through into
|
||||
# node's own argv.
|
||||
#
|
||||
# This is defense in depth, not a trust boundary. `npx`/`uvx` exist to fetch
|
||||
# and run remote code, so an admin can still point one at a package they
|
||||
# published; the boundary remains admin authentication plus not exposing the
|
||||
# Gateway to untrusted networks.
|
||||
_ARBITRARY_EXEC_ARGS = frozenset(
|
||||
{
|
||||
"-c",
|
||||
"--call",
|
||||
"-e",
|
||||
"--eval",
|
||||
"--print",
|
||||
"--shell",
|
||||
"--node-arg",
|
||||
"--node-options",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
# Package launchers parse their own options only until the package name; every
|
||||
# later token is handed to the spawned server's own CLI, where `-c` is commonly
|
||||
# "config" and `-e` "env". Screening those rejected ordinary third-party servers
|
||||
# without covering anything, so the screen is scoped to the option region.
|
||||
#
|
||||
# Finding that region needs each launcher's option *arity*, because a value is
|
||||
# not a positional: `npx -p <pkg> -c '<command>'` runs the command -- `-p` is
|
||||
# exec's `--package`, so `<pkg>` is its value and npm keeps parsing its own
|
||||
# flags. Ending the region at the first non-flag token would walk past it.
|
||||
# (Verified against npm 10.9.4 / uv 0.11.1.)
|
||||
#
|
||||
# The two launchers get opposite defaults for an option neither table lists,
|
||||
# and the reason is the exec set above, not symmetry:
|
||||
#
|
||||
# `npx` really does own exec flags here (`-c`/`--call`), so an unlisted option
|
||||
# must not be able to hide one. Unknown therefore consumes a value, keeping
|
||||
# the region open. npm *errors* on an option it does not define, so this
|
||||
# cannot reject an invocation that would otherwise work; enumerating npm's
|
||||
# booleans (rather than its much larger value-taking set) is what makes the
|
||||
# common `npx -y <pkg> ...` shape land on the package name.
|
||||
#
|
||||
# `uvx` owns no exec flag at all -- uv has no "evaluate this string" option --
|
||||
# so its screen is a tripwire, not a control, and an imprecise region cannot
|
||||
# walk past anything real. Unknown therefore consumes nothing, which keeps
|
||||
# uv's large and growing boolean surface from over-blocking.
|
||||
#
|
||||
# A launcher outside this table is not a package runner and keeps the
|
||||
# conservative whole-args screen below.
|
||||
class _LauncherGrammar(NamedTuple):
|
||||
"""How one package launcher separates its own options from the server's."""
|
||||
|
||||
exec_args: frozenset[str]
|
||||
known_args: frozenset[str]
|
||||
unknown_consumes_value: bool
|
||||
|
||||
def consumes_value(self, flag: str) -> bool:
|
||||
if self.unknown_consumes_value:
|
||||
return flag not in self.known_args
|
||||
return flag in self.known_args
|
||||
|
||||
|
||||
# npm's boolean configs, i.e. the options that do *not* consume the next token.
|
||||
# Generated from `@npmcli/config`'s definitions (npm 10.9.4): every config whose
|
||||
# type is Boolean, plus every nopt shorthand expanding to one of them or to a
|
||||
# complete assignment such as `-d` -> `--loglevel info`. Regenerate against a
|
||||
# newer npm rather than editing by hand. A boolean missing here over-blocks one
|
||||
# invocation and names the flag in the rejection, which is the failure direction
|
||||
# this file prefers.
|
||||
_NPM_BOOLEAN_ARGS = frozenset(
|
||||
{
|
||||
"--all",
|
||||
"--allow-same-version",
|
||||
"--audit",
|
||||
"--bin-links",
|
||||
"--commit-hooks",
|
||||
"--description",
|
||||
"--dev",
|
||||
"--diff-ignore-all-space",
|
||||
"--diff-name-only",
|
||||
"--diff-no-prefix",
|
||||
"--diff-text",
|
||||
"--dry-run",
|
||||
"--engine-strict",
|
||||
"--expect-results",
|
||||
"--force",
|
||||
"--foreground-scripts",
|
||||
"--format-package-lock",
|
||||
"--fund",
|
||||
"--git-tag-version",
|
||||
"--global",
|
||||
"--global-style",
|
||||
"--if-present",
|
||||
"--ignore-scripts",
|
||||
"--include-staged",
|
||||
"--include-workspace-root",
|
||||
"--install-links",
|
||||
"--json",
|
||||
"--legacy-bundling",
|
||||
"--legacy-peer-deps",
|
||||
"--link",
|
||||
"--long",
|
||||
"--offline",
|
||||
"--omit-lockfile-registry-resolved",
|
||||
"--optional",
|
||||
"--package-lock",
|
||||
"--package-lock-only",
|
||||
"--parseable",
|
||||
"--prefer-dedupe",
|
||||
"--prefer-offline",
|
||||
"--prefer-online",
|
||||
"--production",
|
||||
"--progress",
|
||||
"--provenance",
|
||||
"--read-only",
|
||||
"--rebuild-bundle",
|
||||
"--save",
|
||||
"--save-bundle",
|
||||
"--save-dev",
|
||||
"--save-exact",
|
||||
"--save-optional",
|
||||
"--save-peer",
|
||||
"--save-prod",
|
||||
"--shrinkwrap",
|
||||
"--sign-git-commit",
|
||||
"--sign-git-tag",
|
||||
"--strict-peer-deps",
|
||||
"--strict-ssl",
|
||||
"--timing",
|
||||
"--unicode",
|
||||
"--update-notifier",
|
||||
"--usage",
|
||||
"--version",
|
||||
"--versions",
|
||||
"--workspaces",
|
||||
"--workspaces-update",
|
||||
"--yes",
|
||||
"-?",
|
||||
"-B",
|
||||
"-D",
|
||||
"-E",
|
||||
"-H",
|
||||
"-O",
|
||||
"-P",
|
||||
"-S",
|
||||
"-a",
|
||||
"-d",
|
||||
"-dd",
|
||||
"-ddd",
|
||||
"-desc",
|
||||
"-f",
|
||||
"-g",
|
||||
"-h",
|
||||
"-help",
|
||||
"-iwr",
|
||||
"-l",
|
||||
"-local",
|
||||
"-n",
|
||||
"-no",
|
||||
"-porcelain",
|
||||
"-q",
|
||||
"-quiet",
|
||||
"-readonly",
|
||||
"-s",
|
||||
"-silent",
|
||||
"-v",
|
||||
"-verbose",
|
||||
"-ws",
|
||||
"-y",
|
||||
}
|
||||
)
|
||||
|
||||
# `npm exec` overrides the global `-p` shorthand: it is `--package <spec>` there,
|
||||
# not the boolean `--parseable`. Confirmed by running it -- `npx -p . -c '<cmd>'`
|
||||
# executes the command, i.e. `.` was consumed as a value and never ended the
|
||||
# option region. Treating it as boolean is exactly the bypass this table exists
|
||||
# to prevent, so the override is applied explicitly rather than left implicit.
|
||||
_NPX_BOOLEAN_ARGS = _NPM_BOOLEAN_ARGS - {"-p"}
|
||||
|
||||
# uv's value-taking options (`uvx --help`, uv 0.11.1). Everything absent is
|
||||
# treated as boolean; see the unknown-option note above for why that default is
|
||||
# safe here and inverted for npx.
|
||||
_UVX_VALUE_ARGS = frozenset(
|
||||
{
|
||||
"--allow-insecure-host",
|
||||
"--build-constraints",
|
||||
"--cache-dir",
|
||||
"--color",
|
||||
"--config-file",
|
||||
"--config-setting",
|
||||
"--config-settings-package",
|
||||
"--constraints",
|
||||
"--default-index",
|
||||
"--directory",
|
||||
"--env-file",
|
||||
"--exclude-newer",
|
||||
"--exclude-newer-package",
|
||||
"--extra-index-url",
|
||||
"--find-links",
|
||||
"--fork-strategy",
|
||||
"--from",
|
||||
"--index",
|
||||
"--index-strategy",
|
||||
"--index-url",
|
||||
"--keyring-provider",
|
||||
"--link-mode",
|
||||
"--no-binary-package",
|
||||
"--no-build-isolation-package",
|
||||
"--no-build-package",
|
||||
"--no-sources-package",
|
||||
"--overrides",
|
||||
"--prerelease",
|
||||
"--project",
|
||||
"--python",
|
||||
"--python-platform",
|
||||
"--refresh-package",
|
||||
"--reinstall-package",
|
||||
"--resolution",
|
||||
"--torch-backend",
|
||||
"--upgrade-package",
|
||||
"--with",
|
||||
"--with-editable",
|
||||
"--with-requirements",
|
||||
"-C",
|
||||
"-P",
|
||||
"-b",
|
||||
"-c",
|
||||
"-f",
|
||||
"-i",
|
||||
"-p",
|
||||
"-w",
|
||||
}
|
||||
)
|
||||
|
||||
_PACKAGE_LAUNCHERS: dict[str, _LauncherGrammar] = {
|
||||
"npx": _LauncherGrammar(
|
||||
exec_args=_ARBITRARY_EXEC_ARGS,
|
||||
known_args=_NPX_BOOLEAN_ARGS,
|
||||
unknown_consumes_value=True,
|
||||
),
|
||||
# uv spells `-c` `--constraints` and `-p` `--python`, so the short forms are
|
||||
# dropped from its exec set; the long spellings stay as a tripwire in case a
|
||||
# future uv grows one. Derived so a new entry above cannot forget this.
|
||||
"uvx": _LauncherGrammar(
|
||||
exec_args=frozenset(flag for flag in _ARBITRARY_EXEC_ARGS if flag.startswith("--")),
|
||||
known_args=_UVX_VALUE_ARGS,
|
||||
unknown_consumes_value=False,
|
||||
),
|
||||
}
|
||||
|
||||
# `-p` is `--print` (evaluate and print) on node, so exempting it everywhere
|
||||
# left the short and long spellings of one flag disagreeing as soon as an
|
||||
# operator extended the allowlist. It stays scoped to commands outside
|
||||
# `_PACKAGE_LAUNCHERS`, where it is an ordinary selector (`--package` for npx,
|
||||
# `--python` for uv), so the default allowlist is unaffected.
|
||||
_EXEC_ARGS_OUTSIDE_PACKAGE_LAUNCHERS = frozenset({"-p"})
|
||||
|
||||
# Short options combine into one token (`node -pe`, `perl -we`, `python -Ic`),
|
||||
# which whole-token matching does not see. Derived rather than restated so a
|
||||
# new single-letter entry above cannot forget its clustered spelling.
|
||||
_CLUSTERED_EXEC_LETTERS = frozenset(flag[1] for flag in _ARBITRARY_EXEC_ARGS | _EXEC_ARGS_OUTSIDE_PACKAGE_LAUNCHERS if len(flag) == 2 and flag.startswith("-"))
|
||||
|
||||
# Environment variables that inject code into a process at startup, which is
|
||||
# the same bypass as an exec flag by another name.
|
||||
#
|
||||
# `PYTHONPATH` matters most: `site` imports `sitecustomize.py` from any
|
||||
# `sys.path` entry before the tool's entry point runs, so a caller-controlled
|
||||
# directory is code execution under `uvx` -- on the *default* allowlist.
|
||||
# `PYTHONSTARTUP` is inert for the non-interactive launchers in scope and is
|
||||
# kept only as belt-and-braces for an operator who allowlists a REPL.
|
||||
#
|
||||
# Known residual, accepted. Every entry below executes code *unconditionally*
|
||||
# at process startup. Caller-controlled *search paths* are a different, weaker
|
||||
# shape -- they reach code only if the process happens to load a name the
|
||||
# caller can shadow -- and they stay out:
|
||||
#
|
||||
# `LD_LIBRARY_PATH`/`DYLD_LIBRARY_PATH` run a shadowed library's constructor,
|
||||
# and native-dependency servers legitimately set them.
|
||||
#
|
||||
# `NODE_PATH` is narrower still, and not for the reason it first looks like.
|
||||
# Node searches it *after* the local `node_modules` chain -- the resolver
|
||||
# unshifts the requiring module's own paths ahead of it -- so it cannot
|
||||
# shadow an installed dependency, and ESM `import` ignores it entirely. It
|
||||
# can only supply a CJS module that would otherwise fail to resolve, i.e. an
|
||||
# optional `try { require(...) } catch {}` dependency absent from the install.
|
||||
#
|
||||
# Adding them would make the "unconditional" rule above untrue, and a
|
||||
# defense-in-depth list that grows because each entry was cheap is how it ends
|
||||
# up mistaken for a boundary. A denylist is not what makes MCP registration
|
||||
# safe for an untrusted admin anyway.
|
||||
_CODE_INJECTING_ENV_VARS = frozenset(
|
||||
{
|
||||
"BASH_ENV",
|
||||
"DYLD_INSERT_LIBRARIES",
|
||||
"ENV",
|
||||
"LD_AUDIT",
|
||||
"LD_PRELOAD",
|
||||
"NODE_OPTIONS",
|
||||
"PERL5OPT",
|
||||
"PYTHONHOME",
|
||||
"PYTHONPATH",
|
||||
"PYTHONSTARTUP",
|
||||
"RUBYOPT",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class McpOAuthConfigResponse(BaseModel):
|
||||
"""OAuth configuration for an MCP server."""
|
||||
@ -66,6 +377,7 @@ class McpServerConfigResponse(BaseModel):
|
||||
description: str = Field(default="", description="Human-readable description of what this MCP server provides")
|
||||
routing: McpRoutingConfig = Field(default_factory=McpRoutingConfig, description="Soft routing hints for tools from this MCP server")
|
||||
tools: dict[str, McpToolOverride] = Field(default_factory=dict, description="Per-original-tool MCP configuration overrides")
|
||||
tool_name_prefix: bool = Field(default=True, description="Whether to prefix discovered tool names with the MCP server name")
|
||||
tool_call_timeout: float | None = Field(default=None, description="Timeout in seconds for individual stdio MCP tool calls")
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
@ -193,12 +505,91 @@ def _stdio_command_name(command: str | None, *, server_name: str) -> str:
|
||||
return stripped
|
||||
|
||||
|
||||
def _launcher_option_region(args: list[str], *, grammar: _LauncherGrammar) -> list[str]:
|
||||
"""Return the leading args a package launcher parses as its own options.
|
||||
|
||||
The region ends at a bare ``--`` or at the package name -- the first token
|
||||
that is neither a flag nor the value of one. A ``--flag=value`` token
|
||||
carries its own value and never consumes the next one.
|
||||
|
||||
Arity is looked up case-sensitively, because a launcher's short options are:
|
||||
npm reads ``-c`` as ``--call`` but ``-C`` as ``--prefix``, which takes a
|
||||
value.
|
||||
"""
|
||||
region: list[str] = []
|
||||
index = 0
|
||||
while index < len(args):
|
||||
arg = args[index]
|
||||
if not isinstance(arg, str):
|
||||
break
|
||||
token = arg.strip()
|
||||
if token == "--" or token == "-" or not token.startswith("-"):
|
||||
break
|
||||
region.append(token)
|
||||
index += 1
|
||||
if "=" not in token and grammar.consumes_value(token):
|
||||
index += 1
|
||||
return region
|
||||
|
||||
|
||||
def _arbitrary_exec_arg(args: list[str], *, command: str) -> str | None:
|
||||
"""Return the offending flag when an argument makes the launcher eval a string.
|
||||
|
||||
Handles both ``--call value`` and ``--call=value`` spellings.
|
||||
|
||||
For a package launcher (:data:`_PACKAGE_LAUNCHERS`) only the launcher's own
|
||||
option region is screened, because everything from the package name onward
|
||||
is the spawned server's argv -- ``npx -y <pkg> -c config.json`` hands
|
||||
``-c config.json`` to the server, where it is "config", not eval. A bare
|
||||
``--`` ends the region too: only the *first* token after it is the package
|
||||
name, and the rest are that package's arguments.
|
||||
|
||||
Every other command is screened whole, and two extra rules apply because
|
||||
such a command is an interpreter rather than a package runner: ``-p`` is an
|
||||
exec flag (node's ``--print``) instead of a package/python selector, and
|
||||
combined short-option clusters are decomposed so ``-pe`` cannot smuggle
|
||||
past a check that only splits on ``=``.
|
||||
|
||||
Only the normalized flag is returned, never the caller's value, so the
|
||||
rejection message does not echo a payload string back into the response.
|
||||
"""
|
||||
grammar = _PACKAGE_LAUNCHERS.get(command.lower())
|
||||
if grammar is not None:
|
||||
for token in _launcher_option_region(args, grammar=grammar):
|
||||
flag = token.split("=", 1)[0]
|
||||
# Long options are matched case-insensitively as before; a short one
|
||||
# is not, because its case selects a different option -- npm's `-C`
|
||||
# is `--prefix`, and folding it onto `-c` rejected an ordinary flag.
|
||||
flag = flag.lower() if flag.startswith("--") else flag
|
||||
if flag in grammar.exec_args:
|
||||
return flag
|
||||
return None
|
||||
|
||||
denied = _ARBITRARY_EXEC_ARGS | _EXEC_ARGS_OUTSIDE_PACKAGE_LAUNCHERS
|
||||
for arg in args:
|
||||
if not isinstance(arg, str):
|
||||
continue
|
||||
flag = arg.split("=", 1)[0].strip().lower()
|
||||
if flag in denied:
|
||||
return flag
|
||||
if not flag.startswith("-") or flag.startswith("--"):
|
||||
continue
|
||||
for letter in flag[1:]:
|
||||
if letter in _CLUSTERED_EXEC_LETTERS:
|
||||
return f"-{letter}"
|
||||
return None
|
||||
|
||||
|
||||
def _validate_mcp_update_request(request: McpConfigUpdateRequest) -> None:
|
||||
"""Validate API-submitted MCP config before it is persisted.
|
||||
|
||||
Local config files can still express arbitrary advanced setups, but the
|
||||
HTTP API is an untrusted boundary. Restricting stdio commands here reduces
|
||||
the blast radius of a compromised authenticated browser session.
|
||||
|
||||
The command name alone is not a meaningful restriction, so the launcher's
|
||||
``args`` and ``env`` are screened for the flags and variables that turn an
|
||||
allowlisted binary into an arbitrary code evaluator.
|
||||
"""
|
||||
allowed_commands = _allowed_stdio_commands()
|
||||
for name, server in request.mcp_servers.items():
|
||||
@ -214,6 +605,20 @@ def _validate_mcp_update_request(request: McpConfigUpdateRequest) -> None:
|
||||
detail=(f"MCP server '{name}' uses disallowed stdio command '{command_name}'. Allowed commands: {allowed}. Configure {_MCP_STDIO_COMMAND_ALLOWLIST_ENV} to extend this list."),
|
||||
)
|
||||
|
||||
exec_flag = _arbitrary_exec_arg(server.args, command=command_name)
|
||||
if exec_flag is not None:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=(f"MCP server '{name}' passes '{exec_flag}' to '{command_name}', which would run arbitrary code. Point the server at a package or module instead."),
|
||||
)
|
||||
|
||||
for env_name in server.env:
|
||||
if env_name.strip().upper() in _CODE_INJECTING_ENV_VARS:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=(f"MCP server '{name}' sets environment variable '{env_name}', which would run arbitrary code at process startup."),
|
||||
)
|
||||
|
||||
|
||||
def _mask_server_config(server: McpServerConfigResponse) -> McpServerConfigResponse:
|
||||
"""Return a copy of server config with sensitive fields masked.
|
||||
|
||||
@ -116,6 +116,8 @@ def _map_memory_fact_value_error(exc: ValueError) -> HTTPException:
|
||||
detail = "Invalid confidence value; must be between 0 and 1."
|
||||
elif exc.args and exc.args[0] == "agent_name":
|
||||
detail = "An agent name is required for fact operations; user-global memory stores summaries only."
|
||||
elif exc.args and exc.args[0] == "Duplicate fact":
|
||||
return HTTPException(status_code=409, detail="A fact with the same content already exists.")
|
||||
else:
|
||||
detail = "Memory fact content cannot be empty."
|
||||
return HTTPException(status_code=400, detail=detail)
|
||||
|
||||
@ -1,9 +1,19 @@
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
import logging
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.gateway.deps import get_config
|
||||
from app.gateway.authz import (
|
||||
_AuthorizationUnavailable,
|
||||
_is_internal_caller,
|
||||
resolve_model_authorization,
|
||||
)
|
||||
from app.gateway.deps import get_config, get_optional_user_from_request
|
||||
from deerflow.authz.provider import AuthzDecision, AuthzRequest
|
||||
from deerflow.config.app_config import AppConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["models"])
|
||||
|
||||
|
||||
@ -37,12 +47,19 @@ class ModelsListResponse(BaseModel):
|
||||
summary="List All Models",
|
||||
description="Retrieve a list of all available AI models configured in the system.",
|
||||
)
|
||||
async def list_models(config: AppConfig = Depends(get_config)) -> ModelsListResponse:
|
||||
async def list_models(
|
||||
request: Request,
|
||||
config: AppConfig = Depends(get_config),
|
||||
) -> ModelsListResponse:
|
||||
"""List all available models from configuration.
|
||||
|
||||
Returns model information suitable for frontend display,
|
||||
excluding sensitive fields like API keys and internal configuration.
|
||||
|
||||
When ``authorization.enabled`` is true, only models the caller's role may
|
||||
``list`` are returned (filtered via ``provider.filter_resources``). A
|
||||
provider error yields an empty list (fail-closed) or all models (fail-open).
|
||||
|
||||
Returns:
|
||||
A list of all configured models with their metadata and token usage display settings.
|
||||
|
||||
@ -73,6 +90,28 @@ async def list_models(config: AppConfig = Depends(get_config)) -> ModelsListResp
|
||||
}
|
||||
```
|
||||
"""
|
||||
visible_models = config.models
|
||||
fail_closed = config.authorization.fail_closed
|
||||
|
||||
user = await get_optional_user_from_request(request)
|
||||
if user is not None:
|
||||
try:
|
||||
provider, principal = resolve_model_authorization(user, is_internal=_is_internal_caller(request, user))
|
||||
except _AuthorizationUnavailable as exc:
|
||||
if exc.fail_closed:
|
||||
visible_models = []
|
||||
else:
|
||||
if provider is not None and principal is not None:
|
||||
try:
|
||||
allowed_names = provider.filter_resources(principal, "model", [m.name for m in config.models])
|
||||
if not isinstance(allowed_names, list) or any(not isinstance(n, str) for n in allowed_names):
|
||||
raise TypeError("AuthorizationProvider.filter_resources must return list[str]")
|
||||
allowed_set = set(allowed_names)
|
||||
visible_models = [m for m in config.models if m.name in allowed_set]
|
||||
except Exception:
|
||||
logger.warning("Authorization provider failed while filtering models", exc_info=True)
|
||||
visible_models = [] if fail_closed else config.models
|
||||
|
||||
models = [
|
||||
ModelResponse(
|
||||
name=model.name,
|
||||
@ -82,7 +121,7 @@ async def list_models(config: AppConfig = Depends(get_config)) -> ModelsListResp
|
||||
supports_thinking=model.supports_thinking,
|
||||
supports_reasoning_effort=model.supports_reasoning_effort,
|
||||
)
|
||||
for model in config.models
|
||||
for model in visible_models
|
||||
]
|
||||
return ModelsListResponse(
|
||||
models=models,
|
||||
@ -96,7 +135,11 @@ async def list_models(config: AppConfig = Depends(get_config)) -> ModelsListResp
|
||||
summary="Get Model Details",
|
||||
description="Retrieve detailed information about a specific AI model by its name.",
|
||||
)
|
||||
async def get_model(model_name: str, config: AppConfig = Depends(get_config)) -> ModelResponse:
|
||||
async def get_model(
|
||||
model_name: str,
|
||||
request: Request,
|
||||
config: AppConfig = Depends(get_config),
|
||||
) -> ModelResponse:
|
||||
"""Get a specific model by name.
|
||||
|
||||
Args:
|
||||
@ -106,7 +149,10 @@ async def get_model(model_name: str, config: AppConfig = Depends(get_config)) ->
|
||||
Model information if found.
|
||||
|
||||
Raises:
|
||||
HTTPException: 404 if model not found.
|
||||
HTTPException: 404 if model not found; 403 if the caller's role may not
|
||||
``use`` the model (only when ``authorization.enabled`` is true). A
|
||||
provider resolution error yields 403 (fail-closed) or allows the request
|
||||
(fail-open), mirroring ``list_models``'s provider-error semantics.
|
||||
|
||||
Example Response:
|
||||
```json
|
||||
@ -122,6 +168,33 @@ async def get_model(model_name: str, config: AppConfig = Depends(get_config)) ->
|
||||
if model is None:
|
||||
raise HTTPException(status_code=404, detail=f"Model '{model_name}' not found")
|
||||
|
||||
# Phase 3: enforce model:use authorization (deny → 403, not 404, since the
|
||||
# model exists but the role lacks permission to use it).
|
||||
fail_closed = config.authorization.fail_closed
|
||||
user = await get_optional_user_from_request(request)
|
||||
if user is not None:
|
||||
try:
|
||||
provider, principal = resolve_model_authorization(user, is_internal=_is_internal_caller(request, user))
|
||||
except _AuthorizationUnavailable:
|
||||
if fail_closed:
|
||||
raise HTTPException(status_code=403, detail=f"Model '{model_name}' is not available for your role")
|
||||
else:
|
||||
if provider is not None and principal is not None:
|
||||
try:
|
||||
decision = provider.authorize(AuthzRequest(principal=principal, resource="model", action="use", target=model_name))
|
||||
if not isinstance(decision, AuthzDecision):
|
||||
raise TypeError("AuthorizationProvider.authorize must return AuthzDecision")
|
||||
allowed = decision.allow
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Authorization provider failed while checking model:use for %s",
|
||||
model_name,
|
||||
exc_info=True,
|
||||
)
|
||||
allowed = not fail_closed
|
||||
if not allowed:
|
||||
raise HTTPException(status_code=403, detail=f"Model '{model_name}' is not available for your role")
|
||||
|
||||
return ModelResponse(
|
||||
name=model.name,
|
||||
model=model.model,
|
||||
|
||||
@ -8,7 +8,6 @@ is reused so that conversation history is preserved across calls.
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import uuid
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Query, Request
|
||||
from fastapi.responses import StreamingResponse
|
||||
@ -19,6 +18,7 @@ from app.gateway.pagination import trim_run_message_page
|
||||
from app.gateway.run_models import RunCreateRequest
|
||||
from app.gateway.services import build_checkpoint_state_accessor, sse_consumer, start_run, wait_for_run_completion
|
||||
from deerflow.runtime import serialize_channel_values_for_api
|
||||
from deerflow.utils.thread_id import resolve_thread_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
router = APIRouter(prefix="/api/runs", tags=["runs"])
|
||||
@ -27,9 +27,7 @@ router = APIRouter(prefix="/api/runs", tags=["runs"])
|
||||
def _resolve_thread_id(body: RunCreateRequest) -> str:
|
||||
"""Return the thread_id from the request body, or generate a new one."""
|
||||
thread_id = ((body.config or {}).get("configurable") or {}).get("thread_id")
|
||||
if thread_id:
|
||||
return str(thread_id)
|
||||
return str(uuid.uuid4())
|
||||
return resolve_thread_id(thread_id)
|
||||
|
||||
|
||||
@router.post("/stream")
|
||||
|
||||
@ -31,6 +31,7 @@ from pydantic import BaseModel, Field, PlainSerializer
|
||||
|
||||
from deerflow.domain.schedule.commands import ContextChange, CreateScheduledTask, UpdateScheduledTask
|
||||
from deerflow.domain.schedule.model import ScheduledRun, ScheduledTask, ScheduleSpec, ScheduleType
|
||||
from deerflow.utils.thread_id import ThreadId
|
||||
|
||||
UtcTimestamp = Annotated[datetime, PlainSerializer(lambda value: value.isoformat(), return_type=str)]
|
||||
|
||||
@ -68,7 +69,7 @@ class CreateScheduledTaskRequest(BaseModel):
|
||||
through ``to_command``, never read off the wire.
|
||||
"""
|
||||
|
||||
thread_id: str | None = None
|
||||
thread_id: ThreadId | None = None
|
||||
context_mode: str = "fresh_thread_per_run"
|
||||
title: str = Field(min_length=1)
|
||||
prompt: str = Field(min_length=1)
|
||||
@ -115,7 +116,7 @@ class UpdateScheduledTaskRequest(BaseModel):
|
||||
"""
|
||||
|
||||
context_mode: str | None = None
|
||||
thread_id: str | None = None
|
||||
thread_id: ThreadId | None = None
|
||||
title: str | None = Field(default=None, min_length=1)
|
||||
prompt: str | None = Field(default=None, min_length=1)
|
||||
schedule_spec: dict[str, Any] | None = None
|
||||
|
||||
@ -41,6 +41,7 @@ from deerflow.domain.schedule.exceptions import (
|
||||
ThreadNotFoundError,
|
||||
)
|
||||
from deerflow.domain.schedule.model import DispatchOutcome
|
||||
from deerflow.utils.thread_id import ThreadId
|
||||
|
||||
router = APIRouter(prefix="/api", tags=["scheduled-tasks"])
|
||||
|
||||
@ -194,7 +195,7 @@ async def list_scheduled_task_runs(
|
||||
@router.get("/threads/{thread_id}/scheduled-tasks", response_model=list[ScheduledTaskResponse])
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
@_map_domain_errors
|
||||
async def list_thread_scheduled_tasks(thread_id: str, request: Request, service: ScheduleServiceDep):
|
||||
async def list_thread_scheduled_tasks(thread_id: ThreadId, request: Request, service: ScheduleServiceDep):
|
||||
user_id = await _require_user_id(request)
|
||||
tasks = await service.list_tasks_by_thread(user_id, thread_id)
|
||||
return [ScheduledTaskResponse.from_domain(task) for task in tasks]
|
||||
|
||||
@ -31,6 +31,7 @@ from deerflow.skills.security_static_scanner import (
|
||||
)
|
||||
from deerflow.skills.storage import SkillStorage, get_or_new_user_skill_storage
|
||||
from deerflow.skills.types import SKILL_MD_FILE, SkillCategory
|
||||
from deerflow.utils.thread_id import ThreadId
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -65,7 +66,7 @@ class SkillUpdateRequest(BaseModel):
|
||||
class SkillInstallRequest(BaseModel):
|
||||
"""Request model for installing a skill from a .skill file."""
|
||||
|
||||
thread_id: str = Field(..., description="The thread ID where the .skill file is located")
|
||||
thread_id: ThreadId = Field(..., description="The thread ID where the .skill file is located")
|
||||
path: str = Field(..., description="Virtual path to the .skill file (e.g., mnt/user-data/outputs/my-skill.skill)")
|
||||
|
||||
|
||||
@ -273,8 +274,9 @@ async def update_custom_skill(skill_name: str, body: CustomSkillUpdateRequest, r
|
||||
if scan.decision == "block":
|
||||
raise HTTPException(status_code=400, detail=f"Security scan blocked the edit: {scan.reason}")
|
||||
prev_content = storage.read_custom_skill(skill_name)
|
||||
storage.write_custom_skill(skill_name, SKILL_MD_FILE, body.content)
|
||||
storage.append_history(
|
||||
await asyncio.to_thread(storage.write_custom_skill, skill_name, SKILL_MD_FILE, body.content)
|
||||
await asyncio.to_thread(
|
||||
storage.append_history,
|
||||
skill_name,
|
||||
{
|
||||
"action": "human_edit",
|
||||
@ -305,7 +307,8 @@ async def delete_custom_skill(skill_name: str, request: Request, config: AppConf
|
||||
try:
|
||||
skill_name = skill_name.replace("\r\n", "").replace("\n", "")
|
||||
storage = _get_user_skill_storage(config)
|
||||
storage.delete_custom_skill(
|
||||
await asyncio.to_thread(
|
||||
storage.delete_custom_skill,
|
||||
skill_name,
|
||||
history_meta={
|
||||
"action": "human_delete",
|
||||
@ -384,10 +387,10 @@ async def rollback_custom_skill(skill_name: str, body: SkillRollbackRequest, req
|
||||
"scanner": {"decision": scan.decision, "reason": scan.reason, "static_findings": static_findings},
|
||||
}
|
||||
if scan.decision == "block":
|
||||
storage.append_history(skill_name, history_entry)
|
||||
await asyncio.to_thread(storage.append_history, skill_name, history_entry)
|
||||
raise HTTPException(status_code=400, detail=f"Rollback blocked by security scanner: {scan.reason}")
|
||||
storage.write_custom_skill(skill_name, SKILL_MD_FILE, target_content)
|
||||
storage.append_history(skill_name, history_entry)
|
||||
await asyncio.to_thread(storage.write_custom_skill, skill_name, SKILL_MD_FILE, target_content)
|
||||
await asyncio.to_thread(storage.append_history, skill_name, history_entry)
|
||||
await refresh_user_skills_system_prompt_cache_async(get_effective_user_id())
|
||||
return await _read_custom_skill_response(skill_name, config)
|
||||
except HTTPException:
|
||||
@ -426,35 +429,47 @@ async def get_skill(skill_name: str, config: AppConfig = Depends(get_config)) ->
|
||||
raise HTTPException(status_code=500, detail=f"Failed to get skill: {str(e)}")
|
||||
|
||||
|
||||
def _write_extensions_skill_state(skill_name: str, enabled: bool) -> None:
|
||||
def _write_extensions_skill_state(
|
||||
storage: SkillStorage,
|
||||
skill_name: str,
|
||||
enabled: bool,
|
||||
*,
|
||||
rebuild_public_projection: bool,
|
||||
) -> None:
|
||||
"""Read-modify-write a skill's enabled state in the shared extensions_config.json.
|
||||
|
||||
Blocking filesystem IO: always call this via ``asyncio.to_thread``. It takes
|
||||
``extensions_config_write_lock`` itself, so that this router and the MCP
|
||||
router (which performs the same RMW on the same file) cannot interleave and
|
||||
drop each other's change. The lock is held by the worker rather than by the
|
||||
awaiting task, so cancelling the request cannot release it mid-write.
|
||||
the public projection lock before ``extensions_config_write_lock``. The first
|
||||
keeps the enabled-only view synchronized across workers; the second prevents
|
||||
this router and the MCP router from interleaving writes to the shared file.
|
||||
Both locks are held by the worker, so request cancellation cannot release
|
||||
either lock while the write or projection rebuild is still running.
|
||||
"""
|
||||
with extensions_config_write_lock:
|
||||
config_path = ExtensionsConfig.resolve_config_path()
|
||||
if config_path is None:
|
||||
config_path = Path.cwd().parent / "extensions_config.json"
|
||||
logger.info(f"No existing extensions config found. Creating new config at: {config_path}")
|
||||
from contextlib import nullcontext
|
||||
|
||||
# Work on a deep copy rather than the cached singleton: mutating the
|
||||
# singleton in place would publish the new state to readers before it is
|
||||
# durable on disk, and leave it applied even if the write below fails.
|
||||
# to_file_dict() serializes the full extensions_config.json shape (all
|
||||
# top-level keys), so no field is dropped from the file.
|
||||
extensions_config = get_extensions_config().model_copy(deep=True)
|
||||
extensions_config.skills[skill_name] = SkillStateConfig(enabled=enabled)
|
||||
from deerflow.skills.projection import skill_projection_mutation
|
||||
from deerflow.skills.storage.local_skill_storage import LocalSkillStorage
|
||||
|
||||
config_data = extensions_config.to_file_dict()
|
||||
removal_names = (skill_name,) if not enabled else ()
|
||||
projection_update = skill_projection_mutation(storage, "public", remove_names=removal_names) if rebuild_public_projection and isinstance(storage, LocalSkillStorage) else nullcontext()
|
||||
with projection_update:
|
||||
with extensions_config_write_lock:
|
||||
config_path = ExtensionsConfig.resolve_config_path()
|
||||
if config_path is None:
|
||||
config_path = Path.cwd().parent / "extensions_config.json"
|
||||
logger.info(f"No existing extensions config found. Creating new config at: {config_path}")
|
||||
|
||||
atomic_write_extensions_config(config_path, config_data)
|
||||
# The projection lock is cross-process, but the singleton cache is
|
||||
# not. Existing files are therefore re-read under the lock; a new
|
||||
# file starts from a deep snapshot of the cached defaults.
|
||||
extensions_config = ExtensionsConfig.from_file(config_path) if config_path.exists() else get_extensions_config().model_copy(deep=True)
|
||||
extensions_config.skills[skill_name] = SkillStateConfig(enabled=enabled)
|
||||
|
||||
logger.info(f"Skills configuration updated and saved to: {config_path}")
|
||||
reload_extensions_config()
|
||||
config_data = extensions_config.to_file_dict()
|
||||
atomic_write_extensions_config(config_path, config_data)
|
||||
|
||||
logger.info(f"Skills configuration updated and saved to: {config_path}")
|
||||
reload_extensions_config()
|
||||
|
||||
|
||||
@router.put(
|
||||
@ -488,10 +503,13 @@ async def update_skill(skill_name: str, body: SkillUpdateRequest, request: Reque
|
||||
# CUSTOM / LEGACY skills → per-user _skill_states.json (isolated state)
|
||||
# so that two users with same-named custom skills can toggle independently.
|
||||
if skill.category == SkillCategory.PUBLIC:
|
||||
# Shared-file RMW. The worker takes extensions_config_write_lock for
|
||||
# the whole read→write window, so it stays serialized against the MCP
|
||||
# router even if this request is cancelled mid-write.
|
||||
await asyncio.to_thread(_write_extensions_skill_state, skill_name, body.enabled)
|
||||
await asyncio.to_thread(
|
||||
_write_extensions_skill_state,
|
||||
storage,
|
||||
skill_name,
|
||||
body.enabled,
|
||||
rebuild_public_projection=True,
|
||||
)
|
||||
else:
|
||||
# CUSTOM / LEGACY: write per-user state
|
||||
from deerflow.skills.storage.user_scoped_skill_storage import UserScopedSkillStorage
|
||||
@ -500,8 +518,15 @@ async def update_skill(skill_name: str, body: SkillUpdateRequest, request: Reque
|
||||
await asyncio.to_thread(storage.set_skill_enabled_state, skill_name, body.enabled)
|
||||
else:
|
||||
# Fallback for non-user-scoped storage (unlikely in practice):
|
||||
# same shared-file RMW as the PUBLIC branch, same lock.
|
||||
await asyncio.to_thread(_write_extensions_skill_state, skill_name, body.enabled)
|
||||
# same shared-file RMW as the PUBLIC branch, without a public
|
||||
# projection rebuild for this non-public skill.
|
||||
await asyncio.to_thread(
|
||||
_write_extensions_skill_state,
|
||||
storage,
|
||||
skill_name,
|
||||
body.enabled,
|
||||
rebuild_public_projection=False,
|
||||
)
|
||||
|
||||
# PUBLIC skill enabled state lives in the global extensions_config.json
|
||||
# and affects every user, so the prompt cache for ALL users must be
|
||||
|
||||
@ -10,6 +10,7 @@ from app.gateway.deps import get_config
|
||||
from deerflow.config.app_config import AppConfig
|
||||
from deerflow.config.suggestions_config import DEFAULT_MAX_SUGGESTIONS, MAX_SUGGESTIONS_LIMIT
|
||||
from deerflow.utils.oneshot_llm import run_oneshot_llm
|
||||
from deerflow.utils.thread_id import ThreadId
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -102,7 +103,7 @@ async def get_suggestions_config(
|
||||
)
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
async def generate_suggestions(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
body: SuggestionsRequest,
|
||||
request: Request,
|
||||
config: AppConfig = Depends(get_config),
|
||||
|
||||
@ -33,6 +33,7 @@ from app.gateway.checkpoint_lineage import (
|
||||
find_checkpoint_before_message_chronologically,
|
||||
is_duration_only_checkpoint,
|
||||
)
|
||||
from app.gateway.context_usage import build_context_usage
|
||||
from app.gateway.deps import get_current_user, get_feedback_repo, get_run_event_store, get_run_manager, get_run_store, get_stream_bridge
|
||||
from app.gateway.pagination import trim_run_message_page
|
||||
from app.gateway.run_models import RunCreateRequest
|
||||
@ -42,6 +43,7 @@ from deerflow.agents.middlewares.dynamic_context_middleware import strip_injecte
|
||||
from deerflow.runtime import CancelOutcome, RunRecord, RunStatus, serialize_channel_values_for_api
|
||||
from deerflow.runtime.secret_context import redact_config_secrets, redact_metadata_secrets
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY, get_original_user_content_text, message_to_text
|
||||
from deerflow.utils.thread_id import ThreadId
|
||||
from deerflow.workspace_changes import get_workspace_changes_response
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -180,6 +182,12 @@ class ThreadTokenUsageCallerBreakdown(BaseModel):
|
||||
middleware: int = 0
|
||||
|
||||
|
||||
class ThreadContextUsage(BaseModel):
|
||||
token_count: int = 0
|
||||
max_context_tokens: int | None = None
|
||||
percentage: float | None = None
|
||||
|
||||
|
||||
class ThreadTokenUsageResponse(BaseModel):
|
||||
thread_id: str
|
||||
total_tokens: int = 0
|
||||
@ -188,6 +196,7 @@ class ThreadTokenUsageResponse(BaseModel):
|
||||
total_runs: int = 0
|
||||
by_model: dict[str, ThreadTokenUsageModelBreakdown] = Field(default_factory=dict)
|
||||
by_caller: ThreadTokenUsageCallerBreakdown = Field(default_factory=ThreadTokenUsageCallerBreakdown)
|
||||
context_usage: ThreadContextUsage | None = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@ -805,7 +814,7 @@ async def _default_history_hidden_run_ids(run_mgr: Any, thread_id: str, *, user_
|
||||
@router.post("/{thread_id}/runs/regenerate/prepare", response_model=RegeneratePrepareResponse)
|
||||
@require_permission("runs", "create", owner_check=True, require_existing=True)
|
||||
async def prepare_regenerate_run(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
body: RegeneratePrepareRequest,
|
||||
request: Request,
|
||||
) -> RegeneratePrepareResponse:
|
||||
@ -816,7 +825,7 @@ async def prepare_regenerate_run(
|
||||
@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,
|
||||
thread_id: ThreadId,
|
||||
body: EditRegeneratePrepareRequest,
|
||||
request: Request,
|
||||
) -> EditRegeneratePrepareResponse:
|
||||
@ -826,7 +835,7 @@ async def prepare_edit_regenerate_run(
|
||||
|
||||
@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:
|
||||
async def create_run(thread_id: ThreadId, body: RunCreateRequest, request: Request) -> RunResponse:
|
||||
"""Create a background run (returns immediately)."""
|
||||
record = await start_run(body, thread_id, request)
|
||||
return _record_to_response(record)
|
||||
@ -834,7 +843,7 @@ async def create_run(thread_id: str, body: RunCreateRequest, request: Request) -
|
||||
|
||||
@router.post("/{thread_id}/runs/stream")
|
||||
@require_permission("runs", "create", owner_check=True, require_existing=True)
|
||||
async def stream_run(thread_id: str, body: RunCreateRequest, request: Request) -> StreamingResponse:
|
||||
async def stream_run(thread_id: ThreadId, body: RunCreateRequest, request: Request) -> StreamingResponse:
|
||||
"""Create a run and stream events via SSE.
|
||||
|
||||
The response includes a ``Content-Location`` header with the run's
|
||||
@ -862,7 +871,7 @@ async def stream_run(thread_id: str, body: RunCreateRequest, request: Request) -
|
||||
|
||||
@router.post("/{thread_id}/runs/wait", response_model=dict)
|
||||
@require_permission("runs", "create", owner_check=True, require_existing=True)
|
||||
async def wait_run(thread_id: str, body: RunCreateRequest, request: Request) -> dict:
|
||||
async def wait_run(thread_id: ThreadId, body: RunCreateRequest, request: Request) -> dict:
|
||||
"""Create a run and block until it completes, returning the final state."""
|
||||
bridge = get_stream_bridge(request)
|
||||
run_mgr = get_run_manager(request)
|
||||
@ -891,7 +900,7 @@ async def wait_run(thread_id: str, body: RunCreateRequest, request: Request) ->
|
||||
|
||||
@router.get("/{thread_id}/runs", response_model=list[RunResponse])
|
||||
@require_permission("runs", "read", owner_check=True)
|
||||
async def list_runs(thread_id: str, request: Request) -> list[RunResponse]:
|
||||
async def list_runs(thread_id: ThreadId, request: Request) -> list[RunResponse]:
|
||||
"""List all runs for a thread."""
|
||||
run_mgr = get_run_manager(request)
|
||||
user_id = await get_current_user(request)
|
||||
@ -901,7 +910,7 @@ async def list_runs(thread_id: str, request: Request) -> list[RunResponse]:
|
||||
|
||||
@router.get("/{thread_id}/runs/{run_id}", response_model=RunResponse)
|
||||
@require_permission("runs", "read", owner_check=True)
|
||||
async def get_run(thread_id: str, run_id: str, request: Request) -> RunResponse:
|
||||
async def get_run(thread_id: ThreadId, run_id: str, request: Request) -> RunResponse:
|
||||
"""Get details of a specific run."""
|
||||
run_mgr = get_run_manager(request)
|
||||
user_id = await get_current_user(request)
|
||||
@ -914,7 +923,7 @@ async def get_run(thread_id: str, run_id: str, request: Request) -> RunResponse:
|
||||
@router.post("/{thread_id}/runs/{run_id}/cancel")
|
||||
@require_permission("runs", "cancel", owner_check=True, require_existing=True)
|
||||
async def cancel_run(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
run_id: str,
|
||||
request: Request,
|
||||
wait: bool = Query(default=False, description="Block until run completes after cancel"),
|
||||
@ -973,7 +982,7 @@ async def cancel_run(
|
||||
|
||||
@router.get("/{thread_id}/runs/{run_id}/join")
|
||||
@require_permission("runs", "read", owner_check=True)
|
||||
async def join_run(thread_id: str, run_id: str, request: Request) -> StreamingResponse:
|
||||
async def join_run(thread_id: ThreadId, run_id: str, request: Request) -> StreamingResponse:
|
||||
"""Join an existing run's SSE stream."""
|
||||
run_mgr = get_run_manager(request)
|
||||
record = await run_mgr.get(run_id)
|
||||
@ -1002,7 +1011,7 @@ async def join_run(thread_id: str, run_id: str, request: Request) -> StreamingRe
|
||||
@router.post("/{thread_id}/runs/{run_id}/stream", response_model=None)
|
||||
@require_permission("runs", "read", owner_check=True)
|
||||
async def stream_existing_run(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
run_id: str,
|
||||
request: Request,
|
||||
action: Literal["interrupt", "rollback"] | None = Query(default=None, description="Cancel action"),
|
||||
@ -1075,7 +1084,7 @@ async def stream_existing_run(
|
||||
@router.get("/{thread_id}/messages")
|
||||
@require_permission("runs", "read", owner_check=True)
|
||||
async def list_thread_messages(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
request: Request,
|
||||
limit: int = Query(default=50, ge=1, le=200),
|
||||
before_seq: int | None = Query(default=None, ge=1),
|
||||
@ -1329,7 +1338,7 @@ async def _enrich_thread_message_page(
|
||||
@router.get("/{thread_id}/messages/page", response_model=ThreadMessagesPageResponse)
|
||||
@require_permission("runs", "read", owner_check=True)
|
||||
async def list_thread_messages_page(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
request: Request,
|
||||
limit: int = Query(default=50, ge=1, le=200),
|
||||
before_seq: int | None = Query(default=None, ge=1),
|
||||
@ -1357,7 +1366,7 @@ async def list_thread_messages_page(
|
||||
@router.get("/{thread_id}/runs/{run_id}/messages")
|
||||
@require_permission("runs", "read", owner_check=True)
|
||||
async def list_run_messages(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
run_id: str,
|
||||
request: Request,
|
||||
limit: int = Query(default=50, le=200, ge=1),
|
||||
@ -1392,7 +1401,7 @@ async def list_run_messages(
|
||||
@router.get("/{thread_id}/runs/{run_id}/events")
|
||||
@require_permission("runs", "read", owner_check=True)
|
||||
async def list_run_events(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
run_id: str,
|
||||
request: Request,
|
||||
event_types: str | None = Query(default=None),
|
||||
@ -1429,7 +1438,7 @@ async def list_run_events(
|
||||
@router.get("/{thread_id}/runs/{run_id}/workspace-changes")
|
||||
@require_permission("runs", "read", owner_check=True)
|
||||
async def get_run_workspace_changes(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
run_id: str,
|
||||
request: Request,
|
||||
include_files: bool = Query(default=True),
|
||||
@ -1449,7 +1458,7 @@ async def get_run_workspace_changes(
|
||||
@router.get("/{thread_id}/token-usage", response_model=ThreadTokenUsageResponse)
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
async def thread_token_usage(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
request: Request,
|
||||
include_active: bool = Query(default=False, description="Include running run progress snapshots"),
|
||||
) -> ThreadTokenUsageResponse:
|
||||
@ -1459,4 +1468,5 @@ async def thread_token_usage(
|
||||
agg = await run_store.aggregate_tokens_by_thread(thread_id, include_active=True)
|
||||
else:
|
||||
agg = await run_store.aggregate_tokens_by_thread(thread_id)
|
||||
return ThreadTokenUsageResponse(thread_id=thread_id, **agg)
|
||||
context_usage = await build_context_usage(request, thread_id, run_store)
|
||||
return ThreadTokenUsageResponse(thread_id=thread_id, context_usage=context_usage, **agg)
|
||||
|
||||
@ -68,6 +68,7 @@ from deerflow.runtime.runs.worker import valid_duration_entry
|
||||
from deerflow.runtime.secret_context import redact_metadata_secrets
|
||||
from deerflow.runtime.user_context import get_effective_user_id
|
||||
from deerflow.utils.file_io import run_file_io
|
||||
from deerflow.utils.thread_id import ThreadId, resolve_thread_id, validate_thread_id
|
||||
from deerflow.utils.time import coerce_iso, now_iso
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -372,7 +373,7 @@ class ThreadResponse(_MetadataRedactingResponse):
|
||||
class ThreadCreateRequest(BaseModel):
|
||||
"""Request body for creating a thread."""
|
||||
|
||||
thread_id: str | None = Field(default=None, description="Optional thread ID (auto-generated if omitted)")
|
||||
thread_id: ThreadId | None = Field(default=None, description="Optional thread ID (auto-generated if omitted)")
|
||||
assistant_id: str | None = Field(default=None, description="Associate thread with an assistant")
|
||||
metadata: dict[str, Any] = Field(default_factory=dict, description="Initial metadata")
|
||||
|
||||
@ -627,8 +628,18 @@ async def delete_thread_data(thread_id: str, request: Request) -> ThreadDeleteRe
|
||||
"""
|
||||
from app.gateway.deps import get_thread_store
|
||||
|
||||
# Clean local filesystem
|
||||
response = _delete_thread_data(thread_id, user_id=get_effective_user_id())
|
||||
# Legacy IDs may predate the canonical filesystem-safe contract. They can
|
||||
# still be removed from metadata/checkpoint stores, but must never be
|
||||
# interpolated into a host path during cleanup.
|
||||
try:
|
||||
validate_thread_id(thread_id)
|
||||
except ValueError:
|
||||
response = ThreadDeleteResponse(
|
||||
success=True,
|
||||
message="Skipped local data cleanup for legacy thread ID",
|
||||
)
|
||||
else:
|
||||
response = _delete_thread_data(thread_id, user_id=get_effective_user_id())
|
||||
|
||||
# Remove checkpoints (best-effort)
|
||||
checkpointer = getattr(request.app.state, "checkpointer", None)
|
||||
@ -707,7 +718,7 @@ async def create_thread(body: ThreadCreateRequest, request: Request) -> ThreadRe
|
||||
|
||||
checkpointer = get_checkpointer(request)
|
||||
thread_store = get_thread_store(request)
|
||||
thread_id = body.thread_id or str(uuid.uuid4())
|
||||
thread_id = resolve_thread_id(body.thread_id)
|
||||
now = now_iso()
|
||||
thread_owner_user_id = get_trusted_internal_owner_user_id(request)
|
||||
thread_owner_kwargs = {"user_id": thread_owner_user_id} if thread_owner_user_id else {}
|
||||
@ -774,7 +785,7 @@ async def create_thread(body: ThreadCreateRequest, request: Request) -> ThreadRe
|
||||
|
||||
@router.post("/{thread_id}/branches", response_model=ThreadBranchResponse)
|
||||
@require_permission("threads", "write", owner_check=True, require_existing=True)
|
||||
async def branch_thread(thread_id: str, body: ThreadBranchRequest, request: Request) -> ThreadBranchResponse:
|
||||
async def branch_thread(thread_id: ThreadId, body: ThreadBranchRequest, request: Request) -> ThreadBranchResponse:
|
||||
"""Create a new main-thread branch from a completed assistant turn."""
|
||||
from app.gateway.deps import get_thread_store
|
||||
|
||||
@ -977,7 +988,7 @@ async def search_threads(body: ThreadSearchRequest, request: Request) -> list[Th
|
||||
|
||||
@router.patch("/{thread_id}", response_model=ThreadResponse)
|
||||
@require_permission("threads", "write", owner_check=True, require_existing=True)
|
||||
async def patch_thread(thread_id: str, body: ThreadPatchRequest, request: Request) -> ThreadResponse:
|
||||
async def patch_thread(thread_id: ThreadId, body: ThreadPatchRequest, request: Request) -> ThreadResponse:
|
||||
"""Merge metadata into a thread record."""
|
||||
from app.gateway.deps import get_thread_store
|
||||
|
||||
@ -1010,7 +1021,7 @@ async def patch_thread(thread_id: str, body: ThreadPatchRequest, request: Reques
|
||||
|
||||
@router.get("/{thread_id}", response_model=ThreadResponse)
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
async def get_thread(thread_id: str, request: Request) -> ThreadResponse:
|
||||
async def get_thread(thread_id: ThreadId, request: Request) -> ThreadResponse:
|
||||
"""Get thread info from metadata plus the graph's materialized state."""
|
||||
from app.gateway.deps import get_thread_store
|
||||
|
||||
@ -1063,7 +1074,7 @@ async def get_thread(thread_id: str, request: Request) -> ThreadResponse:
|
||||
|
||||
@router.get("/{thread_id}/goal", response_model=ThreadGoalResponse)
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
async def get_thread_goal(thread_id: str, request: Request) -> ThreadGoalResponse:
|
||||
async def get_thread_goal(thread_id: ThreadId, request: Request) -> ThreadGoalResponse:
|
||||
"""Return the active Claude-style goal for a thread, if any."""
|
||||
checkpointer = get_checkpointer(request)
|
||||
try:
|
||||
@ -1076,7 +1087,7 @@ async def get_thread_goal(thread_id: str, request: Request) -> ThreadGoalRespons
|
||||
|
||||
@router.put("/{thread_id}/goal", response_model=ThreadGoalResponse)
|
||||
@require_permission("threads", "write", owner_check=True)
|
||||
async def set_thread_goal(thread_id: str, body: ThreadGoalRequest, request: Request) -> ThreadGoalResponse:
|
||||
async def set_thread_goal(thread_id: ThreadId, body: ThreadGoalRequest, request: Request) -> ThreadGoalResponse:
|
||||
"""Set or replace the active goal for a thread.
|
||||
|
||||
``/chats/new`` pages already hold a generated UUID before the first run, so
|
||||
@ -1102,7 +1113,7 @@ async def set_thread_goal(thread_id: str, body: ThreadGoalRequest, request: Requ
|
||||
|
||||
@router.delete("/{thread_id}/goal", response_model=ThreadGoalResponse)
|
||||
@require_permission("threads", "write", owner_check=True)
|
||||
async def clear_thread_goal(thread_id: str, request: Request) -> ThreadGoalResponse:
|
||||
async def clear_thread_goal(thread_id: ThreadId, request: Request) -> ThreadGoalResponse:
|
||||
"""Clear the active goal for a thread."""
|
||||
checkpointer = get_checkpointer(request)
|
||||
try:
|
||||
@ -1133,7 +1144,7 @@ def _thread_compact_response(result: ThreadCompactionResult) -> ThreadCompactRes
|
||||
|
||||
@router.post("/{thread_id}/compact", response_model=ThreadCompactResponse)
|
||||
@require_permission("threads", "write", owner_check=True, require_existing=True)
|
||||
async def compact_thread(thread_id: str, body: ThreadCompactRequest, request: Request) -> ThreadCompactResponse:
|
||||
async def compact_thread(thread_id: ThreadId, body: ThreadCompactRequest, request: Request) -> ThreadCompactResponse:
|
||||
"""Manually summarize old thread context while preserving the visible history."""
|
||||
# Compaction writes only base-schema channels (messages + summary_text);
|
||||
# every other channel — including middleware-contributed ones — is carried
|
||||
@ -1180,7 +1191,7 @@ async def compact_thread(thread_id: str, body: ThreadCompactRequest, request: Re
|
||||
# ---------------------------------------------------------------------------
|
||||
@router.get("/{thread_id}/state", response_model=ThreadStateResponse)
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
async def get_thread_state(thread_id: str, request: Request) -> ThreadStateResponse:
|
||||
async def get_thread_state(thread_id: ThreadId, request: Request) -> ThreadStateResponse:
|
||||
"""Get the latest materialized graph state for a thread."""
|
||||
# Resolve through the thread's assistant so custom middleware channels
|
||||
# appear in the response instead of being dropped by the default schema.
|
||||
@ -1222,7 +1233,7 @@ async def get_thread_state(thread_id: str, request: Request) -> ThreadStateRespo
|
||||
|
||||
@router.post("/{thread_id}/state", response_model=ThreadStateResponse)
|
||||
@require_permission("threads", "write", owner_check=True, require_existing=True)
|
||||
async def update_thread_state(thread_id: str, body: ThreadStateUpdateRequest, request: Request) -> ThreadStateResponse:
|
||||
async def update_thread_state(thread_id: ThreadId, body: ThreadStateUpdateRequest, request: Request) -> ThreadStateResponse:
|
||||
"""Replace selected thread-state fields through the materialized graph."""
|
||||
from app.gateway.deps import get_thread_store
|
||||
|
||||
@ -1322,7 +1333,7 @@ def _checkpoint_run_durations(metadata: Any) -> dict[str, int]:
|
||||
@router.post("/{thread_id}/history", response_model=list[HistoryEntry])
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
async def get_thread_history(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
body: ThreadHistoryRequest,
|
||||
request: Request,
|
||||
background_tasks: BackgroundTasks,
|
||||
|
||||
@ -35,6 +35,7 @@ from deerflow.uploads.manager import (
|
||||
)
|
||||
from deerflow.utils.file_conversion import CONVERTIBLE_EXTENSIONS, convert_file_to_markdown
|
||||
from deerflow.utils.file_io import run_file_io
|
||||
from deerflow.utils.thread_id import ThreadId
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -299,7 +300,7 @@ def _auto_convert_documents_enabled(app_config: AppConfig) -> bool:
|
||||
@router.post("", response_model=UploadResponse)
|
||||
@require_permission("threads", "write", owner_check=True, require_existing=False)
|
||||
async def upload_files(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
request: Request,
|
||||
files: list[UploadFile] = File(...),
|
||||
config: AppConfig = Depends(get_config),
|
||||
@ -445,7 +446,7 @@ async def upload_files(
|
||||
@router.get("/limits", response_model=UploadLimits)
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
async def get_upload_limits(
|
||||
thread_id: str,
|
||||
thread_id: ThreadId,
|
||||
request: Request,
|
||||
config: AppConfig = Depends(get_config),
|
||||
) -> UploadLimits:
|
||||
@ -455,7 +456,7 @@ async def get_upload_limits(
|
||||
|
||||
@router.get("/list", response_model=UploadListResponse)
|
||||
@require_permission("threads", "read", owner_check=True)
|
||||
async def list_uploaded_files(thread_id: str, request: Request) -> UploadListResponse:
|
||||
async def list_uploaded_files(thread_id: ThreadId, request: Request) -> UploadListResponse:
|
||||
"""List all files in a thread's uploads directory."""
|
||||
try:
|
||||
result = await run_file_io(_list_uploaded_files_for_thread, thread_id, get_effective_user_id())
|
||||
@ -467,7 +468,7 @@ async def list_uploaded_files(thread_id: str, request: Request) -> UploadListRes
|
||||
|
||||
@router.delete("/{filename}")
|
||||
@require_permission("threads", "delete", owner_check=True, require_existing=True)
|
||||
async def delete_uploaded_file(thread_id: str, filename: str, request: Request) -> dict:
|
||||
async def delete_uploaded_file(thread_id: ThreadId, filename: str, request: Request) -> dict:
|
||||
"""Delete a file from a thread's uploads directory."""
|
||||
try:
|
||||
return await run_file_io(_delete_uploaded_file_for_thread, thread_id, filename, get_effective_user_id())
|
||||
|
||||
@ -4,10 +4,11 @@ from __future__ import annotations
|
||||
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_validator, model_validator
|
||||
from pydantic_core import PydanticCustomError
|
||||
|
||||
from deerflow.runtime.stream_modes import RunStreamMode, UnsupportedStreamModeError, normalize_stream_modes
|
||||
from deerflow.utils.thread_id import validate_thread_id
|
||||
|
||||
|
||||
class RunCreateRequest(BaseModel):
|
||||
@ -36,6 +37,19 @@ class RunCreateRequest(BaseModel):
|
||||
if_not_exists: Literal["create"] = Field(default="create", description="Compatibility default; missing threads are created")
|
||||
feedback_keys: None = Field(default=None, description="Compatibility placeholder; feedback key collection is not supported")
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_configurable_thread_id(self) -> RunCreateRequest:
|
||||
"""Validate the stateless-run thread selector inside RunnableConfig."""
|
||||
if not isinstance(self.config, dict):
|
||||
return self
|
||||
configurable = self.config.get("configurable")
|
||||
if not isinstance(configurable, dict) or "thread_id" not in configurable:
|
||||
return self
|
||||
thread_id = configurable["thread_id"]
|
||||
if thread_id is not None:
|
||||
validate_thread_id(thread_id)
|
||||
return self
|
||||
|
||||
@field_validator(
|
||||
"webhook",
|
||||
"on_completion",
|
||||
|
||||
@ -61,6 +61,7 @@ from deerflow.runtime.checkpoint_mode import (
|
||||
)
|
||||
from deerflow.runtime.checkpoint_state import graph_state_schema
|
||||
from deerflow.runtime.goal import goal_thread_lock
|
||||
from deerflow.runtime.journal import build_checkpoint_history_seed_events
|
||||
from deerflow.runtime.runs.naming import resolve_root_run_name
|
||||
from deerflow.runtime.secret_context import (
|
||||
LegacyRunMetadataSecretError,
|
||||
@ -70,6 +71,7 @@ from deerflow.runtime.secret_context import (
|
||||
from deerflow.runtime.stream_modes import normalize_stream_modes
|
||||
from deerflow.runtime.user_context import reset_current_user, set_current_user
|
||||
from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY
|
||||
from deerflow.utils.thread_id import validate_thread_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -983,6 +985,63 @@ async def apply_checkpoint_to_run_config(
|
||||
configurable["checkpoint_map"] = checkpoint_map
|
||||
|
||||
|
||||
async def ensure_checkpoint_history_seeded(
|
||||
request: Request,
|
||||
*,
|
||||
thread_id: str,
|
||||
assistant_id: str | None,
|
||||
) -> None:
|
||||
"""Backfill an empty run-event feed from an existing checkpoint head.
|
||||
|
||||
No-op unless the feed is empty AND a checkpoint head with messages
|
||||
exists — i.e. a legacy checkpoint-only thread facing its first journaled
|
||||
run. This is a migration shim: remove it once pre-journal threads are no
|
||||
longer a supported upgrade source. The info log on a successful seed is
|
||||
the observability hook for that decision — when it stops appearing, the
|
||||
shim is dead.
|
||||
"""
|
||||
event_store = request.app.state.run_event_store
|
||||
# The emptiness check is deliberately thread-scoped, never user-scoped:
|
||||
# seed rows may be stamped with a different principal (NULL for ownerless
|
||||
# seeds, or another user on a shared NULL-owner thread), so a user-scoped
|
||||
# query would miss them and re-seed a duplicate history per principal.
|
||||
# Passing user_id=None also opts out of AUTO resolution explicitly, which
|
||||
# would raise when no user contextvar is set (e.g. the scheduler launch
|
||||
# path for ownerless internal tasks).
|
||||
if await event_store.list_messages(thread_id, limit=1, user_id=None):
|
||||
return
|
||||
|
||||
checkpoint_config = {
|
||||
"configurable": {
|
||||
"thread_id": thread_id,
|
||||
"checkpoint_ns": "",
|
||||
}
|
||||
}
|
||||
if await get_checkpointer(request).aget_tuple(checkpoint_config) is None:
|
||||
return
|
||||
|
||||
accessor, config = build_checkpoint_state_accessor(
|
||||
request,
|
||||
thread_id=thread_id,
|
||||
assistant_id=assistant_id,
|
||||
)
|
||||
snapshot = await accessor.aget(config)
|
||||
values = getattr(snapshot, "values", None)
|
||||
messages = values.get("messages") if isinstance(values, dict) else None
|
||||
if not isinstance(messages, list) or not messages:
|
||||
return
|
||||
|
||||
events = build_checkpoint_history_seed_events(
|
||||
messages,
|
||||
thread_id=thread_id,
|
||||
run_id_prefix=f"checkpoint-seed-{thread_id}",
|
||||
)
|
||||
if not events:
|
||||
return
|
||||
await event_store.put_batch(events)
|
||||
logger.info("Seeded %d checkpoint-history events for thread %s", len(events), thread_id)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Run lifecycle
|
||||
# ---------------------------------------------------------------------------
|
||||
@ -1004,6 +1063,11 @@ async def start_run(
|
||||
request : Request
|
||||
FastAPI request — used to retrieve singletons from ``app.state``.
|
||||
"""
|
||||
try:
|
||||
validate_thread_id(thread_id)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=422, detail=str(exc)) from exc
|
||||
|
||||
body_config = getattr(body, "config", None)
|
||||
config_metadata = body_config.get("metadata") if isinstance(body_config, dict) else None
|
||||
try:
|
||||
@ -1156,6 +1220,11 @@ async def start_run(
|
||||
|
||||
try:
|
||||
async with goal_thread_lock(thread_id):
|
||||
await ensure_checkpoint_history_seeded(
|
||||
request,
|
||||
thread_id=thread_id,
|
||||
assistant_id=body.assistant_id,
|
||||
)
|
||||
record = await run_mgr.create_or_reject(
|
||||
thread_id,
|
||||
body.assistant_id,
|
||||
|
||||
@ -9,7 +9,7 @@ DeerFlow 后端提供了完整的文件上传功能,支持多文件上传,
|
||||
- ✅ 支持多文件同时上传
|
||||
- ✅ 可选地转换文档为 Markdown(PDF、PPT、Excel、Word)
|
||||
- ✅ 文件存储在线程隔离的目录中
|
||||
- ✅ Agent 自动感知已上传的文件
|
||||
- ✅ Agent 自动感知当前消息中附带的文件
|
||||
- ✅ 支持文件列表查询和删除
|
||||
|
||||
## API 端点
|
||||
@ -116,24 +116,30 @@ DELETE /api/threads/{thread_id}/uploads/{filename}
|
||||
|
||||
## Agent 集成
|
||||
|
||||
### 自动文件列举
|
||||
### 当前消息中的文件上下文
|
||||
|
||||
Agent 在每次请求时会自动收到已上传文件的列表,格式如下:
|
||||
发送消息时,前端会把该消息附带的上传文件元数据放入
|
||||
`HumanMessage.additional_kwargs.files`。`UploadsMiddleware` 只把当前消息中的文件
|
||||
注入 Agent 上下文,格式如下:
|
||||
|
||||
```xml
|
||||
<uploaded_files>
|
||||
The following files have been uploaded and are available for use:
|
||||
<current_uploads>
|
||||
The following files were uploaded in this message:
|
||||
|
||||
- document.pdf (1.2 MB)
|
||||
Path: /mnt/user-data/uploads/document.pdf
|
||||
|
||||
- document.md (45.3 KB)
|
||||
Path: /mnt/user-data/uploads/document.md
|
||||
|
||||
You can read these files using the `read_file` tool with the paths shown above.
|
||||
</uploaded_files>
|
||||
To work with these files:
|
||||
- Read from the file first — use the outline line numbers and `read_file` to locate relevant sections.
|
||||
- Use `grep` to search for keywords when you are not sure which section to look at.
|
||||
- Use `glob` to find files by name pattern.
|
||||
</current_uploads>
|
||||
```
|
||||
|
||||
以前轮次上传的文件不会在每次请求中重复注入。Agent 可按需调用
|
||||
`list_uploaded_files` 查询历史上传;如果已知文件名,也可直接使用
|
||||
`read_file` 或 `grep` 访问 `/mnt/user-data/uploads/` 下的文件。
|
||||
|
||||
### 使用上传的文件
|
||||
|
||||
Agent 在沙箱中运行,使用虚拟路径访问文件。Agent 可以直接使用 `read_file` 工具读取上传的文件:
|
||||
@ -240,8 +246,9 @@ backend/.deer-flow/threads/
|
||||
- 使用 markitdown 转换文档
|
||||
|
||||
2. **Uploads Middleware** (`packages/harness/deerflow/agents/middlewares/uploads_middleware.py`)
|
||||
- 在每次 Agent 请求前注入文件列表
|
||||
- 自动生成格式化的文件列表消息
|
||||
- 读取当前消息的 `additional_kwargs.files`
|
||||
- 在 Agent 请求前生成并注入 `<current_uploads>` 文件上下文
|
||||
- 历史上传由 `list_uploaded_files` 按需查询,不会每轮自动注入
|
||||
|
||||
3. **Nginx 配置** (`nginx.conf`)
|
||||
- 路由上传请求到 Gateway API
|
||||
|
||||
@ -208,7 +208,15 @@ The cached value is reused across the blocking (`runs.wait`) and streaming (`_ha
|
||||
|
||||
## IM File Attachment Pipeline
|
||||
|
||||
Inbound files (images, documents) walk through `Channel.receive_file` for materialization, then `_ingest_inbound_files` for owner-bound staging. The agent sees the staged path via the `<uploaded_files>` block injected into its context.
|
||||
Inbound files (images, documents) first pass through `Channel.receive_file` for
|
||||
provider-specific materialization. Attachments that continue through the shared
|
||||
metadata path are staged by `_ingest_inbound_files`; their metadata is placed in
|
||||
`HumanMessage.additional_kwargs.files`, and `UploadsMiddleware` injects a
|
||||
`<current_uploads>` block for the current message. Some providers instead consume
|
||||
their descriptors while downloading and rewrite placeholders or message text with
|
||||
the resulting virtual path (or a failure notice). Historical uploads are not
|
||||
automatically injected on later turns; the agent discovers them with
|
||||
`list_uploaded_files`.
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
@ -218,17 +226,24 @@ sequenceDiagram
|
||||
participant Mgr as ChannelManager
|
||||
participant Ch as Channel impl<br/>.receive_file
|
||||
participant FS as Uploads directory<br/>users/OWNER/.../uploads/
|
||||
participant MW as UploadsMiddleware
|
||||
participant Agent as Agent run
|
||||
|
||||
IM->>Worker: message with file URL/bytes
|
||||
Worker->>Mgr: InboundMessage(files=[...], connection_id, owner_user_id)
|
||||
Mgr->>Mgr: storage_user_id = _channel_storage_user_id(msg)
|
||||
Mgr->>Ch: receive_file(msg, thread_id, user_id=storage_user_id)
|
||||
Note over Ch: provider-specific download<br/>(WeCom: decrypt_file;<br/>WeChat: read_bytes; others: HTTP GET)
|
||||
Ch->>FS: write_upload_file_no_symlink(<br/>uploads/OWNER/.../<br/>, safe_name, data)
|
||||
Ch-->>Mgr: msg with text rewritten to include <uploaded_files>
|
||||
Mgr->>Mgr: _ingest_inbound_files(<br/>thread_id, msg, user_id=storage_user_id)
|
||||
Mgr->>Agent: HumanMessage with <uploaded_files> block<br/>(paths under /mnt/user-data/uploads/)
|
||||
Note over Ch: provider-specific download/decrypt/read;<br/>may persist or hand bytes to manager
|
||||
Ch-->>Mgr: materialized message<br/>(provider may rewrite placeholders/text)
|
||||
alt attachment continues through shared metadata path
|
||||
Mgr->>FS: _ingest_inbound_files(<br/>thread_id, msg, user_id=storage_user_id)
|
||||
FS-->>Mgr: uploaded file metadata
|
||||
Mgr->>MW: HumanMessage with<br/>additional_kwargs.files
|
||||
MW->>Agent: prepend <current_uploads><br/>(paths under /mnt/user-data/uploads/)
|
||||
else provider supplies virtual path in message text
|
||||
Ch->>FS: persist and/or sync attachment
|
||||
Mgr->>Agent: HumanMessage with rewritten path<br/>or failure notice
|
||||
end
|
||||
Agent->>FS: read_file / view_image (sandbox)
|
||||
```
|
||||
|
||||
|
||||
@ -70,6 +70,32 @@ top-level `config.yaml -> tool_search.auto_promote_top_k` setting.
|
||||
- `tool_search.auto_promote_top_k`: global limit for auto-promoted deferred MCP
|
||||
schemas per model call. Default `3`; valid range `1..5`.
|
||||
|
||||
## Tool Name Prefixes
|
||||
|
||||
DeerFlow prefixes discovered MCP tool names with `<server_name>_` by default.
|
||||
This avoids collisions when two enabled servers expose tools with the same
|
||||
name. A server that already namespaces its own tools can opt out:
|
||||
|
||||
```json
|
||||
{
|
||||
"mcpServers": {
|
||||
"semantic-scholar": {
|
||||
"type": "stdio",
|
||||
"command": "uvx",
|
||||
"args": ["s2-mcp-server"],
|
||||
"tool_name_prefix": false
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
With this setting, a server tool named `semantic_scholar_search_papers` keeps
|
||||
that name instead of becoming
|
||||
`semantic-scholar_semantic_scholar_search_papers`. The default is `true` for
|
||||
backward compatibility. Disable it only when every resulting tool name remains
|
||||
unique across the enabled servers. Stdio tools continue to use DeerFlow's
|
||||
persistent per-thread session pool regardless of this setting.
|
||||
|
||||
## Per-Tool Timeout (Stdio MCP Servers)
|
||||
|
||||
For `stdio` MCP servers, set `tool_call_timeout` to limit each individual MCP tool call in seconds:
|
||||
|
||||
@ -26,6 +26,8 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import secrets
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
from langchain.agents import create_agent
|
||||
from langchain.agents.middleware import AgentMiddleware
|
||||
@ -47,6 +49,9 @@ from deerflow.agents.middlewares.token_usage_middleware import TokenUsageMiddlew
|
||||
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.authz.principal import build_principal_from_context
|
||||
from deerflow.authz.provider import AuthzDecision, AuthzRequest
|
||||
from deerflow.authz.runtime import resolve_authorization_provider
|
||||
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
|
||||
@ -135,6 +140,101 @@ def _resolve_model_name(requested_model_name: str | None = None, *, app_config:
|
||||
return default_model_name
|
||||
|
||||
|
||||
def _authorize_model_name(
|
||||
model_name: str,
|
||||
*,
|
||||
context: Mapping[str, Any],
|
||||
app_config: AppConfig,
|
||||
) -> str:
|
||||
"""Enforce ``model:use`` authorization on the resolved model name.
|
||||
|
||||
When ``authorization.enabled`` is false this is a no-op (returns
|
||||
*model_name* unchanged). When enabled, the resolved model is checked
|
||||
against the provider's policy via ``authorize("model", "use")`` so the
|
||||
runtime path and the Gateway ``get_model`` route enforce the same
|
||||
action-scoped contract (matters for custom providers that distinguish
|
||||
``list`` from ``use``). On deny, a graceful fallback to the first
|
||||
``filter_resources``-allowed model is attempted (RFC §9: "fall back to an
|
||||
allowed default, not error, to avoid breaking runs"). If no model is
|
||||
allowed and ``fail_closed`` is true, ``ValueError`` is raised (matching
|
||||
the existing "no models configured" contract); fail-open returns the
|
||||
original name.
|
||||
|
||||
Mirrors the Principal/provider pattern of ``apply_tool_authorization`` so
|
||||
the tool path and the model path share one identity source.
|
||||
"""
|
||||
authz_config = app_config.authorization
|
||||
if authz_config.enabled is not True:
|
||||
return model_name
|
||||
|
||||
provider = resolve_authorization_provider(authz_config)
|
||||
if provider is None:
|
||||
return model_name
|
||||
|
||||
principal = build_principal_from_context(context, default_role=authz_config.default_role)
|
||||
all_names = [m.name for m in app_config.models]
|
||||
|
||||
# Check the resolved model against the action-scoped ``model:use`` policy.
|
||||
# This aligns with the Gateway ``get_model`` route, which also checks
|
||||
# ``authorize("model", "use")``. For the built-in RBAC provider (which
|
||||
# ignores ``action``) this is equivalent to a membership check; for a
|
||||
# custom provider that distinguishes ``list`` from ``use``, it prevents
|
||||
# a model visible via ``filter_resources`` but denied for ``use`` from
|
||||
# being silently selected at runtime.
|
||||
try:
|
||||
decision = provider.authorize(AuthzRequest(principal=principal, resource="model", action="use", target=model_name))
|
||||
if not isinstance(decision, AuthzDecision):
|
||||
raise TypeError("AuthorizationProvider.authorize must return AuthzDecision")
|
||||
if decision.allow:
|
||||
return model_name
|
||||
except Exception:
|
||||
logger.warning("Authorization provider failed while checking model:use for '%s'", model_name, exc_info=True)
|
||||
if authz_config.fail_closed:
|
||||
raise ValueError("No models are authorized for the current role (authorization provider error).")
|
||||
return model_name
|
||||
|
||||
# Denied — graceful fallback: pick the first model that ``filter_resources``
|
||||
# says is visible AND that also passes ``authorize("model", "use")``. For the
|
||||
# built-in RBAC provider (which ignores ``action``) this is equivalent to
|
||||
# picking the first visible name; for a custom provider that distinguishes
|
||||
# ``list`` from ``use``, it ensures the fallback is actually usable.
|
||||
try:
|
||||
allowed_names = provider.filter_resources(principal, "model", all_names)
|
||||
if not isinstance(allowed_names, list) or any(not isinstance(n, str) for n in allowed_names):
|
||||
raise TypeError("AuthorizationProvider.filter_resources must return list[str]")
|
||||
except Exception:
|
||||
logger.warning("Authorization provider failed while resolving allowed models", exc_info=True)
|
||||
if authz_config.fail_closed:
|
||||
raise ValueError("No models are authorized for the current role (authorization provider error).")
|
||||
return model_name
|
||||
|
||||
for candidate in allowed_names:
|
||||
if candidate == model_name:
|
||||
continue # already denied above
|
||||
try:
|
||||
cb_decision = provider.authorize(AuthzRequest(principal=principal, resource="model", action="use", target=candidate))
|
||||
if isinstance(cb_decision, AuthzDecision) and cb_decision.allow:
|
||||
logger.warning(
|
||||
"Model '%s' is not authorized for the current role; fallback to '%s'.",
|
||||
model_name,
|
||||
candidate,
|
||||
)
|
||||
return candidate
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"Authorization provider failed while checking model:use fallback for '%s'",
|
||||
candidate,
|
||||
exc_info=True,
|
||||
)
|
||||
if authz_config.fail_closed:
|
||||
raise ValueError("No models are authorized for the current role (authorization provider error).")
|
||||
return model_name
|
||||
if authz_config.fail_closed:
|
||||
raise ValueError("No models are authorized for the current role.")
|
||||
logger.warning("No models are authorized for the current role; fail_open allows '%s'.", model_name)
|
||||
return model_name
|
||||
|
||||
|
||||
def _create_summarization_middleware(*, app_config: AppConfig | None = None, run_model_name: str | None = None) -> DeerFlowSummarizationMiddleware | None:
|
||||
"""Create and configure the summarization middleware from config.
|
||||
|
||||
@ -581,6 +681,10 @@ def _make_lead_agent(config: RunnableConfig, *, app_config: AppConfig):
|
||||
# Final model name resolution: request → agent config → global default, with fallback for unknown names
|
||||
model_name = _resolve_model_name(requested_model_name or agent_model_name, app_config=resolved_app_config)
|
||||
|
||||
# Phase 3: enforce model:use authorization. On deny, fall back to the first
|
||||
# allowed model (graceful) rather than crashing the run (RFC §9).
|
||||
model_name = _authorize_model_name(model_name, context=cfg, app_config=resolved_app_config)
|
||||
|
||||
model_config = resolved_app_config.get_model_config(model_name)
|
||||
|
||||
if model_config is None:
|
||||
|
||||
@ -15,11 +15,12 @@ template: |-
|
||||
- SKIP: Facts are distinct enough to remain separate.
|
||||
|
||||
Add consolidation decisions to "factsToConsolidate" in your output JSON.
|
||||
Each entry: {{"sourceIds": ["fact_id_1", "fact_id_2"], "consolidated": {{"content": "...", "category": "...", "confidence": 0.9}}}}
|
||||
Each entry: {{"sourceIds": ["fact_id_1", "fact_id_2"], "consolidated": {{"content": "...", "category": "...", "confidence": 0.9, "scope": "user|thread|project", "durability": "durable|temporary", "authority": "descriptive|transactional"}}}}
|
||||
|
||||
Rules:
|
||||
- The consolidated fact must preserve ALL key details from source facts
|
||||
- Only consolidate facts that describe the same aspect of the user
|
||||
- Classify every consolidated fact. It is eligible only when it is user-scoped, durable, descriptive, and safe to inject into unrelated future threads.
|
||||
- Confidence of consolidated fact = max of source confidences
|
||||
- Be conservative - when in doubt, keep facts separate
|
||||
- Maximum {max_groups} consolidation groups per cycle
|
||||
|
||||
@ -12,12 +12,21 @@ messages:
|
||||
|
||||
Before extracting facts, perform a structured reflection on the conversation:
|
||||
1. Error/Retry Detection: Did the agent encounter errors, require retries, or produce incorrect results?
|
||||
If yes, record the root cause and correct approach as a high-confidence fact with category "correction".
|
||||
If yes, record the root cause and correct approach only when it is a durable user-level working pattern.
|
||||
2. User Correction Detection: Did the user correct the agent's direction, understanding, or output?
|
||||
If yes, record the correct interpretation or approach as a high-confidence fact with category "correction".
|
||||
If yes, distinguish a reusable user-level correction from a correction to the current task's facts or files.
|
||||
Include what went wrong in "sourceError" only when category is "correction" and the mistake is explicit in the conversation.
|
||||
3. Project Constraint Discovery: Were any project-specific constraints discovered during the conversation?
|
||||
If yes, record them as facts with the most appropriate category and confidence.
|
||||
If yes, classify them as project-scoped and do not promote them to user memory.
|
||||
|
||||
Scope and Safety Classification:
|
||||
- scope="user": a property, preference, background detail, ongoing personal state, or durable working pattern of the user that is safe and useful to inject into an unrelated future thread, task, project, or repository.
|
||||
- scope="thread": information limited to the current request or conversation, including one-off constraints for a reply, file, email, trip, PR, test, or action.
|
||||
- scope="project": a rule, decision, state, or constraint that is meaningful only inside a particular project or repository, even if it may span several threads.
|
||||
- durability="durable": expected to remain true across future conversations. durability="temporary": current, short-lived, or one-off information.
|
||||
- authority="transactional": an instruction, grant, permission, or authorization to perform an action such as editing, deleting, pushing, closing, publishing, or force-pushing. Transactional content must never become long-term memory.
|
||||
- authority="descriptive": describes the user without granting authority for an action.
|
||||
- When uncertain, use thread/project or temporary; never guess user+durable.
|
||||
|
||||
Memory Section Guidelines:
|
||||
|
||||
@ -43,6 +52,9 @@ messages:
|
||||
Include: Core expertise, longstanding interests, fundamental working style
|
||||
|
||||
**Facts Extraction**:
|
||||
- Every new fact MUST include scope, durability, and authority. These labels are evaluated by a deterministic write gate and are not persisted.
|
||||
- Only facts classified as scope="user", durability="durable", authority="descriptive" are eligible for storage.
|
||||
- Do not create facts for current-task objectives, acceptance criteria, workspace state, exact current file/commit/error state, project-only constraints, or one-time action permissions.
|
||||
- Extract specific, quantifiable details (e.g., "16k+ GitHub stars", "200+ datasets")
|
||||
- Include proper nouns (company names, project names, technology names)
|
||||
- Preserve technical terminology and version numbers
|
||||
@ -85,25 +97,27 @@ messages:
|
||||
Output Format (JSON):
|
||||
{{
|
||||
"user": {{
|
||||
"workContext": {{ "summary": "...", "shouldUpdate": true/false }},
|
||||
"personalContext": {{ "summary": "...", "shouldUpdate": true/false }},
|
||||
"topOfMind": {{ "summary": "...", "shouldUpdate": true/false }}
|
||||
"workContext": {{ "summary": "...", "shouldUpdate": true/false, "scope": "user|thread|project", "authority": "descriptive|transactional" }},
|
||||
"personalContext": {{ "summary": "...", "shouldUpdate": true/false, "scope": "user|thread|project", "authority": "descriptive|transactional" }},
|
||||
"topOfMind": {{ "summary": "...", "shouldUpdate": true/false, "scope": "user|thread|project", "authority": "descriptive|transactional" }}
|
||||
}},
|
||||
"history": {{
|
||||
"recentMonths": {{ "summary": "...", "shouldUpdate": true/false }},
|
||||
"earlierContext": {{ "summary": "...", "shouldUpdate": true/false }},
|
||||
"longTermBackground": {{ "summary": "...", "shouldUpdate": true/false }}
|
||||
"recentMonths": {{ "summary": "...", "shouldUpdate": true/false, "scope": "user|thread|project", "authority": "descriptive|transactional" }},
|
||||
"earlierContext": {{ "summary": "...", "shouldUpdate": true/false, "scope": "user|thread|project", "authority": "descriptive|transactional" }},
|
||||
"longTermBackground": {{ "summary": "...", "shouldUpdate": true/false, "scope": "user|thread|project", "authority": "descriptive|transactional" }}
|
||||
}},
|
||||
"newFacts": [
|
||||
{{ "content": "...", "category": "preference|knowledge|context|behavior|goal|correction", "confidence": 0.0-1.0, "expected_valid_days": 90 }}
|
||||
{{ "content": "...", "category": "preference|knowledge|context|behavior|goal|correction", "confidence": 0.0-1.0, "expected_valid_days": 90, "scope": "user|thread|project", "durability": "durable|temporary", "authority": "descriptive|transactional" }}
|
||||
],
|
||||
"factsToRemove": [
|
||||
{{ "id": "fact_id_1", "scope": "user|thread|project", "reason": "explicit user-level contradiction or retraction", "replacementFactIndex": 0 }}
|
||||
],
|
||||
"factsToRemove": ["fact_id_1", "fact_id_2"],
|
||||
"staleFactsToRemove": [{{ "id": "fact_id", "reason": "brief explanation" }}],
|
||||
"staleFactsToExtend": [{{ "id": "fact_id", "extend_by_days": 365, "reason": "brief explanation" }}],
|
||||
"factsToConsolidate": [
|
||||
{{
|
||||
"sourceIds": ["fact_id_1", "fact_id_2"],
|
||||
"consolidated": {{ "content": "synthesized fact", "category": "knowledge", "confidence": 0.9 }}
|
||||
"consolidated": {{ "content": "synthesized fact", "category": "knowledge", "confidence": 0.9, "scope": "user|thread|project", "durability": "durable|temporary", "authority": "descriptive|transactional" }}
|
||||
}}
|
||||
]
|
||||
}}
|
||||
@ -115,7 +129,9 @@ messages:
|
||||
- Only add facts that are clearly stated (0.9+) or strongly implied (0.7+)
|
||||
- Use category "correction" for explicit agent mistakes or user corrections; assign confidence >= 0.95 when the correction is explicit
|
||||
- Include "sourceError" only for explicit correction facts when the prior mistake or wrong approach is clearly stated; omit it otherwise
|
||||
- Remove facts that are contradicted by new information
|
||||
- Remove an existing fact only when the user explicitly contradicts or retracts it at user scope. A thread/project-local exception does not contradict a user-level fact.
|
||||
- factsToRemove entries MUST include scope and reason. Use replacementFactIndex when a removal depends on a replacement in newFacts; the zero-based index must identify that replacement. Omit replacementFactIndex only for a pure user-level retraction with no replacement.
|
||||
- Every summary with shouldUpdate=true MUST include scope and authority. A summary is user-scoped only when the entire prose block is safe to inject into unrelated future threads; otherwise classify it as thread/project so the write gate rejects it. Any summary containing an instruction, grant, permission, or authorization is transactional and must be rejected even when it describes a user-wide or recurring policy.
|
||||
- When updating topOfMind, integrate new focus areas while removing completed/abandoned ones
|
||||
Keep 3-5 concurrent focus themes that are still active and relevant
|
||||
- For history sections, integrate new information chronologically into appropriate time period
|
||||
|
||||
@ -126,6 +126,52 @@ def _extract_text(content: Any) -> str:
|
||||
|
||||
|
||||
_REQUIRED_MEMORY_UPDATE_TOP_LEVEL_KEYS = frozenset({"user", "history", "newFacts"})
|
||||
_FACT_CLASSIFICATION_FIELDS = ("scope", "durability", "authority")
|
||||
|
||||
|
||||
def _normalize_gate_label(value: Any) -> str | None:
|
||||
"""Normalize a model-produced scope-gate label without validating policy."""
|
||||
if not isinstance(value, str):
|
||||
return None
|
||||
normalized = value.strip().lower()
|
||||
return normalized or None
|
||||
|
||||
|
||||
def _fact_scope_gate_reason(fact: dict[str, Any]) -> str | None:
|
||||
"""Return the deterministic rejection reason for a model-extracted fact."""
|
||||
if any(_normalize_gate_label(fact.get(field)) is None for field in _FACT_CLASSIFICATION_FIELDS):
|
||||
return "missing"
|
||||
if _normalize_gate_label(fact.get("scope")) != "user":
|
||||
return "scope"
|
||||
if _normalize_gate_label(fact.get("durability")) != "durable":
|
||||
return "durability"
|
||||
if _normalize_gate_label(fact.get("authority")) != "descriptive":
|
||||
return "authority"
|
||||
return None
|
||||
|
||||
|
||||
def _summary_scope_gate_reason(section_data: dict[str, Any]) -> str | None:
|
||||
"""Return the deterministic rejection reason for a summary update."""
|
||||
scope = _normalize_gate_label(section_data.get("scope"))
|
||||
authority = _normalize_gate_label(section_data.get("authority"))
|
||||
if scope is None or authority is None:
|
||||
return "missing"
|
||||
if scope != "user":
|
||||
return "scope"
|
||||
if authority != "descriptive":
|
||||
return "authority"
|
||||
return None
|
||||
|
||||
|
||||
def _removal_scope_gate_reason(removal: dict[str, Any]) -> str | None:
|
||||
"""Return the deterministic rejection reason for a contradiction removal."""
|
||||
scope = _normalize_gate_label(removal.get("scope"))
|
||||
reason = removal.get("reason")
|
||||
if scope is None or not isinstance(reason, str) or not reason.strip():
|
||||
return "missing"
|
||||
if scope != "user":
|
||||
return "scope"
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_memory_update_fact(fact: Any) -> dict[str, Any] | None:
|
||||
@ -180,6 +226,14 @@ def _normalize_memory_update_fact(fact: Any) -> dict[str, Any] | None:
|
||||
if evd is not None:
|
||||
normalized_fact["expected_valid_days"] = evd
|
||||
|
||||
# Scope classification is extraction-only metadata. Preserve it through
|
||||
# structural normalization so _apply_updates can fail closed per item, but
|
||||
# never copy it into the persisted fact_entry.
|
||||
for field in _FACT_CLASSIFICATION_FIELDS:
|
||||
normalized_value = _normalize_gate_label(fact.get(field))
|
||||
if normalized_value is not None:
|
||||
normalized_fact[field] = normalized_value
|
||||
|
||||
return normalized_fact
|
||||
|
||||
|
||||
@ -189,7 +243,35 @@ def _normalize_memory_update_data(update_data: dict[str, Any]) -> dict[str, Any]
|
||||
history = update_data.get("history")
|
||||
new_facts = update_data.get("newFacts")
|
||||
facts_to_remove = update_data.get("factsToRemove")
|
||||
normalized_facts_to_remove = [fact_id for fact_id in facts_to_remove if isinstance(fact_id, str)] if isinstance(facts_to_remove, list) else []
|
||||
normalized_facts_to_remove: list[dict[str, Any]] = []
|
||||
if isinstance(facts_to_remove, list):
|
||||
for entry in facts_to_remove:
|
||||
# Preserve the legacy string form as an unclassified removal. The
|
||||
# apply-layer gate will reject it as missing instead of continuing
|
||||
# to allow an unscoped destructive mutation.
|
||||
if isinstance(entry, str):
|
||||
fact_id = entry.strip()
|
||||
if fact_id:
|
||||
normalized_facts_to_remove.append({"id": fact_id})
|
||||
continue
|
||||
if not isinstance(entry, dict):
|
||||
continue
|
||||
raw_id = entry.get("id")
|
||||
if not isinstance(raw_id, str) or not raw_id.strip():
|
||||
continue
|
||||
normalized_removal: dict[str, Any] = {"id": raw_id.strip()}
|
||||
scope = _normalize_gate_label(entry.get("scope"))
|
||||
if scope is not None:
|
||||
normalized_removal["scope"] = scope
|
||||
reason = entry.get("reason")
|
||||
if isinstance(reason, str) and reason.strip():
|
||||
normalized_removal["reason"] = reason.strip()
|
||||
if "replacementFactIndex" in entry:
|
||||
# Preserve invalid values too: the apply layer must reject an
|
||||
# invalid dependency rather than silently treating it as a pure
|
||||
# removal and deleting the old fact.
|
||||
normalized_removal["replacementFactIndex"] = entry.get("replacementFactIndex")
|
||||
normalized_facts_to_remove.append(normalized_removal)
|
||||
normalized_new_facts = []
|
||||
dropped_new_fact = not isinstance(new_facts, list)
|
||||
if isinstance(new_facts, list):
|
||||
@ -290,6 +372,7 @@ def _normalize_memory_update_data(update_data: dict[str, Any]) -> dict[str, Any]
|
||||
"content": content.strip(),
|
||||
"category": _norm_cat,
|
||||
"confidence": _norm_conf,
|
||||
**{field: normalized for field in _FACT_CLASSIFICATION_FIELDS if (normalized := _normalize_gate_label(consolidated.get(field))) is not None},
|
||||
},
|
||||
}
|
||||
)
|
||||
@ -372,6 +455,20 @@ def _fact_content_key(content: Any) -> str | None:
|
||||
return stripped.casefold()
|
||||
|
||||
|
||||
def _raise_if_duplicate_fact_content(memory_data: dict[str, Any], content_key: str | None) -> None:
|
||||
"""Reject a candidate fact whose normalized content already exists.
|
||||
|
||||
Callers must invoke this against the freshest snapshot available inside
|
||||
their read-check-write critical section (i.e. on every revision-conflict
|
||||
retry), so two concurrent creators of the same content cannot both pass
|
||||
the check and store duplicate facts."""
|
||||
if content_key is None:
|
||||
return
|
||||
for fact in memory_data.get("facts", []):
|
||||
if isinstance(fact, dict) and _fact_content_key(fact.get("content")) == content_key:
|
||||
raise ValueError("Duplicate fact")
|
||||
|
||||
|
||||
# ── Staleness review helpers ──────────────────────────────────────────────
|
||||
|
||||
|
||||
@ -815,6 +912,13 @@ class MemoryUpdater:
|
||||
"added" status. This restores both the max_facts cap and the post-trim
|
||||
existence check (upstream's ``create_memory_fact_with_created_fact``),
|
||||
which the vendored copy had dropped together to avoid the dangling id.
|
||||
|
||||
Duplicate rejection is enforced here (not only by callers): the
|
||||
candidate's normalized content key is checked against the fresh
|
||||
memory snapshot inside the revision-conflict retry loop of both
|
||||
storage paths (apply_changes and legacy single-file save), so
|
||||
concurrent creators cannot both store the same content. Raises
|
||||
``ValueError("Duplicate fact")`` on a normalized-content match.
|
||||
"""
|
||||
if agent_name is None:
|
||||
raise ValueError("agent_name")
|
||||
@ -823,6 +927,7 @@ class MemoryUpdater:
|
||||
raise ValueError("content")
|
||||
normalized_category = category.strip() or "context"
|
||||
validated_confidence = _validate_confidence(confidence)
|
||||
candidate_key = _fact_content_key(normalized_content)
|
||||
now = utc_now_iso_z()
|
||||
fact_id = f"fact_{uuid.uuid4().hex[:8]}"
|
||||
candidate = {
|
||||
@ -836,6 +941,12 @@ class MemoryUpdater:
|
||||
if getattr(type(self._storage), "apply_changes", None) is not MemoryStorage.apply_changes:
|
||||
for attempt in range(3):
|
||||
memory_data = self.get_memory_data(agent_name, user_id=user_id) if attempt == 0 else self.reload_memory_data(agent_name, user_id=user_id)
|
||||
# Duplicate rejection lives inside the conflict-retry loop so
|
||||
# it is re-evaluated against the fresh snapshot after every
|
||||
# revision conflict: two concurrent creators of the same
|
||||
# content cannot both store it (the loser reloads, sees the
|
||||
# winner's fact, and is rejected here).
|
||||
_raise_if_duplicate_fact_content(memory_data, candidate_key)
|
||||
updated_memory = dict(memory_data)
|
||||
updated_memory["facts"] = _trim_facts_to_max([*memory_data.get("facts", []), copy.deepcopy(candidate)], self._config.max_facts)
|
||||
kept_ids = {str(fact.get("id")) for fact in updated_memory["facts"]}
|
||||
@ -860,15 +971,24 @@ class MemoryUpdater:
|
||||
raise
|
||||
logger.info("Retrying capped fact creation from a fresh snapshot after a revision conflict")
|
||||
raise AssertionError("bounded create retry did not return or raise")
|
||||
memory_data = self.get_memory_data(agent_name, user_id=user_id)
|
||||
updated_memory = dict(memory_data)
|
||||
updated_memory["facts"] = _trim_facts_to_max([*memory_data.get("facts", []), candidate], self._config.max_facts)
|
||||
if not self._save_memory_to_file(updated_memory, agent_name, user_id=user_id, expected_revision=int(memory_data.get("revision") or 0)):
|
||||
raise OSError("Failed to save memory data after creating fact")
|
||||
# If the cap evicted the just-added (lower-confidence) fact, signal via
|
||||
# None so callers don't report a dangling id as "added".
|
||||
stored = any(f.get("id") == fact_id for f in updated_memory["facts"])
|
||||
return updated_memory, (fact_id if stored else None)
|
||||
# Legacy single-file path: same duplicate-rejection contract as the
|
||||
# apply_changes path above. A revision-conflicted save (False) reloads
|
||||
# the fresh snapshot and re-runs the duplicate check, so a concurrent
|
||||
# creator's commit is rejected with ValueError("Duplicate fact")
|
||||
# instead of surfacing as a generic save failure.
|
||||
for attempt in range(3):
|
||||
memory_data = self.get_memory_data(agent_name, user_id=user_id) if attempt == 0 else self.reload_memory_data(agent_name, user_id=user_id)
|
||||
_raise_if_duplicate_fact_content(memory_data, candidate_key)
|
||||
updated_memory = dict(memory_data)
|
||||
updated_memory["facts"] = _trim_facts_to_max([*memory_data.get("facts", []), copy.deepcopy(candidate)], self._config.max_facts)
|
||||
if self._save_memory_to_file(updated_memory, agent_name, user_id=user_id, expected_revision=int(memory_data.get("revision") or 0)):
|
||||
# If the cap evicted the just-added (lower-confidence) fact,
|
||||
# signal via None so callers don't report a dangling id as
|
||||
# "added".
|
||||
stored = any(f.get("id") == fact_id for f in updated_memory["facts"])
|
||||
return updated_memory, (fact_id if stored else None)
|
||||
logger.info("Retrying capped fact creation from a fresh snapshot after a revision conflict")
|
||||
raise OSError("Failed to save memory data after creating fact")
|
||||
|
||||
def delete_memory_fact(self, fact_id: str, agent_name: str | None = None, *, user_id: str | None = None) -> dict[str, Any]:
|
||||
"""Delete a fact by its id and persist the updated memory data."""
|
||||
@ -984,25 +1104,24 @@ class MemoryUpdater:
|
||||
if "correction" in signals:
|
||||
hints.append(
|
||||
"IMPORTANT: Explicit correction signals were detected in this conversation. "
|
||||
"Pay special attention to what the agent got wrong, what the user corrected, "
|
||||
"and record the correct approach as a fact with category "
|
||||
'"correction" and confidence >= 0.95 when appropriate.'
|
||||
"Record a correction with confidence >= 0.95 only when it describes a durable, user-level "
|
||||
"working preference that is safe to reuse across unrelated tasks. A correction to facts, files, "
|
||||
"directions, or constraints in the current task is thread- or project-scoped and must not be stored."
|
||||
)
|
||||
if "reinforcement" in signals:
|
||||
hints.append(
|
||||
"IMPORTANT: Positive reinforcement signals were detected in this conversation. "
|
||||
"The user explicitly confirmed the agent's approach was correct or helpful. "
|
||||
"Record the confirmed approach, style, or preference as a fact with category "
|
||||
'"preference" or "behavior" and confidence >= 0.9 when appropriate.'
|
||||
"Record the confirmed approach, style, or preference with high confidence only if it is a durable, "
|
||||
"user-level pattern. Approval of the current result or current task is thread-scoped and must not be stored."
|
||||
)
|
||||
if "preference" in signals:
|
||||
hints.append('IMPORTANT: A preference signal was detected. Record the user\'s stated preference or dislike as a fact with category "preference" and high confidence.')
|
||||
hints.append("IMPORTANT: A preference signal was detected. Record it with high confidence only when it is a durable, user-level preference; a one-off choice for the current task is thread-scoped and must not be stored.")
|
||||
if "identity" in signals:
|
||||
hints.append('IMPORTANT: An identity signal was detected. Record the user\'s stated role, profession, or background as a fact with category "identity" and high confidence.')
|
||||
hints.append("IMPORTANT: An identity signal was detected. Record the user's stated role, profession, or background only when it is user-level and durable across tasks.")
|
||||
if "goal" in signals:
|
||||
hints.append('IMPORTANT: A goal signal was detected. Record the user\'s stated objective or intent as a fact with category "goal" and high confidence.')
|
||||
hints.append("IMPORTANT: A goal signal was detected. Record only a durable, user-level goal that remains useful across unrelated tasks; the objective of the current task, sprint, PR, or thread must not be stored.")
|
||||
if "decision" in signals:
|
||||
hints.append('IMPORTANT: A decision signal was detected. Record the user\'s decision or chosen option as a fact with category "decision" and high confidence.')
|
||||
hints.append("IMPORTANT: A decision signal was detected. Record only a durable, user-level decision or working pattern; a choice made for the current task, file, PR, or thread must not be stored.")
|
||||
return "\n".join(hints)
|
||||
|
||||
def _prepare_update_prompt(
|
||||
@ -1497,22 +1616,34 @@ class MemoryUpdater:
|
||||
update_data: Updates from LLM.
|
||||
thread_id: Optional thread ID for tracking.
|
||||
metrics: Optional observability dict. When provided, populated with
|
||||
``facts_passed_confidence`` / ``rejected_low_confidence`` counted
|
||||
at the real confidence-filter site below (the only acceptance
|
||||
gate for new facts), so the metric cannot drift from the actual
|
||||
filter the way a re-derived count in the caller could.
|
||||
confidence and scope-gate counters counted at their real filter
|
||||
sites, so observability cannot drift from actual acceptance.
|
||||
|
||||
Returns:
|
||||
Updated memory data.
|
||||
"""
|
||||
config = self._config
|
||||
now = utc_now_iso_z()
|
||||
scope_gate_rejections: dict[str, dict[str, int]] = {
|
||||
"facts": {"missing": 0, "scope": 0, "durability": 0, "authority": 0},
|
||||
"summaries": {"missing": 0, "scope": 0, "authority": 0},
|
||||
"removals": {"missing": 0, "scope": 0, "replacement": 0},
|
||||
"consolidations": {"missing": 0, "scope": 0, "durability": 0, "authority": 0},
|
||||
}
|
||||
|
||||
def reject_by_scope_gate(kind: str, reason: str) -> None:
|
||||
scope_gate_rejections[kind][reason] += 1
|
||||
|
||||
# Update user sections
|
||||
user_updates = update_data.get("user", {})
|
||||
for section in ["workContext", "personalContext", "topOfMind"]:
|
||||
section_data = user_updates.get(section, {})
|
||||
if section_data.get("shouldUpdate") and section_data.get("summary"):
|
||||
if not isinstance(section_data, dict) or not section_data.get("shouldUpdate") or not section_data.get("summary"):
|
||||
continue
|
||||
rejection_reason = _summary_scope_gate_reason(section_data)
|
||||
if rejection_reason is not None:
|
||||
reject_by_scope_gate("summaries", rejection_reason)
|
||||
else:
|
||||
current_memory["user"][section] = {
|
||||
"summary": section_data["summary"],
|
||||
"updatedAt": now,
|
||||
@ -1522,17 +1653,17 @@ class MemoryUpdater:
|
||||
history_updates = update_data.get("history", {})
|
||||
for section in ["recentMonths", "earlierContext", "longTermBackground"]:
|
||||
section_data = history_updates.get(section, {})
|
||||
if section_data.get("shouldUpdate") and section_data.get("summary"):
|
||||
if not isinstance(section_data, dict) or not section_data.get("shouldUpdate") or not section_data.get("summary"):
|
||||
continue
|
||||
rejection_reason = _summary_scope_gate_reason(section_data)
|
||||
if rejection_reason is not None:
|
||||
reject_by_scope_gate("summaries", rejection_reason)
|
||||
else:
|
||||
current_memory["history"][section] = {
|
||||
"summary": section_data["summary"],
|
||||
"updatedAt": now,
|
||||
}
|
||||
|
||||
# Remove facts (contradiction-based)
|
||||
facts_to_remove = set(update_data.get("factsToRemove", []))
|
||||
if facts_to_remove:
|
||||
current_memory["facts"] = [f for f in current_memory.get("facts", []) if f.get("id") not in facts_to_remove]
|
||||
|
||||
# ── Staleness review: removals + lifetime extensions ──
|
||||
# Both operations share one staleness-candidate guardrail pass and one
|
||||
# candidate_ids set. proposed_remove_ids is hoisted out of the removals
|
||||
@ -1635,63 +1766,103 @@ class MemoryUpdater:
|
||||
# Creation-time lifetime cap shared with the consolidation path below, so
|
||||
# both fact-creation sites apply the identical bound in one place.
|
||||
creation_cap = int(config.staleness_age_days * config.staleness_max_lifetime_multiplier)
|
||||
# Counted at the confidence-gate site (the only real accept filter for new
|
||||
# facts) so the ``facts_passed_confidence`` metric mirrors the actual
|
||||
# filter and cannot drift from it. Facts below the threshold are the
|
||||
# reject count; duplicate / empty / over-cap facts that pass the
|
||||
# threshold are still counted here -- the metric is a confidence-gate
|
||||
# signal (the host's rejection-rate warning monitors confidence
|
||||
# filtering, not dedup / over-cap), not a persisted-fact count.
|
||||
# Two independent accept filters govern new facts: the deterministic
|
||||
# scope gate and this confidence threshold. Each is counted at its own
|
||||
# filter site so neither metric can drift from the filter it reports:
|
||||
# ``facts_passed_confidence`` counts threshold-passers even when the
|
||||
# scope gate rejects them, and the scope-gate counters increment
|
||||
# whether or not the confidence check passes. Duplicate / empty /
|
||||
# over-cap facts that pass the threshold are still counted here -- the
|
||||
# metric is a confidence-gate signal (the host's rejection-rate
|
||||
# warning monitors confidence filtering, not dedup / over-cap), not a
|
||||
# persisted-fact count.
|
||||
passed_threshold = 0
|
||||
for fact in new_facts:
|
||||
replacement_fact_keys: dict[int, str] = {}
|
||||
for fact_index, fact in enumerate(new_facts):
|
||||
confidence = fact.get("confidence", 0.5)
|
||||
if confidence >= config.fact_confidence_threshold:
|
||||
passed_threshold += 1
|
||||
raw_content = fact.get("content", "")
|
||||
if not isinstance(raw_content, str):
|
||||
continue
|
||||
normalized_content = raw_content.strip()
|
||||
fact_key = _fact_content_key(normalized_content)
|
||||
if fact_key is None:
|
||||
# Empty / whitespace-only content: skip it the same way the
|
||||
# non-string guard above does, instead of appending a blank
|
||||
# fact that violates the non-empty-content invariant.
|
||||
continue
|
||||
if fact_key in existing_fact_keys:
|
||||
continue
|
||||
rejection_reason = _fact_scope_gate_reason(fact)
|
||||
if rejection_reason is not None:
|
||||
reject_by_scope_gate("facts", rejection_reason)
|
||||
continue
|
||||
if confidence < config.fact_confidence_threshold:
|
||||
continue
|
||||
raw_content = fact.get("content", "")
|
||||
if not isinstance(raw_content, str):
|
||||
continue
|
||||
normalized_content = raw_content.strip()
|
||||
fact_key = _fact_content_key(normalized_content)
|
||||
if fact_key is None:
|
||||
# Empty / whitespace-only content: skip it the same way the
|
||||
# non-string guard above does, instead of appending a blank
|
||||
# fact that violates the non-empty-content invariant.
|
||||
continue
|
||||
# Remember every eligible replacement's content key even when it is
|
||||
# already present. A paired removal is safe only if the post-trim
|
||||
# memory contains this content under an ID other than its target.
|
||||
replacement_fact_keys[fact_index] = fact_key
|
||||
if fact_key in existing_fact_keys:
|
||||
continue
|
||||
|
||||
fact_entry = {
|
||||
"id": f"fact_{uuid.uuid4().hex[:8]}",
|
||||
"content": normalized_content,
|
||||
"category": fact.get("category", "context"),
|
||||
"confidence": confidence,
|
||||
"createdAt": now,
|
||||
"source": thread_id or "unknown",
|
||||
}
|
||||
source_error = fact.get("sourceError")
|
||||
if isinstance(source_error, str):
|
||||
normalized_source_error = source_error.strip()
|
||||
if normalized_source_error:
|
||||
fact_entry["sourceError"] = normalized_source_error
|
||||
evd = _read_expected_valid_days(fact)
|
||||
if evd is not None:
|
||||
# Apply the creation-time cap so the LLM cannot assign an
|
||||
# unbounded lifetime that defers staleness review indefinitely.
|
||||
# Extensions (staleFactsToExtend) bypass this cap via their own
|
||||
# staleness_max_extension_days ceiling because they represent a
|
||||
# deliberate review decision, not an unchecked initial assignment.
|
||||
fact_entry["expected_valid_days"] = min(evd, creation_cap)
|
||||
current_memory["facts"].append(fact_entry)
|
||||
if fact_key is not None:
|
||||
existing_fact_keys.add(fact_key)
|
||||
|
||||
if metrics is not None:
|
||||
metrics["facts_passed_confidence"] = passed_threshold
|
||||
metrics["rejected_low_confidence"] = len(new_facts) - passed_threshold
|
||||
fact_entry = {
|
||||
"id": f"fact_{uuid.uuid4().hex[:8]}",
|
||||
"content": normalized_content,
|
||||
"category": fact.get("category", "context"),
|
||||
"confidence": confidence,
|
||||
"createdAt": now,
|
||||
"source": thread_id or "unknown",
|
||||
}
|
||||
source_error = fact.get("sourceError")
|
||||
if isinstance(source_error, str):
|
||||
normalized_source_error = source_error.strip()
|
||||
if normalized_source_error:
|
||||
fact_entry["sourceError"] = normalized_source_error
|
||||
evd = _read_expected_valid_days(fact)
|
||||
if evd is not None:
|
||||
# Apply the creation-time cap so the LLM cannot assign an
|
||||
# unbounded lifetime that defers staleness review indefinitely.
|
||||
# Extensions (staleFactsToExtend) bypass this cap via their own
|
||||
# staleness_max_extension_days ceiling because they represent a
|
||||
# deliberate review decision, not an unchecked initial assignment.
|
||||
fact_entry["expected_valid_days"] = min(evd, creation_cap)
|
||||
current_memory["facts"].append(fact_entry)
|
||||
existing_fact_keys.add(fact_key)
|
||||
|
||||
# Enforce max facts limit (coerced confidence -- see _trim_facts_to_max).
|
||||
current_memory["facts"] = _trim_facts_to_max(current_memory["facts"], config.max_facts)
|
||||
|
||||
# Remove contradicted facts only after replacements have passed both
|
||||
# gates and survived deduplication/trimming. Task-local contradictions
|
||||
# cannot delete user memory, and a failed paired replacement cannot
|
||||
# degrade into a delete-only update.
|
||||
fact_ids_to_remove: set[str] = set()
|
||||
for removal in update_data.get("factsToRemove", []):
|
||||
if not isinstance(removal, dict):
|
||||
reject_by_scope_gate("removals", "missing")
|
||||
continue
|
||||
rejection_reason = _removal_scope_gate_reason(removal)
|
||||
if rejection_reason is not None:
|
||||
reject_by_scope_gate("removals", rejection_reason)
|
||||
continue
|
||||
fact_id = removal.get("id")
|
||||
if not isinstance(fact_id, str) or not fact_id:
|
||||
reject_by_scope_gate("removals", "missing")
|
||||
continue
|
||||
if "replacementFactIndex" in removal:
|
||||
replacement_index = removal.get("replacementFactIndex")
|
||||
if not isinstance(replacement_index, int) or isinstance(replacement_index, bool) or replacement_index < 0:
|
||||
reject_by_scope_gate("removals", "replacement")
|
||||
continue
|
||||
replacement_key = replacement_fact_keys.get(replacement_index)
|
||||
if replacement_key is None or not any(fact.get("id") != fact_id and _fact_content_key(fact.get("content")) == replacement_key for fact in current_memory.get("facts", [])):
|
||||
reject_by_scope_gate("removals", "replacement")
|
||||
continue
|
||||
fact_ids_to_remove.add(fact_id)
|
||||
|
||||
if fact_ids_to_remove:
|
||||
current_memory["facts"] = [fact for fact in current_memory.get("facts", []) if fact.get("id") not in fact_ids_to_remove]
|
||||
|
||||
# ── Memory consolidation ──
|
||||
# Runs after the max_facts trim so source facts that were just evicted
|
||||
# (low confidence, pushed out by high-confidence newFacts) are absent
|
||||
@ -1752,6 +1923,10 @@ class MemoryUpdater:
|
||||
content = consolidated.get("content", "")
|
||||
if not isinstance(content, str) or not content.strip():
|
||||
continue
|
||||
rejection_reason = _fact_scope_gate_reason(consolidated)
|
||||
if rejection_reason is not None:
|
||||
reject_by_scope_gate("consolidations", rejection_reason)
|
||||
continue
|
||||
|
||||
source_confidences = [_coerce_source_confidence(fact_index[sid]) for sid in source_ids]
|
||||
# _coerce_source_confidence already clamps each value to [0, 1],
|
||||
@ -1862,4 +2037,11 @@ class MemoryUpdater:
|
||||
current_memory["facts"] = [f for f in current_memory.get("facts", []) if f.get("id") not in ids_consumed]
|
||||
current_memory["facts"].extend(new_consolidated)
|
||||
|
||||
if metrics is not None:
|
||||
metrics["facts_passed_confidence"] = passed_threshold
|
||||
metrics["rejected_low_confidence"] = len(new_facts) - passed_threshold
|
||||
metrics["facts_passed_scope_gate"] = len(new_facts) - sum(scope_gate_rejections["facts"].values())
|
||||
metrics["rejected_by_scope_gate"] = sum(count for reasons in scope_gate_rejections.values() for count in reasons.values())
|
||||
metrics["scope_gate_rejections"] = scope_gate_rejections
|
||||
|
||||
return current_memory
|
||||
|
||||
@ -38,7 +38,8 @@ class Mem0Config:
|
||||
top_k: int = 8
|
||||
#: Minimum relevance score for search() results (mem0 `threshold`, 0-1).
|
||||
score_threshold: float = 0.1
|
||||
#: Hard cap on the injection text returned by get_context.
|
||||
#: Hard cap on the injection text returned by get_context; memories that
|
||||
#: do not fit whole are skipped (truncation happens on entry boundaries).
|
||||
max_injection_chars: int = 12000
|
||||
#: Per-request HTTP timeout in seconds.
|
||||
timeout_seconds: float = 10.0
|
||||
|
||||
@ -196,19 +196,41 @@ class Mem0Manager(MemoryManager):
|
||||
max_items=top_k,
|
||||
),
|
||||
)
|
||||
budget = self._config.max_injection_chars
|
||||
seen: set[str] = set()
|
||||
lines: list[str] = []
|
||||
used = 0
|
||||
shortest_line: int | None = None
|
||||
for record in records:
|
||||
rid = record.get("id")
|
||||
if rid in seen:
|
||||
continue
|
||||
seen.add(rid)
|
||||
text = str(record.get("memory") or "").strip()
|
||||
if text:
|
||||
lines.append(f"- {text}")
|
||||
if not text:
|
||||
continue
|
||||
line = f"- {text}"
|
||||
line_len = len(line)
|
||||
shortest_line = line_len if shortest_line is None else min(shortest_line, line_len)
|
||||
# Truncate on entry boundaries: keep only memories that fit whole
|
||||
# within the remaining budget (+1 for the joining newline), so the
|
||||
# injection never ends mid-entry with a dangling partial line. An
|
||||
# oversized entry is skipped -- a shorter later one may still fit.
|
||||
added = line_len if not lines else line_len + 1
|
||||
if used + added > budget:
|
||||
continue
|
||||
lines.append(line)
|
||||
used += added
|
||||
context = "\n".join(lines)
|
||||
if len(context) > self._config.max_injection_chars:
|
||||
context = context[: self._config.max_injection_chars]
|
||||
if not context and shortest_line is not None:
|
||||
# Every recalled memory was longer than the configured budget.
|
||||
# Keep the entry-boundary guarantee and surface the config problem
|
||||
# with a warning rather than injecting a partial fact.
|
||||
logger.warning(
|
||||
"max_injection_chars=%d is smaller than the shortest recalled memory (%d chars); returning empty context",
|
||||
budget,
|
||||
shortest_line,
|
||||
)
|
||||
return context
|
||||
|
||||
async def aget_context(
|
||||
|
||||
@ -694,6 +694,8 @@ def _host_default_extraction_callback(payload: Any) -> None:
|
||||
extracted = payload.get("facts_extracted")
|
||||
passed_confidence = payload.get("facts_passed_confidence")
|
||||
rejected = payload.get("rejected_low_confidence", 0)
|
||||
rejected_by_scope = payload.get("rejected_by_scope_gate", 0)
|
||||
scope_breakdown = payload.get("scope_gate_rejections")
|
||||
thread_id = payload.get("thread_id")
|
||||
model_name = payload.get("model_name")
|
||||
if isinstance(extracted, int) and isinstance(passed_confidence, int) and extracted > 0:
|
||||
@ -721,6 +723,22 @@ def _host_default_extraction_callback(payload: Any) -> None:
|
||||
payload.get("success"),
|
||||
payload.get("token_usage"),
|
||||
)
|
||||
if isinstance(scope_breakdown, dict):
|
||||
logger.info(
|
||||
"Memory scope-gate metrics: thread=%s model=%s rejected=%s breakdown=%s",
|
||||
thread_id,
|
||||
model_name,
|
||||
rejected_by_scope,
|
||||
scope_breakdown,
|
||||
)
|
||||
fact_breakdown = scope_breakdown.get("facts")
|
||||
fact_scope_rejected = sum(value for value in fact_breakdown.values() if isinstance(value, int)) if isinstance(fact_breakdown, dict) else 0
|
||||
if isinstance(extracted, int) and extracted > 0 and fact_scope_rejected / extracted > 0.6:
|
||||
logger.warning(
|
||||
"Memory fact scope-gate rejection rate %.0f%% exceeds 60%% - review extraction model classification / prompt (thread=%s)",
|
||||
fact_scope_rejected / extracted * 100,
|
||||
thread_id,
|
||||
)
|
||||
|
||||
|
||||
def _collect_host_hooks() -> dict[str, Any]:
|
||||
|
||||
@ -118,9 +118,11 @@ def memory_add_tool(
|
||||
content_key = _memory_content_key(normalized_content)
|
||||
manager = get_memory_manager()
|
||||
existing_facts = manager.get_memory(agent_name=agent_name, user_id=user_id).get("facts", [])
|
||||
# Tool calls normally run one-at-a-time per user turn. If tool-mode
|
||||
# writing broadens to multiple concurrent calls for the same user,
|
||||
# move duplicate rejection into the storage/update critical section.
|
||||
# Fast-path duplicate rejection to spare a write attempt in the common
|
||||
# case. The authoritative check lives in the backend's create critical
|
||||
# section (DeerMem re-checks against a fresh snapshot on every
|
||||
# revision-conflict retry in create_memory_fact), so concurrent tool
|
||||
# calls for the same user cannot both store the same content.
|
||||
if any(_memory_content_key(str(fact.get("content", ""))) == content_key for fact in existing_facts):
|
||||
return json.dumps({"error": "Duplicate fact"})
|
||||
|
||||
|
||||
@ -21,6 +21,31 @@ logger = logging.getLogger(__name__)
|
||||
# Command classification rules
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Executables whose output is dangerous to *execute*. Used by the command
|
||||
# substitution rules below; ``\b`` prevents matching unrelated names that merely
|
||||
# start with one of these words (``shellcheck``, ``shasum``, ``pythonic-tool``).
|
||||
_RISKY_SUBSTITUTION_EXECUTABLES = r"(?:curl|wget|bash|sh|python[\d.]*|ruby|perl|base64)\b"
|
||||
|
||||
# A substitution opening one of those executables, in any of its spellings:
|
||||
# ``$(cmd``, ``<(cmd``, or the backtick form, which has no parenthesis. Sharing
|
||||
# one opener is what keeps ``eval `curl u` `` from slipping past a rule written
|
||||
# only for ``eval $(curl u)``.
|
||||
_RISKY_SUBSTITUTION = rf"(?:[$<]\(\s*|`\s*){_RISKY_SUBSTITUTION_EXECUTABLES}"
|
||||
|
||||
# Interpreters that execute a *code string* handed to them as an argument, and
|
||||
# the flags that receive it: ``-c`` (shells, python), ``-e`` (perl/ruby/node),
|
||||
# ``-p`` (perl/node print loop), ``-r`` (php). Whatever the flag receives is
|
||||
# executed, so a risky substitution there is executed too -- the same class as
|
||||
# ``eval``/``source``, spelled with a flag instead. A here-string (``<<<``)
|
||||
# reaches the same place through stdin.
|
||||
#
|
||||
# These are position-blind on purpose: ``bash -c`` is an execution context
|
||||
# wherever it appears, including as an argument to something else
|
||||
# (``xargs sh -c "$(curl url)"``). The leading-flag repetition is bounded so the
|
||||
# alternation cannot backtrack on long input.
|
||||
_CODE_STRING_INTERPRETERS = r"(?:(?:ba|da|k|z)?sh|python[\d.]*|perl|ruby|node|php)"
|
||||
_LEADING_FLAGS = r"(?:-\w+\s+){0,4}"
|
||||
|
||||
# Each pattern is compiled once at import time.
|
||||
_HIGH_RISK_PATTERNS: list[re.Pattern[str]] = [
|
||||
# --- original rules (retained) ---
|
||||
@ -31,8 +56,11 @@ _HIGH_RISK_PATTERNS: list[re.Pattern[str]] = [
|
||||
re.compile(r">+\s*/etc/"),
|
||||
# --- pipe to sh/bash (generalised, replaces old curl|sh rule) ---
|
||||
re.compile(r"\|\s*(ba)?sh\b"),
|
||||
# --- command substitution (targeted – only dangerous executables) ---
|
||||
re.compile(r"[`$]\(?\s*(curl|wget|bash|sh|python|ruby|perl|base64)"),
|
||||
# --- eval/source execute a substitution regardless of its position ---
|
||||
re.compile(rf"\b(eval|source)\s+[\"']?{_RISKY_SUBSTITUTION}"),
|
||||
# --- an interpreter's code-string flag is an execution context too ---
|
||||
re.compile(rf"\b{_CODE_STRING_INTERPRETERS}\s+{_LEADING_FLAGS}-[cepr]\s+[\"']?{_RISKY_SUBSTITUTION}"),
|
||||
re.compile(rf"\b{_CODE_STRING_INTERPRETERS}\s+{_LEADING_FLAGS}<<<\s*[\"']?{_RISKY_SUBSTITUTION}"),
|
||||
# --- base64 decode piped to execution ---
|
||||
re.compile(r"base64\s+.*-d.*\|"),
|
||||
# --- overwrite system binaries ---
|
||||
@ -50,6 +78,32 @@ _HIGH_RISK_PATTERNS: list[re.Pattern[str]] = [
|
||||
re.compile(r"while\s+true.*&\s*done"), # while true; do bash & done
|
||||
]
|
||||
|
||||
# Command substitution in *command position*: the substitution result becomes the
|
||||
# command that runs, so fetched or interpreted content is executed.
|
||||
#
|
||||
# These are matched anchored against a single sub-command, never against the whole
|
||||
# compound string, because position is what distinguishes the two shapes:
|
||||
#
|
||||
# $(curl url) → executes what was downloaded → block
|
||||
# x=$(curl url) → captures the output into a variable → pass
|
||||
# echo $(curl url) → passes the output as an argument → pass
|
||||
#
|
||||
# The previous unanchored rule could not tell them apart and refused everyday
|
||||
# output capture (issue #4611).
|
||||
#
|
||||
# A command position is not always the first character: POSIX shell allows leading
|
||||
# variable assignments, and exec wrappers keep what follows in command position
|
||||
# (``FOO=1 $(curl url)``, ``env FOO=1 $(curl url)``, ``nohup $(curl url)``). The
|
||||
# assignment branch cannot match ``x=$(curl url)`` because it requires whitespace
|
||||
# between the assignment and the substitution, so value position stays allowed.
|
||||
# The repetition is bounded to keep the alternation from backtracking on long input.
|
||||
_COMMAND_POSITION_PREFIX = r"(?:(?:env|command|builtin|exec|nohup|time|sudo|doas)\s+|\w+=\S*\s+){0,8}"
|
||||
|
||||
_HIGH_RISK_COMMAND_POSITION_PATTERNS: list[re.Pattern[str]] = [
|
||||
re.compile(rf"^{_COMMAND_POSITION_PREFIX}[\"']?\$\(\s*{_RISKY_SUBSTITUTION_EXECUTABLES}"),
|
||||
re.compile(rf"^{_COMMAND_POSITION_PREFIX}[\"']?`\s*{_RISKY_SUBSTITUTION_EXECUTABLES}"),
|
||||
]
|
||||
|
||||
_MEDIUM_RISK_PATTERNS: list[re.Pattern[str]] = [
|
||||
re.compile(r"chmod\s+777"),
|
||||
re.compile(r"pip3?\s+install"),
|
||||
@ -61,7 +115,39 @@ _MEDIUM_RISK_PATTERNS: list[re.Pattern[str]] = [
|
||||
]
|
||||
|
||||
|
||||
def _split_compound_command(command: str) -> list[str]:
|
||||
# A heredoc header and its delimiter: ``<<EOF``, ``<< EOF``, ``<<-EOF``,
|
||||
# ``<<\EOF``, ``<<'EOF'``, ``<<"EOF"``. Both guards are needed to keep ``<<<``
|
||||
# (a here-string, which has no body) from opening one: the lookahead rejects it
|
||||
# at its first ``<``, and the lookbehind stops its trailing ``<<`` from matching
|
||||
# one character later, where ``<<< "text"`` would otherwise read as a heredoc
|
||||
# with delimiter ``text``.
|
||||
_HEREDOC_HEADER = re.compile(r"(?<!<)<<(?!<)-?[ \t]*(?:\\?([A-Za-z_][\w.-]*)|'([^'\n]*)'|\"([^\"\n]*)\")")
|
||||
|
||||
|
||||
def _consume_heredoc_bodies(command: str, pos: int, delimiters: list[str]) -> int:
|
||||
"""Return the index just past the bodies of the *delimiters* opened so far.
|
||||
|
||||
Bodies are consumed in the order their headers appeared, each running until a
|
||||
line whose stripped content equals its delimiter (``<<-`` strips leading tabs,
|
||||
which ``strip()`` covers). An unterminated body consumes the rest of the
|
||||
string: everything after the header genuinely is body, and there is no later
|
||||
statement to find.
|
||||
"""
|
||||
for delimiter in delimiters:
|
||||
while pos < len(command):
|
||||
newline = command.find("\n", pos)
|
||||
if newline == -1:
|
||||
return len(command)
|
||||
line = command[pos:newline]
|
||||
pos = newline + 1
|
||||
if line.strip() == delimiter:
|
||||
break
|
||||
else:
|
||||
return len(command)
|
||||
return pos
|
||||
|
||||
|
||||
def _split_compound_command(command: str, *, split_pipes: bool = False) -> list[str]:
|
||||
"""Split a compound command into sub-commands (quote-aware).
|
||||
|
||||
Scans the raw command string so unquoted shell control operators are
|
||||
@ -70,11 +156,36 @@ def _split_compound_command(command: str) -> list[str]:
|
||||
quotes are ignored. If the command ends with an unclosed quote or a
|
||||
dangling escape, return the whole command unchanged (fail-closed —
|
||||
safer to classify the unsplit string than silently drop parts).
|
||||
|
||||
Sequencing operators (``&&``, ``||``, ``;``) split, and so does an unquoted
|
||||
newline — it separates statements exactly like ``;``, so leaving it joined let
|
||||
``echo hi\\n$(curl url)`` evade the anchored command-position rules that
|
||||
``echo hi; $(curl url)`` triggers, despite identical shell semantics.
|
||||
|
||||
A heredoc body is data, not statements: its newlines and operators are file
|
||||
content. Headers (``<<EOF``, ``<<-EOF``, ``<<'EOF'``) are therefore recorded
|
||||
as they are read and their bodies consumed verbatim at the newline that
|
||||
starts them, so a body line beginning with ``$(curl url)`` is not promoted to
|
||||
command position. ``<<<`` is a here-string, not a heredoc, and does not open
|
||||
one; neither does a ``<<`` inside ``$(( ... ))`` or ``(( ... ))``, where it is
|
||||
a bit shift whose right operand would otherwise read as a delimiter that never
|
||||
appears — swallowing the rest of the command. This is a heuristic, not shell
|
||||
parsing — the goal is only to avoid manufacturing command positions that the
|
||||
shell would never create, and to avoid destroying real ones.
|
||||
|
||||
Pipes do not split by default, because a pipeline is one logical command.
|
||||
Pass ``split_pipes=True`` to also split on ``|``, which is what
|
||||
command-position detection needs — the word after a pipe starts a new
|
||||
command. Rules that span a pipe (``| sh``, ``base64 -d | ...``) are matched by
|
||||
the whole-command scan in :func:`_classify_command`, so they are unaffected by
|
||||
the extra split.
|
||||
"""
|
||||
parts: list[str] = []
|
||||
current: list[str] = []
|
||||
pending_heredocs: list[str] = []
|
||||
in_single_quote = False
|
||||
in_double_quote = False
|
||||
arithmetic_depth = 0
|
||||
escaping = False
|
||||
index = 0
|
||||
|
||||
@ -106,6 +217,47 @@ def _split_compound_command(command: str) -> list[str]:
|
||||
continue
|
||||
|
||||
if not in_single_quote and not in_double_quote:
|
||||
# ``<<`` inside arithmetic is a bit shift, not a redirection, and a
|
||||
# phantom header whose delimiter never appears would swallow the rest
|
||||
# of the command. Both ``$(( ... ))`` and the bare arithmetic command
|
||||
# ``(( ... ))`` are tracked. An unclosed ``((`` leaves the depth
|
||||
# positive, which only disables heredoc detection — newlines keep
|
||||
# splitting, so the failure direction stays towards seeing more
|
||||
# command positions rather than fewer.
|
||||
if char == "(" and command.startswith("((", index):
|
||||
arithmetic_depth += 1
|
||||
current.append("((")
|
||||
index += 2
|
||||
continue
|
||||
if arithmetic_depth and char == ")" and command.startswith("))", index):
|
||||
arithmetic_depth -= 1
|
||||
current.append("))")
|
||||
index += 2
|
||||
continue
|
||||
# A header can only start at ``<``; checking that first keeps the
|
||||
# regex off every other character of a long command.
|
||||
if char == "<" and not arithmetic_depth:
|
||||
heredoc = _HEREDOC_HEADER.match(command, index)
|
||||
if heredoc:
|
||||
pending_heredocs.append(next(group for group in heredoc.groups() if group is not None))
|
||||
current.append(heredoc.group(0))
|
||||
index = heredoc.end()
|
||||
continue
|
||||
if char == "\n":
|
||||
# The newline that follows a heredoc header is the statement
|
||||
# separator, and its body belongs to the statement being closed.
|
||||
if pending_heredocs:
|
||||
body_end = _consume_heredoc_bodies(command, index + 1, pending_heredocs)
|
||||
pending_heredocs = []
|
||||
current.append(command[index:body_end])
|
||||
index = body_end
|
||||
else:
|
||||
index += 1
|
||||
part = "".join(current).strip()
|
||||
if part:
|
||||
parts.append(part)
|
||||
current = []
|
||||
continue
|
||||
if command.startswith("&&", index) or command.startswith("||", index):
|
||||
part = "".join(current).strip()
|
||||
if part:
|
||||
@ -113,6 +265,14 @@ def _split_compound_command(command: str) -> list[str]:
|
||||
current = []
|
||||
index += 2
|
||||
continue
|
||||
# Checked after "||" so a single "|" cannot steal that operator.
|
||||
if split_pipes and char == "|":
|
||||
part = "".join(current).strip()
|
||||
if part:
|
||||
parts.append(part)
|
||||
current = []
|
||||
index += 1
|
||||
continue
|
||||
if char == ";":
|
||||
part = "".join(current).strip()
|
||||
if part:
|
||||
@ -134,21 +294,27 @@ def _split_compound_command(command: str) -> list[str]:
|
||||
return parts if parts else [command]
|
||||
|
||||
|
||||
def _matches_high_risk(candidate: str) -> bool:
|
||||
"""Return True if *candidate* (one sub-command) matches any high-risk rule."""
|
||||
if any(pattern.search(candidate) for pattern in _HIGH_RISK_PATTERNS):
|
||||
return True
|
||||
# Anchored: only meaningful for a single sub-command, not a compound string.
|
||||
return any(pattern.match(candidate) for pattern in _HIGH_RISK_COMMAND_POSITION_PATTERNS)
|
||||
|
||||
|
||||
def _classify_single_command(command: str) -> str:
|
||||
"""Classify a single (non-compound) command. Return 'block', 'warn', or 'pass'."""
|
||||
normalized = " ".join(command.split())
|
||||
|
||||
for pattern in _HIGH_RISK_PATTERNS:
|
||||
if pattern.search(normalized):
|
||||
return "block"
|
||||
if _matches_high_risk(normalized):
|
||||
return "block"
|
||||
|
||||
# Also try shlex-parsed tokens for high-risk detection
|
||||
try:
|
||||
tokens = shlex.split(command)
|
||||
joined = " ".join(tokens)
|
||||
for pattern in _HIGH_RISK_PATTERNS:
|
||||
if pattern.search(joined):
|
||||
return "block"
|
||||
if _matches_high_risk(joined):
|
||||
return "block"
|
||||
except ValueError:
|
||||
# Heredocs and other multiline shell forms may be valid bash but
|
||||
# unparseable by shlex. Raw high-risk patterns were already checked.
|
||||
@ -178,8 +344,9 @@ def _classify_command(command: str) -> str:
|
||||
if pattern.search(normalized):
|
||||
return "block"
|
||||
|
||||
# Pass 2: per-sub-command classification
|
||||
sub_commands = _split_compound_command(command)
|
||||
# Pass 2: per-sub-command classification. Pipes split here too, because the
|
||||
# word after a pipe starts a new command position (``echo hi | $(curl ...)``).
|
||||
sub_commands = _split_compound_command(command, split_pipes=True)
|
||||
worst = "pass"
|
||||
for sub in sub_commands:
|
||||
verdict = _classify_single_command(sub)
|
||||
|
||||
@ -33,7 +33,7 @@ from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from deerflow.agents.lead_agent.agent import build_middlewares
|
||||
from deerflow.agents.lead_agent.agent import _authorize_model_name, 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.authz.principal import build_principal_from_context
|
||||
@ -72,6 +72,7 @@ from deerflow.uploads.manager import (
|
||||
upload_artifact_url,
|
||||
upload_virtual_path,
|
||||
)
|
||||
from deerflow.utils.thread_id import resolve_thread_id, validate_thread_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -289,6 +290,16 @@ class DeerFlowClient:
|
||||
|
||||
thinking_enabled = cfg.get("thinking_enabled", True)
|
||||
model_name = cfg.get("model_name")
|
||||
# Phase 3: enforce model:use authorization on the embedded/library path
|
||||
# too, mirroring the Gateway runtime path in ``_make_lead_agent`` so the
|
||||
# role-scoped model policy cannot be bypassed by constructing the agent
|
||||
# through ``DeerFlowClient``. Resolve the ``None`` default to a concrete
|
||||
# name first (what ``create_chat_model(name=None)`` would pick) so the
|
||||
# policy covers the implicit default model. ``cfg`` already carries the
|
||||
# identity that ``apply_tool_authorization`` reads below.
|
||||
if model_name is None and self._app_config.models:
|
||||
model_name = self._app_config.models[0].name
|
||||
model_name = _authorize_model_name(model_name, context=cfg, app_config=self._app_config)
|
||||
subagent_enabled = cfg.get("subagent_enabled", False)
|
||||
max_concurrent_subagents = cfg.get("max_concurrent_subagents", 3)
|
||||
max_total_subagents = cfg.get("max_total_subagents", self._app_config.subagents.max_total_per_run)
|
||||
@ -525,6 +536,7 @@ class DeerFlowClient:
|
||||
|
||||
def get_goal(self, thread_id: str) -> dict:
|
||||
"""Return the active goal for a thread, if any."""
|
||||
validate_thread_id(thread_id)
|
||||
checkpointer = self._get_thread_checkpointer()
|
||||
goal = _run_async_from_sync(read_thread_goal(checkpointer, thread_id))
|
||||
return {"goal": goal}
|
||||
@ -537,6 +549,7 @@ class DeerFlowClient:
|
||||
max_continuations: int = DEFAULT_MAX_GOAL_CONTINUATIONS,
|
||||
) -> dict:
|
||||
"""Set or replace a thread-scoped goal."""
|
||||
validate_thread_id(thread_id)
|
||||
checkpointer = self._get_thread_checkpointer()
|
||||
goal = build_goal_state(objective, max_continuations=max_continuations)
|
||||
|
||||
@ -549,6 +562,7 @@ class DeerFlowClient:
|
||||
|
||||
def clear_goal(self, thread_id: str) -> dict:
|
||||
"""Clear the active goal for a thread."""
|
||||
validate_thread_id(thread_id)
|
||||
checkpointer = self._get_thread_checkpointer()
|
||||
|
||||
async def _clear_goal() -> None:
|
||||
@ -803,8 +817,7 @@ class DeerFlowClient:
|
||||
Tool results also include ``"artifact"`` when the source ToolMessage has a non-None artifact.
|
||||
- type="end" data={"usage": {"input_tokens": int, "output_tokens": int, "total_tokens": int}}
|
||||
"""
|
||||
if thread_id is None:
|
||||
thread_id = str(uuid.uuid4())
|
||||
thread_id = resolve_thread_id(thread_id)
|
||||
|
||||
config = self._get_runnable_config(thread_id, **kwargs)
|
||||
inject_checkpoint_mode(config, self._checkpoint_channel_mode)
|
||||
@ -1281,13 +1294,19 @@ class DeerFlowClient:
|
||||
if config_path is None:
|
||||
raise FileNotFoundError("Cannot locate extensions_config.json. Set DEER_FLOW_EXTENSIONS_CONFIG_PATH or ensure it exists in the project root.")
|
||||
|
||||
extensions_config = get_extensions_config()
|
||||
extensions_config.skills[name] = SkillStateConfig(enabled=enabled)
|
||||
from deerflow.skills.projection import skill_projection_mutation
|
||||
|
||||
config_data = extensions_config.to_file_dict()
|
||||
removal_names = (name,) if not enabled else ()
|
||||
with skill_projection_mutation(storage, "public", remove_names=removal_names):
|
||||
# The projection lock is cross-process, but the singleton cache
|
||||
# is not. Reload from disk under the lock before this RMW.
|
||||
extensions_config = ExtensionsConfig.from_file(config_path)
|
||||
extensions_config.skills[name] = SkillStateConfig(enabled=enabled)
|
||||
|
||||
self._atomic_write_json(config_path, config_data)
|
||||
reload_extensions_config()
|
||||
config_data = extensions_config.to_file_dict()
|
||||
|
||||
self._atomic_write_json(config_path, config_data)
|
||||
reload_extensions_config()
|
||||
else:
|
||||
# CUSTOM / LEGACY: write per-user state
|
||||
from deerflow.skills.storage.user_scoped_skill_storage import UserScopedSkillStorage
|
||||
@ -1475,6 +1494,7 @@ class DeerFlowClient:
|
||||
FileNotFoundError: If any file does not exist.
|
||||
ValueError: If any supplied path exists but is not a regular file.
|
||||
"""
|
||||
validate_thread_id(thread_id)
|
||||
from deerflow.utils.file_conversion import CONVERTIBLE_EXTENSIONS, convert_file_to_markdown
|
||||
|
||||
# Validate all files upfront to avoid partial uploads.
|
||||
@ -1578,6 +1598,7 @@ class DeerFlowClient:
|
||||
Dict with "files" and "count" keys, matching the Gateway API
|
||||
``list_uploaded_files`` response.
|
||||
"""
|
||||
validate_thread_id(thread_id)
|
||||
uploads_dir = get_uploads_dir(thread_id)
|
||||
result = list_files_in_dir(uploads_dir)
|
||||
return enrich_file_listing(result, thread_id)
|
||||
@ -1597,6 +1618,7 @@ class DeerFlowClient:
|
||||
FileNotFoundError: If the file does not exist.
|
||||
PermissionError: If path traversal is detected.
|
||||
"""
|
||||
validate_thread_id(thread_id)
|
||||
from deerflow.utils.file_conversion import CONVERTIBLE_EXTENSIONS
|
||||
|
||||
uploads_dir = get_uploads_dir(thread_id)
|
||||
@ -1620,6 +1642,7 @@ class DeerFlowClient:
|
||||
FileNotFoundError: If the artifact does not exist.
|
||||
ValueError: If the path is invalid.
|
||||
"""
|
||||
validate_thread_id(thread_id)
|
||||
try:
|
||||
actual = get_paths().resolve_virtual_path(thread_id, path, user_id=get_effective_user_id())
|
||||
except ValueError as exc:
|
||||
|
||||
@ -44,7 +44,6 @@ from deerflow.integrations.lark_cli import LARK_CLI_SANDBOX_CONFIG_DIR, LARK_CLI
|
||||
from deerflow.runtime.user_context import get_effective_user_id
|
||||
from deerflow.sandbox.sandbox import Sandbox
|
||||
from deerflow.sandbox.sandbox_provider import SandboxProvider
|
||||
from deerflow.skills.storage import user_should_see_legacy_skills
|
||||
|
||||
from .aio_sandbox import AioSandbox
|
||||
from .backend import SandboxBackend, wait_for_sandbox_ready, wait_for_sandbox_ready_async
|
||||
@ -868,41 +867,34 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider):
|
||||
that ``Skill.get_container_path()`` category-aware paths resolve
|
||||
correctly inside the sandbox.
|
||||
|
||||
Mount sources use ``DEER_FLOW_HOST_SKILLS_PATH`` and
|
||||
``DEER_FLOW_HOST_BASE_DIR`` when running inside Docker (DooD) so the
|
||||
host Docker daemon can resolve the paths.
|
||||
Mount sources use ``DEER_FLOW_HOST_BASE_DIR`` when running inside
|
||||
Docker (DooD) so the host Docker daemon can resolve the projection
|
||||
paths.
|
||||
"""
|
||||
mounts: list[tuple[str, str, bool]] = []
|
||||
try:
|
||||
config = get_app_config()
|
||||
skills_path = config.skills.get_skills_path()
|
||||
container_path = config.skills.container_path
|
||||
|
||||
# When running inside Docker with DooD, use host-side skills path.
|
||||
host_skills_root = os.environ.get("DEER_FLOW_HOST_SKILLS_PATH") or str(skills_path)
|
||||
effective_user_id = AioSandboxProvider._effective_acquire_user_id(user_id)
|
||||
AioSandboxProvider._ensure_skills_projection(effective_user_id)
|
||||
paths = get_paths()
|
||||
host_base_dir = str(paths.host_base_dir)
|
||||
|
||||
# 1. Public skills: global, read-only — static, shared by all threads
|
||||
public_skills_path = skills_path / "public"
|
||||
if public_skills_path.exists():
|
||||
mounts.append(
|
||||
(
|
||||
join_host_path(host_skills_root, "public"),
|
||||
f"{container_path}/public",
|
||||
True,
|
||||
)
|
||||
mounts.append(
|
||||
(
|
||||
join_host_path(host_base_dir, "skills_view", "public"),
|
||||
f"{container_path}/public",
|
||||
True,
|
||||
)
|
||||
)
|
||||
|
||||
# 2. Per-user custom skills: read-only, per-thread/per-user
|
||||
effective_user_id = AioSandboxProvider._effective_acquire_user_id(user_id)
|
||||
paths = get_paths()
|
||||
user_custom_path = paths.user_custom_skills_dir(effective_user_id)
|
||||
user_custom_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
host_user_custom = join_host_path(
|
||||
str(paths.host_base_dir),
|
||||
host_base_dir,
|
||||
"users",
|
||||
effective_user_id,
|
||||
"skills",
|
||||
"skills_view",
|
||||
"custom",
|
||||
)
|
||||
mounts.append(
|
||||
@ -913,38 +905,66 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider):
|
||||
)
|
||||
)
|
||||
|
||||
# 3. Legacy (pre-migration global-custom) skills: only mount for
|
||||
# users who have no per-user custom skills yet, mirroring
|
||||
# ``UserScopedSkillStorage._iter_skill_files`` visibility rule.
|
||||
legacy_skills_path = skills_path / "custom"
|
||||
if user_should_see_legacy_skills(effective_user_id, host_path=str(skills_path)) and legacy_skills_path.exists():
|
||||
mounts.append(
|
||||
(
|
||||
join_host_path(host_skills_root, "custom"),
|
||||
f"{container_path}/legacy",
|
||||
True,
|
||||
)
|
||||
# 3. Legacy visibility is encoded by projection contents. Keep the
|
||||
# mount stable even when the directory is empty so a later state
|
||||
# change is visible without recreating the sandbox.
|
||||
mounts.append(
|
||||
(
|
||||
join_host_path(host_base_dir, "users", effective_user_id, "skills_view", "legacy"),
|
||||
f"{container_path}/legacy",
|
||||
True,
|
||||
)
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("Could not setup skills mounts: %s", e)
|
||||
|
||||
return mounts
|
||||
|
||||
@staticmethod
|
||||
def _ensure_skills_projection(user_id: str):
|
||||
"""Best-effort: a projection failure must not fail sandbox acquire.
|
||||
|
||||
Called directly (for its side effect) from ``_acquire_internal`` /
|
||||
``_acquire_internal_async`` outside any try/except, as well as from
|
||||
within ``_get_skills_mounts``'s own guarded block — swallowing here
|
||||
keeps both call sites safe without duplicating the guard.
|
||||
"""
|
||||
from deerflow.skills.projection import ensure_skill_projections
|
||||
from deerflow.skills.storage import get_or_new_user_skill_storage
|
||||
|
||||
try:
|
||||
storage = get_or_new_user_skill_storage(user_id, app_config=get_app_config())
|
||||
return ensure_skill_projections(storage)
|
||||
except Exception as exc:
|
||||
logger.warning("Could not ensure skills projection for user %s: %s", user_id, exc, exc_info=True)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_user_skill_mounts(*, user_id: str | None = None) -> list[tuple[str, str, bool]]:
|
||||
"""Mount managed integration skills into AIO sandboxes.
|
||||
"""Mount enabled managed integration skills into AIO sandboxes.
|
||||
|
||||
Per-user custom skills are already mounted by ``_get_skills_mounts``.
|
||||
This helper adds the shared integration skill root so sandbox paths match
|
||||
the skill registry without duplicating ``/mnt/skills/custom``.
|
||||
Integration packages are shared, but their enabled state is per-user, so
|
||||
this helper mounts the user's projection instead of the raw shared root.
|
||||
"""
|
||||
try:
|
||||
config = get_app_config()
|
||||
paths = get_paths()
|
||||
skills_container_path = config.skills.container_path
|
||||
paths.integration_skills_dir().mkdir(parents=True, exist_ok=True)
|
||||
effective_user_id = AioSandboxProvider._effective_acquire_user_id(user_id)
|
||||
AioSandboxProvider._ensure_skills_projection(effective_user_id)
|
||||
return [
|
||||
(paths.host_integration_skills_dir(), f"{skills_container_path}/integrations", True),
|
||||
(
|
||||
join_host_path(
|
||||
str(paths.host_base_dir),
|
||||
"users",
|
||||
effective_user_id,
|
||||
"skills_view",
|
||||
"integrations",
|
||||
),
|
||||
f"{skills_container_path}/integrations",
|
||||
True,
|
||||
),
|
||||
]
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not setup user skill mounts: {e}")
|
||||
@ -1810,6 +1830,7 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider):
|
||||
sandbox_id is deterministic from thread_id so no shared state file
|
||||
is needed — any process can derive the same container name)
|
||||
"""
|
||||
self._ensure_skills_projection(user_id)
|
||||
cached_id = self._reuse_in_process_sandbox(thread_id, user_id=user_id)
|
||||
if cached_id is not None:
|
||||
return cached_id
|
||||
@ -1837,6 +1858,7 @@ class AioSandboxProvider(WarmPoolLifecycleMixin[SandboxInfo], SandboxProvider):
|
||||
|
||||
async def _acquire_internal_async(self, thread_id: str | None, *, user_id: str) -> str:
|
||||
"""Async counterpart to ``_acquire_internal``."""
|
||||
await asyncio.to_thread(self._ensure_skills_projection, user_id)
|
||||
cached_id = await asyncio.to_thread(self._reuse_in_process_sandbox, thread_id, user_id=user_id)
|
||||
if cached_id is not None:
|
||||
return cached_id
|
||||
|
||||
@ -61,7 +61,10 @@ def resolve_ownership_config(config: SandboxOwnershipConfig | None, *, stream_br
|
||||
return SandboxOwnershipConfig()
|
||||
|
||||
|
||||
def _resolve_redis_url(config: SandboxOwnershipConfig) -> str:
|
||||
def resolve_ownership_redis_url(
|
||||
config: SandboxOwnershipConfig,
|
||||
) -> str:
|
||||
"""Resolve the Redis endpoint shared by ownership-adjacent stores."""
|
||||
return config.redis_url or os.getenv(_ENV_OWNERSHIP_REDIS_URL) or os.getenv(_ENV_STREAM_BRIDGE_REDIS_URL) or os.getenv("REDIS_URL") or "redis://localhost:6379/0"
|
||||
|
||||
|
||||
@ -96,7 +99,7 @@ def make_sandbox_ownership_store(config: SandboxOwnershipConfig | None, *, owner
|
||||
if resolved.type == "redis":
|
||||
from .redis import RedisOwnershipStore
|
||||
|
||||
redis_url = _resolve_redis_url(resolved)
|
||||
redis_url = resolve_ownership_redis_url(resolved)
|
||||
logger.info("Sandbox ownership store: redis (ttl=%.1fs, renewal=%.1fs)", ttl, resolved.renewal_interval_seconds)
|
||||
return RedisOwnershipStore(
|
||||
owner_id=effective_owner_id,
|
||||
|
||||
@ -15,7 +15,6 @@ private loop so the core harness installs without it.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import contextlib
|
||||
import logging
|
||||
import os
|
||||
@ -301,7 +300,7 @@ class BrowserSession:
|
||||
# Live screencast state. When streaming, ``_on_frame`` is retained so the
|
||||
# screencast can be re-bound to a new page — login/OAuth flows commonly
|
||||
# open a popup or a fresh tab, and the user must see (and drive) it.
|
||||
self._on_frame: Callable[[str], None] | None = None
|
||||
self._on_frame: Callable[[bytes], None] | None = None
|
||||
# The page the live screencast's CDP session is currently bound to. Frames
|
||||
# are captured from ``self._page`` (the live active page), but the CDP
|
||||
# repaint signal is tied to a specific page; when the active page diverges
|
||||
@ -538,10 +537,9 @@ class BrowserSession:
|
||||
page = await self._ensure_page()
|
||||
return await page.screenshot(full_page=full_page, type="png")
|
||||
|
||||
async def _live_frame(self) -> str:
|
||||
async def _live_frame(self) -> bytes:
|
||||
page = await self._ensure_page()
|
||||
shot = await page.screenshot(type="jpeg", quality=_LIVE_FRAME_JPEG_QUALITY)
|
||||
return base64.b64encode(shot).decode("ascii")
|
||||
return await page.screenshot(type="jpeg", quality=_LIVE_FRAME_JPEG_QUALITY)
|
||||
|
||||
async def _emit_live_frame(self) -> None:
|
||||
if self._on_frame is None:
|
||||
@ -679,7 +677,7 @@ class BrowserSession:
|
||||
self._page_listener_bound = False
|
||||
self._request_guard_bound = False
|
||||
|
||||
async def _start_screencast(self, on_frame: Callable[[str], None]) -> None:
|
||||
async def _start_screencast(self, on_frame: Callable[[bytes], None]) -> None:
|
||||
"""Start Live mode and send an initial JPEG frame.
|
||||
|
||||
This used to attach Chrome's CDP screencast and then turn every repaint
|
||||
@ -698,7 +696,7 @@ class BrowserSession:
|
||||
await self._emit_live_frame()
|
||||
self._schedule_settle_live_frames()
|
||||
|
||||
async def _stop_screencast(self, on_frame: Callable[[str], None] | None = None) -> None:
|
||||
async def _stop_screencast(self, on_frame: Callable[[bytes], None] | None = None) -> None:
|
||||
if on_frame is not None and self._on_frame is not on_frame:
|
||||
return
|
||||
self._on_frame = None
|
||||
@ -792,7 +790,7 @@ class BrowserSession:
|
||||
with self._activity():
|
||||
return await self._loop.run(self._screenshot_bytes(full_page))
|
||||
|
||||
async def live_frame(self) -> str:
|
||||
async def live_frame(self) -> bytes:
|
||||
with self._activity():
|
||||
return await self._loop.run(self._live_frame())
|
||||
|
||||
@ -815,11 +813,11 @@ class BrowserSession:
|
||||
with self._activity():
|
||||
return await self._loop.run(self._tabs())
|
||||
|
||||
async def start_screencast(self, on_frame: Callable[[str], None]) -> None:
|
||||
async def start_screencast(self, on_frame: Callable[[bytes], None]) -> None:
|
||||
with self._activity():
|
||||
await self._loop.run(self._start_screencast(on_frame))
|
||||
|
||||
async def stop_screencast(self, on_frame: Callable[[str], None] | None = None) -> None:
|
||||
async def stop_screencast(self, on_frame: Callable[[bytes], None] | None = None) -> None:
|
||||
with self._activity():
|
||||
await self._loop.run(self._stop_screencast(on_frame))
|
||||
|
||||
|
||||
@ -12,8 +12,8 @@ Configuration example (``config.yaml``)::
|
||||
template: code-interpreter-v1 # e2b template id; defaults to e2b code-interpreter
|
||||
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
|
||||
replicas: 3 # hard capacity shared when ownership is Redis
|
||||
ownership: # multi-worker ownership + capacity coordination
|
||||
type: redis
|
||||
redis_url: $REDIS_URL
|
||||
reconciliation_interval_seconds: 60
|
||||
|
||||
@ -0,0 +1,6 @@
|
||||
"""Redis-backed deployment-wide E2B capacity."""
|
||||
|
||||
from .redis import CapacityBackendError as CapacityBackendError
|
||||
from .redis import RedisE2BCapacityStore as RedisE2BCapacityStore
|
||||
from .redis import ReserveStatus as ReserveStatus
|
||||
from .redis import make_e2b_capacity_store as make_e2b_capacity_store
|
||||
@ -0,0 +1,279 @@
|
||||
"""Atomic Redis Hash ledger for deployment-wide E2B capacity."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import enum
|
||||
import logging
|
||||
|
||||
from deerflow.community.aio_sandbox.ownership.factory import resolve_ownership_redis_url
|
||||
from deerflow.config.sandbox_config import SandboxOwnershipConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_SOCKET_TIMEOUT_SECONDS = 5.0
|
||||
|
||||
_LEDGER_SCRIPT = """
|
||||
local function now_ms()
|
||||
local current = redis.call('TIME')
|
||||
return tonumber(current[1]) * 1000 + math.floor(tonumber(current[2]) / 1000)
|
||||
end
|
||||
|
||||
local function initialize(hard_limit)
|
||||
redis.call('HSET', KEYS[1],
|
||||
'meta:state', 'initializing',
|
||||
'meta:hard_limit', hard_limit,
|
||||
'meta:revision', '0')
|
||||
end
|
||||
|
||||
local function mark_present(sandbox_id)
|
||||
local field = 's:' .. sandbox_id
|
||||
if redis.call('HGET', KEYS[1], field) == '1' then
|
||||
return false
|
||||
end
|
||||
redis.call('HSET', KEYS[1], field, '1')
|
||||
return true
|
||||
end
|
||||
|
||||
local operation = ARGV[1]
|
||||
local hard_limit = ARGV[2]
|
||||
local state = redis.call('HGET', KEYS[1], 'meta:state')
|
||||
local stored_limit = redis.call('HGET', KEYS[1], 'meta:hard_limit')
|
||||
if state ~= false and stored_limit ~= hard_limit then
|
||||
return redis.error_reply(
|
||||
'E2B capacity ledger configuration mismatch: configured hard_limit='
|
||||
.. hard_limit .. ', ledger hard_limit=' .. (stored_limit or '')
|
||||
)
|
||||
end
|
||||
|
||||
if operation == 'revision' then
|
||||
return tonumber(redis.call('HGET', KEYS[1], 'meta:revision') or '0')
|
||||
end
|
||||
|
||||
if operation == 'reserve' then
|
||||
if state ~= 'ready' then
|
||||
return 'NOT_READY'
|
||||
end
|
||||
local field = 'r:' .. ARGV[3]
|
||||
if redis.call('HEXISTS', KEYS[1], field) == 1 then
|
||||
return 'GRANTED'
|
||||
end
|
||||
if redis.call('HLEN', KEYS[1]) - 3 >= tonumber(hard_limit) then
|
||||
return 'FULL'
|
||||
end
|
||||
redis.call('HSET', KEYS[1], field, tostring(now_ms()))
|
||||
redis.call('HINCRBY', KEYS[1], 'meta:revision', 1)
|
||||
return 'GRANTED'
|
||||
end
|
||||
|
||||
if operation == 'release' then
|
||||
if state ~= false and redis.call('HDEL', KEYS[1], 's:' .. ARGV[3]) == 1 then
|
||||
redis.call('HINCRBY', KEYS[1], 'meta:revision', 1)
|
||||
end
|
||||
return 'OK'
|
||||
end
|
||||
|
||||
if state == false then
|
||||
if operation == 'reconcile' and tonumber(ARGV[3]) ~= 0 then
|
||||
return 'STALE'
|
||||
end
|
||||
initialize(hard_limit)
|
||||
state = 'initializing'
|
||||
end
|
||||
|
||||
if operation == 'track' then
|
||||
local changed = false
|
||||
if ARGV[3] ~= '' and redis.call('HDEL', KEYS[1], 'r:' .. ARGV[3]) == 1 then
|
||||
changed = true
|
||||
end
|
||||
if mark_present(ARGV[4]) then
|
||||
changed = true
|
||||
end
|
||||
if changed then
|
||||
redis.call('HINCRBY', KEYS[1], 'meta:revision', 1)
|
||||
end
|
||||
return 'OK'
|
||||
end
|
||||
|
||||
if operation ~= 'reconcile' then
|
||||
return redis.error_reply('unknown E2B capacity operation: ' .. operation)
|
||||
end
|
||||
|
||||
local expected_revision = tonumber(ARGV[3])
|
||||
local revision = tonumber(redis.call('HGET', KEYS[1], 'meta:revision') or '0')
|
||||
if revision ~= expected_revision then
|
||||
return 'STALE'
|
||||
end
|
||||
|
||||
local complete = ARGV[4] == '1'
|
||||
local remote_ids = {}
|
||||
local changed = false
|
||||
for index = 6, #ARGV, 2 do
|
||||
local sandbox_id = ARGV[index]
|
||||
local token = ARGV[index + 1]
|
||||
remote_ids[sandbox_id] = true
|
||||
if token ~= '' and redis.call('HDEL', KEYS[1], 'r:' .. token) == 1 then
|
||||
changed = true
|
||||
end
|
||||
if mark_present(sandbox_id) then
|
||||
changed = true
|
||||
end
|
||||
end
|
||||
|
||||
if complete then
|
||||
local stale_before_ms = now_ms() - tonumber(ARGV[5])
|
||||
for _, field in ipairs(redis.call('HKEYS', KEYS[1])) do
|
||||
local prefix = string.sub(field, 1, 2)
|
||||
local identifier = string.sub(field, 3)
|
||||
if prefix == 's:' and remote_ids[identifier] ~= true then
|
||||
local missing_since_ms = tonumber(string.match(redis.call('HGET', KEYS[1], field) or '', '^m:(%d+)$'))
|
||||
if missing_since_ms == nil then
|
||||
redis.call('HSET', KEYS[1], field, 'm:' .. now_ms())
|
||||
changed = true
|
||||
elseif missing_since_ms <= stale_before_ms then
|
||||
redis.call('HDEL', KEYS[1], field)
|
||||
changed = true
|
||||
end
|
||||
elseif prefix == 'r:' then
|
||||
local created_ms = tonumber(redis.call('HGET', KEYS[1], field))
|
||||
if created_ms ~= nil and created_ms <= stale_before_ms then
|
||||
redis.call('HDEL', KEYS[1], field)
|
||||
changed = true
|
||||
end
|
||||
end
|
||||
end
|
||||
if state ~= 'ready' then
|
||||
redis.call('HSET', KEYS[1], 'meta:state', 'ready')
|
||||
changed = true
|
||||
end
|
||||
end
|
||||
|
||||
if changed then
|
||||
redis.call('HINCRBY', KEYS[1], 'meta:revision', 1)
|
||||
end
|
||||
return 'APPLIED'
|
||||
"""
|
||||
|
||||
|
||||
class CapacityBackendError(RuntimeError):
|
||||
"""Redis could not return a definitive capacity decision."""
|
||||
|
||||
|
||||
class ReserveStatus(enum.StrEnum):
|
||||
GRANTED = "GRANTED"
|
||||
FULL = "FULL"
|
||||
NOT_READY = "NOT_READY"
|
||||
|
||||
|
||||
def _text(value: object) -> str:
|
||||
return value.decode() if isinstance(value, bytes) else str(value)
|
||||
|
||||
|
||||
class RedisE2BCapacityStore:
|
||||
"""One capacity scope stored in one Redis Hash."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
redis_url: str,
|
||||
hard_limit: int,
|
||||
key_prefix: str = "deerflow:sandbox:owner",
|
||||
) -> None:
|
||||
if hard_limit < 1:
|
||||
raise ValueError("hard_limit must be at least 1")
|
||||
try:
|
||||
from redis import Redis
|
||||
from redis.exceptions import RedisError
|
||||
except ImportError: # pragma: no cover - optional extra
|
||||
raise ImportError("Redis E2B capacity requires: cd backend && uv sync --extra redis") from None
|
||||
|
||||
self._hard_limit = hard_limit
|
||||
self._key = f"{key_prefix.rstrip(':')}:e2b-capacity"
|
||||
self._redis_error = RedisError
|
||||
self._redis = Redis.from_url(
|
||||
redis_url,
|
||||
decode_responses=True,
|
||||
socket_timeout=_SOCKET_TIMEOUT_SECONDS,
|
||||
socket_connect_timeout=_SOCKET_TIMEOUT_SECONDS,
|
||||
)
|
||||
self._script = self._redis.register_script(_LEDGER_SCRIPT)
|
||||
|
||||
@property
|
||||
def key(self) -> str:
|
||||
return self._key
|
||||
|
||||
def _run(self, operation: str, *args: object) -> object:
|
||||
try:
|
||||
return self._script(keys=[self._key], args=[operation, self._hard_limit, *args])
|
||||
except self._redis_error as error:
|
||||
raise CapacityBackendError(f"failed to {operation} E2B capacity in Redis: {error}") from error
|
||||
|
||||
def revision(self) -> int:
|
||||
return int(self._run("revision"))
|
||||
|
||||
def reserve(self, token: str) -> ReserveStatus:
|
||||
if not token:
|
||||
raise ValueError("token must not be empty")
|
||||
try:
|
||||
return ReserveStatus(_text(self._run("reserve", token)))
|
||||
except ValueError as error:
|
||||
raise CapacityBackendError("unexpected E2B capacity reserve result") from error
|
||||
|
||||
def track(
|
||||
self,
|
||||
sandbox_id: str,
|
||||
*,
|
||||
reservation_token: str | None = None,
|
||||
) -> None:
|
||||
if not sandbox_id:
|
||||
raise ValueError("sandbox_id must not be empty")
|
||||
self._run("track", reservation_token or "", sandbox_id)
|
||||
|
||||
def release(self, sandbox_id: str) -> None:
|
||||
if sandbox_id:
|
||||
self._run("release", sandbox_id)
|
||||
|
||||
def reconcile(
|
||||
self,
|
||||
*,
|
||||
expected_revision: int,
|
||||
remote_sandboxes: dict[str, str | None],
|
||||
complete: bool,
|
||||
reservation_max_age_ms: int,
|
||||
) -> bool:
|
||||
remote_args = [item for sandbox_id, token in remote_sandboxes.items() for item in (sandbox_id, token or "")]
|
||||
status = _text(
|
||||
self._run(
|
||||
"reconcile",
|
||||
expected_revision,
|
||||
"1" if complete else "0",
|
||||
reservation_max_age_ms,
|
||||
*remote_args,
|
||||
)
|
||||
)
|
||||
if status not in {"APPLIED", "STALE"}:
|
||||
raise CapacityBackendError(f"unexpected E2B capacity reconciliation result: {status}")
|
||||
return status == "APPLIED"
|
||||
|
||||
def close(self) -> None:
|
||||
try:
|
||||
self._redis.close()
|
||||
except Exception as error: # pragma: no cover - teardown best effort
|
||||
logger.warning("Error closing E2B capacity Redis client: %s", error)
|
||||
|
||||
|
||||
def make_e2b_capacity_store(
|
||||
ownership: SandboxOwnershipConfig,
|
||||
*,
|
||||
hard_limit: int,
|
||||
) -> RedisE2BCapacityStore | None:
|
||||
"""Enable the shared ledger only with Redis ownership."""
|
||||
if ownership.type == "memory":
|
||||
return None
|
||||
if ownership.type != "redis":
|
||||
raise ValueError(f"Unknown sandbox ownership type: {ownership.type!r}")
|
||||
logger.info("E2B deployment capacity: redis (key_prefix=%s, hard_limit=%d)", ownership.key_prefix, hard_limit)
|
||||
return RedisE2BCapacityStore(
|
||||
redis_url=resolve_ownership_redis_url(ownership),
|
||||
hard_limit=hard_limit,
|
||||
key_prefix=ownership.key_prefix,
|
||||
)
|
||||
@ -15,6 +15,9 @@ provider fields during startup.
|
||||
overflow_policy: wait # wait | reject | burst (default: wait)
|
||||
acquire_timeout: 30 # seconds for ``wait`` policy (default: 30)
|
||||
burst_limit: 2 # extra slots for ``burst`` policy (default: 0)
|
||||
ownership:
|
||||
type: redis # shares ownership and capacity across Gateways
|
||||
redis_url: redis://redis:6379/0
|
||||
mounts: # one-shot uploads on sandbox start
|
||||
- host_path: /data/skills
|
||||
container_path: /home/user/skills
|
||||
@ -44,6 +47,7 @@ from functools import partial
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from e2b import SandboxQuery
|
||||
from e2b_code_interpreter import Sandbox as E2BClientSandbox
|
||||
|
||||
from deerflow.config import get_app_config
|
||||
@ -56,10 +60,16 @@ from ..aio_sandbox.ownership import (
|
||||
OwnershipBackendError,
|
||||
RenewOutcome,
|
||||
SandboxOwnershipStore,
|
||||
compute_lease_ttl,
|
||||
generate_owner_id,
|
||||
make_sandbox_ownership_store,
|
||||
resolve_ownership_config,
|
||||
)
|
||||
from .capacity import (
|
||||
CapacityBackendError,
|
||||
ReserveStatus,
|
||||
make_e2b_capacity_store,
|
||||
)
|
||||
from .e2b_sandbox import DEFAULT_E2B_HOME_DIR, E2BSandbox, _is_sandbox_gone_error
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -77,6 +87,9 @@ DEFAULT_RECONCILIATION_ORPHAN_TTL_SECONDS = 3600.0
|
||||
DEFAULT_RECONCILIATION_MAX_PAGES = 10
|
||||
DEFAULT_RECONCILIATION_MAX_ITEMS = 200
|
||||
DEFAULT_RECONCILIATION_MAX_SECONDS = 15.0
|
||||
# Twice the E2B SDK's 60-second create request timeout. Short ownership lease
|
||||
# settings must not make an in-flight create look abandoned.
|
||||
MIN_CAPACITY_RESERVATION_SECONDS = 120.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
|
||||
@ -88,6 +101,8 @@ 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_KEY_CAPACITY_LEDGER = "deer_flow_capacity_ledger"
|
||||
META_KEY_CAPACITY_RESERVATION = "deer_flow_capacity_reservation"
|
||||
META_VAL_PROVIDER = "e2b_sandbox_provider"
|
||||
E2B_EXTRA_CONFIG_KEYS = frozenset({"api_key", "domain", "home_dir", "template"})
|
||||
|
||||
@ -168,6 +183,10 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._ownership_config,
|
||||
owner_id=self._owner_id,
|
||||
)
|
||||
self._deployment_capacity = make_e2b_capacity_store(
|
||||
self._ownership_config,
|
||||
hard_limit=self._capacity_limit(),
|
||||
)
|
||||
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.")
|
||||
|
||||
@ -285,6 +304,24 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
def _stable_seed(thread_id: str, user_id: str) -> str:
|
||||
return hashlib.sha256(f"{user_id}:{thread_id}".encode()).hexdigest()[:16]
|
||||
|
||||
def _metadata_matches_capacity_ledger(
|
||||
self,
|
||||
metadata: dict[str, Any],
|
||||
) -> bool:
|
||||
"""Include this ledger and legacy sandboxes that predate the tag."""
|
||||
store = self._deployment_capacity
|
||||
if store is None:
|
||||
return True
|
||||
remote_ledger = metadata.get(META_KEY_CAPACITY_LEDGER)
|
||||
return remote_ledger in (None, "", store.key)
|
||||
|
||||
@staticmethod
|
||||
def _capacity_reservation_from_metadata(
|
||||
metadata: dict[str, Any],
|
||||
) -> str | None:
|
||||
token = metadata.get(META_KEY_CAPACITY_RESERVATION)
|
||||
return token if isinstance(token, str) and token else None
|
||||
|
||||
# ── Signal / shutdown handling ───────────────────────────────────────
|
||||
|
||||
def _register_signal_handlers(self) -> None:
|
||||
@ -390,6 +427,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
sandbox.close()
|
||||
except Exception:
|
||||
pass
|
||||
self._release_deployment_sandbox(sid)
|
||||
return None
|
||||
|
||||
try:
|
||||
@ -496,7 +534,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
"""
|
||||
sandbox_cls = self._get_sandbox_cls()
|
||||
seed = self._stable_seed(thread_id, user_id)
|
||||
entries, _ = self._list_remote_entries(
|
||||
entries, _, _ = self._list_remote_entries(
|
||||
{
|
||||
META_KEY_PROVIDER: META_VAL_PROVIDER,
|
||||
META_KEY_USER: user_id,
|
||||
@ -504,16 +542,21 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
}
|
||||
)
|
||||
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),
|
||||
(
|
||||
(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 and self._metadata_matches_capacity_ledger(metadata)
|
||||
),
|
||||
key=lambda item: (item[1].get(META_KEY_CREATED_AT, ""), item[0]),
|
||||
)
|
||||
for target_id, _metadata in candidates:
|
||||
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,
|
||||
capacity_reservation=(self._capacity_reservation_from_metadata(metadata)),
|
||||
)
|
||||
if adopted is not None:
|
||||
return adopted
|
||||
@ -527,6 +570,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
thread_id: str,
|
||||
user_id: str,
|
||||
seed: str,
|
||||
capacity_reservation: str | None = None,
|
||||
) -> str | None:
|
||||
"""Try to adopt one discovered candidate without harming peer-owned VMs."""
|
||||
|
||||
@ -548,6 +592,10 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
return None
|
||||
|
||||
try:
|
||||
self._track_deployment_sandbox(
|
||||
target_id,
|
||||
reservation_token=capacity_reservation,
|
||||
)
|
||||
self._reserve_capacity(
|
||||
thread_id,
|
||||
user_id,
|
||||
@ -664,6 +712,66 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
with self._lock:
|
||||
self._end_transition_locked()
|
||||
|
||||
def _capacity_reservation_max_age_ms(self) -> int:
|
||||
configured = compute_lease_ttl(self._ownership_config) + float(self._config["reconciliation_grace_seconds"])
|
||||
return int(max(configured, MIN_CAPACITY_RESERVATION_SECONDS) * 1_000)
|
||||
|
||||
def _capacity_error(
|
||||
self,
|
||||
message: str,
|
||||
*,
|
||||
reason: str = "capacity",
|
||||
) -> SandboxCapacityExceededError:
|
||||
with self._lock:
|
||||
return SandboxCapacityExceededError(
|
||||
message,
|
||||
active=len(self._sandboxes),
|
||||
warm=len(self._warm_pool),
|
||||
reserved=self._reserved_slots,
|
||||
replicas=int(self._config["replicas"]),
|
||||
reason=reason,
|
||||
)
|
||||
|
||||
def _track_deployment_sandbox(
|
||||
self,
|
||||
sandbox_id: str,
|
||||
*,
|
||||
reservation_token: str | None = None,
|
||||
required: bool = True,
|
||||
) -> None:
|
||||
store = self._deployment_capacity
|
||||
if store is None:
|
||||
return
|
||||
try:
|
||||
store.track(
|
||||
sandbox_id,
|
||||
reservation_token=reservation_token,
|
||||
)
|
||||
except CapacityBackendError as error:
|
||||
if required:
|
||||
raise self._capacity_error(
|
||||
f"Deployment-wide E2B capacity is unavailable; cannot safely track sandbox {sandbox_id}",
|
||||
reason="capacity_backend",
|
||||
) from error
|
||||
logger.warning(
|
||||
"Could not track E2B sandbox %s in deployment capacity; reconciliation will retry: %s",
|
||||
sandbox_id,
|
||||
error,
|
||||
)
|
||||
|
||||
def _release_deployment_sandbox(self, sandbox_id: str) -> None:
|
||||
store = self._deployment_capacity
|
||||
if store is None:
|
||||
return
|
||||
try:
|
||||
store.release(sandbox_id)
|
||||
except CapacityBackendError as error:
|
||||
logger.warning(
|
||||
"Could not release deployment capacity for destroyed E2B sandbox %s; reconciliation will retry: %s",
|
||||
sandbox_id,
|
||||
error,
|
||||
)
|
||||
|
||||
def _reserve_capacity(
|
||||
self,
|
||||
thread_id: str | None,
|
||||
@ -671,23 +779,71 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
*,
|
||||
remote_id: str | None = None,
|
||||
remote_owned: bool = True,
|
||||
) -> None:
|
||||
"""Acquire a capacity slot, blocking or raising as configured.
|
||||
|
||||
Must be called before ``Sandbox.create()``. The caller MUST call
|
||||
``_commit_capacity()`` on success or ``_release_capacity()`` on
|
||||
failure — otherwise the reserved slot is leaked until shutdown.
|
||||
|
||||
Raises:
|
||||
SandboxCapacityExceededError: when the overflow policy is
|
||||
``reject`` or the wait timeout expires.
|
||||
"""
|
||||
) -> str | None:
|
||||
"""Reserve local capacity, then deployment capacity for a new VM."""
|
||||
store = self._deployment_capacity
|
||||
policy = self._config["overflow_policy"]
|
||||
timeout = float(self._config["acquire_timeout"])
|
||||
deadline = time.monotonic() + timeout
|
||||
token = uuid.uuid4().hex if store is not None and remote_id is None else None
|
||||
|
||||
while True:
|
||||
self._reserve_local_capacity(
|
||||
thread_id,
|
||||
user_id,
|
||||
remote_id=remote_id,
|
||||
remote_owned=remote_owned,
|
||||
deadline=deadline,
|
||||
)
|
||||
if store is None or remote_id is not None:
|
||||
return token
|
||||
|
||||
assert token is not None
|
||||
backend_error = None
|
||||
try:
|
||||
status = store.reserve(token)
|
||||
except CapacityBackendError as error:
|
||||
backend_error = error
|
||||
status = None
|
||||
if status is ReserveStatus.GRANTED:
|
||||
return token
|
||||
self._release_capacity()
|
||||
if status is ReserveStatus.FULL and self._evict_oldest_warm() is not None:
|
||||
continue
|
||||
|
||||
if policy != "wait":
|
||||
if backend_error is not None:
|
||||
raise self._capacity_error(
|
||||
"Deployment-wide E2B capacity is unavailable",
|
||||
reason="capacity_backend",
|
||||
) from backend_error
|
||||
if status is ReserveStatus.NOT_READY:
|
||||
raise self._capacity_error(
|
||||
"Deployment-wide E2B capacity is initializing",
|
||||
reason="capacity_initializing",
|
||||
)
|
||||
raise self._capacity_error("Deployment-wide E2B capacity is full")
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
raise self._capacity_error(f"Timed out after {timeout}s waiting for deployment-wide E2B capacity")
|
||||
with self._capacity_cond:
|
||||
self._capacity_cond.wait(timeout=min(remaining, 1.0))
|
||||
|
||||
def _reserve_local_capacity(
|
||||
self,
|
||||
thread_id: str | None,
|
||||
user_id: str,
|
||||
*,
|
||||
remote_id: str | None = None,
|
||||
remote_owned: bool = True,
|
||||
deadline: float | None = None,
|
||||
) -> None:
|
||||
"""Acquire the existing process-local lifecycle slot."""
|
||||
policy = self._config["overflow_policy"]
|
||||
timeout = float(self._config["acquire_timeout"])
|
||||
deadline = deadline or time.monotonic() + timeout
|
||||
|
||||
while True:
|
||||
# Reject immediately if the provider is shutting down.
|
||||
with self._lock:
|
||||
if self._shutdown_called:
|
||||
raise SandboxCapacityExceededError(
|
||||
@ -697,7 +853,6 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
reason="shutdown",
|
||||
)
|
||||
|
||||
# 1. Try immediate atomic reservation.
|
||||
with self._lock:
|
||||
if self._shutdown_called:
|
||||
raise SandboxCapacityExceededError(
|
||||
@ -714,7 +869,6 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
remote_ops.add(remote_id)
|
||||
return
|
||||
|
||||
# 2. Try evicting a warm entry to free a slot.
|
||||
evicted = self._evict_oldest_warm()
|
||||
if evicted is not None:
|
||||
with self._lock:
|
||||
@ -731,9 +885,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
remote_ops = self._remote_ops_in_progress if remote_owned else self._unowned_remote_ops_in_progress
|
||||
remote_ops.add(remote_id)
|
||||
return
|
||||
# Slot was stolen; fall through to policy / wait.
|
||||
|
||||
# 3. Apply overflow policy.
|
||||
with self._lock:
|
||||
if self._shutdown_called:
|
||||
raise SandboxCapacityExceededError(
|
||||
@ -764,7 +916,6 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
replicas=int(self._config["replicas"]),
|
||||
)
|
||||
|
||||
# policy == "wait": block until a slot frees or timeout.
|
||||
remaining = deadline - time.monotonic()
|
||||
if remaining <= 0:
|
||||
raise SandboxCapacityExceededError(
|
||||
@ -839,7 +990,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
|
||||
Capacity is enforced atomically via :meth:`_reserve_capacity`.
|
||||
"""
|
||||
self._reserve_capacity(thread_id, user_id)
|
||||
reservation_token = self._reserve_capacity(thread_id, user_id)
|
||||
|
||||
sandbox_cls = self._get_sandbox_cls()
|
||||
metadata: dict[str, str] = {
|
||||
@ -847,6 +998,10 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
META_KEY_GATEWAY: self._owner_id,
|
||||
META_KEY_CREATED_AT: str(time.time()),
|
||||
}
|
||||
if self._deployment_capacity is not None:
|
||||
metadata[META_KEY_CAPACITY_LEDGER] = self._deployment_capacity.key
|
||||
if reservation_token is not None:
|
||||
metadata[META_KEY_CAPACITY_RESERVATION] = reservation_token
|
||||
if thread_id:
|
||||
metadata[META_KEY_USER] = user_id
|
||||
metadata[META_KEY_THREAD] = thread_id
|
||||
@ -869,6 +1024,11 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
raise
|
||||
|
||||
sandbox_id: str = getattr(client, "sandbox_id", None) or str(uuid.uuid4())[:8]
|
||||
self._track_deployment_sandbox(
|
||||
sandbox_id,
|
||||
reservation_token=reservation_token,
|
||||
required=False,
|
||||
)
|
||||
if not self._track_reserved_remote_op(sandbox_id):
|
||||
kill_error = self._kill_client(client)
|
||||
cleanup_confirmed = kill_error is None
|
||||
@ -932,7 +1092,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
# One-shot mount uploads. e2b has no host bind-mount, so we copy
|
||||
# files from ``host_path`` into ``container_path`` at sandbox start.
|
||||
try:
|
||||
self._apply_mounts(client)
|
||||
self._apply_mounts(client, user_id=user_id)
|
||||
except Exception as e:
|
||||
logger.warning("Failed to apply some mounts to e2b sandbox %s: %s", sandbox_id, e)
|
||||
|
||||
@ -998,26 +1158,25 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
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."""
|
||||
def _list_remote_entries(
|
||||
self,
|
||||
metadata: dict[str, str],
|
||||
) -> tuple[list[Any], bool, bool]:
|
||||
"""List E2B entries and report budget exhaustion and completeness."""
|
||||
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
|
||||
query = SandboxQuery(metadata=metadata)
|
||||
result = sandbox_cls.list(query=query, **self._common_kwargs()) # type: ignore[attr-defined]
|
||||
except Exception as e:
|
||||
logger.warning("E2B reconciliation list failed: %s", e)
|
||||
return [], False
|
||||
return [], False, 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
|
||||
complete = True
|
||||
|
||||
if hasattr(result, "next_items") and hasattr(result, "has_next"):
|
||||
for page_number in range(max_pages):
|
||||
@ -1028,6 +1187,7 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
page = result.next_items()
|
||||
except Exception as e:
|
||||
logger.warning("E2B reconciliation paginator failed: %s", e)
|
||||
complete = False
|
||||
break
|
||||
if not page:
|
||||
break
|
||||
@ -1045,13 +1205,15 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
all_entries = list(result or [])
|
||||
except TypeError:
|
||||
logger.warning("E2B Sandbox.list returned non-iterable %s", type(result).__name__)
|
||||
return [], False
|
||||
return [], False, False
|
||||
entries = all_entries[:max_items]
|
||||
exhausted = len(all_entries) > max_items
|
||||
|
||||
if time.monotonic() >= deadline:
|
||||
exhausted = True
|
||||
return entries, exhausted
|
||||
if exhausted:
|
||||
complete = False
|
||||
return entries, exhausted, complete
|
||||
|
||||
def _publish_ownership(self, sandbox_id: str) -> None:
|
||||
"""Publish acquire-side ownership before exposing a sandbox locally."""
|
||||
@ -1147,8 +1309,21 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._lease_thread.start()
|
||||
self._reconcile_thread.start()
|
||||
|
||||
def _reserve_reconciliation_capacity(self, sandbox_id: str) -> bool:
|
||||
def _reserve_reconciliation_capacity(
|
||||
self,
|
||||
sandbox_id: str,
|
||||
*,
|
||||
reservation_token: str | None = None,
|
||||
) -> bool:
|
||||
"""Reserve one local slot for adoption without blocking maintenance."""
|
||||
try:
|
||||
self._track_deployment_sandbox(
|
||||
sandbox_id,
|
||||
reservation_token=reservation_token,
|
||||
)
|
||||
except SandboxCapacityExceededError as error:
|
||||
logger.warning("Could not track discovered E2B sandbox %s: %s", sandbox_id, error)
|
||||
return False
|
||||
with self._lock:
|
||||
if self._shutdown_called or self._total_capacity_used_locked() >= self._capacity_limit():
|
||||
return False
|
||||
@ -1162,8 +1337,38 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
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})
|
||||
capacity_revision = None
|
||||
capacity_store = self._deployment_capacity
|
||||
if capacity_store is not None:
|
||||
try:
|
||||
capacity_revision = capacity_store.revision()
|
||||
except CapacityBackendError as error:
|
||||
logger.warning(
|
||||
"Could not read E2B capacity before reconciliation: %s",
|
||||
error,
|
||||
)
|
||||
|
||||
entries, stats.budget_exhausted, inventory_complete = self._list_remote_entries({META_KEY_PROVIDER: META_VAL_PROVIDER})
|
||||
entries = [entry for entry in entries if self._metadata_matches_capacity_ledger(self._entry_metadata(entry))]
|
||||
stats.discovered = len(entries)
|
||||
|
||||
if capacity_store is not None and capacity_revision is not None:
|
||||
records = {sandbox_id: self._capacity_reservation_from_metadata(self._entry_metadata(entry)) for entry in entries if (sandbox_id := self._entry_id(entry))}
|
||||
try:
|
||||
applied = capacity_store.reconcile(
|
||||
expected_revision=capacity_revision,
|
||||
remote_sandboxes=records,
|
||||
complete=inventory_complete,
|
||||
reservation_max_age_ms=self._capacity_reservation_max_age_ms(),
|
||||
)
|
||||
if not applied:
|
||||
logger.debug("E2B capacity inventory became stale during reconciliation; retrying on the next pass")
|
||||
except CapacityBackendError as error:
|
||||
logger.warning(
|
||||
"Could not apply E2B capacity inventory: %s",
|
||||
error,
|
||||
)
|
||||
|
||||
groups: dict[tuple[str, str], list[tuple[str, dict[str, Any]]]] = {}
|
||||
orphans: list[tuple[str, dict[str, Any]]] = []
|
||||
present_ids: set[str] = set()
|
||||
@ -1206,12 +1411,15 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
if not live:
|
||||
continue
|
||||
stats.duplicates += max(0, len(live) - 1)
|
||||
canonical_id, _metadata, canonical_client = live[0]
|
||||
canonical_id, canonical_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):
|
||||
elif not self._reserve_reconciliation_capacity(
|
||||
canonical_id,
|
||||
reservation_token=self._capacity_reservation_from_metadata(canonical_metadata),
|
||||
):
|
||||
self._safe_close_client(canonical_client)
|
||||
stats.deferred += 1
|
||||
elif not self._claim_ownership(canonical_id):
|
||||
@ -1332,10 +1540,16 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
reaped the VM. Closing that host-side client before returning ``None``
|
||||
keeps both acquire paths from leaking a connection.
|
||||
"""
|
||||
client = self._reconnect_client(sandbox_cls, sandbox_id)
|
||||
try:
|
||||
client = self._reconnect_client(sandbox_cls, sandbox_id)
|
||||
except Exception as error:
|
||||
if _is_sandbox_gone_error(error):
|
||||
self._release_deployment_sandbox(sandbox_id)
|
||||
raise
|
||||
if self._client_alive(client):
|
||||
return client
|
||||
self._safe_close_client(client)
|
||||
self._release_deployment_sandbox(sandbox_id)
|
||||
return None
|
||||
|
||||
def _register_connected_sandbox(
|
||||
@ -1503,11 +1717,42 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
if exit_code not in (0, None) or "BOOTSTRAP_OK" not in stdout:
|
||||
raise RuntimeError(f"e2b bootstrap script failed with exit code {exit_code}; stderr={stderr.strip()}")
|
||||
|
||||
def _apply_mounts(self, client: E2BClientSandbox) -> None:
|
||||
mounts = self._config.get("mounts") or []
|
||||
if not mounts:
|
||||
return
|
||||
for mount in mounts:
|
||||
def _skill_projection_mounts(self, user_id: str) -> list[tuple[Path, str, bool]]:
|
||||
"""Best-effort: a projection failure must not drop configured mounts too.
|
||||
|
||||
Unlike Local/AIO's ``_ensure_skills_projection``, this used to raise
|
||||
straight out of ``_apply_mounts`` before the configured-mounts loop
|
||||
ran, so a projection hiccup dropped the operator's own mounts as
|
||||
collateral damage (only caught by ``create()``'s outer warning, with
|
||||
no mounts applied at all). Swallowing here keeps the two mount
|
||||
sources independent, matching the other two providers.
|
||||
"""
|
||||
from deerflow.skills.projection import ensure_skill_projections
|
||||
from deerflow.skills.storage import get_or_new_user_skill_storage
|
||||
|
||||
try:
|
||||
config = get_app_config()
|
||||
storage = get_or_new_user_skill_storage(user_id, app_config=config)
|
||||
projection = ensure_skill_projections(storage)
|
||||
container_root = config.skills.container_path.rstrip("/")
|
||||
return [
|
||||
(projection.public, f"{container_root}/public", True),
|
||||
(projection.custom, f"{container_root}/custom", True),
|
||||
(projection.legacy, f"{container_root}/legacy", True),
|
||||
(projection.integrations, f"{container_root}/integrations", True),
|
||||
]
|
||||
except Exception as exc:
|
||||
logger.warning("Could not ensure skills projection for user %s: %s", user_id, exc, exc_info=True)
|
||||
return []
|
||||
|
||||
def _apply_mounts(self, client: E2BClientSandbox, *, user_id: str | None = None) -> None:
|
||||
effective_user_id = user_id or get_effective_user_id()
|
||||
projection_mounts = self._skill_projection_mounts(effective_user_id)
|
||||
configured_mounts = self._config.get("mounts") or []
|
||||
skills_root = get_app_config().skills.container_path.rstrip("/")
|
||||
|
||||
mounts: list[tuple[Path, str, bool]] = list(projection_mounts)
|
||||
for mount in configured_mounts:
|
||||
try:
|
||||
host_path = Path(getattr(mount, "host_path", "") or "")
|
||||
container_path = (getattr(mount, "container_path", "") or "").rstrip("/")
|
||||
@ -1517,6 +1762,12 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
container_path = (mount.get("container_path", "") or "").rstrip("/")
|
||||
read_only = bool(mount.get("read_only", False))
|
||||
|
||||
if container_path == skills_root or container_path.startswith(skills_root + "/"):
|
||||
logger.warning("Skipping e2b mount that conflicts with managed skills projection: %s", container_path)
|
||||
continue
|
||||
mounts.append((host_path, container_path, read_only))
|
||||
|
||||
for host_path, container_path, read_only in mounts:
|
||||
if not host_path.exists():
|
||||
logger.warning("Skipping e2b mount: host_path %s does not exist", host_path)
|
||||
continue
|
||||
@ -2073,13 +2324,14 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _kill_client(
|
||||
self,
|
||||
client: E2BClientSandbox | None,
|
||||
) -> Exception | None:
|
||||
"""Kill a remote VM and return an exception for the caller to log."""
|
||||
if client is None:
|
||||
return RuntimeError("Cannot confirm remote VM destruction without a client")
|
||||
sandbox_id = getattr(client, "sandbox_id", None)
|
||||
try:
|
||||
kill = getattr(client, "kill", None)
|
||||
if not callable(kill):
|
||||
@ -2087,6 +2339,8 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
kill()
|
||||
except Exception as e:
|
||||
return e
|
||||
if sandbox_id is not None:
|
||||
self._release_deployment_sandbox(sandbox_id)
|
||||
return None
|
||||
|
||||
def reset(self) -> None:
|
||||
@ -2178,3 +2432,11 @@ class E2BSandboxProvider(SandboxProvider):
|
||||
self._ownership.close()
|
||||
except Exception as e:
|
||||
logger.warning("Failed to close E2B ownership store: %s", e)
|
||||
if self._deployment_capacity is not None:
|
||||
try:
|
||||
self._deployment_capacity.close()
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to close E2B deployment capacity store: %s",
|
||||
e,
|
||||
)
|
||||
|
||||
@ -100,6 +100,45 @@ class CheckpointGraphCacheConfig(BaseModel):
|
||||
)
|
||||
|
||||
|
||||
class CheckpointCacheConfig(BaseModel):
|
||||
"""Delta-history cache policy. Performance-only: never frozen, never
|
||||
required to match across processes sharing one checkpoint database.
|
||||
|
||||
Applies only when ``checkpoint_channel_mode`` is ``delta``. ``max_entries``
|
||||
bounds the process-local memory backend; ``0`` disables the cache
|
||||
entirely. The redis backend is bounded by ``ttl_seconds`` and the server's
|
||||
own maxmemory policy.
|
||||
"""
|
||||
|
||||
type: Literal["memory", "redis"] = Field(
|
||||
default="memory",
|
||||
description=("Checkpoint history cache backend. 'memory' = process-local LRU; 'redis' = shared cache for multi-worker deployments (async/Gateway path only; the sync embedded path rejects it)."),
|
||||
)
|
||||
max_entries: int = Field(
|
||||
default=128,
|
||||
ge=0,
|
||||
description="LRU capacity of the memory backend. 0 disables the cache.",
|
||||
)
|
||||
redis_url: str | None = Field(
|
||||
default=None,
|
||||
description=("Redis URL for type=redis. If omitted, DEER_FLOW_CHECKPOINT_CACHE_REDIS_URL, REDIS_URL, or redis://localhost:6379/0 is used."),
|
||||
)
|
||||
ttl_seconds: int = Field(
|
||||
default=86400,
|
||||
ge=0,
|
||||
description=(
|
||||
"Redis entry TTL; a leak safety net, not a correctness mechanism (entries are immutable). "
|
||||
"Thread deletion purges that thread's entries immediately; if the purge fails (redis outage), "
|
||||
"residual copies of the thread's history persist until this TTL expires. "
|
||||
"0 explicitly disables expiry — orphaned keys then rely on the redis maxmemory policy alone."
|
||||
),
|
||||
)
|
||||
key_prefix: str = Field(
|
||||
default="",
|
||||
description="Optional override for the redis key prefix; defaults to a hash of the database identity.",
|
||||
)
|
||||
|
||||
|
||||
class DatabaseConfig(BaseModel):
|
||||
backend: Literal["memory", "sqlite", "postgres"] = Field(
|
||||
default="memory",
|
||||
@ -122,6 +161,10 @@ class DatabaseConfig(BaseModel):
|
||||
default_factory=CheckpointGraphCacheConfig,
|
||||
description="Size caps for the compiled checkpoint graph caches. Hot-reloadable; not restart-required.",
|
||||
)
|
||||
checkpoint_cache: CheckpointCacheConfig = Field(
|
||||
default_factory=CheckpointCacheConfig,
|
||||
description="Delta-mode checkpoint history cache. Performance-only; safe to differ across workers.",
|
||||
)
|
||||
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."),
|
||||
|
||||
@ -98,6 +98,10 @@ class McpServerConfig(BaseModel):
|
||||
description: str = Field(default="", description="Human-readable description of what this MCP server provides")
|
||||
routing: McpRoutingConfig = Field(default_factory=McpRoutingConfig, description="Soft routing hints for tools from this MCP server")
|
||||
tools: dict[str, McpToolOverride] = Field(default_factory=dict, description="Per-original-tool MCP configuration overrides")
|
||||
tool_name_prefix: bool = Field(
|
||||
default=True,
|
||||
description="Whether to prefix discovered tool names with the MCP server name to avoid cross-server collisions",
|
||||
)
|
||||
tool_call_timeout: float | None = Field(
|
||||
default=None,
|
||||
description="Timeout in seconds for individual stdio MCP tool calls. HTTP/SSE servers use transport-level timeouts. None means no timeout.",
|
||||
|
||||
@ -32,6 +32,16 @@ class ModelConfig(BaseModel):
|
||||
description="Extra settings to be passed to the model when thinking is disabled",
|
||||
)
|
||||
supports_vision: bool = Field(default_factory=lambda: False, description="Whether the model supports vision/image inputs")
|
||||
context_window: int | None = Field(
|
||||
default=None,
|
||||
gt=0,
|
||||
description=(
|
||||
"Positive total context window size in tokens (prompt + completion). Used to compute the real-time "
|
||||
"context usage percentage displayed in the chat UI. Distinct from `max_tokens`, which is the "
|
||||
"per-call output cap passed to the provider. Leave unset if unknown; the UI will hide the "
|
||||
"percentage."
|
||||
),
|
||||
)
|
||||
stream_chunk_timeout: float | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
|
||||
@ -6,11 +6,11 @@ import shutil
|
||||
from pathlib import Path, PureWindowsPath
|
||||
|
||||
from deerflow.config.runtime_paths import runtime_home
|
||||
from deerflow.utils.thread_id import validate_thread_id
|
||||
|
||||
# Virtual path prefix seen by agents inside the sandbox
|
||||
VIRTUAL_PATH_PREFIX = "/mnt/user-data"
|
||||
|
||||
_SAFE_THREAD_ID_RE = re.compile(r"^[A-Za-z0-9_\-]+$")
|
||||
_SAFE_USER_ID_RE = re.compile(r"^[A-Za-z0-9_\-]+$")
|
||||
_SAFE_INTEGRATION_ID_RE = re.compile(r"^[A-Za-z0-9_.\-]+$")
|
||||
_UNSAFE_USER_ID_CHAR_RE = re.compile(r"[^A-Za-z0-9_\-]")
|
||||
@ -26,9 +26,7 @@ def _default_local_base_dir() -> Path:
|
||||
|
||||
def _validate_thread_id(thread_id: str) -> str:
|
||||
"""Validate a thread ID before using it in filesystem paths."""
|
||||
if not _SAFE_THREAD_ID_RE.match(thread_id):
|
||||
raise ValueError(f"Invalid thread_id {thread_id!r}: only alphanumeric characters, hyphens, and underscores are allowed.")
|
||||
return thread_id
|
||||
return validate_thread_id(thread_id)
|
||||
|
||||
|
||||
def _validate_user_id(user_id: str) -> str:
|
||||
@ -258,6 +256,32 @@ class Paths:
|
||||
"""
|
||||
return self.base_dir / "integrations" / "skills"
|
||||
|
||||
@property
|
||||
def skills_view_dir(self) -> Path:
|
||||
"""Global sandbox-visible skills projection: ``{base_dir}/skills_view/``."""
|
||||
return self.base_dir / "skills_view"
|
||||
|
||||
@property
|
||||
def public_skills_view_dir(self) -> Path:
|
||||
"""Enabled public skills exposed to sandboxes."""
|
||||
return self.skills_view_dir / "public"
|
||||
|
||||
def user_skills_view_dir(self, user_id: str) -> Path:
|
||||
"""Per-user sandbox-visible skills projection root."""
|
||||
return self.user_dir(user_id) / "skills_view"
|
||||
|
||||
def user_custom_skills_view_dir(self, user_id: str) -> Path:
|
||||
"""Enabled custom skills exposed to one user's sandboxes."""
|
||||
return self.user_skills_view_dir(user_id) / "custom"
|
||||
|
||||
def user_legacy_skills_view_dir(self, user_id: str) -> Path:
|
||||
"""Enabled legacy skills exposed to one user's sandboxes."""
|
||||
return self.user_skills_view_dir(user_id) / "legacy"
|
||||
|
||||
def user_integration_skills_view_dir(self, user_id: str) -> Path:
|
||||
"""Enabled managed integration skills exposed to one user's sandboxes."""
|
||||
return self.user_skills_view_dir(user_id) / "integrations"
|
||||
|
||||
def thread_dir(self, thread_id: str, *, user_id: str | None = None) -> Path:
|
||||
"""
|
||||
Host path for a thread's data.
|
||||
|
||||
@ -76,8 +76,9 @@ class SandboxConfig(BaseModel):
|
||||
|
||||
AioSandboxProvider, BoxliteProvider, and E2BSandboxProvider shared options:
|
||||
image: Sandbox image to use (Docker/AIO image or BoxLite OCI image)
|
||||
replicas: Positive provider capacity per gateway process. Each provider
|
||||
defines which lifecycle states count toward this limit.
|
||||
replicas: Positive provider capacity. E2B shares it across Gateway
|
||||
workers when ownership uses Redis; other modes/providers keep
|
||||
process-local accounting.
|
||||
idle_timeout: Idle timeout in seconds before released warm sandboxes/VMs are stopped (default: 600 = 10 minutes). Set to 0 to disable.
|
||||
environment: Environment variables to inject into the sandbox (values starting with $ are resolved from host env)
|
||||
|
||||
@ -115,7 +116,7 @@ class SandboxConfig(BaseModel):
|
||||
replicas: int | None = Field(
|
||||
default=None,
|
||||
gt=0,
|
||||
description="Positive provider capacity per gateway process. Each provider defines which lifecycle states count toward this limit.",
|
||||
description=("Positive provider capacity. E2B enforces it deployment-wide when sandbox ownership uses Redis; otherwise accounting is per Gateway process. Each provider defines which lifecycle states count."),
|
||||
)
|
||||
overflow_policy: SandboxOverflowPolicy = Field(
|
||||
default="wait",
|
||||
|
||||
@ -431,6 +431,7 @@ def _make_session_pool_tool(
|
||||
connection: dict[str, Any],
|
||||
tool_interceptors: list[Any] | None = None,
|
||||
tool_call_timeout: float | None = None,
|
||||
tool_name_prefix: bool = True,
|
||||
) -> BaseTool:
|
||||
"""Wrap an MCP tool so it reuses a persistent session from the pool.
|
||||
|
||||
@ -442,10 +443,11 @@ def _make_session_pool_tool(
|
||||
The configured ``tool_interceptors`` (OAuth, custom) are preserved and
|
||||
applied on every call before invoking the pooled session.
|
||||
"""
|
||||
# Strip the server-name prefix to recover the original MCP tool name.
|
||||
# Strip only prefixes added by the adapter. An unprefixed server may expose
|
||||
# a tool whose own name happens to start with ``<server_name>_``.
|
||||
original_name = tool.name
|
||||
prefix = f"{server_name}_"
|
||||
if original_name.startswith(prefix):
|
||||
if tool_name_prefix and original_name.startswith(prefix):
|
||||
original_name = original_name[len(prefix) :]
|
||||
|
||||
pool = get_session_pool()
|
||||
@ -582,6 +584,7 @@ async def get_mcp_tools() -> list[BaseTool]:
|
||||
"""
|
||||
try:
|
||||
from langchain_mcp_adapters.client import MultiServerMCPClient
|
||||
from langchain_mcp_adapters.tools import load_mcp_tools
|
||||
except ImportError:
|
||||
logger.warning("langchain-mcp-adapters not installed. Install it to enable MCP tools: pip install langchain-mcp-adapters")
|
||||
return []
|
||||
@ -648,7 +651,18 @@ async def get_mcp_tools() -> list[BaseTool]:
|
||||
|
||||
async def load_server_tools(server_name: str) -> list[BaseTool]:
|
||||
try:
|
||||
return await client.get_tools(server_name=server_name)
|
||||
server_cfg = extensions_config.mcp_servers.get(server_name)
|
||||
tool_name_prefix = server_cfg.tool_name_prefix if server_cfg is not None else True
|
||||
if tool_name_prefix:
|
||||
return await client.get_tools(server_name=server_name)
|
||||
return await load_mcp_tools(
|
||||
None,
|
||||
connection=servers_config[server_name],
|
||||
callbacks=client.callbacks,
|
||||
server_name=server_name,
|
||||
tool_interceptors=client.tool_interceptors,
|
||||
tool_name_prefix=False,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Skipping MCP server '{server_name}' after tool discovery failed: {e}",
|
||||
@ -672,11 +686,11 @@ async def get_mcp_tools() -> list[BaseTool]:
|
||||
# scanning servers_config for a name prefix is ambiguous when one server name is a
|
||||
# prefix of another (e.g. "web" vs "web_scraper" → "web_scraper_search".startswith(
|
||||
# "web_") matches "web" first), which pools the tool under the wrong server. Using the
|
||||
# source grouping makes routing exact; the prefix guard preserves the previous
|
||||
# behavior of leaving unprefixed tools unwrapped.
|
||||
# source grouping makes routing exact even when a server opts out of name prefixing.
|
||||
for source_name, server_tools in zip(servers_config.keys(), tools_by_server, strict=True):
|
||||
transport = servers_config[source_name].get("transport", "stdio")
|
||||
server_cfg = extensions_config.mcp_servers.get(source_name)
|
||||
tool_name_prefix = server_cfg.tool_name_prefix if server_cfg is not None else True
|
||||
for tool in server_tools:
|
||||
if not _VALID_MCP_TOOL_NAME.fullmatch(tool.name or ""):
|
||||
logger.warning(
|
||||
@ -688,13 +702,22 @@ async def get_mcp_tools() -> list[BaseTool]:
|
||||
continue
|
||||
tag_mcp_tool(tool)
|
||||
prefix = f"{source_name}_"
|
||||
original_name = tool.name[len(prefix) :] if tool.name.startswith(prefix) else tool.name
|
||||
original_name = tool.name[len(prefix) :] if tool_name_prefix and tool.name.startswith(prefix) else tool.name
|
||||
routing = resolve_effective_mcp_routing(server_cfg, original_name)
|
||||
if routing.get("mode") != "off":
|
||||
tag_mcp_routing(tool, routing)
|
||||
if tool.name.startswith(f"{source_name}_") and transport == "stdio":
|
||||
if transport == "stdio":
|
||||
_timeout = server_cfg.tool_call_timeout if server_cfg else None
|
||||
wrapped_tools.append(_make_session_pool_tool(tool, source_name, servers_config[source_name], tool_interceptors, tool_call_timeout=_timeout))
|
||||
wrapped_tools.append(
|
||||
_make_session_pool_tool(
|
||||
tool,
|
||||
source_name,
|
||||
servers_config[source_name],
|
||||
tool_interceptors,
|
||||
tool_call_timeout=_timeout,
|
||||
tool_name_prefix=tool_name_prefix,
|
||||
)
|
||||
)
|
||||
else:
|
||||
if transport != "stdio" and server_cfg and server_cfg.tool_call_timeout is not None:
|
||||
logger.warning(
|
||||
|
||||
@ -219,6 +219,9 @@ def create_chat_model(name: str | None = None, thinking_enabled: bool = False, *
|
||||
"when_thinking_disabled",
|
||||
"thinking",
|
||||
"supports_vision",
|
||||
# Runtime/UI metadata used to size the context indicator. Provider
|
||||
# clients do not accept this as a model-constructor argument.
|
||||
"context_window",
|
||||
# Presentation-only metadata (consumed by the console's cost
|
||||
# display) — must never reach the provider client, which would
|
||||
# forward unknown kwargs into the completion request payload.
|
||||
|
||||
@ -0,0 +1,19 @@
|
||||
"""Checkpoint delta-history cache backends (delta mode only)."""
|
||||
|
||||
from deerflow.runtime.checkpoint_cache.base import (
|
||||
CACHE_FORMAT_VERSION,
|
||||
CheckpointCacheStats,
|
||||
CheckpointHistoryCache,
|
||||
SyncCheckpointHistoryCache,
|
||||
make_history_key,
|
||||
)
|
||||
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
||||
|
||||
__all__ = [
|
||||
"CACHE_FORMAT_VERSION",
|
||||
"CheckpointCacheStats",
|
||||
"CheckpointHistoryCache",
|
||||
"MemoryCheckpointHistoryCache",
|
||||
"SyncCheckpointHistoryCache",
|
||||
"make_history_key",
|
||||
]
|
||||
@ -0,0 +1,80 @@
|
||||
"""Cache backend contract for checkpoint delta-history entries.
|
||||
|
||||
Entries are ``DeltaChannelHistory``-shaped dicts (``{"writes": [...], "seed"?}``)
|
||||
keyed by immutable (database, thread, namespace, checkpoint_id, channel)
|
||||
tuples. Checkpoint lineage is append-only and a checkpoint's history excludes
|
||||
its own pending writes, so entries never change once written: correctness
|
||||
never requires invalidation, and a shared backend is coherent across
|
||||
processes without any coordination.
|
||||
|
||||
The only delete API is thread-scoped (``adelete_thread``/``delete_thread``),
|
||||
and it exists purely for data lifecycle, not correctness: when the source
|
||||
checkpoints are erased (thread deletion, tenant offboarding, GDPR-style
|
||||
erasure), the cached history payloads for that thread must go too instead of
|
||||
lingering until LRU eviction or TTL expiry.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Protocol
|
||||
|
||||
CACHE_FORMAT_VERSION = 1
|
||||
|
||||
|
||||
def make_history_key(
|
||||
key_prefix: str,
|
||||
thread_id: str,
|
||||
checkpoint_ns: str,
|
||||
checkpoint_id: str,
|
||||
channel: str,
|
||||
) -> str:
|
||||
"""Build a collision-safe cache key.
|
||||
|
||||
``thread_id`` stays readable for ops debugging; the remaining components
|
||||
are hashed with NUL separators so namespaces containing ':' cannot
|
||||
produce ambiguous keys.
|
||||
"""
|
||||
digest = hashlib.sha256(f"{checkpoint_ns}\x00{checkpoint_id}\x00{channel}".encode()).hexdigest()[:24]
|
||||
return f"{key_prefix}:{thread_id}:{digest}"
|
||||
|
||||
|
||||
def thread_key_stem(key_prefix: str, thread_id: str) -> str:
|
||||
"""Prefix matching every history key of one thread (see make_history_key)."""
|
||||
return f"{key_prefix}:{thread_id}:"
|
||||
|
||||
|
||||
@dataclass
|
||||
class CheckpointCacheStats:
|
||||
hits: int = 0
|
||||
misses: int = 0
|
||||
evictions: int = 0
|
||||
entries: int = 0
|
||||
|
||||
def as_dict(self) -> dict[str, int]:
|
||||
return {
|
||||
"hits": self.hits,
|
||||
"misses": self.misses,
|
||||
"evictions": self.evictions,
|
||||
"entries": self.entries,
|
||||
}
|
||||
|
||||
|
||||
class CheckpointHistoryCache(Protocol):
|
||||
"""Async backend contract. Deletes are thread-scoped lifecycle purges only."""
|
||||
|
||||
async def aget_many(self, keys: list[str]) -> dict[str, dict[str, Any]]: ...
|
||||
async def aset_many(self, entries: dict[str, dict[str, Any]]) -> None: ...
|
||||
async def adelete_thread(self, key_prefix: str, thread_id: str) -> None: ...
|
||||
def stats(self) -> CheckpointCacheStats: ...
|
||||
async def aclose(self) -> None: ...
|
||||
|
||||
|
||||
class SyncCheckpointHistoryCache(Protocol):
|
||||
"""Sync backend contract (embedded/TUI path). Memory backend only."""
|
||||
|
||||
def get_many(self, keys: list[str]) -> dict[str, dict[str, Any]]: ...
|
||||
def set_many(self, entries: dict[str, dict[str, Any]]) -> None: ...
|
||||
def delete_thread(self, key_prefix: str, thread_id: str) -> None: ...
|
||||
def stats(self) -> CheckpointCacheStats: ...
|
||||
@ -0,0 +1,79 @@
|
||||
"""Process-local LRU backend. Zero serialization on the hit path."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import OrderedDict
|
||||
from typing import Any
|
||||
|
||||
from deerflow.runtime.checkpoint_cache.base import CheckpointCacheStats, thread_key_stem
|
||||
|
||||
|
||||
def _copy_entry(entry: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Copy-on-read/write: fresh writes list; seed shared (never mutated in place)."""
|
||||
copied: dict[str, Any] = {"writes": list(entry["writes"])}
|
||||
if "seed" in entry:
|
||||
copied["seed"] = entry["seed"]
|
||||
return copied
|
||||
|
||||
|
||||
class MemoryCheckpointHistoryCache:
|
||||
def __init__(self, max_entries: int = 128) -> None:
|
||||
if max_entries < 0:
|
||||
raise ValueError("max_entries must be >= 0")
|
||||
self._max_entries = max_entries
|
||||
self._data: OrderedDict[str, dict[str, Any]] = OrderedDict()
|
||||
self._hits = 0
|
||||
self._misses = 0
|
||||
self._evictions = 0
|
||||
|
||||
@property
|
||||
def enabled(self) -> bool:
|
||||
return self._max_entries > 0
|
||||
|
||||
def get_many(self, keys: list[str]) -> dict[str, dict[str, Any]]:
|
||||
found: dict[str, dict[str, Any]] = {}
|
||||
for key in keys:
|
||||
entry = self._data.get(key)
|
||||
if entry is None:
|
||||
self._misses += 1
|
||||
continue
|
||||
self._data.move_to_end(key)
|
||||
self._hits += 1
|
||||
found[key] = _copy_entry(entry)
|
||||
return found
|
||||
|
||||
def set_many(self, entries: dict[str, dict[str, Any]]) -> None:
|
||||
if not self.enabled:
|
||||
return
|
||||
for key, entry in entries.items():
|
||||
self._data[key] = _copy_entry(entry)
|
||||
self._data.move_to_end(key)
|
||||
while len(self._data) > self._max_entries:
|
||||
self._data.popitem(last=False)
|
||||
self._evictions += 1
|
||||
|
||||
async def aget_many(self, keys: list[str]) -> dict[str, dict[str, Any]]:
|
||||
return self.get_many(keys)
|
||||
|
||||
async def aset_many(self, entries: dict[str, dict[str, Any]]) -> None:
|
||||
self.set_many(entries)
|
||||
|
||||
def delete_thread(self, key_prefix: str, thread_id: str) -> None:
|
||||
"""Purge every entry of one thread (lifecycle, not invalidation)."""
|
||||
stem = thread_key_stem(key_prefix, thread_id)
|
||||
for key in [k for k in self._data if k.startswith(stem)]:
|
||||
del self._data[key]
|
||||
|
||||
async def adelete_thread(self, key_prefix: str, thread_id: str) -> None:
|
||||
self.delete_thread(key_prefix, thread_id)
|
||||
|
||||
def stats(self) -> CheckpointCacheStats:
|
||||
return CheckpointCacheStats(
|
||||
hits=self._hits,
|
||||
misses=self._misses,
|
||||
evictions=self._evictions,
|
||||
entries=len(self._data),
|
||||
)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
self._data.clear()
|
||||
@ -0,0 +1,101 @@
|
||||
"""Cache factory. Mirrors make_stream_bridge: config -> env fallback -> memory."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import hashlib
|
||||
import logging
|
||||
import os
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any
|
||||
|
||||
from deerflow.config.app_config import AppConfig
|
||||
from deerflow.runtime.checkpoint_cache.base import CACHE_FORMAT_VERSION, CheckpointHistoryCache
|
||||
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_ENV_REDIS_URL = "DEER_FLOW_CHECKPOINT_CACHE_REDIS_URL"
|
||||
|
||||
|
||||
def _resolve_redis_url(config: Any) -> str:
|
||||
return config.redis_url or os.getenv(_ENV_REDIS_URL) or os.getenv("REDIS_URL") or "redis://localhost:6379/0"
|
||||
|
||||
|
||||
def _stable_postgres_identity(postgres_url: str) -> str:
|
||||
"""Credential-free database identity: host/port/database.
|
||||
|
||||
Hashing the raw URL would change the cache namespace on every credential
|
||||
rotation (cold cache + orphaned keys until TTL) even though the database
|
||||
— and thus every cached checkpoint history — is unchanged. Unparseable
|
||||
URLs fall back to the raw string (still stable per deployment).
|
||||
"""
|
||||
if not postgres_url:
|
||||
return ""
|
||||
try:
|
||||
from sqlalchemy.engine.url import make_url
|
||||
|
||||
parsed = make_url(postgres_url)
|
||||
except Exception: # noqa: BLE001 - identity must never fail config load
|
||||
return postgres_url
|
||||
return f"{parsed.host or 'localhost'}:{parsed.port or 5432}/{parsed.database or ''}"
|
||||
|
||||
|
||||
def checkpoint_cache_db_hash(db_config: Any) -> str:
|
||||
"""Deployment-identity hash so two deployments sharing one Redis never collide."""
|
||||
backend = getattr(db_config, "backend", "memory")
|
||||
if backend == "postgres":
|
||||
identity = f"postgres:{_stable_postgres_identity(getattr(db_config, 'postgres_url', ''))}:{getattr(db_config, 'postgres_schema', '')}"
|
||||
elif backend == "sqlite":
|
||||
identity = f"sqlite:{getattr(db_config, 'checkpointer_sqlite_path', '')}"
|
||||
else:
|
||||
identity = "memory"
|
||||
return hashlib.sha256(identity.encode()).hexdigest()[:12]
|
||||
|
||||
|
||||
def checkpoint_cache_key_prefix(app_config: AppConfig) -> str:
|
||||
cache_config = app_config.database.checkpoint_cache
|
||||
if cache_config.key_prefix:
|
||||
return cache_config.key_prefix
|
||||
return f"ckpt-hist:v{CACHE_FORMAT_VERSION}:{checkpoint_cache_db_hash(app_config.database)}"
|
||||
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def make_checkpoint_cache(
|
||||
app_config: AppConfig | None = None,
|
||||
*,
|
||||
serde: Any,
|
||||
) -> AsyncIterator[CheckpointHistoryCache]:
|
||||
"""Yield a history cache for the caller's lifetime.
|
||||
|
||||
``max_entries == 0`` disables the cache uniformly (both types) via a
|
||||
disabled memory backend, so the wrapper never needs a None check.
|
||||
"""
|
||||
config = app_config.database.checkpoint_cache if app_config is not None else None
|
||||
|
||||
if config is None or config.type == "memory" or config.max_entries == 0:
|
||||
max_entries = config.max_entries if config is not None else 128
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=max_entries)
|
||||
logger.info("Checkpoint history cache initialised: memory (max_entries=%d)", max_entries)
|
||||
try:
|
||||
yield cache
|
||||
finally:
|
||||
await cache.aclose()
|
||||
return
|
||||
|
||||
if config.type == "redis":
|
||||
from deerflow.runtime.checkpoint_cache.redis import RedisCheckpointHistoryCache
|
||||
|
||||
cache = RedisCheckpointHistoryCache(
|
||||
_resolve_redis_url(config),
|
||||
serde=serde,
|
||||
ttl_seconds=config.ttl_seconds,
|
||||
)
|
||||
logger.info("Checkpoint history cache initialised: redis (ttl_seconds=%d)", config.ttl_seconds)
|
||||
try:
|
||||
yield cache
|
||||
finally:
|
||||
await cache.aclose()
|
||||
return
|
||||
|
||||
raise ValueError(f"Unknown checkpoint cache type: {config.type!r}")
|
||||
@ -0,0 +1,116 @@
|
||||
"""Shared Redis backend. Entries are immutable, so a multi-worker shared
|
||||
cache needs no invalidation; the TTL is a leak safety net only.
|
||||
|
||||
Thread-scoped purge (``adelete_thread``) exists for data lifecycle: when a
|
||||
thread's checkpoints are deleted, its cached history payloads are removed
|
||||
immediately instead of lingering until TTL expiry.
|
||||
|
||||
The redis import is lazy (module is importable without the optional
|
||||
``redis`` extra), mirroring runtime/stream_bridge/redis.py.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from deerflow.runtime.checkpoint_cache.base import CheckpointCacheStats, thread_key_stem
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
REDIS_INSTALL = "redis is required for the redis checkpoint cache backend. Install it with: uv sync --extra redis"
|
||||
|
||||
_TAG_SEPARATOR = b"\x00"
|
||||
|
||||
|
||||
def _create_client(redis_url: str, *, max_connections: int | None) -> Any:
|
||||
try:
|
||||
import redis.asyncio as redis_async
|
||||
except ImportError as exc:
|
||||
raise ImportError(REDIS_INSTALL) from exc
|
||||
kwargs: dict[str, Any] = {"decode_responses": False}
|
||||
if max_connections is not None:
|
||||
kwargs["max_connections"] = max_connections
|
||||
return redis_async.from_url(redis_url, **kwargs)
|
||||
|
||||
|
||||
def _redis_error() -> type[Exception]:
|
||||
"""Lazy RedisError import, mirroring the lazy client creation above."""
|
||||
try:
|
||||
from redis.exceptions import RedisError
|
||||
except ImportError as exc:
|
||||
raise ImportError(REDIS_INSTALL) from exc
|
||||
return RedisError
|
||||
|
||||
|
||||
class RedisCheckpointHistoryCache:
|
||||
def __init__(
|
||||
self,
|
||||
redis_url: str,
|
||||
*,
|
||||
serde: Any,
|
||||
ttl_seconds: int,
|
||||
max_connections: int | None = None,
|
||||
) -> None:
|
||||
self._client = _create_client(redis_url, max_connections=max_connections)
|
||||
self._serde = serde
|
||||
# ttl_seconds=0 is an explicit opt-out of expiry (no SETEX) — not the
|
||||
# default, and leaked/orphaned keys then rely on redis maxmemory only.
|
||||
self._ttl = ttl_seconds if ttl_seconds > 0 else None
|
||||
self._hits = 0
|
||||
self._misses = 0
|
||||
|
||||
async def aget_many(self, keys: list[str]) -> dict[str, dict[str, Any]]:
|
||||
if not keys:
|
||||
return {}
|
||||
try:
|
||||
raws = await self._client.mget(keys)
|
||||
except _redis_error() as exc:
|
||||
# Performance-only bypass: a redis outage costs hits, never availability.
|
||||
logger.warning("checkpoint history cache mget failed; treating as all-miss: %s", exc)
|
||||
self._misses += len(keys)
|
||||
return {}
|
||||
found: dict[str, dict[str, Any]] = {}
|
||||
for key, raw in zip(keys, raws, strict=True):
|
||||
if raw is None:
|
||||
self._misses += 1
|
||||
continue
|
||||
self._hits += 1
|
||||
tag, payload = raw.split(_TAG_SEPARATOR, 1)
|
||||
found[key] = self._serde.loads_typed((tag.decode(), payload))
|
||||
return found
|
||||
|
||||
async def aset_many(self, entries: dict[str, dict[str, Any]]) -> None:
|
||||
if not entries:
|
||||
return
|
||||
try:
|
||||
pipe = self._client.pipeline(transaction=False)
|
||||
for key, entry in entries.items():
|
||||
tag, data = self._serde.dumps_typed(entry)
|
||||
pipe.set(key, tag.encode() + _TAG_SEPARATOR + data, ex=self._ttl)
|
||||
await pipe.execute()
|
||||
except _redis_error() as exc:
|
||||
# Writes are optional; the next read simply recomputes the history.
|
||||
logger.warning("checkpoint history cache write failed; skipping: %s", exc)
|
||||
|
||||
async def adelete_thread(self, key_prefix: str, thread_id: str) -> None:
|
||||
"""SCAN+UNLINK every entry of one thread. Failure degrades to
|
||||
TTL-bounded residual retention; the source-of-truth delete already
|
||||
happened, so this never raises."""
|
||||
stem = thread_key_stem(key_prefix, thread_id)
|
||||
try:
|
||||
cursor = 0
|
||||
while True:
|
||||
cursor, keys = await self._client.scan(cursor=cursor, match=stem + "*", count=500)
|
||||
if keys:
|
||||
await self._client.unlink(*keys)
|
||||
if cursor == 0:
|
||||
break
|
||||
except _redis_error() as exc:
|
||||
logger.warning("checkpoint history cache thread purge failed; residual entries expire via TTL: %s", exc)
|
||||
|
||||
def stats(self) -> CheckpointCacheStats:
|
||||
return CheckpointCacheStats(hits=self._hits, misses=self._misses)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
await self._client.aclose()
|
||||
@ -185,24 +185,14 @@ async def _async_checkpointer_from_database(db_config) -> AsyncIterator[Checkpoi
|
||||
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def make_checkpointer(app_config: AppConfig | None = None) -> AsyncIterator[Checkpointer]:
|
||||
"""Async context manager that yields a checkpointer for the caller's lifetime.
|
||||
Resources are opened on enter and closed on exit -- no global state::
|
||||
|
||||
async with make_checkpointer(app_config) as checkpointer:
|
||||
app.state.checkpointer = checkpointer
|
||||
|
||||
Yields an ``InMemorySaver`` when no checkpointer is configured in *config.yaml*.
|
||||
async def _select_inner_checkpointer(app_config: AppConfig) -> AsyncIterator[Checkpointer]:
|
||||
"""Yield the raw checkpointer selected by *app_config* (no delta-cache wrapping).
|
||||
|
||||
Priority:
|
||||
1. Legacy ``checkpointer:`` config section (backward compatible)
|
||||
2. Unified ``database:`` config section
|
||||
3. Default InMemorySaver
|
||||
"""
|
||||
|
||||
if app_config is None:
|
||||
app_config = get_app_config()
|
||||
|
||||
# Legacy: standalone checkpointer config takes precedence
|
||||
if app_config.checkpointer is not None:
|
||||
async with _async_checkpointer(app_config.checkpointer) as saver:
|
||||
@ -220,3 +210,44 @@ async def make_checkpointer(app_config: AppConfig | None = None) -> AsyncIterato
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
yield InMemorySaver()
|
||||
|
||||
|
||||
@contextlib.asynccontextmanager
|
||||
async def make_checkpointer(app_config: AppConfig | None = None) -> AsyncIterator[Checkpointer]:
|
||||
"""Async context manager that yields a checkpointer for the caller's lifetime.
|
||||
Resources are opened on enter and closed on exit -- no global state::
|
||||
|
||||
async with make_checkpointer(app_config) as checkpointer:
|
||||
app.state.checkpointer = checkpointer
|
||||
|
||||
Yields an ``InMemorySaver`` when no checkpointer is configured in *config.yaml*.
|
||||
|
||||
Backend selection priority:
|
||||
1. Legacy ``checkpointer:`` config section (backward compatible)
|
||||
2. Unified ``database:`` config section
|
||||
3. Default InMemorySaver
|
||||
|
||||
When the effective checkpoint channel mode is ``delta`` (the process-frozen
|
||||
mode wins, falling back to ``database.checkpoint_channel_mode``), the raw
|
||||
saver is wrapped in a :class:`CachedHistorySaver` backed by a history cache
|
||||
whose lifetime equals this context manager's.
|
||||
"""
|
||||
from deerflow.runtime.checkpoint_mode import frozen_checkpoint_channel_mode
|
||||
|
||||
if app_config is None:
|
||||
app_config = get_app_config()
|
||||
|
||||
async with _select_inner_checkpointer(app_config) as saver:
|
||||
db_config = getattr(app_config, "database", None)
|
||||
mode = frozen_checkpoint_channel_mode() or (db_config.checkpoint_channel_mode if db_config is not None else "full")
|
||||
if mode == "delta":
|
||||
from deerflow.runtime.checkpoint_cache.provider import (
|
||||
checkpoint_cache_key_prefix,
|
||||
make_checkpoint_cache,
|
||||
)
|
||||
from deerflow.runtime.checkpointer.cached_saver import CachedHistorySaver
|
||||
|
||||
async with make_checkpoint_cache(app_config, serde=saver.serde) as cache:
|
||||
yield CachedHistorySaver(saver, cache, key_prefix=checkpoint_cache_key_prefix(app_config))
|
||||
else:
|
||||
yield saver
|
||||
|
||||
@ -0,0 +1,328 @@
|
||||
"""Read-through delta-history cache wrapper for any BaseCheckpointSaver.
|
||||
|
||||
Correctness argument (spec §3): a checkpoint's delta history is a pure
|
||||
function of its sealed ancestor chain — the LangGraph contract excludes the
|
||||
target's own pending writes, parent links are fixed at creation, and an
|
||||
ancestor's writes are sealed once its child exists. Entries keyed by
|
||||
(thread, ns, checkpoint_id, channel) are therefore immutable: no
|
||||
invalidation, and shared backends are coherent across processes.
|
||||
|
||||
The wrapper never caches the "latest checkpoint" resolution; only histories
|
||||
keyed by resolved immutable checkpoint_ids.
|
||||
|
||||
Data lifecycle: thread deletion and prune purge the thread's cached entries
|
||||
(source-of-truth removal must not leave residual history payloads in the
|
||||
cache); run-scoped deletes cannot be mapped to threads cheaply and rely on
|
||||
LRU/TTL bounds.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from collections.abc import Iterator, Sequence
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointTuple, PendingWrite
|
||||
|
||||
from deerflow.runtime.checkpoint_cache.base import make_history_key
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Depth budget for recursive compose before falling back to a chain-warming
|
||||
# walk. Steady-state runs need ~2 (one intermediate checkpoint per step);
|
||||
# deeper cold chains are handled faster by one warming walk than by many
|
||||
# recursive single-tuple fetches.
|
||||
_COMPOSE_MAX_DEPTH = 8
|
||||
|
||||
|
||||
def _checkpoint_ref(tup: CheckpointTuple) -> tuple[str, str, str]:
|
||||
configurable = tup.config["configurable"]
|
||||
return (
|
||||
str(configurable["thread_id"]),
|
||||
str(configurable.get("checkpoint_ns", "")),
|
||||
str(configurable["checkpoint_id"]),
|
||||
)
|
||||
|
||||
|
||||
def _channel_writes(tup: CheckpointTuple, channel: str) -> list[PendingWrite]:
|
||||
"""Writes for one channel, oldest→newest (tuple storage order)."""
|
||||
return [w for w in (tup.pending_writes or []) if w[1] == channel]
|
||||
|
||||
|
||||
class CachedHistorySaver(BaseCheckpointSaver):
|
||||
def __init__(self, inner: BaseCheckpointSaver, cache: Any, *, key_prefix: str) -> None:
|
||||
# Instance attr shadows the base class JsonPlusSerializer default.
|
||||
self.serde = inner.serde
|
||||
self._inner = inner
|
||||
self._cache = cache
|
||||
self._key_prefix = key_prefix
|
||||
self._compose_hits = 0
|
||||
self._full_walks = 0
|
||||
|
||||
def __getattr__(self, name: str) -> Any:
|
||||
# Safety net for saver-specific extras (e.g. AsyncSqliteSaver.setup).
|
||||
# Base-class methods are explicitly delegated below, so this only
|
||||
# fires for attributes BaseCheckpointSaver does not define.
|
||||
inner = self.__dict__.get("_inner")
|
||||
if inner is None:
|
||||
raise AttributeError(name)
|
||||
return getattr(inner, name)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Key building
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _key(self, tup: CheckpointTuple, channel: str) -> str:
|
||||
thread_id, ns, checkpoint_id = _checkpoint_ref(tup)
|
||||
return make_history_key(self._key_prefix, thread_id, ns, checkpoint_id, channel)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Stats
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def stats(self) -> dict[str, int]:
|
||||
backend = self._cache.stats().as_dict()
|
||||
return {**backend, "compose_hits": self._compose_hits, "full_walks": self._full_walks}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Delta history: the only overridden behavior
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
async def aget_delta_channel_history(self, *, config: RunnableConfig, channels: Sequence[str]) -> dict[str, Any]:
|
||||
if not channels:
|
||||
return {}
|
||||
if not getattr(self._cache, "enabled", True):
|
||||
# Disabled cache: pass straight through; composing over all-miss
|
||||
# entries is strictly more work than the raw saver's walk.
|
||||
return await self._walk_inner(config, channels)
|
||||
target = await self._inner.aget_tuple(config)
|
||||
if target is None:
|
||||
return await self._walk_inner(config, channels)
|
||||
|
||||
keys = {ch: self._key(target, ch) for ch in channels}
|
||||
hits = await self._cache.aget_many(list(keys.values()))
|
||||
found: dict[str, dict[str, Any]] = {}
|
||||
missing: list[str] = []
|
||||
for ch in channels:
|
||||
entry = hits.get(keys[ch])
|
||||
if entry is None:
|
||||
missing.append(ch)
|
||||
else:
|
||||
found[ch] = entry
|
||||
|
||||
computed: dict[str, dict[str, Any]] = {}
|
||||
if missing:
|
||||
computed = await self._compose_or_walk(config, target, missing)
|
||||
|
||||
new_entries = {keys[ch]: computed[ch] for ch in missing if ch in computed}
|
||||
if new_entries:
|
||||
await self._cache.aset_many(new_entries)
|
||||
|
||||
return {ch: found.get(ch) or computed.get(ch) or {"writes": []} for ch in channels}
|
||||
|
||||
async def _compose_or_walk(self, config: RunnableConfig, target: CheckpointTuple, missing: list[str]) -> dict[str, dict[str, Any]]:
|
||||
return {ch: await self._aresolve(target, ch, _COMPOSE_MAX_DEPTH) for ch in missing}
|
||||
|
||||
async def _aresolve(self, tup: CheckpointTuple, channel: str, depth: int) -> dict[str, Any]:
|
||||
"""Recursively compose history(tup) from the nearest warm ancestor.
|
||||
|
||||
Real runs create several checkpoints per super-step and only some are
|
||||
ever materialized as targets, so the parent is usually an unwarmed
|
||||
intermediate checkpoint (measured: 0 cache hits with single-level
|
||||
compose on a 500-step sqlite run). Recursing one level per
|
||||
intermediate lands on a warmed ancestor within ~2 levels in steady
|
||||
state; each composed level is cached, so the warm frontier follows
|
||||
the run. At depth 0 on a cold chain it delegates one inner fast-path
|
||||
walk (2 SQL) rather than crawling ancestors tuple-by-tuple.
|
||||
"""
|
||||
parent_config = tup.parent_config
|
||||
if parent_config is None:
|
||||
return {"writes": []}
|
||||
parent = await self._inner.aget_tuple(parent_config)
|
||||
if parent is None:
|
||||
return {"writes": []}
|
||||
|
||||
channel_values = parent.checkpoint.get("channel_values") or {}
|
||||
writes = _channel_writes(parent, channel)
|
||||
if channel in channel_values:
|
||||
self._compose_hits += 1
|
||||
return {"writes": writes, "seed": channel_values[channel]}
|
||||
|
||||
key = self._key(parent, channel)
|
||||
hits = await self._cache.aget_many([key])
|
||||
parent_history = hits.get(key)
|
||||
if parent_history is None:
|
||||
if depth > 0:
|
||||
parent_history = await self._aresolve(parent, channel, depth - 1)
|
||||
else:
|
||||
# Depth budget exhausted on a cold chain: delegate ONE inner
|
||||
# fast-path walk (2 SQL total) for this level instead of
|
||||
# fetching every ancestor tuple individually. Ancestors below
|
||||
# stay cold; resolving them later recurses up to the nearest
|
||||
# warm level, so the frontier still follows the run.
|
||||
self._full_walks += 1
|
||||
walked = await self._inner.aget_delta_channel_history(config=parent.config, channels=[channel])
|
||||
parent_history = walked.get(channel) or {"writes": []}
|
||||
if parent_history is not None:
|
||||
await self._cache.aset_many({key: parent_history})
|
||||
|
||||
self._compose_hits += 1
|
||||
entry: dict[str, Any] = {"writes": list(parent_history["writes"]) + writes}
|
||||
if "seed" in parent_history:
|
||||
entry["seed"] = parent_history["seed"]
|
||||
return entry
|
||||
|
||||
async def _walk_inner(self, config: RunnableConfig, channels: Sequence[str]) -> dict[str, Any]:
|
||||
self._full_walks += 1
|
||||
return dict(await self._inner.aget_delta_channel_history(config=config, channels=channels))
|
||||
|
||||
def get_delta_channel_history(self, *, config: RunnableConfig, channels: Sequence[str]) -> dict[str, Any]:
|
||||
if not channels:
|
||||
return {}
|
||||
if not getattr(self._cache, "enabled", True):
|
||||
return self._walk_inner_sync(config, channels)
|
||||
get_many = getattr(self._cache, "get_many", None)
|
||||
set_many = getattr(self._cache, "set_many", None)
|
||||
if get_many is None or set_many is None:
|
||||
raise TypeError("sync get_delta_channel_history requires a SyncCheckpointHistoryCache (memory backend)")
|
||||
target = self._inner.get_tuple(config)
|
||||
if target is None:
|
||||
return self._walk_inner_sync(config, channels)
|
||||
|
||||
keys = {ch: self._key(target, ch) for ch in channels}
|
||||
hits = get_many(list(keys.values()))
|
||||
found: dict[str, dict[str, Any]] = {}
|
||||
missing: list[str] = []
|
||||
for ch in channels:
|
||||
entry = hits.get(keys[ch])
|
||||
if entry is None:
|
||||
missing.append(ch)
|
||||
else:
|
||||
found[ch] = entry
|
||||
|
||||
computed: dict[str, dict[str, Any]] = {}
|
||||
if missing:
|
||||
computed = {ch: self._resolve_sync(target, ch, _COMPOSE_MAX_DEPTH) for ch in missing}
|
||||
|
||||
new_entries = {keys[ch]: computed[ch] for ch in missing if ch in computed}
|
||||
if new_entries:
|
||||
set_many(new_entries)
|
||||
return {ch: found.get(ch) or computed.get(ch) or {"writes": []} for ch in channels}
|
||||
|
||||
def _resolve_sync(self, tup: CheckpointTuple, channel: str, depth: int) -> dict[str, Any]:
|
||||
"""Sync twin of _aresolve (recursive compose, see its docstring)."""
|
||||
parent_config = tup.parent_config
|
||||
if parent_config is None:
|
||||
return {"writes": []}
|
||||
parent = self._inner.get_tuple(parent_config)
|
||||
if parent is None:
|
||||
return {"writes": []}
|
||||
|
||||
channel_values = parent.checkpoint.get("channel_values") or {}
|
||||
writes = _channel_writes(parent, channel)
|
||||
if channel in channel_values:
|
||||
self._compose_hits += 1
|
||||
return {"writes": writes, "seed": channel_values[channel]}
|
||||
|
||||
key = self._key(parent, channel)
|
||||
parent_history = self._cache.get_many([key]).get(key)
|
||||
if parent_history is None:
|
||||
if depth > 0:
|
||||
parent_history = self._resolve_sync(parent, channel, depth - 1)
|
||||
else:
|
||||
# See _aresolve: one inner fast-path walk, no per-tuple crawl.
|
||||
self._full_walks += 1
|
||||
parent_history = self._inner.get_delta_channel_history(config=parent.config, channels=[channel]).get(channel) or {"writes": []}
|
||||
if parent_history is not None:
|
||||
self._cache.set_many({key: parent_history})
|
||||
|
||||
self._compose_hits += 1
|
||||
entry: dict[str, Any] = {"writes": list(parent_history["writes"]) + writes}
|
||||
if "seed" in parent_history:
|
||||
entry["seed"] = parent_history["seed"]
|
||||
return entry
|
||||
|
||||
def _walk_inner_sync(self, config: RunnableConfig, channels: Sequence[str]) -> dict[str, Any]:
|
||||
self._full_walks += 1
|
||||
return dict(self._inner.get_delta_channel_history(config=config, channels=channels))
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Explicit delegation (BaseCheckpointSaver defines these concretely,
|
||||
# so __getattr__ never fires for them)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
return self._inner.get_tuple(config)
|
||||
|
||||
def list(self, config: RunnableConfig | None, *, filter: dict[str, Any] | None = None, before: RunnableConfig | None = None, limit: int | None = None) -> Iterator[CheckpointTuple]:
|
||||
return self._inner.list(config, filter=filter, before=before, limit=limit)
|
||||
|
||||
def put(self, config: RunnableConfig, checkpoint: dict[str, Any], metadata: dict[str, Any], new_versions: dict[str, Any]) -> RunnableConfig:
|
||||
return self._inner.put(config, checkpoint, metadata, new_versions)
|
||||
|
||||
def put_writes(self, config: RunnableConfig, writes: Sequence[tuple[str, str, Any]], task_id: str, task_path: str = "") -> None:
|
||||
self._inner.put_writes(config, writes, task_id, task_path)
|
||||
|
||||
def delete_thread(self, thread_id: str) -> None:
|
||||
self._inner.delete_thread(thread_id)
|
||||
self._purge_thread_sync(thread_id)
|
||||
|
||||
def delete_for_runs(self, run_ids: Sequence[str]) -> None:
|
||||
# Run-scoped deletes cannot be mapped back to threads without an extra
|
||||
# query, so cached entries are left in place: they stay *correct* (the
|
||||
# sealed-chain argument is unaffected by other chains) and residual
|
||||
# retention is bounded by LRU/TTL. No in-tree callers today.
|
||||
self._inner.delete_for_runs(run_ids)
|
||||
|
||||
def _purge_thread_sync(self, thread_id: str) -> None:
|
||||
delete = getattr(self._cache, "delete_thread", None)
|
||||
if delete is not None:
|
||||
delete(self._key_prefix, thread_id)
|
||||
|
||||
def copy_thread(self, source_thread_id: str, target_thread_id: str) -> None:
|
||||
self._inner.copy_thread(source_thread_id, target_thread_id)
|
||||
|
||||
def prune(self, thread_ids: Sequence[str], *, strategy: str = "keep_latest") -> None:
|
||||
self._inner.prune(thread_ids, strategy=strategy)
|
||||
# Pruning rewrites these threads' chains: purge so no cached history
|
||||
# references a deleted ancestor (retention) or its pre-prune chain.
|
||||
for thread_id in thread_ids:
|
||||
self._purge_thread_sync(thread_id)
|
||||
|
||||
async def aget_tuple(self, config: RunnableConfig) -> CheckpointTuple | None:
|
||||
return await self._inner.aget_tuple(config)
|
||||
|
||||
def alist(self, config: RunnableConfig | None, *, filter: dict[str, Any] | None = None, before: RunnableConfig | None = None, limit: int | None = None) -> Any:
|
||||
return self._inner.alist(config, filter=filter, before=before, limit=limit)
|
||||
|
||||
async def aput(self, config: RunnableConfig, checkpoint: dict[str, Any], metadata: dict[str, Any], new_versions: dict[str, Any]) -> RunnableConfig:
|
||||
return await self._inner.aput(config, checkpoint, metadata, new_versions)
|
||||
|
||||
async def aput_writes(self, config: RunnableConfig, writes: Sequence[tuple[str, str, Any]], task_id: str, task_path: str = "") -> None:
|
||||
await self._inner.aput_writes(config, writes, task_id, task_path)
|
||||
|
||||
async def adelete_thread(self, thread_id: str) -> None:
|
||||
await self._inner.adelete_thread(thread_id)
|
||||
await self._apurge_thread(thread_id)
|
||||
|
||||
async def adelete_for_runs(self, run_ids: Sequence[str]) -> None:
|
||||
# See delete_for_runs: unscoped residual retention bounded by LRU/TTL.
|
||||
await self._inner.adelete_for_runs(run_ids)
|
||||
|
||||
async def _apurge_thread(self, thread_id: str) -> None:
|
||||
delete = getattr(self._cache, "adelete_thread", None)
|
||||
if delete is not None:
|
||||
await delete(self._key_prefix, thread_id)
|
||||
|
||||
async def acopy_thread(self, source_thread_id: str, target_thread_id: str) -> None:
|
||||
await self._inner.acopy_thread(source_thread_id, target_thread_id)
|
||||
|
||||
async def aprune(self, thread_ids: Sequence[str], *, strategy: str = "keep_latest") -> None:
|
||||
await self._inner.aprune(thread_ids, strategy=strategy)
|
||||
# See prune: rewritten chains must not keep pre-prune cached histories.
|
||||
for thread_id in thread_ids:
|
||||
await self._apurge_thread(thread_id)
|
||||
|
||||
def get_next_version(self, current: Any, channel: Any) -> Any:
|
||||
return self._inner.get_next_version(current, channel)
|
||||
@ -29,6 +29,7 @@ from langgraph.types import Checkpointer
|
||||
from deerflow.config.app_config import AppConfig, get_app_config
|
||||
from deerflow.config.checkpointer_config import CheckpointerConfig, ensure_config_loaded, get_checkpointer_config
|
||||
from deerflow.persistence.postgres_schema import dsn_with_search_path, ensure_postgres_schema
|
||||
from deerflow.runtime.checkpoint_mode import frozen_checkpoint_channel_mode
|
||||
from deerflow.runtime.store._sqlite_utils import ensure_sqlite_parent_dir, resolve_sqlite_conn_str
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -155,6 +156,41 @@ def _sync_checkpointer_cm(config: CheckpointerConfig) -> Iterator[Checkpointer]:
|
||||
_checkpointer: Checkpointer | None = None
|
||||
_checkpointer_ctx = None # open context manager keeping the connection alive
|
||||
_checkpointer_lock = threading.Lock()
|
||||
_checkpointer_cache = None # MemoryCheckpointHistoryCache singleton shared by wrapped sync savers
|
||||
_checkpointer_cache_prefix: str | None = None # key prefix the singleton was built for
|
||||
|
||||
|
||||
def _wrap_sync_if_delta(saver: Checkpointer, app_config: AppConfig) -> Checkpointer:
|
||||
"""Wrap *saver* in a delta-history cache when the effective mode is ``delta``.
|
||||
|
||||
The process-frozen mode wins; ``database.checkpoint_channel_mode`` is the
|
||||
fallback when nothing is frozen yet. Only the memory cache backend is
|
||||
supported on the sync path (TUI/embedded) — it is process-local anyway.
|
||||
"""
|
||||
global _checkpointer_cache, _checkpointer_cache_prefix
|
||||
# The ``_checkpointer_cache`` singleton is reassigned here without holding
|
||||
# ``_checkpointer_lock`` on the ``checkpointer_context()`` path (and under
|
||||
# the lock on the ``get_checkpointer()`` path). The race is intentional
|
||||
# and benign: worst case two wrappers get their own fresh memory cache —
|
||||
# last writer wins, and the cache is performance-only.
|
||||
db_config = getattr(app_config, "database", None)
|
||||
mode = frozen_checkpoint_channel_mode() or (db_config.checkpoint_channel_mode if db_config is not None else "full")
|
||||
if mode != "delta":
|
||||
return saver
|
||||
cache_config = app_config.database.checkpoint_cache
|
||||
if cache_config.type == "redis":
|
||||
raise ValueError("database.checkpoint_cache.type 'redis' is not supported on the sync checkpointer path (TUI/embedded); use 'memory'.")
|
||||
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
||||
from deerflow.runtime.checkpoint_cache.provider import checkpoint_cache_key_prefix
|
||||
from deerflow.runtime.checkpointer.cached_saver import CachedHistorySaver
|
||||
|
||||
key_prefix = checkpoint_cache_key_prefix(app_config)
|
||||
# Recreate on capacity OR namespace change: entries under a stale prefix
|
||||
# would be unreachable and no longer covered by thread purges.
|
||||
if _checkpointer_cache is None or _checkpointer_cache._max_entries != cache_config.max_entries or _checkpointer_cache_prefix != key_prefix:
|
||||
_checkpointer_cache = MemoryCheckpointHistoryCache(max_entries=cache_config.max_entries)
|
||||
_checkpointer_cache_prefix = key_prefix
|
||||
return CachedHistorySaver(saver, _checkpointer_cache, key_prefix=key_prefix)
|
||||
|
||||
|
||||
def get_checkpointer() -> Checkpointer:
|
||||
@ -177,12 +213,27 @@ def get_checkpointer() -> Checkpointer:
|
||||
# config outside this provider lock to avoid cross-provider lock-order inversion.
|
||||
config = _get_checkpointer_config()
|
||||
|
||||
# ``get_app_config()`` can trigger a config reload whose
|
||||
# ``_apply_singleton_configs`` calls ``reset_checkpointer()`` — which takes
|
||||
# ``_checkpointer_lock``. Resolve it (non-reentrant lock) BEFORE acquiring
|
||||
# the lock below, exactly like ``_get_checkpointer_config()`` above.
|
||||
try:
|
||||
app_config = get_app_config()
|
||||
except FileNotFoundError:
|
||||
app_config = None
|
||||
|
||||
with _checkpointer_lock:
|
||||
if _checkpointer is not None:
|
||||
return _checkpointer
|
||||
|
||||
checkpointer_ctx = _sync_checkpointer_cm(config)
|
||||
checkpointer = checkpointer_ctx.__enter__()
|
||||
try:
|
||||
if app_config is not None:
|
||||
checkpointer = _wrap_sync_if_delta(checkpointer, app_config)
|
||||
except Exception:
|
||||
checkpointer_ctx.__exit__(None, None, None)
|
||||
raise
|
||||
_checkpointer_ctx = checkpointer_ctx
|
||||
_checkpointer = checkpointer
|
||||
|
||||
@ -195,7 +246,7 @@ def reset_checkpointer() -> None:
|
||||
Closes any open backend connections and clears the cached instance.
|
||||
Useful in tests or after a configuration change.
|
||||
"""
|
||||
global _checkpointer, _checkpointer_ctx
|
||||
global _checkpointer, _checkpointer_ctx, _checkpointer_cache, _checkpointer_cache_prefix
|
||||
with _checkpointer_lock:
|
||||
if _checkpointer_ctx is not None:
|
||||
try:
|
||||
@ -204,6 +255,8 @@ def reset_checkpointer() -> None:
|
||||
logger.warning("Error during checkpointer cleanup", exc_info=True)
|
||||
_checkpointer_ctx = None
|
||||
_checkpointer = None
|
||||
_checkpointer_cache = None
|
||||
_checkpointer_cache_prefix = None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@ -227,6 +280,7 @@ def checkpointer_context() -> Iterator[Checkpointer]:
|
||||
``InMemorySaver`` when neither selects a persistent backend.
|
||||
"""
|
||||
|
||||
config = _resolve_checkpointer_config(get_app_config())
|
||||
app_config = get_app_config()
|
||||
config = _resolve_checkpointer_config(app_config)
|
||||
with _sync_checkpointer_cm(config) as saver:
|
||||
yield saver
|
||||
yield _wrap_sync_if_delta(saver, app_config)
|
||||
|
||||
@ -32,6 +32,7 @@ from typing import Any
|
||||
|
||||
from deerflow.runtime.events.store.base import RunEventStore
|
||||
from deerflow.runtime.user_context import AUTO, _AutoSentinel
|
||||
from deerflow.utils.thread_id import validate_thread_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -56,7 +57,7 @@ class JsonlRunEventStore(RunEventStore):
|
||||
return value
|
||||
|
||||
def _thread_dir(self, thread_id: str) -> Path:
|
||||
self._validate_id(thread_id, "thread_id")
|
||||
validate_thread_id(thread_id)
|
||||
return self._base_dir / "threads" / thread_id / "runs"
|
||||
|
||||
def _run_file(self, thread_id: str, run_id: str) -> Path:
|
||||
|
||||
@ -87,24 +87,16 @@ def _coerce_seed_message(message: Any) -> Any:
|
||||
return message
|
||||
|
||||
|
||||
def build_branch_history_seed_events(
|
||||
def _build_history_seed_events(
|
||||
messages: Sequence[Any],
|
||||
*,
|
||||
thread_id: str,
|
||||
run_id_prefix: str,
|
||||
parent_thread_id: str,
|
||||
seed_metadata: Mapping[str, Any],
|
||||
) -> list[dict]:
|
||||
"""Serialize a branch checkpoint's messages into run-event message rows.
|
||||
"""Serialize checkpoint messages into run-event rows.
|
||||
|
||||
Thread branching copies checkpoint state, but the thread feed
|
||||
(``list_messages`` / ``GET /threads/{id}/messages/page``) reads the
|
||||
run-event store — which a fresh branch has no rows in, so the inherited
|
||||
history vanishes from the UI as soon as the branch's first run refreshes
|
||||
the feed (#4380). Seeding the branch's run_events from the same
|
||||
checkpoint snapshot the branch was created from keeps the feed
|
||||
consistent with what the branch actually contains.
|
||||
|
||||
Rows are grouped into one synthetic run per inherited turn
|
||||
Rows are grouped into one synthetic run per checkpoint turn
|
||||
(``{run_id_prefix}-{n}``), a new turn starting at every persisted human
|
||||
message — the same boundary a real run has, since a run begins with a
|
||||
human input (including the allowlisted hidden ``ask_clarification``
|
||||
@ -117,9 +109,9 @@ def build_branch_history_seed_events(
|
||||
id per turn confines the drop to the turn actually regenerated.
|
||||
|
||||
Mirrors RunJournal's message-event contract so seeded rows are
|
||||
indistinguishable from journaled ones except by the ``branch_seed``
|
||||
marker: same event types, ``category="message"``, ``content=
|
||||
message.model_dump()``, the human-input persistence rule
|
||||
indistinguishable from journaled ones except by the supplied seed metadata:
|
||||
same event types, ``category="message"``, ``content=message.model_dump()``,
|
||||
the human-input persistence rule
|
||||
(``_should_persist_human_input_message``), the original-user-text
|
||||
restoration, and the same treatment of ``hide_from_ui`` AI/tool rows —
|
||||
RunJournal persists them (``on_llm_end`` / ``_persist_tool_result_message``
|
||||
@ -136,7 +128,6 @@ def build_branch_history_seed_events(
|
||||
"""
|
||||
events: list[dict] = []
|
||||
created_at = datetime.now(UTC).isoformat()
|
||||
seed_metadata = {"branch_seed": True, "branch_parent_thread_id": parent_thread_id}
|
||||
# Messages ahead of the first human turn (none in practice) stay in turn 0.
|
||||
turn_index = 0
|
||||
for raw_message in messages:
|
||||
@ -175,6 +166,45 @@ def build_branch_history_seed_events(
|
||||
return events
|
||||
|
||||
|
||||
def build_branch_history_seed_events(
|
||||
messages: Sequence[Any],
|
||||
*,
|
||||
thread_id: str,
|
||||
run_id_prefix: str,
|
||||
parent_thread_id: str,
|
||||
) -> list[dict]:
|
||||
"""Serialize inherited branch history into the branch's empty event feed."""
|
||||
return _build_history_seed_events(
|
||||
messages,
|
||||
thread_id=thread_id,
|
||||
run_id_prefix=run_id_prefix,
|
||||
seed_metadata={
|
||||
"branch_seed": True,
|
||||
"branch_parent_thread_id": parent_thread_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def build_checkpoint_history_seed_events(
|
||||
messages: Sequence[Any],
|
||||
*,
|
||||
thread_id: str,
|
||||
run_id_prefix: str,
|
||||
) -> list[dict]:
|
||||
"""Serialize legacy checkpoint history for a thread's empty event feed.
|
||||
|
||||
Reuse the branch seed's message normalization and per-turn synthetic run
|
||||
grouping, but stamp migration-specific metadata so these rows are not
|
||||
misidentified as history inherited from another thread.
|
||||
"""
|
||||
return _build_history_seed_events(
|
||||
messages,
|
||||
thread_id=thread_id,
|
||||
run_id_prefix=run_id_prefix,
|
||||
seed_metadata={"checkpoint_history_seed": True},
|
||||
)
|
||||
|
||||
|
||||
class RunJournal(BaseCallbackHandler):
|
||||
"""LangChain callback handler that captures events to RunEventStore."""
|
||||
|
||||
|
||||
@ -8,6 +8,7 @@ class ThreadOperationKind(StrEnum):
|
||||
|
||||
run = "run"
|
||||
checkpoint_write = "checkpoint_write"
|
||||
artifact_write = "artifact_write"
|
||||
|
||||
|
||||
class RunStatus(StrEnum):
|
||||
|
||||
@ -6,7 +6,6 @@ from pathlib import Path
|
||||
from deerflow.sandbox.local.local_sandbox import LocalSandbox, PathMapping
|
||||
from deerflow.sandbox.sandbox import Sandbox
|
||||
from deerflow.sandbox.sandbox_provider import SandboxProvider
|
||||
from deerflow.skills.storage import user_should_see_legacy_skills
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -101,11 +100,11 @@ class LocalSandboxProvider(SandboxProvider):
|
||||
from deerflow.config import get_app_config
|
||||
|
||||
config = get_app_config()
|
||||
skills_path = config.skills.get_skills_path()
|
||||
container_path = config.skills.container_path
|
||||
projection = self._ensure_skills_projection()
|
||||
|
||||
# Public skills: global, read-only — static, shared by all threads
|
||||
public_skills_path = skills_path / "public"
|
||||
public_skills_path = projection.public
|
||||
if public_skills_path.exists():
|
||||
mappings.append(
|
||||
PathMapping(
|
||||
@ -217,6 +216,52 @@ class LocalSandboxProvider(SandboxProvider):
|
||||
def _thread_key(thread_id: str, user_id: str) -> tuple[str, str]:
|
||||
return (user_id, thread_id)
|
||||
|
||||
@staticmethod
|
||||
def _ensure_skills_projection(user_id: str | None = None):
|
||||
"""Best-effort: a projection failure must not fail sandbox acquire.
|
||||
|
||||
Mirrors the surrounding skill-mount setup, which has always logged
|
||||
and continued rather than failing the whole acquire (e.g. missing
|
||||
config.yaml in a test double). Callers see ``None`` and skip the
|
||||
skill mounts for this acquire; the projection self-heals on a later
|
||||
acquire once the underlying condition clears.
|
||||
"""
|
||||
from deerflow.config import get_app_config
|
||||
from deerflow.skills.projection import ensure_skill_projections
|
||||
from deerflow.skills.storage import get_or_new_skill_storage, get_or_new_user_skill_storage
|
||||
|
||||
try:
|
||||
config = get_app_config()
|
||||
if user_id is None:
|
||||
storage = get_or_new_skill_storage(app_config=config)
|
||||
else:
|
||||
storage = get_or_new_user_skill_storage(user_id, app_config=config)
|
||||
return ensure_skill_projections(storage)
|
||||
except Exception as exc:
|
||||
logger.warning("Could not ensure skills projection for user %s: %s", user_id, exc, exc_info=True)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _append_public_skill_mapping(mappings: list[PathMapping], projection) -> None:
|
||||
if projection is None:
|
||||
return
|
||||
try:
|
||||
from deerflow.config import get_app_config
|
||||
|
||||
container_path = get_app_config().skills.container_path.rstrip("/")
|
||||
public_container_path = f"{container_path}/public"
|
||||
if any(mapping.container_path.rstrip("/") == public_container_path for mapping in mappings):
|
||||
return
|
||||
mappings.append(
|
||||
PathMapping(
|
||||
container_path=public_container_path,
|
||||
local_path=str(projection.public),
|
||||
read_only=True,
|
||||
)
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Could not append public skill mapping: %s", exc, exc_info=True)
|
||||
|
||||
@staticmethod
|
||||
def _sandbox_id_for_thread(thread_id: str, user_id: str) -> str:
|
||||
return f"local:{user_id}:{thread_id}"
|
||||
@ -232,7 +277,7 @@ class LocalSandboxProvider(SandboxProvider):
|
||||
return (user_id, thread_id)
|
||||
|
||||
@staticmethod
|
||||
def _build_thread_path_mappings(thread_id: str, *, user_id: str | None = None) -> list[PathMapping]:
|
||||
def _build_thread_path_mappings(thread_id: str, *, user_id: str | None = None, skill_projection=None) -> list[PathMapping]:
|
||||
"""Build per-thread path mappings for /mnt/user-data, /mnt/acp-workspace,
|
||||
and /mnt/skills/custom.
|
||||
|
||||
@ -281,56 +326,35 @@ class LocalSandboxProvider(SandboxProvider):
|
||||
),
|
||||
]
|
||||
|
||||
# Per-user custom skills mount (read-only). This must be per-thread
|
||||
# because ``/mnt/skills/custom`` resolves to different host directories
|
||||
# for different users.
|
||||
# Per-user category mounts stay present for the sandbox lifetime. Their
|
||||
# enabled-only contents change beneath these stable roots.
|
||||
try:
|
||||
config = get_app_config()
|
||||
skills_container_path = config.skills.container_path
|
||||
user_custom_path = paths.user_custom_skills_dir(effective_user_id)
|
||||
integrations_path = paths.integration_skills_dir()
|
||||
user_custom_path.mkdir(parents=True, exist_ok=True)
|
||||
integrations_path.mkdir(parents=True, exist_ok=True)
|
||||
projection = skill_projection if skill_projection is not None else LocalSandboxProvider._ensure_skills_projection(effective_user_id)
|
||||
|
||||
mappings.append(
|
||||
PathMapping(
|
||||
container_path=f"{skills_container_path}/custom",
|
||||
local_path=str(user_custom_path),
|
||||
read_only=True,
|
||||
)
|
||||
)
|
||||
mappings.append(
|
||||
PathMapping(
|
||||
container_path=f"{skills_container_path}/integrations",
|
||||
local_path=str(integrations_path),
|
||||
read_only=True,
|
||||
)
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Could not setup per-thread custom skills mount: %s", exc, exc_info=True)
|
||||
|
||||
# Legacy (pre-migration global-custom) skills: only mount for users
|
||||
# who have no per-user custom skills yet, mirroring the
|
||||
# ``UserScopedSkillStorage._iter_skill_files`` visibility rule. Users
|
||||
# with their own per-user custom skills cannot see LEGACY in the
|
||||
# listing/prompt and must not be able to read it via the sandbox
|
||||
# either — otherwise the listing layer and the sandbox layer disagree
|
||||
# about visibility, and the sandbox layer is the more permissive one.
|
||||
try:
|
||||
config = get_app_config()
|
||||
skills_container_path = config.skills.container_path
|
||||
user_custom_path = paths.user_custom_skills_dir(effective_user_id)
|
||||
legacy_skills_path = config.skills.get_skills_path() / "custom"
|
||||
if user_should_see_legacy_skills(effective_user_id, host_path=str(config.skills.get_skills_path())) and legacy_skills_path.exists():
|
||||
mappings.append(
|
||||
PathMapping(
|
||||
container_path=f"{skills_container_path}/legacy",
|
||||
local_path=str(legacy_skills_path),
|
||||
read_only=True,
|
||||
)
|
||||
if projection is not None:
|
||||
mappings.extend(
|
||||
[
|
||||
PathMapping(
|
||||
container_path=f"{skills_container_path}/custom",
|
||||
local_path=str(projection.custom),
|
||||
read_only=True,
|
||||
),
|
||||
PathMapping(
|
||||
container_path=f"{skills_container_path}/legacy",
|
||||
local_path=str(projection.legacy),
|
||||
read_only=True,
|
||||
),
|
||||
PathMapping(
|
||||
container_path=f"{skills_container_path}/integrations",
|
||||
local_path=str(projection.integrations),
|
||||
read_only=True,
|
||||
),
|
||||
]
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Could not setup per-thread legacy skills mount: %s", exc, exc_info=True)
|
||||
logger.warning("Could not setup per-thread skills projection mounts: %s", exc, exc_info=True)
|
||||
|
||||
return mappings
|
||||
|
||||
@ -350,13 +374,23 @@ class LocalSandboxProvider(SandboxProvider):
|
||||
global _singleton
|
||||
|
||||
if thread_id is None:
|
||||
skill_projection = self._ensure_skills_projection()
|
||||
with self._lock:
|
||||
if self._generic_sandbox is None:
|
||||
self._generic_sandbox = LocalSandbox("local", path_mappings=list(self._path_mappings))
|
||||
mappings = list(self._path_mappings)
|
||||
self._append_public_skill_mapping(mappings, skill_projection)
|
||||
self._generic_sandbox = LocalSandbox("local", path_mappings=mappings)
|
||||
_singleton = self._generic_sandbox
|
||||
return self._generic_sandbox.id
|
||||
|
||||
effective_user_id = self._effective_acquire_user_id(user_id)
|
||||
# Runs on every acquire, including cache hits, to self-heal drift —
|
||||
# cheap (~3-4 ms metadata walk) when the manifest is fresh. If another
|
||||
# worker mutated this user's skills since the last check, this
|
||||
# triggers a full rebuild (~400 ms measured locally) under the
|
||||
# cross-process projection lock, serializing concurrent acquires and
|
||||
# mutations for that user. Acceptable for an editing-frequency event.
|
||||
skill_projection = self._ensure_skills_projection(effective_user_id)
|
||||
key = self._thread_key(thread_id, effective_user_id)
|
||||
|
||||
# Fast path under lock.
|
||||
@ -366,11 +400,18 @@ class LocalSandboxProvider(SandboxProvider):
|
||||
# Mark as most-recently used so frequently-touched threads
|
||||
# survive eviction.
|
||||
self._thread_sandboxes.move_to_end(key)
|
||||
return cached.id
|
||||
if cached is not None:
|
||||
return cached.id
|
||||
|
||||
# ``_build_thread_path_mappings`` touches the filesystem
|
||||
# (``ensure_thread_dirs``); release the lock during I/O.
|
||||
new_mappings = list(self._path_mappings) + self._build_thread_path_mappings(thread_id, user_id=effective_user_id)
|
||||
new_mappings = list(self._path_mappings)
|
||||
self._append_public_skill_mapping(new_mappings, skill_projection)
|
||||
new_mappings += self._build_thread_path_mappings(
|
||||
thread_id,
|
||||
user_id=effective_user_id,
|
||||
skill_projection=skill_projection,
|
||||
)
|
||||
|
||||
with self._lock:
|
||||
# Re-check after the lock-free I/O: another caller may have
|
||||
|
||||
552
backend/packages/harness/deerflow/skills/projection.py
Normal file
552
backend/packages/harness/deerflow/skills/projection.py
Normal file
@ -0,0 +1,552 @@
|
||||
"""Materialize enabled-only skill trees for sandbox filesystem exposure."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import errno
|
||||
import hashlib
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
import threading
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from deerflow.skills.parser import parse_skill_file
|
||||
from deerflow.skills.types import SKILL_MD_FILE, Skill, SkillCategory
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from deerflow.skills.storage.skill_storage import SkillStorage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
import fcntl
|
||||
except ImportError: # pragma: no cover - Windows
|
||||
fcntl = None # type: ignore[assignment]
|
||||
import msvcrt
|
||||
|
||||
_locks_guard = threading.Lock()
|
||||
_process_locks: dict[Path, threading.RLock] = {}
|
||||
_MANIFEST_VERSION = 1
|
||||
_MAX_REBUILD_ATTEMPTS = 2
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SkillProjectionPaths:
|
||||
"""Stable category roots mounted or uploaded by sandbox providers."""
|
||||
|
||||
public: Path
|
||||
custom: Path
|
||||
legacy: Path
|
||||
integrations: Path
|
||||
|
||||
|
||||
def get_skill_projection_paths(storage: SkillStorage) -> SkillProjectionPaths:
|
||||
from deerflow.config.paths import get_paths
|
||||
|
||||
paths = getattr(storage, "_paths", None) or get_paths()
|
||||
user_id = getattr(storage, "user_id", None)
|
||||
if user_id is None:
|
||||
return SkillProjectionPaths(
|
||||
public=paths.public_skills_view_dir,
|
||||
custom=paths.skills_view_dir / "custom",
|
||||
legacy=paths.skills_view_dir / "legacy",
|
||||
integrations=paths.skills_view_dir / "integrations",
|
||||
)
|
||||
return SkillProjectionPaths(
|
||||
public=paths.public_skills_view_dir,
|
||||
custom=paths.user_custom_skills_view_dir(user_id),
|
||||
legacy=paths.user_legacy_skills_view_dir(user_id),
|
||||
integrations=paths.user_integration_skills_view_dir(user_id),
|
||||
)
|
||||
|
||||
|
||||
def _lock_for(path: Path) -> threading.RLock:
|
||||
resolved = path.resolve()
|
||||
with _locks_guard:
|
||||
return _process_locks.setdefault(resolved, threading.RLock())
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _projection_lock(root: Path) -> Iterator[None]:
|
||||
"""Serialize projection replacement in-process and across POSIX workers."""
|
||||
lock_path = root.parent / f".{root.name}.projection.lock"
|
||||
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
process_lock = _lock_for(lock_path)
|
||||
with process_lock, lock_path.open("a", encoding="utf-8") as lock_file:
|
||||
if fcntl is not None:
|
||||
fcntl.flock(lock_file, fcntl.LOCK_EX)
|
||||
else: # pragma: no cover - Windows
|
||||
lock_file.seek(0)
|
||||
msvcrt.locking(lock_file.fileno(), msvcrt.LK_LOCK, 1)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
if fcntl is not None:
|
||||
fcntl.flock(lock_file, fcntl.LOCK_UN)
|
||||
else: # pragma: no cover - Windows
|
||||
lock_file.seek(0)
|
||||
msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1)
|
||||
|
||||
|
||||
def _link_or_copy(source: str, target: str, *, follow_symlinks: bool = True) -> str:
|
||||
# Hardlinks share the source inode and provide no write isolation. Any
|
||||
# read-only guarantee must come from the consuming sandbox or mount.
|
||||
try:
|
||||
os.link(source, target, follow_symlinks=follow_symlinks)
|
||||
except OSError as exc:
|
||||
if exc.errno not in {errno.EXDEV, errno.EPERM, errno.EACCES, errno.ENOTSUP}:
|
||||
raise
|
||||
shutil.copy2(source, target, follow_symlinks=follow_symlinks)
|
||||
return target
|
||||
|
||||
|
||||
def _stage_skill(source: Path, target: Path, nested_skill_roots: set[Path]) -> None:
|
||||
def _exclude_nested_skills(current: str, names: list[str]) -> list[str]:
|
||||
relative_root = Path(current).relative_to(source)
|
||||
return [name for name in names if relative_root / name in nested_skill_roots]
|
||||
|
||||
shutil.copytree(
|
||||
source,
|
||||
target,
|
||||
copy_function=_link_or_copy,
|
||||
symlinks=True,
|
||||
ignore=_exclude_nested_skills,
|
||||
dirs_exist_ok=True,
|
||||
)
|
||||
|
||||
|
||||
def _path_kind(path: Path) -> str:
|
||||
if path.is_symlink():
|
||||
return "symlink"
|
||||
if path.is_dir():
|
||||
return "directory"
|
||||
return "file"
|
||||
|
||||
|
||||
def _tree_entries(root: Path) -> dict[Path, str]:
|
||||
entries: dict[Path, str] = {}
|
||||
for current_root, dir_names, file_names in os.walk(root, followlinks=False):
|
||||
current = Path(current_root)
|
||||
for name in dir_names:
|
||||
path = current / name
|
||||
entries[path.relative_to(root)] = _path_kind(path)
|
||||
dir_names[:] = [name for name in dir_names if not (current / name).is_symlink()]
|
||||
for name in file_names:
|
||||
path = current / name
|
||||
entries[path.relative_to(root)] = _path_kind(path)
|
||||
return entries
|
||||
|
||||
|
||||
def _remove_projection_entry(path: Path) -> None:
|
||||
if path.is_dir() and not path.is_symlink():
|
||||
shutil.rmtree(path)
|
||||
else:
|
||||
path.unlink(missing_ok=True)
|
||||
|
||||
|
||||
def _validate_projection_relative_path(relative_path: Path) -> None:
|
||||
if relative_path.is_absolute() or not relative_path.parts or any(part in {"", ".", ".."} for part in relative_path.parts):
|
||||
raise ValueError("Projection removal path must identify a package within its category root")
|
||||
|
||||
|
||||
def _remove_projection_relative(root: Path, relative_path: Path) -> None:
|
||||
"""Remove a projected package without following a drifted namespace symlink."""
|
||||
current = root
|
||||
for part in relative_path.parts:
|
||||
current /= part
|
||||
if current.is_symlink():
|
||||
current.unlink()
|
||||
return
|
||||
_remove_projection_entry(current)
|
||||
|
||||
|
||||
def _sync_staged_category(root: Path, staging: Path) -> None:
|
||||
desired = _tree_entries(staging)
|
||||
live = _tree_entries(root)
|
||||
|
||||
for relative_path, live_kind in sorted(live.items(), key=lambda item: len(item[0].parts), reverse=True):
|
||||
if desired.get(relative_path) != live_kind:
|
||||
_remove_projection_entry(root / relative_path)
|
||||
|
||||
for relative_path, kind in sorted(desired.items(), key=lambda item: len(item[0].parts)):
|
||||
if kind == "directory":
|
||||
(root / relative_path).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for relative_path, kind in desired.items():
|
||||
if kind == "directory":
|
||||
continue
|
||||
target = root / relative_path
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
(staging / relative_path).replace(target)
|
||||
|
||||
|
||||
def _replace_category(root: Path, desired: dict[Path, Skill], skill_boundaries: set[Path]) -> None:
|
||||
"""Reconcile entries beneath a stable category root without blanking it."""
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
with tempfile.TemporaryDirectory(prefix=f".{root.name}.projection-", dir=root.parent) as staging_dir:
|
||||
staging = Path(staging_dir)
|
||||
for relative_path, skill in desired.items():
|
||||
nested_roots = {boundary.relative_to(relative_path) for boundary in skill_boundaries if boundary != relative_path and boundary.is_relative_to(relative_path)}
|
||||
_stage_skill(skill.skill_dir, staging / relative_path, nested_roots)
|
||||
_sync_staged_category(root, staging)
|
||||
|
||||
|
||||
def _clear_category(root: Path) -> None:
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
for path in root.iterdir():
|
||||
if path.is_dir() and not path.is_symlink():
|
||||
shutil.rmtree(path)
|
||||
else:
|
||||
path.unlink()
|
||||
|
||||
|
||||
def _clear_projection_scope(scope_root: Path, *category_roots: Path) -> None:
|
||||
for category_root in category_roots:
|
||||
_clear_category(category_root)
|
||||
_manifest_path(scope_root).unlink(missing_ok=True)
|
||||
|
||||
|
||||
def _update_tree_digest(digest, root: Path, label: str) -> None:
|
||||
"""Hash directory metadata (inode/mode/size/mtime), not file contents.
|
||||
|
||||
Trade-off: fast enough to run on every sandbox acquire (O(files), no
|
||||
reads), but an external edit that preserves inode+size+mtime — unlikely,
|
||||
not zero-probability — is invisible to this signature and leaves the
|
||||
projection stale until the next explicit rebuild. Runtime writes through
|
||||
this codebase are covered regardless: the mutation path rebuilds under
|
||||
lock, and atomic-rename always changes the inode.
|
||||
"""
|
||||
digest.update(f"root:{label}\0".encode())
|
||||
if not root.exists():
|
||||
digest.update(b"absent\0")
|
||||
return
|
||||
|
||||
stack = [(root, Path("."))]
|
||||
while stack:
|
||||
current, relative_root = stack.pop()
|
||||
with os.scandir(current) as entries:
|
||||
ordered = sorted(entries, key=lambda entry: entry.name)
|
||||
child_dirs: list[tuple[Path, Path]] = []
|
||||
for entry in ordered:
|
||||
relative = relative_root / entry.name
|
||||
metadata = entry.stat(follow_symlinks=False)
|
||||
if entry.is_symlink():
|
||||
kind = "link"
|
||||
elif entry.is_dir(follow_symlinks=False):
|
||||
kind = "dir"
|
||||
child_dirs.append((Path(entry.path), relative))
|
||||
else:
|
||||
kind = "file"
|
||||
digest.update((f"{label}:{relative.as_posix()}:{kind}:{metadata.st_ino}:{metadata.st_mode}:{metadata.st_size}:{metadata.st_mtime_ns}\0").encode())
|
||||
stack.extend(reversed(child_dirs))
|
||||
|
||||
|
||||
def _extensions_state() -> dict:
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
|
||||
config = ExtensionsConfig.from_file()
|
||||
return {name: state.model_dump(mode="json") for name, state in config.skills.items()}
|
||||
|
||||
|
||||
def _source_signature(storage: SkillStorage, scope: str) -> str:
|
||||
digest = hashlib.sha256()
|
||||
host_root = storage.get_skills_root_path()
|
||||
if scope == "public":
|
||||
_update_tree_digest(digest, host_root / SkillCategory.PUBLIC.value, "public")
|
||||
state = {"extensions": _extensions_state()}
|
||||
elif scope == "user":
|
||||
user_custom_root = storage.get_user_custom_root()
|
||||
integration_root = storage.get_user_integrations_root()
|
||||
_update_tree_digest(digest, user_custom_root, "custom")
|
||||
_update_tree_digest(digest, host_root / SkillCategory.CUSTOM.value, "legacy")
|
||||
_update_tree_digest(digest, integration_root, "integrations")
|
||||
# CUSTOM/LEGACY/INTEGRATION visibility is the intersection of the
|
||||
# per-user state and the global extensions default, so both belong in
|
||||
# this signature.
|
||||
state = {
|
||||
"extensions": _extensions_state(),
|
||||
"user": storage._read_skill_states(),
|
||||
}
|
||||
else: # pragma: no cover - internal invariant
|
||||
raise ValueError(f"Unknown skill projection scope: {scope}")
|
||||
digest.update(json.dumps(state, sort_keys=True, separators=(",", ":")).encode())
|
||||
return digest.hexdigest()
|
||||
|
||||
|
||||
def _manifest_path(scope_root: Path) -> Path:
|
||||
return scope_root / ".projection-manifest.json"
|
||||
|
||||
|
||||
def _read_manifest(scope_root: Path) -> dict | None:
|
||||
try:
|
||||
value = json.loads(_manifest_path(scope_root).read_text(encoding="utf-8"))
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return None
|
||||
return value if isinstance(value, dict) else None
|
||||
|
||||
|
||||
def _write_manifest(scope_root: Path, source_signature: str) -> None:
|
||||
scope_root.mkdir(parents=True, exist_ok=True)
|
||||
target = _manifest_path(scope_root)
|
||||
fd, temporary_name = tempfile.mkstemp(prefix=".projection-manifest-", suffix=".tmp", dir=scope_root)
|
||||
temporary = Path(temporary_name)
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as stream:
|
||||
json.dump({"version": _MANIFEST_VERSION, "source_signature": source_signature}, stream, sort_keys=True)
|
||||
temporary.replace(target)
|
||||
except Exception:
|
||||
temporary.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
|
||||
def _load_public_skills(storage: SkillStorage, *, enabled_only: bool) -> list[Skill]:
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
|
||||
public_root = storage.get_skills_root_path() / SkillCategory.PUBLIC.value
|
||||
if not public_root.is_dir():
|
||||
return []
|
||||
extensions = ExtensionsConfig.from_file()
|
||||
skills: list[Skill] = []
|
||||
for current_root, dir_names, file_names in os.walk(public_root, followlinks=True):
|
||||
dir_names[:] = sorted(name for name in dir_names if not name.startswith("."))
|
||||
if SKILL_MD_FILE not in file_names:
|
||||
continue
|
||||
# Match the runtime loader: nested SKILL.md files inside a package are
|
||||
# support data, not independently configurable skills.
|
||||
dir_names.clear()
|
||||
skill_file = Path(current_root) / SKILL_MD_FILE
|
||||
skill = parse_skill_file(
|
||||
skill_file,
|
||||
category=SkillCategory.PUBLIC,
|
||||
relative_path=skill_file.parent.relative_to(public_root),
|
||||
)
|
||||
if skill is None:
|
||||
continue
|
||||
enabled = extensions.is_skill_enabled(skill.name, SkillCategory.PUBLIC.value)
|
||||
if not enabled_only or enabled:
|
||||
skills.append(skill)
|
||||
return skills
|
||||
|
||||
|
||||
def _by_relative_path(skills: list[Skill], category: SkillCategory) -> dict[Path, Skill]:
|
||||
return {skill.relative_path: skill for skill in skills if skill.category == category}
|
||||
|
||||
|
||||
def _category_boundaries(skills: list[Skill], category: SkillCategory) -> set[Path]:
|
||||
return {skill.relative_path for skill in skills if skill.category == category}
|
||||
|
||||
|
||||
def _rebuild_public_locked(storage: SkillStorage, paths: SkillProjectionPaths) -> None:
|
||||
scope_root = paths.public.parent
|
||||
try:
|
||||
for _attempt in range(_MAX_REBUILD_ATTEMPTS):
|
||||
before = _source_signature(storage, "public")
|
||||
all_public_skills = _load_public_skills(storage, enabled_only=False)
|
||||
enabled_public_skills = _load_public_skills(storage, enabled_only=True)
|
||||
_replace_category(
|
||||
paths.public,
|
||||
_by_relative_path(enabled_public_skills, SkillCategory.PUBLIC),
|
||||
_category_boundaries(all_public_skills, SkillCategory.PUBLIC),
|
||||
)
|
||||
after = _source_signature(storage, "public")
|
||||
if before == after:
|
||||
_write_manifest(scope_root, after)
|
||||
return
|
||||
raise RuntimeError("Public skills changed repeatedly while rebuilding the sandbox projection")
|
||||
except Exception:
|
||||
_clear_projection_scope(scope_root, paths.public)
|
||||
raise
|
||||
|
||||
|
||||
def _rebuild_user_locked(storage: SkillStorage, paths: SkillProjectionPaths) -> None:
|
||||
scope_root = paths.custom.parent
|
||||
try:
|
||||
for _attempt in range(_MAX_REBUILD_ATTEMPTS):
|
||||
before = _source_signature(storage, "user")
|
||||
all_user_skills = storage.load_skills(enabled_only=False)
|
||||
enabled_user_skills = [skill for skill in all_user_skills if skill.enabled]
|
||||
_replace_category(
|
||||
paths.custom,
|
||||
_by_relative_path(enabled_user_skills, SkillCategory.CUSTOM),
|
||||
_category_boundaries(all_user_skills, SkillCategory.CUSTOM),
|
||||
)
|
||||
_replace_category(
|
||||
paths.legacy,
|
||||
_by_relative_path(enabled_user_skills, SkillCategory.LEGACY),
|
||||
_category_boundaries(all_user_skills, SkillCategory.LEGACY),
|
||||
)
|
||||
_replace_category(
|
||||
paths.integrations,
|
||||
_by_relative_path(enabled_user_skills, SkillCategory.INTEGRATION),
|
||||
_category_boundaries(all_user_skills, SkillCategory.INTEGRATION),
|
||||
)
|
||||
after = _source_signature(storage, "user")
|
||||
if before == after:
|
||||
_write_manifest(scope_root, after)
|
||||
return
|
||||
raise RuntimeError("User skills changed repeatedly while rebuilding the sandbox projection")
|
||||
except Exception:
|
||||
_clear_projection_scope(scope_root, paths.custom, paths.legacy, paths.integrations)
|
||||
raise
|
||||
|
||||
|
||||
def rebuild_skill_projections(
|
||||
storage: SkillStorage,
|
||||
*,
|
||||
include_public: bool = True,
|
||||
include_user: bool = True,
|
||||
) -> SkillProjectionPaths:
|
||||
"""Rebuild enabled-only projection scopes visible through ``storage``."""
|
||||
paths = get_skill_projection_paths(storage)
|
||||
user_id = getattr(storage, "user_id", None)
|
||||
if include_public:
|
||||
with _projection_lock(paths.public.parent):
|
||||
_rebuild_public_locked(storage, paths)
|
||||
if include_user and user_id is not None:
|
||||
with _projection_lock(paths.custom.parent):
|
||||
_rebuild_user_locked(storage, paths)
|
||||
|
||||
return paths
|
||||
|
||||
|
||||
def _public_projection_is_fresh(storage: SkillStorage, paths: SkillProjectionPaths) -> bool:
|
||||
if not paths.public.is_dir():
|
||||
return False
|
||||
manifest_before = _read_manifest(paths.public.parent)
|
||||
if manifest_before is None or manifest_before.get("version") != _MANIFEST_VERSION:
|
||||
return False
|
||||
signature = _source_signature(storage, "public")
|
||||
manifest_after = _read_manifest(paths.public.parent)
|
||||
return manifest_before == manifest_after and manifest_before.get("source_signature") == signature
|
||||
|
||||
|
||||
def ensure_skill_projections(storage: SkillStorage) -> SkillProjectionPaths:
|
||||
"""Repair stale projection scopes, otherwise leave their inodes untouched."""
|
||||
paths = get_skill_projection_paths(storage)
|
||||
|
||||
try:
|
||||
public_is_fresh = _public_projection_is_fresh(storage, paths)
|
||||
except Exception:
|
||||
# Re-check under the mutation lock before failing closed. A concurrent
|
||||
# writer may have exposed a transient source/manifest state.
|
||||
public_is_fresh = False
|
||||
if not public_is_fresh:
|
||||
with _projection_lock(paths.public.parent):
|
||||
try:
|
||||
if not _public_projection_is_fresh(storage, paths):
|
||||
_rebuild_public_locked(storage, paths)
|
||||
except Exception:
|
||||
_clear_projection_scope(paths.public.parent, paths.public)
|
||||
raise
|
||||
|
||||
if getattr(storage, "user_id", None) is not None:
|
||||
with _projection_lock(paths.custom.parent):
|
||||
try:
|
||||
manifest = _read_manifest(paths.custom.parent)
|
||||
signature = _source_signature(storage, "user")
|
||||
if not paths.custom.is_dir() or not paths.legacy.is_dir() or not paths.integrations.is_dir() or manifest is None or manifest.get("version") != _MANIFEST_VERSION or manifest.get("source_signature") != signature:
|
||||
_rebuild_user_locked(storage, paths)
|
||||
except Exception:
|
||||
_clear_projection_scope(paths.custom.parent, paths.custom, paths.legacy, paths.integrations)
|
||||
raise
|
||||
return paths
|
||||
|
||||
|
||||
@contextmanager
|
||||
def skill_projection_mutation(
|
||||
storage: SkillStorage,
|
||||
scope: str,
|
||||
*,
|
||||
remove: tuple[tuple[SkillCategory, Path], ...] = (),
|
||||
remove_names: tuple[str, ...] = (),
|
||||
) -> Iterator[None]:
|
||||
"""Hold a projection scope lock across a source/state mutation."""
|
||||
if not isinstance(storage.get_skills_root_path(), Path):
|
||||
# Lightweight unit-test doubles sometimes return MagicMock here. The
|
||||
# SkillStorage contract requires a Path; real storage implementations
|
||||
# therefore never take this compatibility branch.
|
||||
yield
|
||||
return
|
||||
paths = get_skill_projection_paths(storage)
|
||||
if scope == "public":
|
||||
scope_root = paths.public.parent
|
||||
category_roots = {SkillCategory.PUBLIC: paths.public}
|
||||
|
||||
def rebuild() -> None:
|
||||
_rebuild_public_locked(storage, paths)
|
||||
|
||||
elif scope == "user":
|
||||
scope_root = paths.custom.parent
|
||||
category_roots = {
|
||||
SkillCategory.CUSTOM: paths.custom,
|
||||
SkillCategory.LEGACY: paths.legacy,
|
||||
SkillCategory.INTEGRATION: paths.integrations,
|
||||
}
|
||||
|
||||
def rebuild() -> None:
|
||||
_rebuild_user_locked(storage, paths)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown skill projection scope: {scope}")
|
||||
|
||||
removals: set[tuple[Path, Path]] = set()
|
||||
for category, relative_path in remove:
|
||||
root = category_roots.get(category)
|
||||
if root is None:
|
||||
raise ValueError(f"Skill category {category.value!r} does not belong to projection scope {scope!r}")
|
||||
_validate_projection_relative_path(relative_path)
|
||||
removals.add((root, relative_path))
|
||||
|
||||
with _projection_lock(scope_root):
|
||||
if remove_names:
|
||||
names = set(remove_names)
|
||||
skills = _load_public_skills(storage, enabled_only=False) if scope == "public" else storage.load_skills(enabled_only=False)
|
||||
for skill in skills:
|
||||
root = category_roots.get(skill.category)
|
||||
if skill.name not in names or root is None:
|
||||
continue
|
||||
_validate_projection_relative_path(skill.relative_path)
|
||||
removals.add((root, skill.relative_path))
|
||||
|
||||
try:
|
||||
_manifest_path(scope_root).unlink(missing_ok=True)
|
||||
for root, relative_path in removals:
|
||||
_remove_projection_relative(root, relative_path)
|
||||
yield
|
||||
rebuild()
|
||||
except Exception:
|
||||
_clear_projection_scope(scope_root, *category_roots.values())
|
||||
raise
|
||||
|
||||
|
||||
def ensure_public_skill_projection(*, app_config=None) -> bool:
|
||||
"""Ensure the global public view during boot without scanning user data.
|
||||
|
||||
User projections are repaired lazily by sandbox acquire. Eagerly rebuilding
|
||||
every historical user would make gateway readiness scale with tenant count,
|
||||
while providing no additional safety before that user's next acquire.
|
||||
"""
|
||||
from deerflow.config import get_app_config
|
||||
from deerflow.config.paths import get_paths
|
||||
from deerflow.skills.storage import get_or_new_skill_storage
|
||||
|
||||
try:
|
||||
config = app_config or get_app_config()
|
||||
public_storage = get_or_new_skill_storage(app_config=config)
|
||||
ensure_skill_projections(public_storage)
|
||||
except Exception:
|
||||
logger.warning("Failed to ensure the public skill projection during boot; clearing it until a sandbox acquire self-heals it", exc_info=True)
|
||||
try:
|
||||
paths = get_paths()
|
||||
with _projection_lock(paths.public_skills_view_dir.parent):
|
||||
_clear_projection_scope(paths.public_skills_view_dir.parent, paths.public_skills_view_dir)
|
||||
except Exception:
|
||||
logger.error("Failed to clear the public skill projection after a boot-time error", exc_info=True)
|
||||
return False
|
||||
return True
|
||||
@ -10,6 +10,7 @@ import os
|
||||
import shutil
|
||||
import tempfile
|
||||
from collections.abc import Iterable
|
||||
from contextlib import nullcontext
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
|
||||
@ -107,8 +108,18 @@ class LocalSkillStorage(SkillStorage):
|
||||
) as tmp_file:
|
||||
tmp_file.write(content)
|
||||
tmp_path = Path(tmp_file.name)
|
||||
tmp_path.replace(target)
|
||||
make_skill_written_path_sandbox_readable(self.get_custom_skill_dir(name), target)
|
||||
try:
|
||||
with self._skill_projection_mutation():
|
||||
tmp_path.replace(target)
|
||||
make_skill_written_path_sandbox_readable(self.get_custom_skill_dir(name), target)
|
||||
except Exception:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
def remove_custom_skill_file(self, name: str, relative_path: str) -> str:
|
||||
removal = ((SkillCategory.CUSTOM, Path(name)),)
|
||||
with self._skill_projection_mutation(remove=removal):
|
||||
return super().remove_custom_skill_file(name, relative_path)
|
||||
|
||||
async def ainstall_skill_from_archive(self, archive_path: str | Path) -> dict:
|
||||
from deerflow.skills.installer import _scan_skill_archive_contents_or_raise
|
||||
@ -200,11 +211,12 @@ class LocalSkillStorage(SkillStorage):
|
||||
"""Stage and move the validated skill into place (blocking; runs off the event loop)."""
|
||||
from deerflow.skills.installer import _move_staged_skill_into_reserved_target
|
||||
|
||||
with tempfile.TemporaryDirectory(prefix=f".installing-{skill_name}-", dir=custom_dir) as staging_root:
|
||||
staging_target = Path(staging_root) / skill_name
|
||||
shutil.copytree(skill_dir, staging_target)
|
||||
_move_staged_skill_into_reserved_target(staging_target, target)
|
||||
make_skill_written_path_sandbox_readable(custom_dir, target)
|
||||
with self._skill_projection_mutation():
|
||||
with tempfile.TemporaryDirectory(prefix=f".installing-{skill_name}-", dir=custom_dir) as staging_root:
|
||||
staging_target = Path(staging_root) / skill_name
|
||||
shutil.copytree(skill_dir, staging_target)
|
||||
_move_staged_skill_into_reserved_target(staging_target, target)
|
||||
make_skill_written_path_sandbox_readable(custom_dir, target)
|
||||
|
||||
def delete_custom_skill(self, name: str, *, history_meta: dict | None = None) -> None:
|
||||
self.validate_skill_name(name)
|
||||
@ -222,8 +234,22 @@ class LocalSkillStorage(SkillStorage):
|
||||
name,
|
||||
e,
|
||||
)
|
||||
if target.exists():
|
||||
shutil.rmtree(target)
|
||||
removal = ((SkillCategory.CUSTOM, Path(name)),)
|
||||
with self._skill_projection_mutation(remove=removal):
|
||||
if target.exists():
|
||||
shutil.rmtree(target)
|
||||
|
||||
def _skill_projection_mutation(
|
||||
self,
|
||||
*,
|
||||
remove: tuple[tuple[SkillCategory, Path], ...] = (),
|
||||
remove_names: tuple[str, ...] = (),
|
||||
):
|
||||
if getattr(self, "user_id", None) is None:
|
||||
return nullcontext()
|
||||
from deerflow.skills.projection import skill_projection_mutation
|
||||
|
||||
return skill_projection_mutation(self, "user", remove=remove, remove_names=remove_names)
|
||||
|
||||
def append_history(self, name: str, record: dict) -> None:
|
||||
self.validate_skill_name(name)
|
||||
|
||||
@ -155,6 +155,15 @@ class SkillStorage(ABC):
|
||||
Origin: ``deerflow.skills.manager.atomic_write``.
|
||||
"""
|
||||
|
||||
def remove_custom_skill_file(self, name: str, relative_path: str) -> str:
|
||||
"""Remove a supporting file and return its previous text content."""
|
||||
target = self.ensure_safe_support_path(name, relative_path)
|
||||
if not target.exists():
|
||||
raise FileNotFoundError(f"Supporting file '{relative_path}' not found for skill '{name}'.")
|
||||
previous_content = target.read_text(encoding="utf-8")
|
||||
target.unlink()
|
||||
return previous_content
|
||||
|
||||
@abstractmethod
|
||||
async def ainstall_skill_from_archive(self, archive_path: str | Path) -> dict:
|
||||
"""Async install of a skill from a ``.skill`` ZIP archive.
|
||||
|
||||
@ -82,6 +82,7 @@ class UserScopedSkillStorage(LocalSkillStorage):
|
||||
|
||||
self._user_id = _validate_user_id(user_id)
|
||||
paths = get_paths()
|
||||
self._paths = paths
|
||||
self._user_custom_root: Path = paths.user_custom_skills_dir(self._user_id)
|
||||
self._integrations_root: Path = paths.integration_skills_dir()
|
||||
self._user_skills_root: Path = paths.user_skills_dir(self._user_id)
|
||||
@ -153,9 +154,11 @@ class UserScopedSkillStorage(LocalSkillStorage):
|
||||
|
||||
def set_skill_enabled_state(self, skill_name: str, enabled: bool) -> None:
|
||||
"""Set the enabled state for a custom/legacy skill and persist."""
|
||||
states = self._read_skill_states()
|
||||
states[skill_name] = {"enabled": enabled}
|
||||
self._write_skill_states(states)
|
||||
removal_names = (skill_name,) if not enabled else ()
|
||||
with self._skill_projection_mutation(remove_names=removal_names):
|
||||
states = self._read_skill_states()
|
||||
states[skill_name] = {"enabled": enabled}
|
||||
self._write_skill_states(states)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Path helpers — redirect custom skill paths to user directory
|
||||
@ -204,10 +207,12 @@ class UserScopedSkillStorage(LocalSkillStorage):
|
||||
# being silently re-enabled by an absent per-user entry, while still
|
||||
# letting the per-user state override the global default when both
|
||||
# are present. PUBLIC skill state remains governed solely by
|
||||
# extensions_config (handled by ``super().load_skills`` above).
|
||||
from deerflow.config.extensions_config import get_extensions_config
|
||||
# extensions_config (handled by ``super().load_skills`` above). Re-read
|
||||
# from disk here too so another worker's update cannot be masked by
|
||||
# this process's singleton cache while rebuilding a user projection.
|
||||
from deerflow.config.extensions_config import ExtensionsConfig
|
||||
|
||||
extensions_config = get_extensions_config()
|
||||
extensions_config = ExtensionsConfig.from_file()
|
||||
skills = [
|
||||
dataclasses.replace(s, enabled=self.get_skill_enabled_state(s.name) and extensions_config.is_skill_enabled(s.name, s.category.value if hasattr(s.category, "value") else s.category))
|
||||
if dataclasses.is_dataclass(s) and not isinstance(s, type) and (s.category.value if hasattr(s.category, "value") else s.category) != SkillCategory.PUBLIC.value
|
||||
@ -369,8 +374,13 @@ class UserScopedSkillStorage(LocalSkillStorage):
|
||||
) as tmp_file:
|
||||
tmp_file.write(content)
|
||||
tmp_path = Path(tmp_file.name)
|
||||
tmp_path.replace(target)
|
||||
make_skill_written_path_sandbox_readable(self.get_custom_skill_dir(name), target)
|
||||
try:
|
||||
with self._skill_projection_mutation():
|
||||
tmp_path.replace(target)
|
||||
make_skill_written_path_sandbox_readable(self.get_custom_skill_dir(name), target)
|
||||
except Exception:
|
||||
tmp_path.unlink(missing_ok=True)
|
||||
raise
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Public helpers
|
||||
|
||||
@ -242,11 +242,7 @@ async def _skill_manage_impl(
|
||||
await _to_thread(skill_storage.ensure_custom_skill_is_editable, name)
|
||||
if path is None:
|
||||
raise ValueError("path is required for remove_file.")
|
||||
target = await _to_thread(skill_storage.ensure_safe_support_path, name, path)
|
||||
if not await _to_thread(target.exists):
|
||||
raise FileNotFoundError(f"Supporting file '{path}' not found for skill '{name}'.")
|
||||
prev_content = await _to_thread(target.read_text, encoding="utf-8")
|
||||
await _to_thread(target.unlink)
|
||||
prev_content = await _to_thread(skill_storage.remove_custom_skill_file, name, path)
|
||||
await _to_thread(
|
||||
skill_storage.append_history,
|
||||
name,
|
||||
|
||||
@ -39,18 +39,29 @@ class Session:
|
||||
|
||||
Matches an existing thread by id first, then by exact title. Falls back to
|
||||
the literal ref (treated as an id) when nothing matches, so an unknown id
|
||||
still continues/creates that namespace.
|
||||
still continues/creates that namespace — provided it satisfies the
|
||||
canonical thread ID contract.
|
||||
"""
|
||||
try:
|
||||
threads = self.client.list_threads(limit=100).get("thread_list", [])
|
||||
except Exception: # noqa: BLE001 - resolution is best-effort
|
||||
return ref
|
||||
return self._validated_literal_ref(ref)
|
||||
if any(t.get("thread_id") == ref for t in threads):
|
||||
return ref
|
||||
for thread in threads:
|
||||
if (thread.get("title") or "") == ref:
|
||||
return thread.get("thread_id") or ref
|
||||
return ref
|
||||
return thread.get("thread_id") or self._validated_literal_ref(ref)
|
||||
return self._validated_literal_ref(ref)
|
||||
|
||||
@staticmethod
|
||||
def _validated_literal_ref(ref: str) -> str:
|
||||
"""Validate a literal ref before it is adopted as a thread id."""
|
||||
from deerflow.utils.thread_id import validate_thread_id
|
||||
|
||||
try:
|
||||
return validate_thread_id(ref)
|
||||
except ValueError as exc:
|
||||
raise ValueError(f"Thread reference {ref!r} matches no existing thread and is not a valid thread id (expected 1-64 ASCII letters, digits, hyphens, or underscores).") from exc
|
||||
|
||||
def recent_threads(self, limit: int = 20) -> list[dict]:
|
||||
return self.client.list_threads(limit=limit).get("thread_list", [])
|
||||
|
||||
@ -7,13 +7,13 @@ Both Gateway and Client delegate to these functions.
|
||||
import errno
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import stat
|
||||
from pathlib import Path
|
||||
from urllib.parse import quote
|
||||
|
||||
from deerflow.config.paths import VIRTUAL_PATH_PREFIX, get_paths
|
||||
from deerflow.runtime.user_context import get_effective_user_id
|
||||
from deerflow.utils.thread_id import validate_thread_id
|
||||
|
||||
|
||||
class PathTraversalError(ValueError):
|
||||
@ -26,22 +26,10 @@ class UnsafeUploadPathError(ValueError):
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# thread_id must be alphanumeric, hyphens, underscores, or dots only.
|
||||
_SAFE_THREAD_ID = re.compile(r"^[a-zA-Z0-9._-]+$")
|
||||
UPLOAD_STAGING_PREFIX = ".upload-"
|
||||
UPLOAD_STAGING_SUFFIX = ".part"
|
||||
|
||||
|
||||
def validate_thread_id(thread_id: str) -> None:
|
||||
"""Reject thread IDs containing characters unsafe for filesystem paths.
|
||||
|
||||
Raises:
|
||||
ValueError: If thread_id is empty or contains unsafe characters.
|
||||
"""
|
||||
if not thread_id or not _SAFE_THREAD_ID.match(thread_id):
|
||||
raise ValueError(f"Invalid thread_id: {thread_id!r}")
|
||||
|
||||
|
||||
def get_uploads_dir(thread_id: str, *, user_id: str | None = None) -> Path:
|
||||
"""Return the uploads directory path for a thread (no side effects)."""
|
||||
validate_thread_id(thread_id)
|
||||
|
||||
37
backend/packages/harness/deerflow/utils/thread_id.py
Normal file
37
backend/packages/harness/deerflow/utils/thread_id.py
Normal file
@ -0,0 +1,37 @@
|
||||
"""Canonical thread identifier validation shared across DeerFlow backends."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import uuid
|
||||
from typing import Annotated
|
||||
|
||||
from pydantic import AfterValidator, StringConstraints
|
||||
|
||||
THREAD_ID_PATTERN = r"^[A-Za-z0-9_-]{1,64}$"
|
||||
_THREAD_ID_RE = re.compile(THREAD_ID_PATTERN)
|
||||
|
||||
|
||||
def validate_thread_id(thread_id: str) -> str:
|
||||
"""Return a valid thread ID or raise ``ValueError``.
|
||||
|
||||
Thread IDs are caller-defined opaque identifiers, not necessarily UUIDs,
|
||||
but they must be safe for every persistence and filesystem backend.
|
||||
"""
|
||||
if not isinstance(thread_id, str) or _THREAD_ID_RE.fullmatch(thread_id) is None:
|
||||
raise ValueError("Invalid thread_id: expected 1-64 ASCII letters, digits, hyphens, or underscores")
|
||||
return thread_id
|
||||
|
||||
|
||||
def resolve_thread_id(thread_id: str | None) -> str:
|
||||
"""Validate a supplied ID, generating a UUID only when it is ``None``."""
|
||||
if thread_id is None:
|
||||
return str(uuid.uuid4())
|
||||
return validate_thread_id(thread_id)
|
||||
|
||||
|
||||
ThreadId = Annotated[
|
||||
str,
|
||||
StringConstraints(min_length=1, max_length=64, pattern=THREAD_ID_PATTERN),
|
||||
AfterValidator(validate_thread_id),
|
||||
]
|
||||
@ -350,13 +350,41 @@ def _validate_materialized(case: BenchmarkCase, expected: list[BaseMessage], war
|
||||
return len(cold), cold_digest
|
||||
|
||||
|
||||
_HISTORY_CACHE_ENV = "DEERFLOW_CHECKPOINT_BENCH_HISTORY_CACHE"
|
||||
|
||||
|
||||
def _wrap_history_cache(saver: Any) -> Any:
|
||||
"""Wrap *saver* in a CachedHistorySaver with a fresh, unbounded memory cache.
|
||||
|
||||
Opt-in via DEERFLOW_CHECKPOINT_BENCH_HISTORY_CACHE=1 so default rows are
|
||||
byte-identical to the pre-cache benchmark. A fresh wrapper per phase keeps
|
||||
the cold read genuinely cold: the write-phase cache is discarded, mirroring
|
||||
a process restart (cache lifetime == checkpointer CM lifetime).
|
||||
"""
|
||||
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
||||
from deerflow.runtime.checkpointer.cached_saver import CachedHistorySaver
|
||||
|
||||
return CachedHistorySaver(
|
||||
saver,
|
||||
MemoryCheckpointHistoryCache(max_entries=1_000_000),
|
||||
key_prefix="bench:v1:checkpoint-bench",
|
||||
)
|
||||
|
||||
|
||||
def _history_cache_stats(wrapper: Any, prefix: str) -> dict[str, Any]:
|
||||
return {f"{prefix}{key}": value for key, value in wrapper.stats().items()}
|
||||
|
||||
|
||||
def _run_memory_case(case: BenchmarkCase, messages: list[BaseMessage]) -> dict[str, Any]:
|
||||
cache_opt_in = os.environ.get(_HISTORY_CACHE_ENV) == "1"
|
||||
saver = InMemorySaver()
|
||||
metrics, warm = _write_and_read(case, saver, messages)
|
||||
write_saver = _wrap_history_cache(saver) if cache_opt_in else saver
|
||||
metrics, warm = _write_and_read(case, write_saver, messages)
|
||||
stats = _collect_storage_stats(lambda: _memory_storage_stats(saver, _config(case)["configurable"]["thread_id"]))
|
||||
cold_read_ms, cold = _cold_read(case, saver)
|
||||
cold_saver = _wrap_history_cache(saver) if cache_opt_in else saver
|
||||
cold_read_ms, cold = _cold_read(case, cold_saver)
|
||||
actual_count, digest = _validate_materialized(case, messages, warm, cold)
|
||||
return {
|
||||
result = {
|
||||
**metrics,
|
||||
**stats,
|
||||
"cold_read_ms": cold_read_ms,
|
||||
@ -371,12 +399,19 @@ def _run_memory_case(case: BenchmarkCase, messages: list[BaseMessage]) -> dict[s
|
||||
"actual_message_count": actual_count,
|
||||
"content_sha256": digest,
|
||||
}
|
||||
if cache_opt_in:
|
||||
result["history_cache_enabled"] = True
|
||||
result.update(_history_cache_stats(write_saver, "history_cache_write_"))
|
||||
result.update(_history_cache_stats(cold_saver, "history_cache_cold_"))
|
||||
return result
|
||||
|
||||
|
||||
def _run_sqlite_case(case: BenchmarkCase, messages: list[BaseMessage], db_path: Path) -> dict[str, Any]:
|
||||
cache_opt_in = os.environ.get(_HISTORY_CACHE_ENV) == "1"
|
||||
with SqliteSaver.from_conn_string(str(db_path)) as saver:
|
||||
saver.setup()
|
||||
metrics, warm = _write_and_read(case, saver, messages)
|
||||
write_saver = _wrap_history_cache(saver) if cache_opt_in else saver
|
||||
metrics, warm = _write_and_read(case, write_saver, messages)
|
||||
stats = _collect_storage_stats(lambda: _sqlite_storage_stats(saver, _config(case)["configurable"]["thread_id"]))
|
||||
db_bytes = _file_size(db_path)
|
||||
wal_bytes = _file_size(Path(f"{db_path}-wal"))
|
||||
@ -387,10 +422,11 @@ def _run_sqlite_case(case: BenchmarkCase, messages: list[BaseMessage], db_path:
|
||||
with SqliteSaver.from_conn_string(str(db_path)) as reopened:
|
||||
reopened.setup()
|
||||
saver_reopen_ms = (time.perf_counter() - reopen_start) * 1000
|
||||
cold_read_ms, cold = _cold_read(case, reopened)
|
||||
cold_saver = _wrap_history_cache(reopened) if cache_opt_in else reopened
|
||||
cold_read_ms, cold = _cold_read(case, cold_saver)
|
||||
|
||||
actual_count, digest = _validate_materialized(case, messages, warm, cold)
|
||||
return {
|
||||
result = {
|
||||
**metrics,
|
||||
**stats,
|
||||
"cold_read_ms": cold_read_ms,
|
||||
@ -403,6 +439,11 @@ def _run_sqlite_case(case: BenchmarkCase, messages: list[BaseMessage], db_path:
|
||||
"actual_message_count": actual_count,
|
||||
"content_sha256": digest,
|
||||
}
|
||||
if cache_opt_in:
|
||||
result["history_cache_enabled"] = True
|
||||
result.update(_history_cache_stats(write_saver, "history_cache_write_"))
|
||||
result.update(_history_cache_stats(cold_saver, "history_cache_cold_"))
|
||||
return result
|
||||
|
||||
|
||||
def _run_case(case: BenchmarkCase, *, work_dir: Path) -> dict[str, Any]:
|
||||
|
||||
@ -1,7 +1,7 @@
|
||||
"""Regression anchor: serving artifacts must not block the event loop.
|
||||
|
||||
``get_artifact`` probes the artifact path (``exists`` / ``is_file``), reads
|
||||
text content (``read_text``), sniffs text-ness (``is_text_file_by_content``),
|
||||
``get_artifact`` probes the artifact path (``exists`` / ``is_file``), sniffs
|
||||
text-ness (``is_text_file_by_content``),
|
||||
and extracts ``.skill`` archive members — all blocking filesystem IO. The
|
||||
handler offloads each branch's IO via ``asyncio.to_thread``; if any regresses
|
||||
back onto the event loop, the strict Blockbuster gate raises ``BlockingError``
|
||||
@ -26,19 +26,23 @@ sit at module top so any import-time IO runs at collection, outside the gate.
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import zipfile
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from starlette.responses import FileResponse
|
||||
|
||||
import app.gateway.routers.artifacts as artifacts_router
|
||||
from app.gateway.path_utils import resolve_thread_virtual_path
|
||||
from app.gateway.routers.artifacts import get_artifact
|
||||
from app.gateway.routers.artifacts import ArtifactUpdateRequest, get_artifact, update_artifact
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
# The undecorated coroutine (``require_permission`` uses ``functools.wraps``).
|
||||
_get_artifact = get_artifact.__wrapped__
|
||||
_update_artifact = update_artifact.__wrapped__
|
||||
|
||||
|
||||
async def _seed(tmp_path: Path, monkeypatch, thread_id: str, virtual_path: str) -> Path:
|
||||
@ -61,8 +65,10 @@ async def test_get_artifact_text_does_not_block_event_loop(tmp_path: Path, monke
|
||||
|
||||
resp = await _get_artifact("t1", vpath, request=None, download=False)
|
||||
|
||||
assert isinstance(resp, FileResponse)
|
||||
assert resp.status_code == 200
|
||||
assert resp.body == b"hello world"
|
||||
assert Path(resp.path) == target
|
||||
assert resp.headers.get("content-disposition", "").startswith("inline;")
|
||||
|
||||
|
||||
async def test_get_artifact_binary_does_not_block_event_loop(tmp_path: Path, monkeypatch) -> None:
|
||||
@ -96,3 +102,33 @@ async def test_get_artifact_skill_archive_member_does_not_block_event_loop(tmp_p
|
||||
|
||||
assert resp.status_code == 200
|
||||
assert b"# demo skill" in resp.body
|
||||
|
||||
|
||||
async def test_update_artifact_does_not_block_event_loop(tmp_path: Path, monkeypatch) -> None:
|
||||
vpath = "/mnt/user-data/outputs/notes.txt"
|
||||
target = await _seed(tmp_path, monkeypatch, "t1", vpath)
|
||||
original = b"hello world"
|
||||
await asyncio.to_thread(target.write_bytes, original)
|
||||
|
||||
@asynccontextmanager
|
||||
async def allow_write(*_args, **_kwargs):
|
||||
yield
|
||||
|
||||
class MountedProvider:
|
||||
uses_thread_data_mounts = True
|
||||
|
||||
monkeypatch.setattr(artifacts_router, "reserve_artifact_write", allow_write)
|
||||
monkeypatch.setattr(artifacts_router, "get_sandbox_provider", lambda: MountedProvider())
|
||||
|
||||
result = await _update_artifact(
|
||||
"t1",
|
||||
vpath,
|
||||
ArtifactUpdateRequest(
|
||||
content="updated",
|
||||
expected_sha256=hashlib.sha256(original).hexdigest(),
|
||||
),
|
||||
request=None,
|
||||
)
|
||||
|
||||
assert result.sha256 == hashlib.sha256(b"updated").hexdigest()
|
||||
assert await asyncio.to_thread(target.read_bytes) == b"updated"
|
||||
|
||||
192
backend/tests/blocking_io/test_channel_outbound_files.py
Normal file
192
backend/tests/blocking_io/test_channel_outbound_files.py
Normal file
@ -0,0 +1,192 @@
|
||||
"""Regression anchors for outbound IM attachment file IO.
|
||||
|
||||
Feishu, Telegram, and WeCom send attachments from async channel handlers. File
|
||||
open/read/hash work must run off the event loop; otherwise a large outbound
|
||||
artifact stalls every channel and Gateway coroutine on that worker.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import builtins
|
||||
import hashlib
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from app.channels.feishu import FeishuChannel
|
||||
from app.channels.message_bus import MessageBus, OutboundMessage, ResolvedAttachment
|
||||
from app.channels.telegram import TelegramChannel
|
||||
from app.channels.wecom import WeComChannel
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
def _attachment(path: Path, *, is_image: bool) -> ResolvedAttachment:
|
||||
return ResolvedAttachment(
|
||||
virtual_path=f"/mnt/user-data/outputs/{path.name}",
|
||||
actual_path=path,
|
||||
filename=path.name,
|
||||
mime_type="image/png" if is_image else "application/octet-stream",
|
||||
size=path.stat().st_size,
|
||||
is_image=is_image,
|
||||
)
|
||||
|
||||
|
||||
def _outbound(channel_name: str) -> OutboundMessage:
|
||||
return OutboundMessage(
|
||||
channel_name=channel_name,
|
||||
chat_id="123",
|
||||
thread_id="thread-1",
|
||||
text="attachment",
|
||||
)
|
||||
|
||||
|
||||
def _builder(*, captured_file: dict[str, object] | None = None) -> MagicMock:
|
||||
builder = MagicMock()
|
||||
for method_name in ("request_body", "image_type", "file_type", "file_name"):
|
||||
getattr(builder, method_name).return_value = builder
|
||||
if captured_file is not None:
|
||||
|
||||
def _capture(file_obj):
|
||||
captured_file["file"] = file_obj
|
||||
return builder
|
||||
|
||||
builder.image.side_effect = _capture
|
||||
builder.file.side_effect = _capture
|
||||
builder.build.return_value = object()
|
||||
return builder
|
||||
|
||||
|
||||
async def test_feishu_outbound_uploads_do_not_block_event_loop(tmp_path: Path, monkeypatch) -> None:
|
||||
image_path = tmp_path / "chart.png"
|
||||
file_path = tmp_path / "report.pdf"
|
||||
await asyncio.to_thread(image_path.write_bytes, b"image-bytes")
|
||||
await asyncio.to_thread(file_path.write_bytes, b"file-bytes")
|
||||
|
||||
channel = FeishuChannel(MessageBus(), {})
|
||||
channel._api_client = MagicMock()
|
||||
|
||||
real_open = builtins.open
|
||||
opened_on_threads: list[int] = []
|
||||
tracked_paths = {str(image_path), str(file_path)}
|
||||
|
||||
def _tracked_open(file, *args, **kwargs):
|
||||
if str(file) in tracked_paths:
|
||||
opened_on_threads.append(threading.get_ident())
|
||||
return real_open(file, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(builtins, "open", _tracked_open)
|
||||
|
||||
image_capture: dict[str, object] = {}
|
||||
channel._CreateImageRequest = MagicMock()
|
||||
channel._CreateImageRequest.builder.return_value = _builder()
|
||||
channel._CreateImageRequestBody = MagicMock()
|
||||
channel._CreateImageRequestBody.builder.return_value = _builder(captured_file=image_capture)
|
||||
|
||||
image_response = MagicMock()
|
||||
image_response.success.return_value = True
|
||||
image_response.data.image_key = "image-key"
|
||||
image_worker_thread: int | None = None
|
||||
|
||||
def _create_image(_request):
|
||||
nonlocal image_worker_thread
|
||||
image_worker_thread = threading.get_ident()
|
||||
file_obj = image_capture["file"]
|
||||
assert not file_obj.closed
|
||||
assert file_obj.read() == b"image-bytes"
|
||||
return image_response
|
||||
|
||||
channel._api_client.im.v1.image.create.side_effect = _create_image
|
||||
|
||||
file_capture: dict[str, object] = {}
|
||||
channel._CreateFileRequest = MagicMock()
|
||||
channel._CreateFileRequest.builder.return_value = _builder()
|
||||
channel._CreateFileRequestBody = MagicMock()
|
||||
channel._CreateFileRequestBody.builder.return_value = _builder(captured_file=file_capture)
|
||||
|
||||
file_response = MagicMock()
|
||||
file_response.success.return_value = True
|
||||
file_response.data.file_key = "file-key"
|
||||
file_worker_thread: int | None = None
|
||||
|
||||
def _create_file(_request):
|
||||
nonlocal file_worker_thread
|
||||
file_worker_thread = threading.get_ident()
|
||||
file_obj = file_capture["file"]
|
||||
assert not file_obj.closed
|
||||
assert file_obj.read() == b"file-bytes"
|
||||
return file_response
|
||||
|
||||
channel._api_client.im.v1.file.create.side_effect = _create_file
|
||||
|
||||
event_loop_thread = threading.get_ident()
|
||||
assert await channel._upload_image(image_path) == "image-key"
|
||||
assert await channel._upload_file(file_path, file_path.name) == "file-key"
|
||||
assert len(opened_on_threads) == 2
|
||||
assert all(thread_id != event_loop_thread for thread_id in opened_on_threads)
|
||||
assert image_worker_thread != event_loop_thread
|
||||
assert file_worker_thread != event_loop_thread
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_image", [False, True])
|
||||
async def test_telegram_outbound_upload_does_not_block_event_loop(tmp_path: Path, is_image: bool) -> None:
|
||||
path = tmp_path / ("chart.png" if is_image else "report.bin")
|
||||
payload = b"telegram-payload"
|
||||
await asyncio.to_thread(path.write_bytes, payload)
|
||||
|
||||
sent_file = None
|
||||
|
||||
class _Bot:
|
||||
async def send_document(self, **kwargs):
|
||||
nonlocal sent_file
|
||||
sent_file = kwargs["document"]
|
||||
return SimpleNamespace(message_id=7)
|
||||
|
||||
async def send_photo(self, **kwargs):
|
||||
nonlocal sent_file
|
||||
sent_file = kwargs["photo"]
|
||||
if hasattr(sent_file, "read"):
|
||||
sent_file.read()
|
||||
return SimpleNamespace(message_id=7)
|
||||
|
||||
channel = TelegramChannel(MessageBus(), {})
|
||||
channel._application = SimpleNamespace(bot=_Bot())
|
||||
|
||||
assert await channel.send_file(_outbound("telegram"), _attachment(path, is_image=is_image))
|
||||
assert sent_file.input_file_content == payload
|
||||
assert sent_file.filename == path.name
|
||||
|
||||
|
||||
async def test_wecom_outbound_upload_does_not_block_event_loop(tmp_path: Path) -> None:
|
||||
path = tmp_path / "report.bin"
|
||||
payload = b"x" * (512 * 1024 + 17)
|
||||
await asyncio.to_thread(path.write_bytes, payload)
|
||||
|
||||
channel = WeComChannel(MessageBus(), {})
|
||||
channel._ws_client = object()
|
||||
channel._send_ws_upload_command = AsyncMock(
|
||||
side_effect=[
|
||||
{"body": {"upload_id": "upload-1"}},
|
||||
{"body": {}},
|
||||
{"body": {}},
|
||||
{"body": {"media_id": "media-1"}},
|
||||
]
|
||||
)
|
||||
|
||||
result = await channel._upload_media_ws(
|
||||
media_type="file",
|
||||
filename=path.name,
|
||||
path=str(path),
|
||||
size=len(payload),
|
||||
)
|
||||
|
||||
assert result == "media-1"
|
||||
calls = channel._send_ws_upload_command.await_args_list
|
||||
assert calls[0].args[1]["md5"] == hashlib.md5(payload).hexdigest()
|
||||
encoded_chunks = [call.args[1]["base64_data"] for call in calls[1:-1]]
|
||||
assert b"".join(base64.b64decode(chunk) for chunk in encoded_chunks) == payload
|
||||
126
backend/tests/blocking_io/test_feishu_receive_file.py
Normal file
126
backend/tests/blocking_io/test_feishu_receive_file.py
Normal file
@ -0,0 +1,126 @@
|
||||
"""Regression anchors: Feishu ``receive_file`` must not block the event loop."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from io import BytesIO
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = pytest.mark.asyncio
|
||||
|
||||
|
||||
def _channel_with_file(content: bytes = b"DATA", filename: str = "report.pdf"):
|
||||
from app.channels.feishu import FeishuChannel
|
||||
from app.channels.message_bus import MessageBus
|
||||
|
||||
channel = FeishuChannel(MessageBus(), {"app_id": "test", "app_secret": "test"})
|
||||
channel._GetMessageResourceRequest = MagicMock()
|
||||
builder = MagicMock()
|
||||
builder.message_id.return_value = builder
|
||||
builder.file_key.return_value = builder
|
||||
builder.type.return_value = builder
|
||||
builder.build.return_value = object()
|
||||
channel._GetMessageResourceRequest.builder.return_value = builder
|
||||
|
||||
response = MagicMock()
|
||||
response.success.return_value = True
|
||||
response.file = BytesIO(content)
|
||||
response.file_name = filename
|
||||
channel._api_client = MagicMock()
|
||||
channel._api_client.im.v1.message_resource.get.return_value = response
|
||||
return channel
|
||||
|
||||
|
||||
class _RemoteSandbox:
|
||||
def __init__(self) -> None:
|
||||
self.updates: list[tuple[str, bytes]] = []
|
||||
self.update_thread_id: int | None = None
|
||||
|
||||
def update_file(self, path: str, content: bytes) -> None:
|
||||
self.update_thread_id = threading.get_ident()
|
||||
self.updates.append((path, content))
|
||||
|
||||
|
||||
class _RemoteProvider:
|
||||
uses_thread_data_mounts = False
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.sandbox = _RemoteSandbox()
|
||||
self.acquire_async_calls: list[tuple[str | None, str | None]] = []
|
||||
|
||||
def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str:
|
||||
raise AssertionError("Feishu receive_file must use acquire_async")
|
||||
|
||||
async def acquire_async(self, thread_id: str | None = None, *, user_id: str | None = None) -> str:
|
||||
self.acquire_async_calls.append((thread_id, user_id))
|
||||
return "remote-sandbox"
|
||||
|
||||
def get(self, sandbox_id: str):
|
||||
return self.sandbox if sandbox_id == "remote-sandbox" else None
|
||||
|
||||
|
||||
class _MountedProvider:
|
||||
uses_thread_data_mounts = True
|
||||
|
||||
def acquire(self, thread_id: str | None = None, *, user_id: str | None = None) -> str:
|
||||
raise AssertionError("mounted uploads must not acquire a sandbox")
|
||||
|
||||
async def acquire_async(self, thread_id: str | None = None, *, user_id: str | None = None) -> str:
|
||||
raise AssertionError("mounted uploads must not acquire a sandbox")
|
||||
|
||||
def get(self, sandbox_id: str):
|
||||
raise AssertionError("mounted uploads must not look up a sandbox")
|
||||
|
||||
|
||||
async def test_receive_file_remote_sandbox_does_not_block_event_loop(tmp_path, monkeypatch) -> None:
|
||||
from deerflow.config.paths import Paths
|
||||
|
||||
paths = await asyncio.to_thread(Paths, str(tmp_path))
|
||||
provider = _RemoteProvider()
|
||||
provider_lookup_thread_id = None
|
||||
|
||||
def _get_provider():
|
||||
nonlocal provider_lookup_thread_id
|
||||
provider_lookup_thread_id = threading.get_ident()
|
||||
return provider
|
||||
|
||||
monkeypatch.setattr("app.channels.feishu.get_paths", lambda: paths)
|
||||
monkeypatch.setattr("app.channels.feishu.get_sandbox_provider", _get_provider)
|
||||
|
||||
loop_thread_id = threading.get_ident()
|
||||
result = await _channel_with_file()._receive_single_file(
|
||||
"message-1",
|
||||
"file-key",
|
||||
"file",
|
||||
"thread-1",
|
||||
user_id="ou-user",
|
||||
)
|
||||
|
||||
assert result == "/mnt/user-data/uploads/report.pdf"
|
||||
assert provider.acquire_async_calls == [("thread-1", "ou-user")]
|
||||
assert provider.sandbox.updates == [(result, b"DATA")]
|
||||
assert provider_lookup_thread_id != loop_thread_id
|
||||
assert provider.sandbox.update_thread_id != loop_thread_id
|
||||
|
||||
|
||||
async def test_receive_file_mounted_sandbox_skips_redundant_sync(tmp_path, monkeypatch) -> None:
|
||||
from deerflow.config.paths import Paths
|
||||
|
||||
paths = await asyncio.to_thread(Paths, str(tmp_path))
|
||||
monkeypatch.setattr("app.channels.feishu.get_paths", lambda: paths)
|
||||
monkeypatch.setattr("app.channels.feishu.get_sandbox_provider", lambda: _MountedProvider())
|
||||
|
||||
result = await _channel_with_file()._receive_single_file(
|
||||
"message-1",
|
||||
"file-key",
|
||||
"file",
|
||||
"thread-1",
|
||||
user_id="ou-user",
|
||||
)
|
||||
|
||||
assert result == "/mnt/user-data/uploads/report.pdf"
|
||||
uploaded = tmp_path / "users" / "ou-user" / "threads" / "thread-1" / "user-data" / "uploads" / "report.pdf"
|
||||
assert await asyncio.to_thread(uploaded.read_bytes) == b"DATA"
|
||||
@ -119,6 +119,22 @@ async def test_update_skill_writes_from_snapshot_without_mutating_singleton(tmp_
|
||||
assert "middlewares" in written
|
||||
|
||||
|
||||
async def test_update_skill_persists_state_when_source_omits_skills(tmp_path: Path, monkeypatch) -> None:
|
||||
config_path = tmp_path / "extensions_config.json"
|
||||
await asyncio.to_thread(
|
||||
config_path.write_text,
|
||||
json.dumps({"mcpServers": {}, "middlewares": ["pkg:Middleware"]}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
_patch_config_infra(monkeypatch, config_path)
|
||||
|
||||
await update_skill("demo-skill", SkillUpdateRequest(enabled=False), _admin_request(), SimpleNamespace())
|
||||
|
||||
written = json.loads(await asyncio.to_thread(config_path.read_text, encoding="utf-8"))
|
||||
assert written["skills"] == {"demo-skill": {"enabled": False}}
|
||||
assert written["middlewares"] == ["pkg:Middleware"]
|
||||
|
||||
|
||||
@pytest.mark.allow_blocking_io # gate-exempt: needs real worker-thread overlap to observe serialization
|
||||
async def test_update_skill_serializes_concurrent_writes(tmp_path: Path, monkeypatch) -> None:
|
||||
state_lock = threading.Lock()
|
||||
|
||||
@ -219,8 +219,9 @@ def test_get_user_skill_mounts_mounts_only_global_integrations(tmp_path, monkeyp
|
||||
|
||||
assert set(alice) == {"/mnt/skills/integrations"}
|
||||
assert set(bob) == {"/mnt/skills/integrations"}
|
||||
assert alice["/mnt/skills/integrations"] == bob["/mnt/skills/integrations"]
|
||||
assert alice["/mnt/skills/integrations"] == str(tmp_path / "home" / "integrations" / "skills")
|
||||
assert alice["/mnt/skills/integrations"] != bob["/mnt/skills/integrations"]
|
||||
assert alice["/mnt/skills/integrations"] == str(tmp_path / "home" / "users" / "alice" / "skills_view" / "integrations")
|
||||
assert bob["/mnt/skills/integrations"] == str(tmp_path / "home" / "users" / "bob" / "skills_view" / "integrations")
|
||||
|
||||
|
||||
def test_get_extra_mounts_provisioner_payload_has_unique_container_paths(tmp_path, monkeypatch, provisioner_module):
|
||||
@ -243,7 +244,7 @@ def test_get_extra_mounts_provisioner_payload_has_unique_container_paths(tmp_pat
|
||||
monkeypatch.setattr(aio_mod, "get_app_config", lambda: config)
|
||||
monkeypatch.setattr(aio_mod, "get_paths", lambda: Paths(base_dir=home))
|
||||
monkeypatch.setattr(aio_mod, "get_effective_user_id", lambda: "default")
|
||||
monkeypatch.setattr(aio_mod, "user_should_see_legacy_skills", lambda *_args, **_kwargs: False)
|
||||
monkeypatch.setattr(remote_backend, "user_should_see_legacy_skills", lambda *_args, **_kwargs: False)
|
||||
|
||||
provider = _make_provider(tmp_path)
|
||||
mounts = provider._get_extra_mounts("thread-1", user_id="alice")
|
||||
@ -526,7 +527,10 @@ async def test_acquire_internal_async_offloads_cached_reuse_health_check(tmp_pat
|
||||
sandbox_id = await provider._acquire_internal_async("thread-cached-async", user_id="default")
|
||||
|
||||
assert sandbox_id == "sandbox-cached-async"
|
||||
assert to_thread_calls == [(provider._reuse_in_process_sandbox, ("thread-cached-async",))]
|
||||
assert to_thread_calls == [
|
||||
(provider._ensure_skills_projection, ("default",)),
|
||||
(provider._reuse_in_process_sandbox, ("thread-cached-async",)),
|
||||
]
|
||||
|
||||
|
||||
def test_remote_backend_create_forwards_effective_user_id(monkeypatch):
|
||||
@ -548,7 +552,7 @@ def test_remote_backend_create_forwards_effective_user_id(monkeypatch):
|
||||
return _Response()
|
||||
|
||||
monkeypatch.setattr(remote_mod.requests, "post", _post)
|
||||
monkeypatch.setattr(remote_mod, "user_should_see_legacy_skills", lambda user_id: True)
|
||||
monkeypatch.setattr(remote_mod, "user_should_see_legacy_skills", lambda _user_id: True)
|
||||
|
||||
try:
|
||||
backend.create("thread-42", "sandbox-42")
|
||||
@ -585,7 +589,7 @@ def test_remote_backend_create_prefers_explicit_user_id(monkeypatch):
|
||||
|
||||
monkeypatch.setattr(remote_mod.requests, "post", _post)
|
||||
monkeypatch.setattr(remote_mod, "get_effective_user_id", lambda: "default")
|
||||
monkeypatch.setattr(remote_mod, "user_should_see_legacy_skills", lambda user_id: False)
|
||||
monkeypatch.setattr(remote_mod, "user_should_see_legacy_skills", lambda _user_id: False)
|
||||
|
||||
backend.create("thread-42", "sandbox-42", user_id="ou-user")
|
||||
|
||||
|
||||
@ -1,5 +1,8 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import stat
|
||||
import zipfile
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
@ -32,18 +35,355 @@ def test_get_artifact_reads_utf8_text_file_on_windows_locale(tmp_path, monkeypat
|
||||
|
||||
original_read_text = Path.read_text
|
||||
|
||||
def read_text_with_gbk_default(self, *args, **kwargs):
|
||||
kwargs.setdefault("encoding", "gbk")
|
||||
def reject_artifact_read_text(self, *args, **kwargs):
|
||||
if self == artifact_path:
|
||||
pytest.fail("text files must stream")
|
||||
return original_read_text(self, *args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(Path, "read_text", read_text_with_gbk_default)
|
||||
monkeypatch.setattr(Path, "read_text", reject_artifact_read_text)
|
||||
monkeypatch.setattr(artifacts_router, "resolve_thread_virtual_path", lambda _thread_id, _path, user_id=None: artifact_path)
|
||||
|
||||
request = _make_request()
|
||||
response = asyncio.run(call_unwrapped(artifacts_router.get_artifact, "thread-1", "mnt/user-data/outputs/note.txt", request))
|
||||
app = make_authed_test_app()
|
||||
app.include_router(artifacts_router.router)
|
||||
with TestClient(app) as client:
|
||||
response = client.get("/api/threads/thread-1/artifacts/mnt/user-data/outputs/note.txt")
|
||||
|
||||
assert bytes(response.body).decode("utf-8") == text
|
||||
assert response.media_type == "text/plain"
|
||||
assert response.text == text
|
||||
assert response.headers["content-type"].startswith("text/plain")
|
||||
assert response.headers["accept-ranges"] == "bytes"
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def _allow_artifact_write(*_args, **_kwargs):
|
||||
yield
|
||||
|
||||
|
||||
class _MountedSandboxProvider:
|
||||
uses_thread_data_mounts = True
|
||||
|
||||
|
||||
class _RemoteSandbox:
|
||||
def __init__(self, *, fail_next_update: bool = False) -> None:
|
||||
self.updates: list[tuple[str, bytes]] = []
|
||||
self.fail_next_update = fail_next_update
|
||||
|
||||
def update_file(self, path: str, content: bytes) -> None:
|
||||
if self.fail_next_update:
|
||||
self.fail_next_update = False
|
||||
raise RuntimeError("sandbox sync failed")
|
||||
self.updates.append((path, content))
|
||||
|
||||
|
||||
class _RemoteSandboxProvider:
|
||||
uses_thread_data_mounts = False
|
||||
|
||||
def __init__(self, *, fail_next_update: bool = False) -> None:
|
||||
self.sandbox = _RemoteSandbox(fail_next_update=fail_next_update)
|
||||
self.released: list[str] = []
|
||||
|
||||
async def acquire_async(self, _thread_id: str, *, user_id: str | None = None) -> str:
|
||||
return "sandbox-1"
|
||||
|
||||
def get(self, sandbox_id: str):
|
||||
assert sandbox_id == "sandbox-1"
|
||||
return self.sandbox
|
||||
|
||||
def release(self, sandbox_id: str) -> None:
|
||||
self.released.append(sandbox_id)
|
||||
|
||||
|
||||
def _artifact_sha256(content: str) -> str:
|
||||
return hashlib.sha256(content.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _patch_artifact_update_dependencies(monkeypatch, artifact_path: Path, provider=None) -> None:
|
||||
monkeypatch.setattr(artifacts_router, "resolve_thread_virtual_path", lambda _thread_id, _path, user_id=None: artifact_path)
|
||||
monkeypatch.setattr(artifacts_router, "reserve_artifact_write", _allow_artifact_write)
|
||||
monkeypatch.setattr(artifacts_router, "get_sandbox_provider", lambda: provider or _MountedSandboxProvider())
|
||||
|
||||
|
||||
def test_update_artifact_replaces_utf8_text_atomically(tmp_path, monkeypatch) -> None:
|
||||
artifact_path = tmp_path / "note.txt"
|
||||
artifact_path.write_text("before", encoding="utf-8")
|
||||
artifact_path.chmod(0o600)
|
||||
_patch_artifact_update_dependencies(monkeypatch, artifact_path)
|
||||
|
||||
response = asyncio.run(
|
||||
call_unwrapped(
|
||||
artifacts_router.update_artifact,
|
||||
"thread-1",
|
||||
"mnt/user-data/outputs/note.txt",
|
||||
artifacts_router.ArtifactUpdateRequest(content="after", expected_sha256=_artifact_sha256("before")),
|
||||
_make_request(),
|
||||
)
|
||||
)
|
||||
|
||||
assert artifact_path.read_text(encoding="utf-8") == "after"
|
||||
assert response.path == "/mnt/user-data/outputs/note.txt"
|
||||
assert response.sha256 == _artifact_sha256("after")
|
||||
assert response.size == len(b"after")
|
||||
if hasattr(artifacts_router.os, "fchmod"):
|
||||
replacement_mode = stat.S_IMODE(artifact_path.stat().st_mode)
|
||||
assert replacement_mode == 0o660
|
||||
assert not replacement_mode & stat.S_IWOTH
|
||||
|
||||
|
||||
def test_update_artifact_replaces_when_fchmod_is_unavailable(tmp_path, monkeypatch) -> None:
|
||||
artifact_path = tmp_path / "note.txt"
|
||||
artifact_path.write_text("before", encoding="utf-8")
|
||||
_patch_artifact_update_dependencies(monkeypatch, artifact_path)
|
||||
monkeypatch.delattr(artifacts_router.os, "fchmod", raising=False)
|
||||
|
||||
response = asyncio.run(
|
||||
call_unwrapped(
|
||||
artifacts_router.update_artifact,
|
||||
"thread-1",
|
||||
"mnt/user-data/outputs/note.txt",
|
||||
artifacts_router.ArtifactUpdateRequest(content="after", expected_sha256=_artifact_sha256("before")),
|
||||
_make_request(),
|
||||
)
|
||||
)
|
||||
|
||||
assert artifact_path.read_text(encoding="utf-8") == "after"
|
||||
assert response.sha256 == _artifact_sha256("after")
|
||||
|
||||
|
||||
def test_update_artifact_rejects_stale_revision_without_changing_file(tmp_path, monkeypatch) -> None:
|
||||
artifact_path = tmp_path / "note.txt"
|
||||
artifact_path.write_text("agent version", encoding="utf-8")
|
||||
_patch_artifact_update_dependencies(monkeypatch, artifact_path)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
asyncio.run(
|
||||
call_unwrapped(
|
||||
artifacts_router.update_artifact,
|
||||
"thread-1",
|
||||
"mnt/user-data/outputs/note.txt",
|
||||
artifacts_router.ArtifactUpdateRequest(content="user version", expected_sha256=_artifact_sha256("old version")),
|
||||
_make_request(),
|
||||
)
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 412
|
||||
assert artifact_path.read_text(encoding="utf-8") == "agent version"
|
||||
|
||||
|
||||
def test_update_artifact_rejects_non_output_path(tmp_path, monkeypatch) -> None:
|
||||
artifact_path = tmp_path / "note.txt"
|
||||
artifact_path.write_text("before", encoding="utf-8")
|
||||
_patch_artifact_update_dependencies(monkeypatch, artifact_path)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
asyncio.run(
|
||||
call_unwrapped(
|
||||
artifacts_router.update_artifact,
|
||||
"thread-1",
|
||||
"mnt/user-data/workspace/note.txt",
|
||||
artifacts_router.ArtifactUpdateRequest(content="after", expected_sha256=_artifact_sha256("before")),
|
||||
_make_request(),
|
||||
)
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert artifact_path.read_text(encoding="utf-8") == "before"
|
||||
|
||||
|
||||
def test_update_artifact_rejects_binary_file(tmp_path, monkeypatch) -> None:
|
||||
artifact_path = tmp_path / "blob.bin"
|
||||
artifact_path.write_bytes(b"before\x00binary")
|
||||
_patch_artifact_update_dependencies(monkeypatch, artifact_path)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
asyncio.run(
|
||||
call_unwrapped(
|
||||
artifacts_router.update_artifact,
|
||||
"thread-1",
|
||||
"mnt/user-data/outputs/blob.bin",
|
||||
artifacts_router.ArtifactUpdateRequest(content="after", expected_sha256=hashlib.sha256(b"before\x00binary").hexdigest()),
|
||||
_make_request(),
|
||||
)
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 415
|
||||
|
||||
|
||||
def test_update_artifact_syncs_non_mounted_sandbox(tmp_path, monkeypatch) -> None:
|
||||
artifact_path = tmp_path / "note.txt"
|
||||
artifact_path.write_text("before", encoding="utf-8")
|
||||
provider = _RemoteSandboxProvider()
|
||||
_patch_artifact_update_dependencies(monkeypatch, artifact_path, provider)
|
||||
|
||||
asyncio.run(
|
||||
call_unwrapped(
|
||||
artifacts_router.update_artifact,
|
||||
"thread-1",
|
||||
"mnt/user-data/outputs/note.txt",
|
||||
artifacts_router.ArtifactUpdateRequest(content="after", expected_sha256=_artifact_sha256("before")),
|
||||
_make_request(),
|
||||
)
|
||||
)
|
||||
|
||||
assert provider.sandbox.updates == [("/mnt/user-data/outputs/note.txt", b"after")]
|
||||
assert provider.released == ["sandbox-1"]
|
||||
assert artifact_path.read_text(encoding="utf-8") == "after"
|
||||
|
||||
|
||||
def test_update_artifact_releases_sandbox_when_initial_sync_fails(tmp_path, monkeypatch) -> None:
|
||||
artifact_path = tmp_path / "note.txt"
|
||||
artifact_path.write_text("before", encoding="utf-8")
|
||||
provider = _RemoteSandboxProvider(fail_next_update=True)
|
||||
_patch_artifact_update_dependencies(monkeypatch, artifact_path, provider)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
asyncio.run(
|
||||
call_unwrapped(
|
||||
artifacts_router.update_artifact,
|
||||
"thread-1",
|
||||
"mnt/user-data/outputs/note.txt",
|
||||
artifacts_router.ArtifactUpdateRequest(content="after", expected_sha256=_artifact_sha256("before")),
|
||||
_make_request(),
|
||||
)
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert provider.released == ["sandbox-1"]
|
||||
assert provider.sandbox.updates == [("/mnt/user-data/outputs/note.txt", b"before")]
|
||||
assert artifact_path.read_text(encoding="utf-8") == "before"
|
||||
|
||||
|
||||
def test_update_artifact_rolls_back_remote_when_local_replace_fails(tmp_path, monkeypatch) -> None:
|
||||
artifact_path = tmp_path / "note.txt"
|
||||
artifact_path.write_text("before", encoding="utf-8")
|
||||
provider = _RemoteSandboxProvider()
|
||||
_patch_artifact_update_dependencies(monkeypatch, artifact_path, provider)
|
||||
|
||||
def fail_replace(*_args, **_kwargs) -> None:
|
||||
raise OSError("replace failed")
|
||||
|
||||
monkeypatch.setattr(artifacts_router, "_replace_artifact_atomically", fail_replace)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
asyncio.run(
|
||||
call_unwrapped(
|
||||
artifacts_router.update_artifact,
|
||||
"thread-1",
|
||||
"mnt/user-data/outputs/note.txt",
|
||||
artifacts_router.ArtifactUpdateRequest(content="after", expected_sha256=_artifact_sha256("before")),
|
||||
_make_request(),
|
||||
)
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 500
|
||||
assert provider.sandbox.updates == [
|
||||
("/mnt/user-data/outputs/note.txt", b"after"),
|
||||
("/mnt/user-data/outputs/note.txt", b"before"),
|
||||
]
|
||||
assert provider.released == ["sandbox-1"]
|
||||
assert artifact_path.read_text(encoding="utf-8") == "before"
|
||||
|
||||
|
||||
def test_update_artifact_rejects_oversized_content(tmp_path, monkeypatch) -> None:
|
||||
artifact_path = tmp_path / "note.txt"
|
||||
artifact_path.write_text("before", encoding="utf-8")
|
||||
_patch_artifact_update_dependencies(monkeypatch, artifact_path)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
asyncio.run(
|
||||
call_unwrapped(
|
||||
artifacts_router.update_artifact,
|
||||
"thread-1",
|
||||
"mnt/user-data/outputs/note.txt",
|
||||
artifacts_router.ArtifactUpdateRequest(
|
||||
content="x" * (artifacts_router.MAX_EDITABLE_ARTIFACT_BYTES + 1),
|
||||
expected_sha256=_artifact_sha256("before"),
|
||||
),
|
||||
_make_request(),
|
||||
)
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 413
|
||||
assert artifact_path.read_text(encoding="utf-8") == "before"
|
||||
|
||||
|
||||
def test_update_artifact_reports_active_run_conflict(tmp_path, monkeypatch) -> None:
|
||||
artifact_path = tmp_path / "note.txt"
|
||||
artifact_path.write_text("before", encoding="utf-8")
|
||||
monkeypatch.setattr(artifacts_router, "resolve_thread_virtual_path", lambda _thread_id, _path, user_id=None: artifact_path)
|
||||
|
||||
@asynccontextmanager
|
||||
async def reject_artifact_write(*_args, **_kwargs):
|
||||
raise artifacts_router.ConflictError("active run")
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(artifacts_router, "reserve_artifact_write", reject_artifact_write)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
asyncio.run(
|
||||
call_unwrapped(
|
||||
artifacts_router.update_artifact,
|
||||
"thread-1",
|
||||
"mnt/user-data/outputs/note.txt",
|
||||
artifacts_router.ArtifactUpdateRequest(content="after", expected_sha256=_artifact_sha256("before")),
|
||||
_make_request(),
|
||||
)
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 409
|
||||
assert artifact_path.read_text(encoding="utf-8") == "before"
|
||||
|
||||
|
||||
def test_get_artifact_text_preview_supports_bounded_range_requests(tmp_path, monkeypatch) -> None:
|
||||
payload = ("0123456789abcdef" * 131_072).encode()
|
||||
artifact_path = tmp_path / "large.txt"
|
||||
artifact_path.write_bytes(payload)
|
||||
monkeypatch.setattr(artifacts_router, "resolve_thread_virtual_path", lambda _thread_id, _path, user_id=None: artifact_path)
|
||||
|
||||
app = make_authed_test_app()
|
||||
app.include_router(artifacts_router.router)
|
||||
with TestClient(app) as client:
|
||||
preview = client.get(
|
||||
"/api/threads/thread-1/artifacts/mnt/user-data/outputs/large.txt",
|
||||
headers={"Range": "bytes=0-1048575"},
|
||||
)
|
||||
invalid = client.get(
|
||||
"/api/threads/thread-1/artifacts/mnt/user-data/outputs/large.txt",
|
||||
headers={"Range": f"bytes={len(payload)}-"},
|
||||
)
|
||||
|
||||
assert preview.status_code == 206
|
||||
assert preview.content == payload[:1_048_576]
|
||||
assert preview.headers["content-range"] == f"bytes 0-1048575/{len(payload)}"
|
||||
assert preview.headers["content-disposition"].startswith("inline;")
|
||||
assert invalid.status_code == 416
|
||||
assert invalid.headers["content-range"] == f"bytes */{len(payload)}"
|
||||
|
||||
|
||||
def test_get_skill_archive_preview_supports_bounded_range_requests(tmp_path, monkeypatch) -> None:
|
||||
payload = ("skill preview \u4e2d\u6587\n" * 100_000).encode()
|
||||
skill_path = tmp_path / "sample.skill"
|
||||
with zipfile.ZipFile(skill_path, "w", compression=zipfile.ZIP_DEFLATED) as zip_ref:
|
||||
zip_ref.writestr("SKILL.md", payload)
|
||||
|
||||
monkeypatch.setattr(artifacts_router, "resolve_thread_virtual_path", lambda _thread_id, _path, user_id=None: skill_path)
|
||||
|
||||
app = make_authed_test_app()
|
||||
app.include_router(artifacts_router.router)
|
||||
with TestClient(app) as client:
|
||||
preview = client.get(
|
||||
"/api/threads/thread-1/artifacts/mnt/user-data/outputs/sample.skill/SKILL.md",
|
||||
headers={"Range": "bytes=0-1048575"},
|
||||
)
|
||||
invalid = client.get(
|
||||
"/api/threads/thread-1/artifacts/mnt/user-data/outputs/sample.skill/SKILL.md",
|
||||
headers={"Range": f"bytes={len(payload)}-"},
|
||||
)
|
||||
|
||||
assert preview.status_code == 206
|
||||
assert preview.content == payload[:1_048_576]
|
||||
assert preview.headers["accept-ranges"] == "bytes"
|
||||
assert preview.headers["content-range"] == f"bytes 0-1048575/{len(payload)}"
|
||||
assert invalid.status_code == 416
|
||||
assert invalid.headers["content-range"] == f"bytes */{len(payload)}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("filename", "content"), ACTIVE_ARTIFACT_CASES)
|
||||
@ -87,7 +427,7 @@ def test_get_artifact_download_false_does_not_force_attachment(tmp_path, monkeyp
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.text == "hello"
|
||||
assert "content-disposition" not in response.headers
|
||||
assert response.headers["content-disposition"].startswith("inline;")
|
||||
|
||||
|
||||
def test_get_artifact_binary_preview_is_inline_file_response(tmp_path, monkeypatch) -> None:
|
||||
|
||||
@ -301,7 +301,7 @@ async def test_wheel_input_falls_back_to_native_wheel_when_js_scroll_fails():
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_live_frame_returns_base64_jpeg_screenshot():
|
||||
async def test_live_frame_returns_jpeg_bytes_without_base64_expansion():
|
||||
session = BrowserSession(
|
||||
MagicMock(),
|
||||
headless=True,
|
||||
@ -314,7 +314,7 @@ async def test_live_frame_returns_base64_jpeg_screenshot():
|
||||
|
||||
frame = await session._live_frame()
|
||||
|
||||
assert frame == "/9hq cGVnLWJ5dGVz".replace(" ", "")
|
||||
assert frame == b"\xff\xd8jpeg-bytes"
|
||||
page.screenshot.assert_awaited_once_with(type="jpeg", quality=_LIVE_FRAME_JPEG_QUALITY)
|
||||
|
||||
|
||||
@ -390,7 +390,7 @@ async def test_continuous_inputs_refresh_before_input_stops(monkeypatch):
|
||||
page.keyboard.press = AsyncMock()
|
||||
session._ensure_page = AsyncMock(return_value=page)
|
||||
session._on_frame = MagicMock()
|
||||
session._live_frame = AsyncMock(return_value="frame")
|
||||
session._live_frame = AsyncMock(return_value=b"frame")
|
||||
session._schedule_settle_live_frames = MagicMock()
|
||||
|
||||
stop = asyncio.Event()
|
||||
@ -633,7 +633,7 @@ async def test_live_frame_screenshots_current_active_page_after_switch():
|
||||
frame = await session._live_frame()
|
||||
|
||||
new_page.screenshot.assert_awaited_once_with(type="jpeg", quality=_LIVE_FRAME_JPEG_QUALITY)
|
||||
assert frame # base64 payload of the new page
|
||||
assert frame # JPEG bytes from the new page
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@ -646,10 +646,10 @@ async def test_stop_screencast_ignores_stale_connection_callback():
|
||||
viewport={"width": 1000, "height": 500},
|
||||
)
|
||||
|
||||
def old_frame(_data: str) -> None:
|
||||
def old_frame(_data: bytes) -> None:
|
||||
pass
|
||||
|
||||
def new_frame(_data: str) -> None:
|
||||
def new_frame(_data: bytes) -> None:
|
||||
pass
|
||||
|
||||
session._on_frame = new_frame
|
||||
@ -675,10 +675,10 @@ async def test_start_screencast_rejects_second_live_viewer():
|
||||
page = MagicMock()
|
||||
session._ensure_page = AsyncMock(return_value=page)
|
||||
|
||||
def old_frame(_data: str) -> None:
|
||||
def old_frame(_data: bytes) -> None:
|
||||
pass
|
||||
|
||||
def new_frame(_data: str) -> None:
|
||||
def new_frame(_data: bytes) -> None:
|
||||
pass
|
||||
|
||||
session._on_frame = old_frame
|
||||
|
||||
@ -1,3 +1,4 @@
|
||||
import json
|
||||
import logging
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
@ -10,7 +11,12 @@ from starlette.websockets import WebSocketDisconnect
|
||||
|
||||
from app.gateway.auth.models import User
|
||||
from app.gateway.routers import browser as browser_router
|
||||
from app.gateway.routers.browser import _should_apply_browser_seed, _ws_origin_allowed
|
||||
from app.gateway.routers.browser import (
|
||||
_negotiate_browser_frame_format,
|
||||
_send_browser_frame,
|
||||
_should_apply_browser_seed,
|
||||
_ws_origin_allowed,
|
||||
)
|
||||
|
||||
|
||||
class _FakeWebSocket:
|
||||
@ -207,6 +213,64 @@ def test_ws_origin_allowed_same_origin_host():
|
||||
assert _ws_origin_allowed(ws) is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_browser_frame_uses_binary_websocket_message_when_requested():
|
||||
websocket = MagicMock()
|
||||
websocket.send_bytes = AsyncMock()
|
||||
websocket.send_text = AsyncMock()
|
||||
|
||||
await _send_browser_frame(websocket, b"\xff\xd8jpeg", binary=True)
|
||||
|
||||
websocket.send_bytes.assert_awaited_once_with(b"\xff\xd8jpeg")
|
||||
websocket.send_text.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_send_browser_frame_keeps_legacy_base64_json_protocol():
|
||||
websocket = MagicMock()
|
||||
websocket.send_bytes = AsyncMock()
|
||||
websocket.send_text = AsyncMock()
|
||||
|
||||
await _send_browser_frame(websocket, b"\xff\xd8jpeg", binary=False)
|
||||
|
||||
websocket.send_bytes.assert_not_awaited()
|
||||
payload = json.loads(websocket.send_text.await_args.args[0])
|
||||
assert payload == {"type": "frame", "data": "/9hq cGVn".replace(" ", "")}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_browser_frame_format_rejects_unknown_capability():
|
||||
websocket = MagicMock()
|
||||
websocket.query_params = {"frame_format": "avif"}
|
||||
websocket.accept = AsyncMock()
|
||||
websocket.send_text = AsyncMock()
|
||||
websocket.close = AsyncMock()
|
||||
|
||||
result = await _negotiate_browser_frame_format(websocket)
|
||||
|
||||
assert result is None
|
||||
websocket.accept.assert_awaited_once()
|
||||
payload = json.loads(websocket.send_text.await_args.args[0])
|
||||
assert payload["type"] == "error"
|
||||
assert "frame_format" in payload["message"]
|
||||
websocket.close.assert_awaited_once_with(code=1008)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("value", "expected"), [(None, False), ("binary", True)])
|
||||
async def test_browser_frame_format_accepts_legacy_and_binary(value, expected):
|
||||
websocket = MagicMock()
|
||||
websocket.query_params = {} if value is None else {"frame_format": value}
|
||||
websocket.accept = AsyncMock()
|
||||
websocket.send_text = AsyncMock()
|
||||
websocket.close = AsyncMock()
|
||||
|
||||
assert await _negotiate_browser_frame_format(websocket) is expected
|
||||
websocket.accept.assert_awaited_once()
|
||||
websocket.send_text.assert_not_awaited()
|
||||
websocket.close.assert_not_awaited()
|
||||
|
||||
|
||||
def test_ws_origin_allowed_rejects_cross_origin():
|
||||
ws = _FakeWebSocket({"origin": "https://evil.example.com", "host": "app.example.com"})
|
||||
assert _ws_origin_allowed(ws) is False
|
||||
|
||||
398
backend/tests/test_cached_history_saver.py
Normal file
398
backend/tests/test_cached_history_saver.py
Normal file
@ -0,0 +1,398 @@
|
||||
"""CachedHistorySaver composition vs. the saver's own full walk."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointTuple
|
||||
|
||||
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
||||
from deerflow.runtime.checkpointer.cached_saver import CachedHistorySaver
|
||||
|
||||
PREFIX = "ckpt-hist:v1:testdb"
|
||||
|
||||
|
||||
class _DictSaver(BaseCheckpointSaver):
|
||||
"""Minimal in-memory saver. Deliberately does NOT override the delta
|
||||
history methods, so the base class full parent-chain walk is the
|
||||
differential oracle."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.checkpoints: dict[str, tuple[CheckpointTuple, ...]] = {}
|
||||
self.tuple_reads = 0
|
||||
self.history_walks = 0
|
||||
|
||||
def put_tuple(self, tup: CheckpointTuple) -> None:
|
||||
cid = tup.config["configurable"]["checkpoint_id"]
|
||||
self.checkpoints[cid] = (tup,)
|
||||
|
||||
def get_tuple(self, config):
|
||||
self.tuple_reads += 1
|
||||
configurable = config["configurable"]
|
||||
cid = configurable.get("checkpoint_id")
|
||||
if cid is None:
|
||||
cid = next(reversed(self.checkpoints), None)
|
||||
if cid is None:
|
||||
return None
|
||||
stored = self.checkpoints.get(cid)
|
||||
return stored[0] if stored else None
|
||||
|
||||
async def aget_tuple(self, config):
|
||||
return self.get_tuple(config)
|
||||
|
||||
def get_delta_channel_history(self, *, config, channels):
|
||||
self.history_walks += 1
|
||||
return super().get_delta_channel_history(config=config, channels=channels)
|
||||
|
||||
async def aget_delta_channel_history(self, *, config, channels):
|
||||
self.history_walks += 1
|
||||
return await super().aget_delta_channel_history(config=config, channels=channels)
|
||||
|
||||
# Unused abstract surface.
|
||||
def list(self, config, *, filter=None, before=None, limit=None):
|
||||
yield from ()
|
||||
|
||||
def put(self, config, checkpoint, metadata, new_versions):
|
||||
raise NotImplementedError
|
||||
|
||||
def put_writes(self, config, writes, task_id, task_path=""):
|
||||
raise NotImplementedError
|
||||
|
||||
def delete_thread(self, thread_id):
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def _cfg(thread: str, cid: str | None) -> dict:
|
||||
configurable: dict[str, Any] = {"thread_id": thread, "checkpoint_ns": ""}
|
||||
if cid is not None:
|
||||
configurable["checkpoint_id"] = cid
|
||||
return {"configurable": configurable}
|
||||
|
||||
|
||||
def _tup(thread: str, cid: str, parent: str | None, *, channel_values: dict, writes: list) -> CheckpointTuple:
|
||||
config = _cfg(thread, cid)
|
||||
parent_config = _cfg(thread, parent) if parent else None
|
||||
checkpoint = {"v": 1, "id": cid, "channel_values": channel_values, "channel_versions": {}, "versions_seen": {}, "updated_at": None}
|
||||
return CheckpointTuple(config=config, checkpoint=checkpoint, metadata={}, parent_config=parent_config, pending_writes=list(writes))
|
||||
|
||||
|
||||
def _chain(saver: _DictSaver, thread: str = "t1") -> list[str]:
|
||||
"""c0(seed snapshot) -> c1(writes w1) -> c2(writes w2)."""
|
||||
saver.put_tuple(_tup(thread, "c0", None, channel_values={"messages": ["seed-msg"]}, writes=[]))
|
||||
saver.put_tuple(_tup(thread, "c1", "c0", channel_values={}, writes=[("task1", "messages", "w1")]))
|
||||
saver.put_tuple(_tup(thread, "c2", "c1", channel_values={}, writes=[("task2", "messages", "w2")]))
|
||||
return ["c0", "c1", "c2"]
|
||||
|
||||
|
||||
def _wrap(saver: _DictSaver, cache) -> CachedHistorySaver:
|
||||
return CachedHistorySaver(saver, cache, key_prefix=PREFIX)
|
||||
|
||||
|
||||
class _DeletableSaver(_DictSaver):
|
||||
"""Functional delete/prune so purge behavior is observable."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.deleted_threads: list[str] = []
|
||||
self.deleted_run_ids: list[list[str]] = []
|
||||
self.pruned_threads: list[list[str]] = []
|
||||
|
||||
def delete_for_runs(self, run_ids):
|
||||
self.deleted_run_ids.append(list(run_ids))
|
||||
|
||||
async def adelete_for_runs(self, run_ids):
|
||||
self.deleted_run_ids.append(list(run_ids))
|
||||
|
||||
def _drop(self, thread_id: str) -> None:
|
||||
self.deleted_threads.append(thread_id)
|
||||
self.checkpoints = {cid: stored for cid, stored in self.checkpoints.items() if stored[0].config["configurable"]["thread_id"] != thread_id}
|
||||
|
||||
def delete_thread(self, thread_id):
|
||||
self._drop(thread_id)
|
||||
|
||||
async def adelete_thread(self, thread_id):
|
||||
self._drop(thread_id)
|
||||
|
||||
def prune(self, thread_ids, *, strategy="keep_latest"):
|
||||
self.pruned_threads.append(list(thread_ids))
|
||||
|
||||
async def aprune(self, thread_ids, *, strategy="keep_latest"):
|
||||
self.pruned_threads.append(list(thread_ids))
|
||||
|
||||
|
||||
def _thread_entries(cache: MemoryCheckpointHistoryCache, thread_id: str) -> int:
|
||||
stem = f"{PREFIX}:{thread_id}:"
|
||||
return sum(1 for key in cache._data if key.startswith(stem))
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_adelete_thread_purges_only_that_threads_cache_entries():
|
||||
inner = _DeletableSaver()
|
||||
_chain(inner, "t1")
|
||||
# _DictSaver keys tuples by checkpoint_id alone: t2 needs distinct cids.
|
||||
inner.put_tuple(_tup("t2", "u0", None, channel_values={"messages": ["seed2"]}, writes=[]))
|
||||
inner.put_tuple(_tup("t2", "u1", "u0", channel_values={}, writes=[("task1", "messages", "x1")]))
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=32)
|
||||
saver = _wrap(inner, cache)
|
||||
await saver.aget_delta_channel_history(config=_cfg("t1", "c2"), channels=["messages"])
|
||||
await saver.aget_delta_channel_history(config=_cfg("t2", "u1"), channels=["messages"])
|
||||
assert _thread_entries(cache, "t1") > 0
|
||||
assert _thread_entries(cache, "t2") > 0
|
||||
|
||||
await saver.adelete_thread("t1")
|
||||
|
||||
assert inner.deleted_threads == ["t1"]
|
||||
assert inner.get_tuple(_cfg("t1", "c2")) is None # source of truth gone
|
||||
assert _thread_entries(cache, "t1") == 0 # residual history payloads purged
|
||||
assert _thread_entries(cache, "t2") > 0 # other threads untouched
|
||||
|
||||
|
||||
def test_sync_delete_thread_purges_cache_entries():
|
||||
inner = _DeletableSaver()
|
||||
_chain(inner, "t1")
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=32)
|
||||
saver = _wrap(inner, cache)
|
||||
saver.get_delta_channel_history(config=_cfg("t1", "c2"), channels=["messages"])
|
||||
assert _thread_entries(cache, "t1") > 0
|
||||
|
||||
saver.delete_thread("t1")
|
||||
|
||||
assert inner.deleted_threads == ["t1"]
|
||||
assert _thread_entries(cache, "t1") == 0
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_prune_purges_rewritten_threads_cache_entries():
|
||||
inner = _DeletableSaver()
|
||||
_chain(inner, "t1")
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=32)
|
||||
saver = _wrap(inner, cache)
|
||||
await saver.aget_delta_channel_history(config=_cfg("t1", "c2"), channels=["messages"])
|
||||
assert _thread_entries(cache, "t1") > 0
|
||||
|
||||
await saver.aprune(["t1"], strategy="keep_latest")
|
||||
|
||||
assert inner.pruned_threads == [["t1"]]
|
||||
# The chain was rewritten: pre-prune histories must not linger.
|
||||
assert _thread_entries(cache, "t1") == 0
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_delete_for_runs_delegates_without_cache_purge():
|
||||
"""Run-scoped deletes cannot be mapped to threads cheaply; documented
|
||||
behavior is delegation with LRU/TTL-bounded residual retention."""
|
||||
inner = _DeletableSaver()
|
||||
_chain(inner, "t1")
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=32)
|
||||
saver = _wrap(inner, cache)
|
||||
await saver.aget_delta_channel_history(config=_cfg("t1", "c2"), channels=["messages"])
|
||||
entries_before = cache.stats().entries
|
||||
|
||||
await saver.adelete_for_runs(["run-1"]) # base-class no-op on _DictSaver lineage
|
||||
|
||||
assert inner.deleted_run_ids == [["run-1"]]
|
||||
assert cache.stats().entries == entries_before
|
||||
|
||||
|
||||
class _RecordingSaver(_DictSaver):
|
||||
"""Captures the config passed into each fallback walk."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.walk_configs: list[dict] = []
|
||||
|
||||
def get_delta_channel_history(self, *, config, channels):
|
||||
self.walk_configs.append(config)
|
||||
return super().get_delta_channel_history(config=config, channels=channels)
|
||||
|
||||
async def aget_delta_channel_history(self, *, config, channels):
|
||||
self.walk_configs.append(config)
|
||||
return await super().aget_delta_channel_history(config=config, channels=channels)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_composition_matches_full_walk_and_avoids_it():
|
||||
inner = _DictSaver()
|
||||
_chain(inner)
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=16)
|
||||
saver = _wrap(inner, cache)
|
||||
|
||||
# Cold: c1 composes from snapshot parent c0 (channel_values hit) — no walk.
|
||||
h1 = await saver.aget_delta_channel_history(config=_cfg("t1", "c1"), channels=["messages"])
|
||||
assert h1["messages"]["writes"] == []
|
||||
assert h1["messages"]["seed"] == ["seed-msg"]
|
||||
assert inner.history_walks == 0
|
||||
|
||||
# c2 composes from cached history(c1) + c1's pending writes: no inner walk.
|
||||
# NOTE: history(c2) excludes c2's OWN pending writes (they belong to the
|
||||
# next super-step per the LangGraph contract) — on-path writes are c1's.
|
||||
h2 = await saver.aget_delta_channel_history(config=_cfg("t1", "c2"), channels=["messages"])
|
||||
assert h2["messages"]["writes"] == [("task1", "messages", "w1")]
|
||||
assert h2["messages"]["seed"] == ["seed-msg"]
|
||||
assert inner.history_walks == 0 # unchanged: composition, not a walk
|
||||
|
||||
# Differential oracle: identical to the inner saver's own full walk.
|
||||
oracle = _DictSaver()
|
||||
_chain(oracle)
|
||||
for cid in ("c0", "c1", "c2"):
|
||||
expected = await oracle.aget_delta_channel_history(config=_cfg("t1", cid), channels=["messages"])
|
||||
actual = await saver.aget_delta_channel_history(config=_cfg("t1", cid), channels=["messages"])
|
||||
assert actual == expected, cid
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_snapshot_parent_composes_without_parent_history_lookup():
|
||||
inner = _DictSaver()
|
||||
_chain(inner)
|
||||
# c3 has c0-style snapshot directly at parent c1: rewrite c1 with channel_values.
|
||||
inner.put_tuple(_tup("t1", "c1b", "c0", channel_values={"messages": ["snap"]}, writes=[("t", "messages", "wx")]))
|
||||
inner.put_tuple(_tup("t1", "c2b", "c1b", channel_values={}, writes=[("t2", "messages", "wy")]))
|
||||
saver = _wrap(inner, MemoryCheckpointHistoryCache(max_entries=16))
|
||||
h = await saver.aget_delta_channel_history(config=_cfg("t1", "c2b"), channels=["messages"])
|
||||
assert h["messages"] == {"writes": [("t", "messages", "wx")], "seed": ["snap"]}
|
||||
assert inner.history_walks == 0
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_root_checkpoint_history_is_empty_writes():
|
||||
inner = _DictSaver()
|
||||
_chain(inner)
|
||||
saver = _wrap(inner, MemoryCheckpointHistoryCache(max_entries=16))
|
||||
h = await saver.aget_delta_channel_history(config=_cfg("t1", "c0"), channels=["messages"])
|
||||
assert h["messages"] == {"writes": []}
|
||||
assert "seed" not in h["messages"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_latest_config_caches_under_resolved_checkpoint_id():
|
||||
inner = _DictSaver()
|
||||
_chain(inner)
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=16)
|
||||
saver = _wrap(inner, cache)
|
||||
await saver.aget_delta_channel_history(config=_cfg("t1", None), channels=["messages"])
|
||||
reads_after_first = inner.tuple_reads
|
||||
await saver.aget_delta_channel_history(config=_cfg("t1", "c2"), channels=["messages"])
|
||||
# Second call resolves by id from cache: target tuple only, no parent refetch.
|
||||
assert inner.tuple_reads == reads_after_first + 1
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_eviction_falls_back_to_walk_but_stays_correct():
|
||||
inner = _DictSaver()
|
||||
_chain(inner)
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=1)
|
||||
saver = _wrap(inner, cache)
|
||||
await saver.aget_delta_channel_history(config=_cfg("t1", "c2"), channels=["messages"]) # cold: fallback walk
|
||||
await saver.aget_delta_channel_history(config=_cfg("t1", "c1"), channels=["messages"]) # compose; evicts c2 entry
|
||||
h = await saver.aget_delta_channel_history(config=_cfg("t1", "c2"), channels=["messages"]) # recomposes via cached history(c1)
|
||||
assert h["messages"]["writes"] == [("task1", "messages", "w1")]
|
||||
assert h["messages"]["seed"] == ["seed-msg"]
|
||||
assert saver.stats()["full_walks"] == 0 # cold read resolved by recursive compose
|
||||
assert inner.history_walks == 0 # resolution never delegates to the inner walk
|
||||
assert cache.stats().evictions >= 1
|
||||
|
||||
|
||||
def test_sync_path_matches_async():
|
||||
inner = _DictSaver()
|
||||
_chain(inner)
|
||||
saver = _wrap(inner, MemoryCheckpointHistoryCache(max_entries=16))
|
||||
h = saver.get_delta_channel_history(config=_cfg("t1", "c2"), channels=["messages"])
|
||||
assert h["messages"]["writes"] == [("task1", "messages", "w1")]
|
||||
assert h["messages"]["seed"] == ["seed-msg"]
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cold_resolve_never_delegates_to_inner_history():
|
||||
"""Cold reads compose recursively (or walk themselves via aget_tuple):
|
||||
the inner saver's own history method is never called, so a 'latest'
|
||||
config cannot be re-resolved mid-resolution (the old pinned race is gone
|
||||
by construction)."""
|
||||
inner = _RecordingSaver()
|
||||
_chain(inner)
|
||||
saver = _wrap(inner, MemoryCheckpointHistoryCache(max_entries=16))
|
||||
h = await saver.aget_delta_channel_history(config=_cfg("t1", None), channels=["messages"])
|
||||
assert h["messages"]["writes"] == [("task1", "messages", "w1")]
|
||||
assert h["messages"]["seed"] == ["seed-msg"]
|
||||
assert inner.walk_configs == [], "resolution must not delegate to the inner history walk"
|
||||
|
||||
|
||||
def test_sync_cold_resolve_never_delegates_to_inner_history():
|
||||
inner = _RecordingSaver()
|
||||
_chain(inner)
|
||||
saver = _wrap(inner, MemoryCheckpointHistoryCache(max_entries=16))
|
||||
h = saver.get_delta_channel_history(config=_cfg("t1", None), channels=["messages"])
|
||||
assert h["messages"]["writes"] == [("task1", "messages", "w1")]
|
||||
assert inner.walk_configs == []
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_recursive_resolve_caches_intermediate_levels():
|
||||
"""Cold short chains resolve by recursive compose: every intermediate
|
||||
level is computed once and cached, later reads hit directly."""
|
||||
inner = _DictSaver()
|
||||
_chain(inner)
|
||||
inner.put_tuple(_tup("t1", "c3", "c2", channel_values={}, writes=[("task3", "messages", "w3")]))
|
||||
oracle = _DictSaver()
|
||||
_chain(oracle)
|
||||
oracle.put_tuple(_tup("t1", "c3", "c2", channel_values={}, writes=[("task3", "messages", "w3")]))
|
||||
saver = _wrap(inner, MemoryCheckpointHistoryCache(max_entries=16))
|
||||
|
||||
await saver.aget_delta_channel_history(config=_cfg("t1", "c3"), channels=["messages"])
|
||||
reads_after_cold = inner.tuple_reads
|
||||
|
||||
for cid in ("c0", "c1", "c2", "c3"):
|
||||
expected = await oracle.aget_delta_channel_history(config=_cfg("t1", cid), channels=["messages"])
|
||||
actual = await saver.aget_delta_channel_history(config=_cfg("t1", cid), channels=["messages"])
|
||||
assert actual == expected, cid
|
||||
# Every level is now warm: each read costs exactly its target tuple fetch.
|
||||
assert inner.tuple_reads - reads_after_cold == 4
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_deep_cold_chain_delegates_one_inner_walk_at_depth_limit():
|
||||
"""A cold chain deeper than the compose budget resolves the deepest
|
||||
reached level with ONE inner fast-path walk (2 SQL), caches every level
|
||||
above it, and leaves deeper ancestors cold until asked."""
|
||||
inner = _DictSaver()
|
||||
# 12-deep chain: c0(seed) <- c1 <- ... <- c11, deeper than the budget.
|
||||
inner.put_tuple(_tup("t1", "c0", None, channel_values={"messages": ["seed-msg"]}, writes=[]))
|
||||
for i in range(1, 12):
|
||||
inner.put_tuple(_tup("t1", f"c{i}", f"c{i - 1}", channel_values={}, writes=[("task", "messages", f"w{i}")]))
|
||||
oracle = _DictSaver()
|
||||
oracle.put_tuple(_tup("t1", "c0", None, channel_values={"messages": ["seed-msg"]}, writes=[]))
|
||||
for i in range(1, 12):
|
||||
oracle.put_tuple(_tup("t1", f"c{i}", f"c{i - 1}", channel_values={}, writes=[("task", "messages", f"w{i}")]))
|
||||
saver = _wrap(inner, MemoryCheckpointHistoryCache(max_entries=64))
|
||||
|
||||
h = await saver.aget_delta_channel_history(config=_cfg("t1", "c11"), channels=["messages"])
|
||||
assert saver.stats()["full_walks"] == 1
|
||||
assert inner.history_walks == 1 # exactly one delegated fast-path walk
|
||||
assert h["messages"]["writes"] == [("task", "messages", f"w{i}") for i in range(1, 11)]
|
||||
assert h["messages"]["seed"] == ["seed-msg"]
|
||||
|
||||
reads_after_cold = inner.tuple_reads
|
||||
for cid in [f"c{i}" for i in range(12)]:
|
||||
expected = await oracle.aget_delta_channel_history(config=_cfg("t1", cid), channels=["messages"])
|
||||
actual = await saver.aget_delta_channel_history(config=_cfg("t1", cid), channels=["messages"])
|
||||
assert actual == expected, cid
|
||||
# Warm after the cold read: c2..c11. Reads: c0 = 1 (root, no parent);
|
||||
# c1 = 2 (target + snapshot parent c0); c2..c11 = 1 each.
|
||||
assert inner.tuple_reads - reads_after_cold == 13
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_stats_expose_composition_counters():
|
||||
inner = _DictSaver()
|
||||
_chain(inner)
|
||||
saver = _wrap(inner, MemoryCheckpointHistoryCache(max_entries=16))
|
||||
await saver.aget_delta_channel_history(config=_cfg("t1", "c2"), channels=["messages"]) # cold: recursive compose ×2, warms c1+c2
|
||||
await saver.aget_delta_channel_history(config=_cfg("t1", "c1"), channels=["messages"]) # direct hit (warmed)
|
||||
# Compose: add c3 (parent c2, not a snapshot) AFTER the cold read — c3 is
|
||||
# not cached but history(c2) is, so c3 composes without a walk.
|
||||
inner.put_tuple(_tup("t1", "c3", "c2", channel_values={}, writes=[("task3", "messages", "w3")]))
|
||||
h3 = await saver.aget_delta_channel_history(config=_cfg("t1", "c3"), channels=["messages"])
|
||||
assert h3["messages"]["writes"] == [("task1", "messages", "w1"), ("task2", "messages", "w2")]
|
||||
stats = saver.stats()
|
||||
assert stats["full_walks"] == 0
|
||||
assert stats["compose_hits"] == 3 # c1-level + c2-level (cold) + c3
|
||||
assert stats["hits"] >= 1
|
||||
364
backend/tests/test_cached_history_saver_integration.py
Normal file
364
backend/tests/test_cached_history_saver_integration.py
Normal file
@ -0,0 +1,364 @@
|
||||
"""Behavioral integration tests for CachedHistorySaver on REAL LangGraph execution.
|
||||
|
||||
Unlike tests/test_cached_history_saver.py (fake saver, hand-built chains), these
|
||||
tests drive compiled StateGraphs through pregel in delta mode
|
||||
(``DeltaChannel(merge_message_writes, snapshot_frequency=2)``) and verify the
|
||||
cache against a differential oracle: the identical scenario executed on a raw
|
||||
``InMemorySaver`` in a fresh thread. Digests are (type, content, id) triples of
|
||||
the materialized ``messages`` channel, so any history corruption shows up as a
|
||||
digest mismatch.
|
||||
|
||||
Observed pregel call pattern on langgraph 1.2.9 (5-step linear graph, 7
|
||||
checkpoints, snapshot cadence 2): one ``aget_delta_channel_history`` per run
|
||||
start (empty-thread load), none for snapshot checkpoints, one per materialized
|
||||
non-snapshot checkpoint on the raw saver. The cached saver pays the run-start
|
||||
walk plus one cold fallback walk; every other materialization composes from a
|
||||
parent snapshot seed or a cached parent history, and a second identical read
|
||||
pass costs zero inner walks.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Annotated, Any, TypedDict
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
|
||||
from langgraph.channels import DeltaChannel
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.types import Command, interrupt
|
||||
|
||||
from deerflow.agents.thread_state import merge_message_writes
|
||||
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
||||
from deerflow.runtime.checkpoint_state import CheckpointStateAccessor
|
||||
from deerflow.runtime.checkpointer.cached_saver import CachedHistorySaver
|
||||
|
||||
STEPS = 5
|
||||
SNAPSHOT_FREQUENCY = 2
|
||||
|
||||
|
||||
class _CountingInMemorySaver(InMemorySaver):
|
||||
"""InMemorySaver that counts full delta-history walks.
|
||||
|
||||
Placed under ``CachedHistorySaver`` it records exactly the walks the cache
|
||||
could not serve; used bare it is the uncached oracle's walk counter.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.history_walks = 0
|
||||
|
||||
def get_delta_channel_history(self, *, config: Any, channels: Any) -> Any:
|
||||
self.history_walks += 1
|
||||
return super().get_delta_channel_history(config=config, channels=channels)
|
||||
|
||||
async def aget_delta_channel_history(self, *, config: Any, channels: Any) -> Any:
|
||||
self.history_walks += 1
|
||||
return await super().aget_delta_channel_history(config=config, channels=channels)
|
||||
|
||||
|
||||
def _state_schema() -> type:
|
||||
class State(TypedDict):
|
||||
messages: Annotated[
|
||||
list[AnyMessage],
|
||||
DeltaChannel(merge_message_writes, snapshot_frequency=SNAPSHOT_FREQUENCY),
|
||||
]
|
||||
|
||||
return State
|
||||
|
||||
|
||||
def _make_step_node(n: int) -> Any:
|
||||
def node(state: dict) -> dict:
|
||||
return {"messages": [AIMessage(content=f"step-{n}", id=f"ai-{n}")]}
|
||||
|
||||
return node
|
||||
|
||||
|
||||
def _build_graph(saver: Any, steps: int = STEPS) -> Any:
|
||||
builder = StateGraph(_state_schema())
|
||||
for i in range(steps):
|
||||
builder.add_node(f"step{i}", _make_step_node(i))
|
||||
builder.set_entry_point("step0")
|
||||
for i in range(steps - 1):
|
||||
builder.add_edge(f"step{i}", f"step{i + 1}")
|
||||
builder.set_finish_point(f"step{steps - 1}")
|
||||
return builder.compile(checkpointer=saver)
|
||||
|
||||
|
||||
def _build_interrupt_graph(saver: Any) -> Any:
|
||||
"""step0 -> step1 -> pause(interrupt) -> step2 -> step3."""
|
||||
|
||||
def pause(state: dict) -> dict:
|
||||
answer = interrupt({"question": "continue?"})
|
||||
return {"messages": [AIMessage(content=f"resumed:{answer}", id="ai-resume")]}
|
||||
|
||||
builder = StateGraph(_state_schema())
|
||||
builder.add_node("step0", _make_step_node(0))
|
||||
builder.add_node("step1", _make_step_node(1))
|
||||
builder.add_node("pause", pause)
|
||||
builder.add_node("step2", _make_step_node(2))
|
||||
builder.add_node("step3", _make_step_node(3))
|
||||
builder.set_entry_point("step0")
|
||||
builder.add_edge("step0", "step1")
|
||||
builder.add_edge("step1", "pause")
|
||||
builder.add_edge("pause", "step2")
|
||||
builder.add_edge("step2", "step3")
|
||||
builder.set_finish_point("step3")
|
||||
return builder.compile(checkpointer=saver)
|
||||
|
||||
|
||||
def _config() -> dict[str, Any]:
|
||||
return {"configurable": {"thread_id": f"cache-itest-{uuid4().hex}"}}
|
||||
|
||||
|
||||
def _input() -> dict[str, Any]:
|
||||
return {"messages": [HumanMessage(content="kickoff", id="h-0")]}
|
||||
|
||||
|
||||
def _digest(values: dict[str, Any]) -> list[tuple[str, str, str | None]]:
|
||||
return [(m.type, m.content, m.id) for m in values["messages"]]
|
||||
|
||||
|
||||
def _history_digests(snapshots: list[Any]) -> list[list[tuple[str, str, str | None]]]:
|
||||
return [_digest(s.values) for s in snapshots]
|
||||
|
||||
|
||||
def _expected_final_digest(steps: int = STEPS) -> list[tuple[str, str, str | None]]:
|
||||
return [("human", "kickoff", "h-0"), *[("ai", f"step-{n}", f"ai-{n}") for n in range(steps)]]
|
||||
|
||||
|
||||
def _make_cached_stack(max_entries: int = 128) -> tuple[_CountingInMemorySaver, CachedHistorySaver, Any, CheckpointStateAccessor]:
|
||||
inner = _CountingInMemorySaver()
|
||||
saver = CachedHistorySaver(inner, MemoryCheckpointHistoryCache(max_entries), key_prefix=f"itest-{uuid4().hex}")
|
||||
graph = _build_graph(saver)
|
||||
accessor = CheckpointStateAccessor.bind(graph, saver, mode="delta")
|
||||
return inner, saver, graph, accessor
|
||||
|
||||
|
||||
def _make_oracle_stack() -> tuple[_CountingInMemorySaver, Any, CheckpointStateAccessor]:
|
||||
inner = _CountingInMemorySaver()
|
||||
graph = _build_graph(inner)
|
||||
accessor = CheckpointStateAccessor.bind(graph, inner, mode="delta")
|
||||
return inner, graph, accessor
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_sequential_run_composes_without_inner_walks() -> None:
|
||||
"""A cached run must serve warm reads with strictly fewer inner history
|
||||
walks than the identical uncached run, composing histories instead."""
|
||||
inner, saver, graph, accessor = _make_cached_stack()
|
||||
config = _config()
|
||||
await graph.ainvoke(_input(), config)
|
||||
|
||||
final = await accessor.aget(config)
|
||||
assert _digest(final.values) == _expected_final_digest()
|
||||
|
||||
# Cold pass: materialize every checkpoint in the thread.
|
||||
cold = await accessor.ahistory(config)
|
||||
cold_walks = inner.history_walks
|
||||
|
||||
# Warm pass: identical reads must be served entirely from the cache.
|
||||
warm = await accessor.ahistory(config)
|
||||
assert inner.history_walks == cold_walks, "warm re-read triggered an inner walk"
|
||||
assert _history_digests(warm) == _history_digests(cold)
|
||||
|
||||
# Differential oracle: same run through the raw saver.
|
||||
oracle_inner, oracle_graph, oracle_accessor = _make_oracle_stack()
|
||||
oracle_config = _config()
|
||||
await oracle_graph.ainvoke(_input(), oracle_config)
|
||||
await oracle_accessor.aget(oracle_config)
|
||||
oracle_history = await oracle_accessor.ahistory(oracle_config)
|
||||
|
||||
assert _history_digests(cold) == _history_digests(oracle_history)
|
||||
assert oracle_inner.history_walks > 0
|
||||
assert cold_walks < oracle_inner.history_walks
|
||||
assert saver.stats()["compose_hits"] > 0
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_cache_disabled_parity() -> None:
|
||||
"""A zero-entry cache must behave exactly like the raw inner saver."""
|
||||
inner, saver, graph, accessor = _make_cached_stack(max_entries=0)
|
||||
config = _config()
|
||||
await graph.ainvoke(_input(), config)
|
||||
final = await accessor.aget(config)
|
||||
history = await accessor.ahistory(config)
|
||||
|
||||
oracle_inner, oracle_graph, oracle_accessor = _make_oracle_stack()
|
||||
oracle_config = _config()
|
||||
await oracle_graph.ainvoke(_input(), oracle_config)
|
||||
oracle_final = await oracle_accessor.aget(oracle_config)
|
||||
oracle_history = await oracle_accessor.ahistory(oracle_config)
|
||||
|
||||
assert _digest(final.values) == _expected_final_digest()
|
||||
assert _digest(final.values) == _digest(oracle_final.values)
|
||||
assert _history_digests(history) == _history_digests(oracle_history)
|
||||
assert saver.stats()["hits"] == 0
|
||||
|
||||
|
||||
async def _run_branch_scenario(
|
||||
graph: Any,
|
||||
accessor: CheckpointStateAccessor,
|
||||
config: dict[str, Any],
|
||||
*,
|
||||
fork_next: str,
|
||||
as_node: str,
|
||||
branch_id: str,
|
||||
) -> dict[tuple[str, ...], Any]:
|
||||
"""Run to completion, fork at the checkpoint whose next node is
|
||||
``fork_next``, then resume the branch to completion (the branch head is
|
||||
then the thread's latest checkpoint). Returns the original chain's
|
||||
snapshots keyed by their ``next`` tuple for pinned re-reads.
|
||||
|
||||
``fork_next`` must name a NON-snapshot checkpoint: on langgraph 1.2.9 a
|
||||
``aupdate_state`` fork at a snapshot checkpoint silently drops the update
|
||||
(verified against a raw InMemorySaver - upstream behavior, not the cache).
|
||||
"""
|
||||
await graph.ainvoke(_input(), config)
|
||||
history = await accessor.ahistory(config)
|
||||
by_next = {s.next: s for s in history}
|
||||
base = by_next[(fork_next,)]
|
||||
branch_config = await accessor.aupdate(
|
||||
base.config,
|
||||
{"messages": [AIMessage(content=branch_id, id=f"ai-{branch_id}")]},
|
||||
as_node=as_node,
|
||||
)
|
||||
await graph.ainvoke(None, branch_config)
|
||||
return by_next
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_branch_divergence_no_cross_contamination() -> None:
|
||||
"""A forked branch and the original head must each materialize their own
|
||||
distinct history through the SAME cached saver."""
|
||||
inner, saver, graph, accessor = _make_cached_stack()
|
||||
config = _config()
|
||||
by_next = await _run_branch_scenario(graph, accessor, config, fork_next="step2", as_node="step2", branch_id="branch")
|
||||
|
||||
# Oracle: identical branch scenario on the raw saver.
|
||||
oracle_inner, oracle_graph, oracle_accessor = _make_oracle_stack()
|
||||
oracle_config = _config()
|
||||
oracle_by_next = await _run_branch_scenario(oracle_graph, oracle_accessor, oracle_config, fork_next="step2", as_node="step2", branch_id="branch")
|
||||
|
||||
# Branch head = thread's latest checkpoint after the forked resume.
|
||||
branch_head = await accessor.aget(config)
|
||||
oracle_branch_head = await oracle_accessor.aget(oracle_config)
|
||||
expected_branch = [
|
||||
("human", "kickoff", "h-0"),
|
||||
("ai", "step-0", "ai-0"),
|
||||
("ai", "step-1", "ai-1"),
|
||||
("ai", "branch", "ai-branch"),
|
||||
("ai", "step-3", "ai-3"),
|
||||
("ai", "step-4", "ai-4"),
|
||||
]
|
||||
assert _digest(branch_head.values) == expected_branch
|
||||
assert _digest(branch_head.values) == _digest(oracle_branch_head.values)
|
||||
|
||||
# The original chain still materializes its own un-branched history: the
|
||||
# snapshot head and a non-snapshot mid checkpoint (cache-exercising read).
|
||||
for next_key, expected_len in [((), 6), (("step4",), 5)]:
|
||||
reread_original = await accessor.aget(by_next[next_key].config)
|
||||
oracle_reread_original = await oracle_accessor.aget(oracle_by_next[next_key].config)
|
||||
assert _digest(reread_original.values) == _expected_final_digest()[:expected_len]
|
||||
assert _digest(reread_original.values) == _digest(oracle_reread_original.values)
|
||||
assert _digest(reread_original.values) != _digest(branch_head.values)
|
||||
|
||||
|
||||
async def _run_interrupt_scenario(graph: Any, accessor: CheckpointStateAccessor, config: dict[str, Any]) -> None:
|
||||
result = await graph.ainvoke(_input(), config)
|
||||
assert "__interrupt__" in result
|
||||
await graph.ainvoke(Command(resume="yes"), config)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_interrupt_resume_appended_head_writes() -> None:
|
||||
"""Resume appends writes under the interrupted head checkpoint; the cached
|
||||
final state must equal the no-cache reference."""
|
||||
inner = _CountingInMemorySaver()
|
||||
saver = CachedHistorySaver(inner, MemoryCheckpointHistoryCache(128), key_prefix=f"itest-{uuid4().hex}")
|
||||
graph = _build_interrupt_graph(saver)
|
||||
accessor = CheckpointStateAccessor.bind(graph, saver, mode="delta")
|
||||
config = _config()
|
||||
await _run_interrupt_scenario(graph, accessor, config)
|
||||
|
||||
oracle_inner = _CountingInMemorySaver()
|
||||
oracle_graph = _build_interrupt_graph(oracle_inner)
|
||||
oracle_accessor = CheckpointStateAccessor.bind(oracle_graph, oracle_inner, mode="delta")
|
||||
oracle_config = _config()
|
||||
await _run_interrupt_scenario(oracle_graph, oracle_accessor, oracle_config)
|
||||
|
||||
final = await accessor.aget(config)
|
||||
oracle_final = await oracle_accessor.aget(oracle_config)
|
||||
expected = [
|
||||
("human", "kickoff", "h-0"),
|
||||
("ai", "step-0", "ai-0"),
|
||||
("ai", "step-1", "ai-1"),
|
||||
("ai", "resumed:yes", "ai-resume"),
|
||||
("ai", "step-2", "ai-2"),
|
||||
("ai", "step-3", "ai-3"),
|
||||
]
|
||||
assert _digest(final.values) == expected
|
||||
assert _digest(final.values) == _digest(oracle_final.values)
|
||||
|
||||
# Every checkpoint along the resumed thread matches the oracle, including
|
||||
# the interrupted head whose pending writes grew at resume time.
|
||||
history = await accessor.ahistory(config)
|
||||
oracle_history = await oracle_accessor.ahistory(oracle_config)
|
||||
assert _history_digests(history) == _history_digests(oracle_history)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_eviction_only_costs_performance() -> None:
|
||||
"""A 1-entry LRU thrashes on every read but must stay correct."""
|
||||
inner, saver, graph, accessor = _make_cached_stack(max_entries=1)
|
||||
config = _config()
|
||||
await graph.ainvoke(_input(), config)
|
||||
|
||||
first_pass = await accessor.ahistory(config)
|
||||
second_pass = await accessor.ahistory(config)
|
||||
|
||||
oracle_inner, oracle_graph, oracle_accessor = _make_oracle_stack()
|
||||
oracle_config = _config()
|
||||
await oracle_graph.ainvoke(_input(), oracle_config)
|
||||
oracle_history = await oracle_accessor.ahistory(oracle_config)
|
||||
|
||||
assert _history_digests(first_pass) == _history_digests(oracle_history)
|
||||
assert _history_digests(second_pass) == _history_digests(oracle_history)
|
||||
assert saver.stats()["evictions"] > 0
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_rollback_supersede_does_not_pollute() -> None:
|
||||
"""Re-running from an early checkpoint supersedes the head; the original
|
||||
head's cached history must remain intact and re-readable."""
|
||||
inner, saver, graph, accessor = _make_cached_stack()
|
||||
config = _config()
|
||||
by_next = await _run_branch_scenario(graph, accessor, config, fork_next="step0", as_node="step0", branch_id="rollback")
|
||||
|
||||
oracle_inner, oracle_graph, oracle_accessor = _make_oracle_stack()
|
||||
oracle_config = _config()
|
||||
oracle_by_next = await _run_branch_scenario(oracle_graph, oracle_accessor, oracle_config, fork_next="step0", as_node="step0", branch_id="rollback")
|
||||
|
||||
# New head equals the reference run's new head.
|
||||
new_head = await accessor.aget(config)
|
||||
oracle_new_head = await oracle_accessor.aget(oracle_config)
|
||||
expected_new_head = [
|
||||
("human", "kickoff", "h-0"),
|
||||
("ai", "rollback", "ai-rollback"),
|
||||
("ai", "step-1", "ai-1"),
|
||||
("ai", "step-2", "ai-2"),
|
||||
("ai", "step-3", "ai-3"),
|
||||
("ai", "step-4", "ai-4"),
|
||||
]
|
||||
assert _digest(new_head.values) == expected_new_head
|
||||
assert _digest(new_head.values) == _digest(oracle_new_head.values)
|
||||
|
||||
# Re-reading ORIGINAL chain checkpoints (pinned by checkpoint_id) returns
|
||||
# their own original histories - their cached entries predate the fork and
|
||||
# must be untouched by the superseding branch.
|
||||
for next_key, expected_len in [((), 6), (("step4",), 5), (("step2",), 3)]:
|
||||
reread = await accessor.aget(by_next[next_key].config)
|
||||
oracle_reread = await oracle_accessor.aget(oracle_by_next[next_key].config)
|
||||
assert _digest(reread.values) == _expected_final_digest()[:expected_len]
|
||||
assert _digest(reread.values) == _digest(oracle_reread.values)
|
||||
46
backend/tests/test_checkpoint_cache_config.py
Normal file
46
backend/tests/test_checkpoint_cache_config.py
Normal file
@ -0,0 +1,46 @@
|
||||
"""Config parsing for database.checkpoint_cache."""
|
||||
|
||||
from deerflow.config.database_config import CheckpointCacheConfig, DatabaseConfig
|
||||
|
||||
|
||||
def test_checkpoint_cache_defaults():
|
||||
cfg = DatabaseConfig()
|
||||
assert cfg.checkpoint_cache.type == "memory"
|
||||
assert cfg.checkpoint_cache.max_entries == 128
|
||||
assert cfg.checkpoint_cache.redis_url is None
|
||||
assert cfg.checkpoint_cache.ttl_seconds == 86400
|
||||
assert cfg.checkpoint_cache.key_prefix == ""
|
||||
|
||||
|
||||
def test_checkpoint_cache_from_dict_redis():
|
||||
cfg = DatabaseConfig.model_validate(
|
||||
{
|
||||
"backend": "postgres",
|
||||
"postgres_url": "postgresql://u:p@h/db",
|
||||
"checkpoint_cache": {
|
||||
"type": "redis",
|
||||
"max_entries": 256,
|
||||
"redis_url": "redis://localhost:6379/3",
|
||||
"ttl_seconds": 3600,
|
||||
"key_prefix": "prod:",
|
||||
},
|
||||
}
|
||||
)
|
||||
assert cfg.checkpoint_cache.type == "redis"
|
||||
assert cfg.checkpoint_cache.max_entries == 256
|
||||
assert cfg.checkpoint_cache.redis_url == "redis://localhost:6379/3"
|
||||
assert cfg.checkpoint_cache.ttl_seconds == 3600
|
||||
assert cfg.checkpoint_cache.key_prefix == "prod:"
|
||||
|
||||
|
||||
def test_checkpoint_cache_zero_max_entries_means_disabled():
|
||||
cfg = CheckpointCacheConfig(max_entries=0)
|
||||
assert cfg.max_entries == 0
|
||||
|
||||
|
||||
def test_checkpoint_cache_rejects_negative_max_entries():
|
||||
import pydantic
|
||||
import pytest
|
||||
|
||||
with pytest.raises(pydantic.ValidationError):
|
||||
CheckpointCacheConfig(max_entries=-1)
|
||||
130
backend/tests/test_checkpoint_cache_memory.py
Normal file
130
backend/tests/test_checkpoint_cache_memory.py
Normal file
@ -0,0 +1,130 @@
|
||||
"""Memory LRU backend for the checkpoint history cache."""
|
||||
|
||||
import pytest
|
||||
|
||||
from deerflow.runtime.checkpoint_cache.base import (
|
||||
CACHE_FORMAT_VERSION,
|
||||
CheckpointCacheStats,
|
||||
make_history_key,
|
||||
thread_key_stem,
|
||||
)
|
||||
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
||||
|
||||
|
||||
def _entry(tag: str) -> dict:
|
||||
return {"writes": [("task-1", "messages", tag)], "seed": f"seed-{tag}"}
|
||||
|
||||
|
||||
def test_make_history_key_is_stable_and_scoped():
|
||||
k1 = make_history_key("ckpt-hist:v1:db0", "t1", "", "c1", "messages")
|
||||
k2 = make_history_key("ckpt-hist:v1:db0", "t1", "", "c1", "messages")
|
||||
assert k1 == k2
|
||||
assert k1.startswith("ckpt-hist:v1:db0:t1:")
|
||||
# ns / checkpoint / channel each change the key
|
||||
assert k1 != make_history_key("ckpt-hist:v1:db0", "t1", "sub", "c1", "messages")
|
||||
assert k1 != make_history_key("ckpt-hist:v1:db0", "t1", "", "c2", "messages")
|
||||
assert k1 != make_history_key("ckpt-hist:v1:db0", "t1", "", "c1", "todos")
|
||||
assert k1 != make_history_key("ckpt-hist:v1:db9", "t1", "", "c1", "messages")
|
||||
assert CACHE_FORMAT_VERSION == 1
|
||||
|
||||
|
||||
def test_get_many_miss_then_hit():
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=4)
|
||||
assert cache.get_many(["a"]) == {}
|
||||
assert cache.stats().misses == 1
|
||||
cache.set_many({"a": _entry("x")})
|
||||
hit = cache.get_many(["a"])
|
||||
assert hit["a"]["writes"] == [("task-1", "messages", "x")]
|
||||
assert hit["a"]["seed"] == "seed-x"
|
||||
assert cache.stats().hits == 1
|
||||
|
||||
|
||||
def test_entry_without_seed_roundtrips_without_seed_key():
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=4)
|
||||
cache.set_many({"a": {"writes": []}})
|
||||
hit = cache.get_many(["a"])
|
||||
assert hit["a"] == {"writes": []}
|
||||
assert "seed" not in hit["a"]
|
||||
|
||||
|
||||
def test_copy_on_read_returns_fresh_writes_list():
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=4)
|
||||
cache.set_many({"a": _entry("x")})
|
||||
first = cache.get_many(["a"])["a"]
|
||||
first["writes"].append(("task-2", "messages", "MUTATION"))
|
||||
second = cache.get_many(["a"])["a"]
|
||||
assert second["writes"] == [("task-1", "messages", "x")]
|
||||
|
||||
|
||||
def test_caller_mutation_after_set_does_not_leak():
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=4)
|
||||
entry = _entry("x")
|
||||
cache.set_many({"a": entry})
|
||||
entry["writes"].append(("task-2", "messages", "MUTATION"))
|
||||
assert cache.get_many(["a"])["a"]["writes"] == [("task-1", "messages", "x")]
|
||||
|
||||
|
||||
def test_lru_evicts_oldest_and_counts():
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=2)
|
||||
cache.set_many({"a": _entry("a"), "b": _entry("b")})
|
||||
cache.get_many(["a"]) # refresh a
|
||||
cache.set_many({"c": _entry("c")}) # evicts b
|
||||
assert cache.get_many(["b"]) == {}
|
||||
assert cache.get_many(["a"]) != {}
|
||||
assert cache.stats().evictions == 1
|
||||
assert cache.stats().entries == 2
|
||||
|
||||
|
||||
def test_zero_max_entries_disables():
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=0)
|
||||
assert cache.enabled is False
|
||||
cache.set_many({"a": _entry("x")})
|
||||
assert cache.get_many(["a"]) == {}
|
||||
assert cache.stats().entries == 0
|
||||
|
||||
|
||||
def test_delete_thread_purges_only_that_thread():
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=16)
|
||||
prefix = "ckpt-hist:v1:db0"
|
||||
t1_keys = [make_history_key(prefix, "t1", "", f"c{i}", "messages") for i in range(3)]
|
||||
t2_key = make_history_key(prefix, "t2", "", "c0", "messages")
|
||||
# A thread_id that is a prefix of another must not over-match: the stem
|
||||
# ends with ':' so "t1" never matches "t10"'s keys.
|
||||
t10_key = make_history_key(prefix, "t10", "", "c0", "messages")
|
||||
cache.set_many({k: _entry(k) for k in [*t1_keys, t2_key, t10_key]})
|
||||
|
||||
cache.delete_thread(prefix, "t1")
|
||||
|
||||
assert cache.stats().entries == 2
|
||||
assert all(cache.get_many([k]) == {} for k in t1_keys)
|
||||
assert cache.get_many([t2_key]) != {}
|
||||
assert cache.get_many([t10_key]) != {}
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_adelete_thread_matches_sync():
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=4)
|
||||
prefix = "ckpt-hist:v1:db0"
|
||||
key = make_history_key(prefix, "t1", "", "c0", "messages")
|
||||
await cache.aset_many({key: _entry("x")})
|
||||
await cache.adelete_thread(prefix, "t1")
|
||||
assert cache.get_many([key]) == {}
|
||||
|
||||
|
||||
def test_thread_key_stem_matches_make_history_key_layout():
|
||||
key = make_history_key("p", "t1", "ns", "c1", "messages")
|
||||
assert key.startswith(thread_key_stem("p", "t1"))
|
||||
assert not key.startswith(thread_key_stem("p", "t"))
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_protocol_matches_sync():
|
||||
cache = MemoryCheckpointHistoryCache(max_entries=4)
|
||||
await cache.aset_many({"a": _entry("x")})
|
||||
hit = await cache.aget_many(["a"])
|
||||
assert hit["a"]["seed"] == "seed-x"
|
||||
stats = cache.stats()
|
||||
assert isinstance(stats, CheckpointCacheStats)
|
||||
assert stats.as_dict()["hits"] == 1
|
||||
await cache.aclose()
|
||||
assert cache.get_many(["a"]) == {}
|
||||
96
backend/tests/test_checkpoint_cache_provider.py
Normal file
96
backend/tests/test_checkpoint_cache_provider.py
Normal file
@ -0,0 +1,96 @@
|
||||
"""Provider wiring: mode-gated wrapping in async and sync checkpointer factories."""
|
||||
|
||||
import pytest
|
||||
|
||||
from deerflow.config.app_config import AppConfig, set_app_config
|
||||
from deerflow.runtime.checkpoint_mode import freeze_checkpoint_channel_mode
|
||||
from deerflow.runtime.checkpointer.async_provider import make_checkpointer
|
||||
from deerflow.runtime.checkpointer.cached_saver import CachedHistorySaver
|
||||
from deerflow.runtime.checkpointer.provider import checkpointer_context, reset_checkpointer
|
||||
|
||||
|
||||
# AppConfig requires the sandbox section (no default); the rest of the config
|
||||
# is optional. Mirrors test_checkpoint_cache_redis.py's construction pattern.
|
||||
def _app_config(mode: str, cache: dict | None = None) -> AppConfig:
|
||||
database: dict = {"backend": "memory", "checkpoint_channel_mode": mode}
|
||||
if cache is not None:
|
||||
database["checkpoint_cache"] = cache
|
||||
return AppConfig.model_validate(
|
||||
{
|
||||
"sandbox": {"use": "deerflow.sandbox.local.provider:LocalSandboxProvider"},
|
||||
"database": database,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_delta_mode_wraps_with_cached_saver():
|
||||
set_app_config(_app_config("delta"))
|
||||
freeze_checkpoint_channel_mode("delta")
|
||||
async with make_checkpointer() as saver:
|
||||
assert isinstance(saver, CachedHistorySaver)
|
||||
assert saver.stats()["entries"] == 0
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_full_mode_yields_raw_saver():
|
||||
set_app_config(_app_config("full"))
|
||||
freeze_checkpoint_channel_mode("full")
|
||||
async with make_checkpointer() as saver:
|
||||
assert not isinstance(saver, CachedHistorySaver)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_zero_max_entries_disables_but_still_wraps():
|
||||
set_app_config(_app_config("delta", {"max_entries": 0}))
|
||||
freeze_checkpoint_channel_mode("delta")
|
||||
async with make_checkpointer() as saver:
|
||||
assert isinstance(saver, CachedHistorySaver)
|
||||
# Disabled cache -> every history call is a full walk on the inner saver.
|
||||
assert saver._cache.enabled is False
|
||||
|
||||
|
||||
def test_sync_delta_mode_wraps_memory():
|
||||
set_app_config(_app_config("delta"))
|
||||
freeze_checkpoint_channel_mode("delta")
|
||||
reset_checkpointer()
|
||||
with checkpointer_context() as saver:
|
||||
assert isinstance(saver, CachedHistorySaver)
|
||||
|
||||
|
||||
def test_sync_redis_cache_type_is_config_error():
|
||||
set_app_config(_app_config("delta", {"type": "redis"}))
|
||||
freeze_checkpoint_channel_mode("delta")
|
||||
reset_checkpointer()
|
||||
with pytest.raises(ValueError, match="redis"):
|
||||
with checkpointer_context():
|
||||
pass
|
||||
|
||||
|
||||
def test_sync_full_mode_unwrapped():
|
||||
set_app_config(_app_config("full"))
|
||||
freeze_checkpoint_channel_mode("full")
|
||||
reset_checkpointer()
|
||||
with checkpointer_context() as saver:
|
||||
assert not isinstance(saver, CachedHistorySaver)
|
||||
|
||||
|
||||
def test_sync_cache_recreated_when_key_prefix_changes():
|
||||
"""The singleton must not outlive its namespace: a prefix change without
|
||||
a process restart leaves old-prefix entries unreachable and unpurgeable."""
|
||||
reset_checkpointer()
|
||||
set_app_config(_app_config("delta", {"key_prefix": "ns-a"}))
|
||||
freeze_checkpoint_channel_mode("delta")
|
||||
with checkpointer_context() as saver:
|
||||
saver._cache.set_many({"ns-a:t1:x": {"writes": []}})
|
||||
first_cache = saver._cache
|
||||
# Same prefix: singleton is reused (warm across wrappers).
|
||||
with checkpointer_context() as saver:
|
||||
assert saver._cache is first_cache
|
||||
assert saver._cache.stats().entries == 1
|
||||
# Prefix change: fresh cache, stale namespace gone with the old instance.
|
||||
set_app_config(_app_config("delta", {"key_prefix": "ns-b"}))
|
||||
with checkpointer_context() as saver:
|
||||
assert saver._cache is not first_cache
|
||||
assert saver._cache.stats().entries == 0
|
||||
reset_checkpointer()
|
||||
245
backend/tests/test_checkpoint_cache_redis.py
Normal file
245
backend/tests/test_checkpoint_cache_redis.py
Normal file
@ -0,0 +1,245 @@
|
||||
"""Redis backend and provider factory for the checkpoint history cache."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
||||
|
||||
from deerflow.config.app_config import AppConfig
|
||||
from deerflow.runtime.checkpoint_cache.provider import (
|
||||
checkpoint_cache_db_hash,
|
||||
checkpoint_cache_key_prefix,
|
||||
make_checkpoint_cache,
|
||||
)
|
||||
|
||||
|
||||
class _FakeRedis:
|
||||
"""Minimal async redis stand-in: mget / set / pipeline / scan / unlink."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.store: dict[str, bytes] = {}
|
||||
self.ttls: dict[str, int | None] = {}
|
||||
self.unlinked: list[tuple[str, ...]] = []
|
||||
|
||||
async def mget(self, keys: list[str]) -> list[bytes | None]:
|
||||
return [self.store.get(k) for k in keys]
|
||||
|
||||
def set(self, key: str, value: bytes, ex: int | None = None) -> None:
|
||||
self.store[key] = value
|
||||
self.ttls[key] = ex
|
||||
|
||||
async def scan(self, cursor: int = 0, match: str | None = None, count: int = 500) -> tuple[int, list[str]]:
|
||||
import fnmatch
|
||||
|
||||
keys = sorted(self.store)
|
||||
batch = keys[cursor : cursor + count]
|
||||
if match is not None:
|
||||
batch = [k for k in batch if fnmatch.fnmatchcase(k, match)]
|
||||
next_cursor = cursor + count
|
||||
return (0 if next_cursor >= len(keys) else next_cursor), batch
|
||||
|
||||
async def unlink(self, *keys: str) -> int:
|
||||
self.unlinked.append(tuple(keys))
|
||||
removed = 0
|
||||
for key in keys:
|
||||
removed += self.store.pop(key, None) is not None
|
||||
return removed
|
||||
|
||||
def pipeline(self, transaction: bool = False) -> "_FakePipeline":
|
||||
return _FakePipeline(self)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class _FakePipeline:
|
||||
def __init__(self, client: _FakeRedis) -> None:
|
||||
self._client = client
|
||||
|
||||
def set(self, key: str, value: bytes, ex: int | None = None) -> "_FakePipeline":
|
||||
self._client.set(key, value, ex=ex)
|
||||
return self
|
||||
|
||||
async def execute(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class _FailingRedis(_FakeRedis):
|
||||
"""Simulates a redis outage: every operation raises RedisError."""
|
||||
|
||||
async def mget(self, keys: list[str]) -> list[bytes | None]:
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
raise RedisError("connection refused")
|
||||
|
||||
async def scan(self, cursor: int = 0, match: str | None = None, count: int = 500) -> tuple[int, list[str]]:
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
raise RedisError("connection refused")
|
||||
|
||||
def pipeline(self, transaction: bool = False) -> "_FakePipeline":
|
||||
from redis.exceptions import RedisError
|
||||
|
||||
raise RedisError("connection refused")
|
||||
|
||||
|
||||
def _make_cache(monkeypatch: pytest.MonkeyPatch, fake: _FakeRedis, ttl_seconds: int = 60, **kwargs: Any):
|
||||
import deerflow.runtime.checkpoint_cache.redis as redis_mod
|
||||
|
||||
monkeypatch.setattr(redis_mod, "_create_client", lambda *a, **k: fake)
|
||||
return redis_mod.RedisCheckpointHistoryCache("redis://unused", serde=JsonPlusSerializer(), ttl_seconds=ttl_seconds, **kwargs)
|
||||
|
||||
|
||||
def _entry(i: int) -> dict:
|
||||
# Real message-like payloads to prove serde fidelity beyond plain dicts.
|
||||
from langchain_core.messages import AIMessage
|
||||
|
||||
return {"writes": [("task-1", "messages", AIMessage(content=f"m{i}", id=f"ai-{i}"))], "seed": [AIMessage(content="s", id="ai-s")]}
|
||||
|
||||
|
||||
# AppConfig requires the sandbox section (no default); the rest of the config
|
||||
# is optional. Mirrors test_checkpoint_mode.py's construction pattern.
|
||||
def _app_config(database: dict) -> AppConfig:
|
||||
return AppConfig.model_validate(
|
||||
{
|
||||
"sandbox": {"use": "deerflow.sandbox.local.provider:LocalSandboxProvider"},
|
||||
"database": database,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_roundtrip_preserves_types(monkeypatch: pytest.MonkeyPatch):
|
||||
fake = _FakeRedis()
|
||||
cache = _make_cache(monkeypatch, fake)
|
||||
await cache.aset_many({"k1": _entry(1), "k2": {"writes": []}})
|
||||
hit = await cache.aget_many(["k1", "k2", "k3"])
|
||||
assert set(hit) == {"k1", "k2"}
|
||||
msg = hit["k1"]["writes"][0][2]
|
||||
assert msg.content == "m1" and msg.id == "ai-1" and msg.type == "ai"
|
||||
assert "seed" not in hit["k2"]
|
||||
assert cache.stats().hits == 2 and cache.stats().misses == 1
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_keys_land_verbatim_and_ttl_set(monkeypatch: pytest.MonkeyPatch):
|
||||
fake = _FakeRedis()
|
||||
cache = _make_cache(monkeypatch, fake)
|
||||
await cache.aset_many({"k1": _entry(1)})
|
||||
assert list(fake.store) == ["k1"]
|
||||
assert fake.ttls["k1"] == 60
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_outage_degrades_to_all_miss(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture):
|
||||
fake = _FailingRedis()
|
||||
cache = _make_cache(monkeypatch, fake)
|
||||
with caplog.at_level("WARNING"):
|
||||
assert await cache.aget_many(["k1", "k2"]) == {}
|
||||
assert cache.stats().misses == 2 and cache.stats().hits == 0
|
||||
assert "mget failed" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_redis_outage_skips_write(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture):
|
||||
fake = _FailingRedis()
|
||||
cache = _make_cache(monkeypatch, fake)
|
||||
with caplog.at_level("WARNING"):
|
||||
await cache.aset_many({"k1": _entry(1)}) # must not raise
|
||||
assert fake.store == {}
|
||||
assert "write failed" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_zero_ttl_disables_expiry_explicitly(monkeypatch: pytest.MonkeyPatch):
|
||||
fake = _FakeRedis()
|
||||
cache = _make_cache(monkeypatch, fake, ttl_seconds=0)
|
||||
await cache.aset_many({"k1": _entry(1)})
|
||||
assert cache._ttl is None
|
||||
assert fake.ttls["k1"] is None # SET without EX: redis maxmemory policy only
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_adelete_thread_purges_matching_keys_only(monkeypatch: pytest.MonkeyPatch):
|
||||
fake = _FakeRedis()
|
||||
cache = _make_cache(monkeypatch, fake)
|
||||
prefix = "ckpt-hist:v1:db0"
|
||||
await cache.aset_many(
|
||||
{
|
||||
f"{prefix}:t1:aaa": _entry(1),
|
||||
f"{prefix}:t1:bbb": _entry(2),
|
||||
f"{prefix}:t10:ccc": _entry(3), # 't1' stem must not over-match 't10'
|
||||
f"{prefix}:t2:ddd": _entry(4),
|
||||
}
|
||||
)
|
||||
|
||||
await cache.adelete_thread(prefix, "t1")
|
||||
|
||||
assert sorted(fake.store) == [f"{prefix}:t10:ccc", f"{prefix}:t2:ddd"]
|
||||
assert fake.unlinked # UNLINK, not DEL: non-blocking on big histories
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_adelete_thread_outage_degrades_without_raising(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture):
|
||||
fake = _FailingRedis()
|
||||
cache = _make_cache(monkeypatch, fake)
|
||||
with caplog.at_level("WARNING"):
|
||||
await cache.adelete_thread("p", "t1") # must not raise
|
||||
assert "thread purge failed" in caplog.text
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_provider_memory_default():
|
||||
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
||||
|
||||
async with make_checkpoint_cache(_app_config({"backend": "sqlite"}), serde=JsonPlusSerializer()) as cache:
|
||||
assert isinstance(cache, MemoryCheckpointHistoryCache)
|
||||
assert cache.enabled is True
|
||||
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_provider_zero_max_entries_disables_any_type():
|
||||
app_config = _app_config({"backend": "sqlite", "checkpoint_cache": {"type": "redis", "max_entries": 0}})
|
||||
from deerflow.runtime.checkpoint_cache.memory import MemoryCheckpointHistoryCache
|
||||
|
||||
async with make_checkpoint_cache(app_config, serde=JsonPlusSerializer()) as cache:
|
||||
assert isinstance(cache, MemoryCheckpointHistoryCache)
|
||||
assert cache.enabled is False
|
||||
|
||||
|
||||
def test_db_hash_distinguishes_backends_and_targets():
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
|
||||
sqlite_cfg = DatabaseConfig.model_validate({"backend": "sqlite", "sqlite_dir": "/tmp/a"})
|
||||
pg_cfg = DatabaseConfig.model_validate({"backend": "postgres", "postgres_url": "postgresql://u:p@h/db"})
|
||||
pg_cfg2 = DatabaseConfig.model_validate({"backend": "postgres", "postgres_url": "postgresql://u:p@h/other"})
|
||||
assert checkpoint_cache_db_hash(sqlite_cfg) != checkpoint_cache_db_hash(pg_cfg)
|
||||
assert checkpoint_cache_db_hash(pg_cfg) != checkpoint_cache_db_hash(pg_cfg2)
|
||||
assert len(checkpoint_cache_db_hash(pg_cfg)) == 12
|
||||
|
||||
|
||||
def test_db_hash_stable_across_credential_rotation():
|
||||
"""Same database, rotated user/password -> same cache namespace."""
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
|
||||
before = DatabaseConfig.model_validate({"backend": "postgres", "postgres_url": "postgresql://alice:secret1@pg.internal:5432/deerflow"})
|
||||
rotated = DatabaseConfig.model_validate({"backend": "postgres", "postgres_url": "postgresql://bob:secret2@pg.internal:5432/deerflow"})
|
||||
driver_suffix = DatabaseConfig.model_validate({"backend": "postgres", "postgres_url": "postgresql+asyncpg://alice:secret1@pg.internal:5432/deerflow"})
|
||||
other_db = DatabaseConfig.model_validate({"backend": "postgres", "postgres_url": "postgresql://alice:secret1@pg.internal:5432/other"})
|
||||
assert checkpoint_cache_db_hash(before) == checkpoint_cache_db_hash(rotated)
|
||||
assert checkpoint_cache_db_hash(before) == checkpoint_cache_db_hash(driver_suffix)
|
||||
assert checkpoint_cache_db_hash(before) != checkpoint_cache_db_hash(other_db)
|
||||
|
||||
|
||||
def test_db_hash_unparseable_url_falls_back_to_raw():
|
||||
from deerflow.config.database_config import DatabaseConfig
|
||||
|
||||
cfg = DatabaseConfig.model_validate({"backend": "postgres", "postgres_url": "not-a-url"})
|
||||
assert len(checkpoint_cache_db_hash(cfg)) == 12 # stable, never raises
|
||||
|
||||
|
||||
def test_key_prefix_override_wins():
|
||||
app_config = _app_config({"backend": "sqlite", "checkpoint_cache": {"key_prefix": "custom:"}})
|
||||
assert checkpoint_cache_key_prefix(app_config) == "custom:"
|
||||
default = checkpoint_cache_key_prefix(_app_config({"backend": "sqlite"}))
|
||||
assert default.startswith("ckpt-hist:v1:")
|
||||
@ -1088,6 +1088,8 @@ class TestExtractText:
|
||||
|
||||
class TestEnsureAgent:
|
||||
def test_authorization_filters_framework_tools_and_reuses_provider(self, client, mock_app_config):
|
||||
from deerflow.authz.provider import AuthzDecision, AuthzReason
|
||||
|
||||
class Provider:
|
||||
name = "test"
|
||||
|
||||
@ -1095,10 +1097,13 @@ class TestEnsureAgent:
|
||||
return [name for name in candidates if name == "safe_tool"]
|
||||
|
||||
def authorize(self, request):
|
||||
raise AssertionError("not called while assembling")
|
||||
# Phase 3: model:use is now checked during assembly; allow it so
|
||||
# the model name passes through. Tool-level authorize is still
|
||||
# not invoked here (filter_resources drives tool assembly).
|
||||
return AuthzDecision(allow=True, reasons=[AuthzReason(code="authz.allowed")])
|
||||
|
||||
async def aauthorize(self, request):
|
||||
raise AssertionError("not called while assembling")
|
||||
return self.authorize(request)
|
||||
|
||||
provider = Provider()
|
||||
mock_app_config.authorization = AuthorizationConfig(
|
||||
@ -1121,6 +1126,7 @@ class TestEnsureAgent:
|
||||
patch("deerflow.client.build_skill_search_setup", return_value=SimpleNamespace(describe_skill_tool=describe_tool, skill_names=frozenset({"example"}))),
|
||||
patch.object(client, "_get_tools", return_value=[safe_tool, denied_tool]),
|
||||
patch("deerflow.authz.tool_filter.resolve_authorization_provider", return_value=provider),
|
||||
patch("deerflow.agents.lead_agent.agent.resolve_authorization_provider", return_value=provider),
|
||||
patch("deerflow.runtime.checkpointer.get_checkpointer", return_value=None),
|
||||
):
|
||||
client._ensure_agent(client._get_runnable_config("t1"), context={"user_role": "user"})
|
||||
@ -1861,10 +1867,8 @@ class TestSkillsManagement:
|
||||
skill = self._make_skill(enabled=True)
|
||||
updated_skill = self._make_skill(enabled=False)
|
||||
|
||||
ext_config = ExtensionsConfig()
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
|
||||
json.dump({}, f)
|
||||
json.dump({"mcpServers": {}, "skills": {"untouched-skill": {"enabled": False}}}, f)
|
||||
tmp_path = Path(f.name)
|
||||
|
||||
try:
|
||||
@ -1881,12 +1885,37 @@ class TestSkillsManagement:
|
||||
side_effect=[[skill], [skill], [updated_skill], [updated_skill]],
|
||||
),
|
||||
patch("deerflow.client.ExtensionsConfig.resolve_config_path", return_value=tmp_path),
|
||||
patch("deerflow.client.get_extensions_config", return_value=ext_config),
|
||||
patch("deerflow.client.reload_extensions_config"),
|
||||
):
|
||||
result = client.update_skill("test-skill", enabled=False)
|
||||
assert result["enabled"] is False
|
||||
assert client._agent is None # M2: agent invalidated
|
||||
persisted = json.loads(tmp_path.read_text(encoding="utf-8"))
|
||||
assert persisted["skills"]["untouched-skill"] == {"enabled": False}
|
||||
finally:
|
||||
tmp_path.unlink()
|
||||
|
||||
def test_update_skill_persists_state_when_source_omits_skills(self, client):
|
||||
skill = self._make_skill(enabled=True)
|
||||
updated_skill = self._make_skill(enabled=False)
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
|
||||
json.dump({"mcpServers": {}}, f)
|
||||
tmp_path = Path(f.name)
|
||||
|
||||
try:
|
||||
with (
|
||||
patch(
|
||||
"deerflow.skills.storage.local_skill_storage.LocalSkillStorage.load_skills",
|
||||
side_effect=[[skill], [skill], [updated_skill], [updated_skill]],
|
||||
),
|
||||
patch("deerflow.client.ExtensionsConfig.resolve_config_path", return_value=tmp_path),
|
||||
patch("deerflow.client.reload_extensions_config"),
|
||||
):
|
||||
client.update_skill("test-skill", enabled=False)
|
||||
|
||||
persisted = json.loads(tmp_path.read_text(encoding="utf-8"))
|
||||
assert persisted["skills"] == {"test-skill": {"enabled": False}}
|
||||
finally:
|
||||
tmp_path.unlink()
|
||||
|
||||
|
||||
112
backend/tests/test_compose_default_bind_host.py
Normal file
112
backend/tests/test_compose_default_bind_host.py
Normal file
@ -0,0 +1,112 @@
|
||||
"""Regression test for the Docker Compose default published bind address.
|
||||
|
||||
``README.md`` documents DeerFlow as being deployed by default "in a local
|
||||
trusted environment (accessible only via the 127.0.0.1 loopback interface)",
|
||||
but the shipped compose files published the nginx entry as
|
||||
``"${PORT:-2026}:2026"``, which Docker binds to ``0.0.0.0`` (and ``[::]``). The
|
||||
shipped artifact therefore did not match its own documented default, and an
|
||||
operator running it on a LAN or cloud host got a wider surface than the docs
|
||||
implied without changing anything.
|
||||
|
||||
The Gateway itself binds ``0.0.0.0`` inside the container on purpose (nginx has
|
||||
to reach it over the compose network) and its port is deliberately not
|
||||
published, so the published nginx port is the whole external surface. This test
|
||||
pins the loopback default there while keeping it overridable for operators who
|
||||
intentionally expose the stack behind their own TLS/auth front door.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
COMPOSE_PATHS = {
|
||||
"prod": REPO_ROOT / "docker" / "docker-compose.yaml",
|
||||
"dev": REPO_ROOT / "docker" / "docker-compose-dev.yaml",
|
||||
}
|
||||
|
||||
EXPECTED_NGINX_PORT_MAPPING = "${BIND_HOST:-127.0.0.1}:${PORT:-2026}:2026"
|
||||
|
||||
|
||||
def _published_ports(compose_path: Path) -> dict[str, list[str]]:
|
||||
"""Return {service_name: [port mapping, ...]} for every published port."""
|
||||
compose = yaml.safe_load(compose_path.read_text(encoding="utf-8"))
|
||||
published: dict[str, list[str]] = {}
|
||||
for service_name, service in (compose.get("services") or {}).items():
|
||||
ports = service.get("ports") if isinstance(service, dict) else None
|
||||
if not ports:
|
||||
continue
|
||||
published[service_name] = [str(entry) for entry in ports]
|
||||
return published
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", sorted(COMPOSE_PATHS))
|
||||
def test_nginx_entry_defaults_to_loopback(variant: str):
|
||||
"""With BIND_HOST unset, the entry port must bind 127.0.0.1, not 0.0.0.0."""
|
||||
published = _published_ports(COMPOSE_PATHS[variant])
|
||||
|
||||
assert published.get("nginx") == [EXPECTED_NGINX_PORT_MAPPING], f"{variant} compose must publish nginx as {EXPECTED_NGINX_PORT_MAPPING!r}; got: {published.get('nginx')!r}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", sorted(COMPOSE_PATHS))
|
||||
def test_no_service_publishes_on_all_interfaces(variant: str):
|
||||
"""No compose service may publish a port without an explicit bind address.
|
||||
|
||||
A bare ``"HOST:CONTAINER"`` mapping binds every interface. Any port added
|
||||
later must either stay internal to the compose network or opt in to the
|
||||
same ``BIND_HOST`` default.
|
||||
"""
|
||||
offenders: list[str] = []
|
||||
for service_name, mappings in _published_ports(COMPOSE_PATHS[variant]).items():
|
||||
for mapping in mappings:
|
||||
# A bind address is present only when the mapping has three
|
||||
# colon-separated parts (``ADDR:HOST:CONTAINER``). Variable
|
||||
# substitutions such as ``${PORT:-2026}`` also contain colons, so
|
||||
# count separators outside ``${...}`` instead of splitting naively.
|
||||
if _bind_address(mapping) is None:
|
||||
offenders.append(f"{service_name}: {mapping}")
|
||||
|
||||
assert not offenders, f"{variant} compose publishes ports on all interfaces (add a bind address): {offenders}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("variant", sorted(COMPOSE_PATHS))
|
||||
def test_bind_address_remains_overridable(variant: str):
|
||||
"""Operators fronting the stack themselves must be able to widen the bind."""
|
||||
mapping = _published_ports(COMPOSE_PATHS[variant])["nginx"][0]
|
||||
|
||||
assert _bind_address(mapping) == "${BIND_HOST:-127.0.0.1}", f"{variant} compose must keep the bind address overridable via BIND_HOST; got: {mapping!r}"
|
||||
|
||||
|
||||
def _bind_address(mapping: str) -> str | None:
|
||||
"""Return the bind-address segment of a compose port mapping, if any.
|
||||
|
||||
Splits on ``:`` at nesting depth zero so ``${PORT:-2026}`` is treated as a
|
||||
single segment rather than two.
|
||||
"""
|
||||
segments: list[str] = []
|
||||
current: list[str] = []
|
||||
depth = 0
|
||||
index = 0
|
||||
while index < len(mapping):
|
||||
char = mapping[index]
|
||||
if mapping.startswith("${", index):
|
||||
depth += 1
|
||||
current.append("${")
|
||||
index += 2
|
||||
continue
|
||||
if char == "}" and depth > 0:
|
||||
depth -= 1
|
||||
elif char == ":" and depth == 0:
|
||||
segments.append("".join(current))
|
||||
current = []
|
||||
index += 1
|
||||
continue
|
||||
current.append(char)
|
||||
index += 1
|
||||
segments.append("".join(current))
|
||||
|
||||
# ADDR:HOST:CONTAINER -> bound; HOST:CONTAINER or CONTAINER -> unbound.
|
||||
return segments[0] if len(segments) >= 3 else None
|
||||
@ -3,6 +3,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
@ -513,6 +514,20 @@ class TestMemoryFilePath:
|
||||
# ===========================================================================
|
||||
|
||||
|
||||
# Model names the agents API tests may send in a create/update payload. The
|
||||
# router validates `model` against the app config, so the fixture pins a stub
|
||||
# config exposing exactly these instead of leaving the assertion at the mercy
|
||||
# of whatever `config.yaml` happens to sit in the repo root: CI has none (the
|
||||
# validation is skipped and the request passes), a real dev checkout does (an
|
||||
# unlisted model name yields 422 and the test fails).
|
||||
_KNOWN_TEST_MODELS = frozenset({"deepseek-v3"})
|
||||
|
||||
|
||||
def _stub_app_config():
|
||||
"""App config that knows only ``_KNOWN_TEST_MODELS``."""
|
||||
return SimpleNamespace(get_model_config=lambda name: SimpleNamespace(name=name) if name in _KNOWN_TEST_MODELS else None)
|
||||
|
||||
|
||||
def _make_test_app(tmp_path: Path):
|
||||
"""Create a FastAPI app with the agents router, patching paths to tmp_path."""
|
||||
from fastapi import FastAPI
|
||||
@ -532,7 +547,11 @@ def agent_client(tmp_path):
|
||||
paths_instance = _make_paths(tmp_path)
|
||||
previous_config = AgentsApiConfig(**get_agents_api_config().model_dump())
|
||||
|
||||
with patch("deerflow.config.agents_config.get_paths", return_value=paths_instance), patch.object(agents_router, "get_paths", return_value=paths_instance):
|
||||
with (
|
||||
patch("deerflow.config.agents_config.get_paths", return_value=paths_instance),
|
||||
patch.object(agents_router, "get_paths", return_value=paths_instance),
|
||||
patch.object(agents_router, "get_app_config", _stub_app_config),
|
||||
):
|
||||
set_agents_api_config(AgentsApiConfig(enabled=True))
|
||||
try:
|
||||
app = _make_test_app(tmp_path)
|
||||
|
||||
@ -240,7 +240,7 @@ def test_trace_id_threads_through_to_callbacks(deermem_data_dir):
|
||||
|
||||
|
||||
def test_default_passive_update_persists_fact_in_reserved_default_bucket(deermem_data_dir):
|
||||
dm = _deermem_with_fake_llm(payload='{"user":{},"history":{},"newFacts":[{"content":"Default agent fact","category":"context","confidence":0.9}],"factsToRemove":[]}')
|
||||
dm = _deermem_with_fake_llm(payload='{"user":{},"history":{},"newFacts":[{"content":"Default agent fact","category":"context","confidence":0.9,"scope":"user","durability":"durable","authority":"descriptive"}],"factsToRemove":[]}')
|
||||
|
||||
dm.add(
|
||||
thread_id="default-thread",
|
||||
|
||||
178
backend/tests/test_e2b_capacity_store_redis.py
Normal file
178
backend/tests/test_e2b_capacity_store_redis.py
Normal file
@ -0,0 +1,178 @@
|
||||
"""Integration tests for deployment-wide E2B admission."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import threading
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
import pytest
|
||||
|
||||
from deerflow.community.e2b_sandbox.capacity import (
|
||||
CapacityBackendError,
|
||||
RedisE2BCapacityStore,
|
||||
ReserveStatus,
|
||||
make_e2b_capacity_store,
|
||||
)
|
||||
from deerflow.config.sandbox_config import SandboxOwnershipConfig
|
||||
|
||||
REDIS_URL = os.environ.get("DEER_FLOW_TEST_REDIS_URL", "redis://localhost:6379/15")
|
||||
pytestmark = pytest.mark.integration
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def make_store():
|
||||
redis = pytest.importorskip("redis")
|
||||
probe = redis.Redis.from_url(REDIS_URL, decode_responses=True, socket_connect_timeout=0.5)
|
||||
try:
|
||||
probe.ping()
|
||||
except Exception:
|
||||
probe.close()
|
||||
pytest.skip(f"Redis not reachable at {REDIS_URL}")
|
||||
|
||||
prefix = f"deerflow:test:{uuid.uuid4().hex}"
|
||||
stores = []
|
||||
|
||||
def make(hard_limit: int = 1):
|
||||
store = RedisE2BCapacityStore(
|
||||
redis_url=REDIS_URL,
|
||||
hard_limit=hard_limit,
|
||||
key_prefix=prefix,
|
||||
)
|
||||
stores.append(store)
|
||||
return store
|
||||
|
||||
try:
|
||||
yield make
|
||||
finally:
|
||||
probe.delete(f"{prefix}:e2b-capacity")
|
||||
probe.close()
|
||||
for store in stores:
|
||||
store.close()
|
||||
|
||||
|
||||
def _initialize(store) -> None:
|
||||
assert store.reconcile(
|
||||
expected_revision=store.revision(),
|
||||
remote_sandboxes={},
|
||||
complete=True,
|
||||
reservation_max_age_ms=0,
|
||||
)
|
||||
|
||||
|
||||
def _counts(store) -> tuple[int, int]:
|
||||
fields = store._redis.hkeys(store.key)
|
||||
return (
|
||||
sum(field.startswith("s:") for field in fields),
|
||||
sum(field.startswith("r:") for field in fields),
|
||||
)
|
||||
|
||||
|
||||
def test_factory_is_lazy_and_backend_errors_fail_closed() -> None:
|
||||
assert make_e2b_capacity_store(SandboxOwnershipConfig(type="memory"), hard_limit=3) is None
|
||||
store = make_e2b_capacity_store(
|
||||
SandboxOwnershipConfig(type="redis", redis_url="redis://127.0.0.1:1/0", key_prefix="test"),
|
||||
hard_limit=3,
|
||||
)
|
||||
assert store is not None and store.key == "test:e2b-capacity"
|
||||
try:
|
||||
with pytest.raises(CapacityBackendError):
|
||||
store.reserve("reservation")
|
||||
finally:
|
||||
store.close()
|
||||
|
||||
|
||||
def test_two_gateways_atomically_share_one_hash(make_store) -> None:
|
||||
gateway_a, gateway_b = make_store(), make_store()
|
||||
assert gateway_a.reserve("not-ready") is ReserveStatus.NOT_READY
|
||||
_initialize(gateway_a)
|
||||
barrier = threading.Barrier(2)
|
||||
|
||||
def reserve(args):
|
||||
store, token = args
|
||||
barrier.wait()
|
||||
return store.reserve(token)
|
||||
|
||||
tokens = ["reservation-a", "reservation-b"]
|
||||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||||
results = list(executor.map(reserve, zip((gateway_a, gateway_b), tokens)))
|
||||
|
||||
assert results.count(ReserveStatus.GRANTED) == 1
|
||||
assert results.count(ReserveStatus.FULL) == 1
|
||||
assert list(gateway_a._redis.scan_iter(f"{gateway_a.key}")) == [gateway_a.key]
|
||||
winner = tokens[results.index(ReserveStatus.GRANTED)]
|
||||
gateway_a.track("sandbox-a", reservation_token=winner)
|
||||
gateway_a.track("sandbox-a", reservation_token=winner)
|
||||
assert _counts(gateway_a) == (1, 0)
|
||||
# A successful but stale list must not release a just-tracked slot.
|
||||
assert gateway_b.reconcile(
|
||||
expected_revision=gateway_b.revision(),
|
||||
remote_sandboxes={},
|
||||
complete=True,
|
||||
reservation_max_age_ms=120_000,
|
||||
)
|
||||
assert gateway_b.reserve("stale-inventory") is ReserveStatus.FULL
|
||||
gateway_b.release("sandbox-a")
|
||||
gateway_b.release("sandbox-a")
|
||||
assert _counts(gateway_a) == (0, 0)
|
||||
|
||||
|
||||
def test_reconcile_repairs_crashes_without_erasing_concurrent_changes(make_store) -> None:
|
||||
gateway_a, gateway_b = make_store(2), make_store(2)
|
||||
_initialize(gateway_a)
|
||||
assert gateway_a.reserve("crashed-create") is ReserveStatus.GRANTED
|
||||
assert gateway_b.reconcile(
|
||||
expected_revision=gateway_b.revision(),
|
||||
remote_sandboxes={"sandbox-a": "crashed-create"},
|
||||
complete=True,
|
||||
reservation_max_age_ms=0,
|
||||
)
|
||||
stale_revision = gateway_a.revision()
|
||||
assert gateway_b.reserve("concurrent") is ReserveStatus.GRANTED
|
||||
assert not gateway_a.reconcile(
|
||||
expected_revision=stale_revision,
|
||||
remote_sandboxes={},
|
||||
complete=True,
|
||||
reservation_max_age_ms=0,
|
||||
)
|
||||
assert _counts(gateway_a) == (1, 1)
|
||||
|
||||
|
||||
def test_reconcile_keeps_incomplete_inventory_and_fresh_reservations(make_store) -> None:
|
||||
store = make_store(2)
|
||||
_initialize(store)
|
||||
store.track("sandbox-a")
|
||||
assert store.reserve("creating") is ReserveStatus.GRANTED
|
||||
|
||||
assert store.reconcile(
|
||||
expected_revision=store.revision(),
|
||||
remote_sandboxes={},
|
||||
complete=False,
|
||||
reservation_max_age_ms=0,
|
||||
)
|
||||
assert _counts(store) == (1, 1)
|
||||
assert store.reconcile(
|
||||
expected_revision=store.revision(),
|
||||
remote_sandboxes={"sandbox-a": None},
|
||||
complete=True,
|
||||
reservation_max_age_ms=120_000,
|
||||
)
|
||||
assert _counts(store) == (1, 1)
|
||||
store._redis.hset(store.key, "s:sandbox-a", "m:0")
|
||||
assert store.reconcile(
|
||||
expected_revision=store.revision(),
|
||||
remote_sandboxes={},
|
||||
complete=True,
|
||||
reservation_max_age_ms=0,
|
||||
)
|
||||
assert _counts(store) == (0, 0)
|
||||
|
||||
|
||||
def test_mismatched_hard_limits_fail_closed(make_store) -> None:
|
||||
gateway_a, gateway_b = make_store(), make_store(2)
|
||||
_initialize(gateway_a)
|
||||
with pytest.raises(CapacityBackendError, match="configuration mismatch"):
|
||||
gateway_b.revision()
|
||||
with pytest.raises(CapacityBackendError, match="configuration mismatch"):
|
||||
gateway_b.reserve("reservation")
|
||||
@ -10,6 +10,7 @@ import threading
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
@ -17,6 +18,10 @@ from unittest.mock import MagicMock
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from deerflow.community.e2b_sandbox.capacity import (
|
||||
CapacityBackendError,
|
||||
ReserveStatus,
|
||||
)
|
||||
from deerflow.config.paths import Paths
|
||||
from deerflow.config.sandbox_config import SandboxConfig
|
||||
from deerflow.sandbox.exceptions import SandboxCapacityExceededError
|
||||
@ -253,7 +258,12 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_poli
|
||||
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._ownership_config = SimpleNamespace(
|
||||
renewal_interval_seconds=60.0,
|
||||
ttl_multiplier=4.0,
|
||||
key_prefix="deerflow:test",
|
||||
)
|
||||
provider._deployment_capacity = None
|
||||
provider._owned_sandbox_ids = set()
|
||||
provider._acquire_inflight = set()
|
||||
provider._orphan_first_seen = {}
|
||||
@ -282,12 +292,122 @@ def _make_provider(*, replicas: int = 3, idle_timeout: int = 1800, overflow_poli
|
||||
return provider
|
||||
|
||||
|
||||
def _install_shared_deployment_capacity(
|
||||
*providers,
|
||||
reserve_results: list[ReserveStatus] | None = None,
|
||||
) -> MagicMock:
|
||||
store = MagicMock()
|
||||
store.key = "deerflow:test:e2b-capacity"
|
||||
store.revision.return_value = 0
|
||||
store.reserve.return_value = ReserveStatus.GRANTED
|
||||
store.reconcile.return_value = True
|
||||
if reserve_results is not None:
|
||||
store.reserve.side_effect = reserve_results
|
||||
for provider in providers:
|
||||
provider._deployment_capacity = store
|
||||
return store
|
||||
|
||||
|
||||
def _install_fake_sdk(monkeypatch, provider) -> FakeSandboxClass:
|
||||
fake_cls = FakeSandboxClass()
|
||||
monkeypatch.setattr(provider, "_get_sandbox_cls", lambda: fake_cls)
|
||||
return fake_cls
|
||||
|
||||
|
||||
def _write_skill(root: Path, name: str) -> None:
|
||||
target = root / name / "SKILL.md"
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
target.write_text(f"---\nname: {name}\ndescription: test\n---\n", encoding="utf-8")
|
||||
|
||||
|
||||
def test_apply_mounts_uploads_only_enabled_skill_projection(monkeypatch, tmp_path):
|
||||
from deerflow.config.extensions_config import ExtensionsConfig, SkillStateConfig
|
||||
|
||||
mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider")
|
||||
paths = Paths(base_dir=tmp_path)
|
||||
skills_root = tmp_path / "skills"
|
||||
_write_skill(skills_root / "public", "enabled-skill")
|
||||
_write_skill(skills_root / "public", "disabled-skill")
|
||||
(skills_root / "custom").mkdir()
|
||||
_write_skill(paths.integration_skills_dir() / "lark-cli", "enabled-integration")
|
||||
_write_skill(paths.integration_skills_dir() / "lark-cli", "disabled-integration")
|
||||
user_skills_root = paths.user_skills_dir("user-1")
|
||||
user_skills_root.mkdir(parents=True, exist_ok=True)
|
||||
(user_skills_root / "_skill_states.json").write_text(
|
||||
json.dumps({"disabled-integration": {"enabled": False}}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
extensions = ExtensionsConfig(skills={"disabled-skill": SkillStateConfig(enabled=False)})
|
||||
config = SimpleNamespace(
|
||||
skills=SimpleNamespace(
|
||||
get_skills_path=lambda: skills_root,
|
||||
container_path="/mnt/skills",
|
||||
use="deerflow.skills.storage.local_skill_storage:LocalSkillStorage",
|
||||
)
|
||||
)
|
||||
monkeypatch.setattr(mod, "get_app_config", lambda: config)
|
||||
monkeypatch.setattr("deerflow.config.paths.get_paths", lambda: paths)
|
||||
monkeypatch.setattr("deerflow.config.extensions_config.ExtensionsConfig.from_file", lambda *_args, **_kwargs: extensions)
|
||||
monkeypatch.setattr("deerflow.config.extensions_config.get_extensions_config", lambda: extensions)
|
||||
|
||||
provider = _make_provider()
|
||||
client = FakeClient()
|
||||
provider._apply_mounts(client, user_id="user-1")
|
||||
|
||||
uploaded_paths = {path for path, _content in client.files.write_calls}
|
||||
assert "/mnt/skills/public/enabled-skill/SKILL.md" in uploaded_paths
|
||||
assert "/mnt/skills/public/disabled-skill/SKILL.md" not in uploaded_paths
|
||||
assert "/mnt/skills/integrations/lark-cli/enabled-integration/SKILL.md" in uploaded_paths
|
||||
assert "/mnt/skills/integrations/lark-cli/disabled-integration/SKILL.md" not in uploaded_paths
|
||||
|
||||
|
||||
def test_skill_projection_mounts_swallows_projection_failure(monkeypatch):
|
||||
"""``_skill_projection_mounts`` must not raise — a projection failure used
|
||||
to propagate out of ``_apply_mounts`` before the configured-mounts loop
|
||||
ran, dropping the operator's own configured mounts as collateral (#4107
|
||||
review)."""
|
||||
mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider")
|
||||
|
||||
config = SimpleNamespace(skills=SimpleNamespace(container_path="/mnt/skills"))
|
||||
monkeypatch.setattr(mod, "get_app_config", lambda: config)
|
||||
monkeypatch.setattr(
|
||||
"deerflow.skills.projection.ensure_skill_projections",
|
||||
lambda storage: (_ for _ in ()).throw(RuntimeError("simulated projection failure")),
|
||||
)
|
||||
|
||||
provider = _make_provider()
|
||||
|
||||
assert provider._skill_projection_mounts("user-1") == []
|
||||
|
||||
|
||||
def test_apply_mounts_keeps_configured_mounts_when_projection_fails(monkeypatch, tmp_path):
|
||||
"""End-to-end: a skills-projection failure must not drop the operator's
|
||||
own configured mounts too — the two mount sources are independent."""
|
||||
mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox_provider")
|
||||
|
||||
host_dir = tmp_path / "operator-mount"
|
||||
host_dir.mkdir()
|
||||
(host_dir / "notes.txt").write_text("hello", encoding="utf-8")
|
||||
|
||||
config = SimpleNamespace(skills=SimpleNamespace(container_path="/mnt/skills"))
|
||||
monkeypatch.setattr(mod, "get_app_config", lambda: config)
|
||||
monkeypatch.setattr(
|
||||
"deerflow.skills.projection.ensure_skill_projections",
|
||||
lambda storage: (_ for _ in ()).throw(RuntimeError("simulated projection failure")),
|
||||
)
|
||||
|
||||
provider = _make_provider()
|
||||
provider._config["mounts"] = [
|
||||
SimpleNamespace(host_path=str(host_dir), container_path="/mnt/operator", read_only=True),
|
||||
]
|
||||
|
||||
client = FakeClient()
|
||||
provider._apply_mounts(client, user_id="user-1")
|
||||
|
||||
uploaded_paths = {path for path, _content in client.files.write_calls}
|
||||
assert "/mnt/operator/notes.txt" in uploaded_paths
|
||||
|
||||
|
||||
def _make_sandbox(client: FakeClient, *, sandbox_id: str | None = None) -> Any:
|
||||
mod = importlib.import_module("deerflow.community.e2b_sandbox.e2b_sandbox")
|
||||
return mod.E2BSandbox(
|
||||
@ -1081,11 +1201,17 @@ def test_bootstrap_failure_does_not_kill_without_destroy_lease(monkeypatch):
|
||||
|
||||
def test_kill_client_returns_exception_without_raising():
|
||||
p = _make_provider()
|
||||
client = FakeClient()
|
||||
store = _install_shared_deployment_capacity(p)
|
||||
failed_client = FakeClient()
|
||||
error = RuntimeError("already gone")
|
||||
client.kill = MagicMock(side_effect=error)
|
||||
failed_client.kill = MagicMock(side_effect=error)
|
||||
|
||||
assert p._kill_client(client) is error
|
||||
assert p._kill_client(failed_client) is error
|
||||
store.release.assert_not_called()
|
||||
|
||||
client = FakeClient()
|
||||
assert p._kill_client(client) is None
|
||||
store.release.assert_called_once_with(client.sandbox_id)
|
||||
|
||||
|
||||
def test_kill_client_reports_uncertain_cleanup_without_callable_kill():
|
||||
@ -2107,9 +2233,132 @@ def test_grep_single_file_path_with_matching_glob():
|
||||
assert truncated is False
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Capacity enforcement tests (#4339)
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_deployment_capacity_reserves_commits_and_rejects_globally(monkeypatch) -> None:
|
||||
gateway_a = _make_provider(replicas=1, overflow_policy="reject")
|
||||
gateway_b = _make_provider(replicas=1, overflow_policy="reject")
|
||||
store = _install_shared_deployment_capacity(
|
||||
gateway_a,
|
||||
gateway_b,
|
||||
reserve_results=[ReserveStatus.GRANTED, ReserveStatus.FULL],
|
||||
)
|
||||
sdk_a = _install_fake_sdk(monkeypatch, gateway_a)
|
||||
sdk_b = FakeSandboxClass()
|
||||
monkeypatch.setattr(gateway_b, "_get_sandbox_cls", lambda: sdk_b)
|
||||
|
||||
sandbox_id = gateway_a.acquire("thread-a", user_id="user-a")
|
||||
with pytest.raises(SandboxCapacityExceededError):
|
||||
gateway_b.acquire("thread-b", user_id="user-b")
|
||||
|
||||
metadata = sdk_a.create_calls[0]["metadata"]
|
||||
assert metadata["deer_flow_capacity_ledger"] == store.key
|
||||
assert metadata["deer_flow_capacity_reservation"]
|
||||
store.track.assert_called_once_with(
|
||||
sandbox_id,
|
||||
reservation_token=metadata["deer_flow_capacity_reservation"],
|
||||
)
|
||||
assert len(sdk_a.create_calls) == 1
|
||||
assert sdk_b.create_calls == []
|
||||
assert store.reserve.call_count == 2
|
||||
|
||||
|
||||
def test_ambiguous_create_failure_retains_deployment_reservation(monkeypatch) -> None:
|
||||
provider = _make_provider(replicas=1, overflow_policy="reject")
|
||||
store = _install_shared_deployment_capacity(provider)
|
||||
sdk = _install_fake_sdk(monkeypatch, provider)
|
||||
sdk.create_factory = lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("control-plane timeout"))
|
||||
|
||||
with pytest.raises(RuntimeError, match="control-plane timeout"):
|
||||
provider.acquire("thread-a", user_id="user-a")
|
||||
|
||||
assert provider._reserved_slots == 0
|
||||
store.reserve.assert_called_once()
|
||||
store.track.assert_not_called()
|
||||
store.release.assert_not_called()
|
||||
|
||||
|
||||
def test_discovery_uses_sdk_query_and_tracks_without_reserving(monkeypatch) -> None:
|
||||
provider = _make_provider(replicas=1, overflow_policy="reject")
|
||||
store = _install_shared_deployment_capacity(provider)
|
||||
sdk = _install_fake_sdk(monkeypatch, provider)
|
||||
entry = SimpleNamespace(
|
||||
sandbox_id="sandbox-existing",
|
||||
metadata={
|
||||
"deer_flow_provider": "e2b_sandbox_provider",
|
||||
"deer_flow_user": "user-a",
|
||||
"deer_flow_thread": "thread-a",
|
||||
"deer_flow_capacity_ledger": store.key,
|
||||
},
|
||||
)
|
||||
expected_query = {key: entry.metadata[key] for key in ("deer_flow_provider", "deer_flow_user", "deer_flow_thread")}
|
||||
sdk.list_return = SimpleNamespace(
|
||||
has_next=False,
|
||||
next_items=lambda: [entry] if sdk.list_calls[-1]["query"].metadata == expected_query else [],
|
||||
)
|
||||
assert provider.acquire("thread-a", user_id="user-a") == entry.sandbox_id
|
||||
assert sdk.create_calls == []
|
||||
store.reserve.assert_not_called()
|
||||
store.track.assert_called_once_with(entry.sandbox_id, reservation_token=None)
|
||||
|
||||
|
||||
def test_reconciliation_repairs_crash_and_uses_safe_reservation_age(monkeypatch) -> None:
|
||||
provider = _make_provider(replicas=1, overflow_policy="reject")
|
||||
provider._ownership_config.renewal_interval_seconds = 1.0
|
||||
provider._ownership_config.ttl_multiplier = 2.0
|
||||
provider._config["reconciliation_grace_seconds"] = 0.0
|
||||
store = _install_shared_deployment_capacity(provider)
|
||||
sdk = _install_fake_sdk(monkeypatch, provider)
|
||||
sdk.list_return = [
|
||||
{
|
||||
"sandbox_id": "sandbox-existing",
|
||||
"metadata": {
|
||||
"deer_flow_provider": "e2b_sandbox_provider",
|
||||
"deer_flow_capacity_ledger": store.key,
|
||||
"deer_flow_capacity_reservation": "reservation-crashed",
|
||||
},
|
||||
},
|
||||
{
|
||||
"sandbox_id": "sandbox-other-deployment",
|
||||
"metadata": {
|
||||
"deer_flow_provider": "e2b_sandbox_provider",
|
||||
"deer_flow_capacity_ledger": "deerflow:other:e2b-capacity",
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
stats = provider._reconcile_remote_sandboxes(now=100.0)
|
||||
args = store.reconcile.call_args.kwargs
|
||||
assert stats.discovered == 1
|
||||
assert args["remote_sandboxes"] == {"sandbox-existing": "reservation-crashed"}
|
||||
assert args["complete"] is True
|
||||
assert args["reservation_max_age_ms"] == 120_000
|
||||
|
||||
|
||||
def test_failed_inventory_and_redis_error_both_prevent_create(monkeypatch) -> None:
|
||||
provider = _make_provider(replicas=1, overflow_policy="reject")
|
||||
store = _install_shared_deployment_capacity(
|
||||
provider,
|
||||
reserve_results=[ReserveStatus.NOT_READY],
|
||||
)
|
||||
sdk = _install_fake_sdk(monkeypatch, provider)
|
||||
sdk.list = MagicMock(side_effect=RuntimeError("E2B unavailable"))
|
||||
|
||||
provider._reconcile_remote_sandboxes(now=100.0)
|
||||
|
||||
reconcile_args = store.reconcile.call_args.kwargs
|
||||
assert reconcile_args["complete"] is False
|
||||
assert reconcile_args["remote_sandboxes"] == {}
|
||||
with pytest.raises(SandboxCapacityExceededError):
|
||||
provider._create_sandbox("thread-a", user_id="user-a")
|
||||
store.reserve.side_effect = CapacityBackendError("Redis unavailable")
|
||||
with pytest.raises(SandboxCapacityExceededError) as error:
|
||||
provider._create_sandbox("thread-a", user_id="user-a")
|
||||
|
||||
assert error.value.reason == "capacity_backend"
|
||||
assert sdk.create_calls == []
|
||||
assert provider._reserved_slots == 0
|
||||
|
||||
|
||||
def test_capacity_reject_policy_raises_when_full(monkeypatch):
|
||||
@ -2146,7 +2395,11 @@ def test_capacity_reject_frees_slot_on_release(monkeypatch):
|
||||
|
||||
def test_capacity_reject_evicts_other_thread_warm_entry_before_create(monkeypatch):
|
||||
"""Reject policy can evict one warm VM before it rejects new capacity."""
|
||||
p = _make_provider(replicas=1, overflow_policy="reject")
|
||||
p = _make_provider(replicas=3, overflow_policy="reject")
|
||||
store = _install_shared_deployment_capacity(
|
||||
p,
|
||||
reserve_results=[ReserveStatus.GRANTED, ReserveStatus.FULL, ReserveStatus.GRANTED],
|
||||
)
|
||||
fake_cls = _install_fake_sdk(monkeypatch, p)
|
||||
|
||||
sid1 = p.acquire("t1", user_id="u1")
|
||||
@ -2158,6 +2411,7 @@ def test_capacity_reject_evicts_other_thread_warm_entry_before_create(monkeypatc
|
||||
assert sid2 != sid1
|
||||
assert len(p._warm_pool) == 0
|
||||
assert len(fake_cls.create_calls) == 2
|
||||
store.release.assert_called_once_with(sid1)
|
||||
|
||||
|
||||
def test_capacity_wait_policy_times_out(monkeypatch):
|
||||
@ -3330,8 +3584,14 @@ def test_shutdown_during_discovery_does_not_kill_unowned_vm(monkeypatch):
|
||||
allow_commit = threading.Event()
|
||||
reserve_capacity = p._reserve_capacity
|
||||
|
||||
def pause_after_reserve(thread_id, user_id, *, remote_id=None, remote_owned=True):
|
||||
reserve_capacity(
|
||||
def pause_after_reserve(
|
||||
thread_id,
|
||||
user_id,
|
||||
*,
|
||||
remote_id=None,
|
||||
remote_owned=True,
|
||||
):
|
||||
reservation = reserve_capacity(
|
||||
thread_id,
|
||||
user_id,
|
||||
remote_id=remote_id,
|
||||
@ -3339,6 +3599,7 @@ def test_shutdown_during_discovery_does_not_kill_unowned_vm(monkeypatch):
|
||||
)
|
||||
reserved.set()
|
||||
assert allow_commit.wait(timeout=2)
|
||||
return reservation
|
||||
|
||||
monkeypatch.setattr(p, "_reserve_capacity", pause_after_reserve)
|
||||
result: list[str | None] = []
|
||||
|
||||
@ -177,7 +177,9 @@ def test_feishu_receive_file_syncs_sandbox_with_explicit_user_id(tmp_path, monke
|
||||
channel._api_client.im.v1.message_resource.get.return_value = response
|
||||
|
||||
provider = MagicMock()
|
||||
provider.acquire.return_value = "aio-1"
|
||||
provider.uses_thread_data_mounts = False
|
||||
provider.acquire.side_effect = AssertionError("receive_file must use acquire_async")
|
||||
provider.acquire_async = AsyncMock(return_value="aio-1")
|
||||
sandbox = MagicMock()
|
||||
provider.get.return_value = sandbox
|
||||
|
||||
@ -189,7 +191,7 @@ def test_feishu_receive_file_syncs_sandbox_with_explicit_user_id(tmp_path, monke
|
||||
|
||||
assert virtual_path == "/mnt/user-data/uploads/report.md"
|
||||
assert (tmp_path / "users" / "ou-user" / "threads" / "thread-1" / "user-data" / "uploads" / "report.md").read_bytes() == b"file-bytes"
|
||||
provider.acquire.assert_called_once_with("thread-1", user_id="ou-user")
|
||||
provider.acquire_async.assert_awaited_once_with("thread-1", user_id="ou-user")
|
||||
sandbox.update_file.assert_called_once_with("/mnt/user-data/uploads/report.md", b"file-bytes")
|
||||
|
||||
_run(go())
|
||||
|
||||
@ -1,7 +1,8 @@
|
||||
"""Dual-mode (full/delta) parity for the gateway thread-state endpoints.
|
||||
|
||||
Drives ``GET /api/threads/{id}``, ``GET /api/threads/{id}/state`` and
|
||||
``POST /api/threads/{id}/history`` through the real route stack
|
||||
Drives ``GET /api/threads/{id}``, ``GET /api/threads/{id}/state``,
|
||||
``POST /api/threads/{id}/history``, and the context-usage checkpoint reader
|
||||
through the real materialization stack
|
||||
(``build_thread_checkpoint_state_accessor`` -> factory-built graph ->
|
||||
``CheckpointStateAccessor``) against a real ``InMemorySaver``, once per
|
||||
checkpoint channel mode, and asserts the wire responses are identical apart
|
||||
@ -24,6 +25,7 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.store.memory import InMemoryStore
|
||||
|
||||
from app.gateway import context_usage
|
||||
from app.gateway import services as gateway_services
|
||||
from app.gateway.routers import threads
|
||||
from deerflow.agents.thread_state import get_thread_state_schema
|
||||
@ -113,6 +115,46 @@ def test_thread_state_endpoints_are_mode_invariant(_stub_app_config, monkeypatch
|
||||
assert any(full["history_messages"]), "expected history snapshots with messages"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["full", "delta"])
|
||||
def test_context_usage_reads_materialized_messages_in_both_modes(
|
||||
mode: str,
|
||||
_stub_app_config,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Context usage must not read raw delta ``channel_values``."""
|
||||
app = make_authed_test_app()
|
||||
store = InMemoryStore()
|
||||
checkpointer = InMemorySaver()
|
||||
app.state.store = store
|
||||
app.state.checkpointer = checkpointer
|
||||
app.state.thread_store = MemoryThreadMetaStore(store)
|
||||
app.state.checkpoint_channel_mode = mode
|
||||
app.state.run_event_store = SimpleNamespace()
|
||||
|
||||
graph = _build_reply_graph(mode, checkpointer)
|
||||
monkeypatch.setattr(
|
||||
gateway_services,
|
||||
"resolve_agent_factory",
|
||||
lambda assistant_id=None: lambda config: graph,
|
||||
)
|
||||
|
||||
config: dict[str, Any] = {"configurable": {"thread_id": _THREAD_ID}}
|
||||
inject_checkpoint_mode(config, mode)
|
||||
for i in range(2):
|
||||
asyncio.run(graph.ainvoke({"messages": [HumanMessage(content=f"question-{i}", id=f"h{i}")]}, config))
|
||||
|
||||
request = SimpleNamespace(app=app)
|
||||
accessor, read_config = asyncio.run(gateway_services.build_thread_checkpoint_state_accessor(request, thread_id=_THREAD_ID))
|
||||
messages = asyncio.run(context_usage._load_checkpoint_messages(accessor, read_config))
|
||||
|
||||
assert [(message.type, message.content, message.id) for message in messages] == [
|
||||
("human", "question-0", "h0"),
|
||||
("ai", "answer-1", "a1"),
|
||||
("human", "question-1", "h1"),
|
||||
("ai", "answer-3", "a3"),
|
||||
]
|
||||
|
||||
|
||||
def test_full_mode_gateway_rejects_delta_thread_with_409(_stub_app_config, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Fail-closed gate at the HTTP boundary, against a real checkpointer.
|
||||
|
||||
|
||||
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