Compare commits

...
99 changed files with 6459 additions and 950 deletions
@@ -0,0 +1,127 @@
# Nullable Home Snapshot and Backend Default Home — Implementation Note
## Result
Home Snapshot is now an optional immutable checkpoint rather than a prerequisite
for creating an Agent.
An Agent config whose `home_snapshot_id` is `NULL` creates its first Execution
Binding from the selected backend's deployment-default Home. No logical
`AgentHomeSnapshot` row or physical immutable snapshot is created for that
default. A real snapshot is still created when Build Draft Apply checkpoints the
participant's current Materialized Home.
## Major additions
### Nullable API data flow
- `AgentConfigDraft.home_snapshot_id`,
`AgentConfigSnapshot.home_snapshot_id`, and
`AgentWorkspaceBinding.base_home_snapshot_id` are nullable in the ORM and
database schema.
- Agent creation, backing-Agent creation, Composer recovery, Workflow-only
creation, DSL import, and clone paths seed config versions with
`home_snapshot_id=None` without contacting Dify Agent.
- Publish and config-version propagation preserve `None`. Explicit logical
snapshot ids are still resolved against the tenant, Agent owner, and active
ledger state.
- `AgentWorkspaceService` passes `home_snapshot_ref=None` to Dify Agent when the
config has no checkpoint. It does not query the snapshot ledger or manufacture
a placeholder id.
- Agent App, Workflow Agent, preview, and Build Draft callers carry the nullable
generation through their session and Binding paths.
- Build Draft Apply accepts a Binding whose base snapshot is `None`, creates an
immutable checkpoint from that exact Binding, records the resulting
`AgentHomeSnapshot`, and assigns its logical id to the normal draft.
### Backend-default Home contract
- `CreateExecutionBindingRequest.home_snapshot_ref` and
`ExecutionBindingCreateSpec.home_snapshot_ref` accept `None`.
- The backend protocol documents that `None` means an independent mutable
deployment-default Home, while a non-empty ref must be materialized exactly
and must never fall back.
- Local creates an empty per-Binding Home for `None`; explicit snapshots are
validated before lease allocation and before Workspace/Home writes.
- E2B creates the Sandbox from the configured E2B template for `None`, and from
the explicit snapshot ref otherwise. The template is owned by the Execution
Binding backend rather than the snapshot backend.
- Enterprise creates a default Sandbox through `POST /v1/sandboxes`, initializes
the canonical Home and empty Workspace through shellctl, and returns the
Sandbox id as the Binding and Workspace refs. Explicit immutable snapshots and
shared Workspace attachment remain fail-fast unsupported.
### Linear schema convergence
- Historical revision `2f39536b3feb` now adds the draft, snapshot, and
intermediate runtime-session fields as nullable, so a database with existing
rows can reach later revisions.
- Historical revision `f6e4c5686857` now creates the Binding base snapshot field
as nullable and recreates the intermediate field as nullable on downgrade.
- Generated linear revision `e4708db55c1d` alters the three current fields to
nullable for databases that already executed the former NOT NULL revisions.
On a fresh corrected chain, the repeated alters are intentional schema
convergence.
- No rows are backfilled and no migration calls an external runtime backend.
## Major deletions
- Removed `AgentHomeSnapshotService.create_initial()`.
- Removed `InitializeHomeSnapshotSpec`,
`InitializeHomeSnapshotRequest`, `HomeSnapshotBackend.initialize()`, the
`/home-snapshots/initialize` route and server method, sync/async client
methods, backend implementations, exports, fixtures, and tests.
- Removed the temporary E2B initialization-Sandbox flow.
- Removed unused Enterprise HomeSnapshot Gateway configuration and unreachable
exception branches left behind by the old initialization design.
- Removed documentation that described an initial/baseline snapshot or a
temporary Home initialization resource.
## Review-driven refinements
- Local validates an explicit snapshot before allocating an operation lease or
mutating Home/Workspace resources.
- Tests independently cover owner routing and nullable generation instead of
coupling the two parameter dimensions.
- The migration suite includes both a fresh pre-`2f39536b3feb` upgrade and an
old-`f6e4c5686857` NOT NULL convergence case.
- Publish, full Build Draft Apply, API explicit-snapshot fail-fast, and Local
validation-before-write behavior have direct regression coverage.
- Dead snapshot fixtures, duplicate assertions, incidental call-order
assertions, unused backend configuration, and unreachable guards were
removed.
## Differences from the proposal
There is no intentional product or architecture deviation from the proposal.
The implementation keeps the proposed limits:
- `NULL` does not create a logical or physical snapshot.
- Invalid explicit snapshot ids and refs fail fast without fallback.
- No new lifecycle states, compatibility layer, sentinel ids, or snapshot
backfill were introduced.
- Enterprise supports the required default-Home Binding path but does not claim
immutable snapshot support.
Generated API/frontend artifacts were not changed because the removed
Home-Snapshot initialization surface is private to the Python Dify Agent client
and server and has no repository generation target.
## Verification
- Alembic: one linear head, `e4708db55c1d`.
- Final focused API regression suite: `231 passed`.
- Final focused Dify Agent protocol/backend suite: `103 passed`.
- Dify Agent backend/profile review fixes: `64 passed`; `make check` passed.
- Focused changed-production-module type checks passed during review.
- `git diff --check` passed.
Not run locally:
- CI-only PostgreSQL/MySQL migration jobs.
- Live E2B template integration.
- Live Enterprise Gateway integration.
The full Dify Agent typecheck still reports repository baseline and missing
optional-dependency errors outside the changed production modules; those were
not expanded into this implementation.
@@ -0,0 +1,453 @@
# Nullable Home Snapshot and Backend Default Home
## 1. 背景
当前 Agent 创建流程会立即调用 Dify Agent 创建一个空的 Home Snapshot,并把其逻辑
`home_snapshot_id` 写入首个 `AgentConfigSnapshot`。这使一个尚未运行、也没有 Home
内容需要保存的 Agent 在创建时就依赖外部 runtime backend,并产生一份没有产品价值的
物理快照。
同时,现有 Alembic revision `2f39536b3feb` 向已有的
`agent_config_drafts``agent_config_snapshots`
`agent_runtime_sessions` 直接添加 `nullable=False`
`home_snapshot_id`。只要这些表已有数据,数据库就无法为历史行构造合法值,migration
会在到达后续 revision 之前失败。
本方案让 Home Snapshot 成为真正按需产生的不可变 checkpoint
- 创建 Agent 和配置版本时可以没有 Home Snapshot
- 新建 Execution Binding 时,如果配置没有明确的 Home Snapshot,由物理后端物化其部署默认
Home
- 只有需要保存实际 Home 内容时,例如 Build Draft Apply,才创建
`agent_home_snapshots` 记录和物理快照;
- Local、E2B 和 Enterprise 都必须支持“无明确 Home Snapshot”这一 Binding 创建方式。
## 2. 目标
1. `AgentConfigDraft.home_snapshot_id`
`AgentConfigSnapshot.home_snapshot_id` 支持 `NULL`
2. `AgentWorkspaceBinding.base_home_snapshot_id` 支持 `NULL`,准确记录该参与者是从 backend
default Home 物化,而不是从某个逻辑 Home Snapshot 物化。
3. Agent 创建、导入、复制和 Workflow-only Agent 创建不再创建初始 Home Snapshot。
4. Execution Binding 协议显式支持 `home_snapshot_ref = null`
5. 所有 runtime backend 都实现无 snapshot 的默认 Home 物化。
6. Build Draft Apply 保持现有 checkpoint 语义:从确切 Binding 创建不可变 Home Snapshot
并把新逻辑 id 写入 normal draft。
7. 修复 migration 链,使带有历史 Agent 数据的数据库可以从旧 revision 升级,并使已经执行过
当前 head 的数据库收敛到 nullable schema。
8. Alembic 继续保持单一线性 head,不创建 merge revision。
## 3. 非目标
- 不创建“默认 Home Snapshot”逻辑记录或全局默认 snapshot id。
- 不用空字符串、固定 UUID 或其他 sentinel 代替 `NULL`
- 不在首次运行时懒创建并持久化 `agent_home_snapshots` 记录。
- 不为无效、已退休或不属于当前 Agent 的非空 snapshot 自动回退到默认 Home。
- 不改变 RuntimeLease、Workspace、Binding 的既有所有权和 retire/collect 模型。
- 不解决跨 backend 的 snapshot 可移植性。
- 不在本方案中为 Enterprise Gateway 增加不可变 snapshot/checkpoint 能力;本方案只要求
Enterprise 支持无 snapshot 的默认 Home Binding。
## 4. 核心语义
### 4.1 `NULL` 的含义
`home_snapshot_id = NULL` 表示该配置没有绑定任何逻辑 Home Snapshot。它不指向
`agent_home_snapshots`,也不产生可 retire 或 collect 的 snapshot 资源。
创建新 Binding 时:
- `home_snapshot_id is None`Dify API 向 Dify Agent 传递
`home_snapshot_ref=None`backend 为该 Binding 创建独立、可变的部署默认 Home;
- `home_snapshot_id is not None`Dify API 必须解析到当前 tenant、Agent 所拥有且状态为
`ACTIVE``AgentHomeSnapshot`,并传递其准确 `snapshot_ref`
- 非空 id 无法解析时直接失败,不回退到默认 Home。
### 4.2 默认 Home 不是不可变 snapshot
默认 Home 只在创建 Binding 时物化。Dify API 不为它写
`agent_home_snapshots`Dify Agent 也不返回隐式 snapshot ref。
默认 Home 由部署配置决定,因此同一配置版本在不同时刻创建的新 Binding 可能看到不同版本的
backend default Home。这个行为是 `home_snapshot_id=NULL` 的明确语义。已经存在的 Binding
继续使用自己的 Materialized Home,不受后续 backend template 变化影响。
### 4.3 Binding generation
`AgentWorkspaceBinding.base_home_snapshot_id` 保存物化来源:
- 非空:来自该逻辑 Home Snapshot
- `NULL`:来自 backend default Home。
generation 校验继续比较
`(base_home_snapshot_id, agent_config_version_id, agent_config_version_kind)`
`None == None` 是有效匹配;`None` 与任意具体 snapshot id 不匹配。Binding 不会因为配置后来
获得 snapshot 而被隐式迁移或重新物化。
## 5. 数据模型
| 表/模型 | 字段 | 新类型 | 含义 |
| --- | --- | --- | --- |
| `agent_config_drafts` / `AgentConfigDraft` | `home_snapshot_id` | `str \| None`, nullable | Draft 可使用 backend default Home |
| `agent_config_snapshots` / `AgentConfigSnapshot` | `home_snapshot_id` | `str \| None`, nullable | 不可变配置版本可以不绑定 Home Snapshot |
| `agent_workspace_bindings` / `AgentWorkspaceBinding` | `base_home_snapshot_id` | `str \| None`, nullable | Binding 的 Home 来源;`NULL` 表示 backend default |
`agent_home_snapshots` 不变:
- 只有实际存在的不可变 Home checkpoint 才有记录;
- `snapshot_ref` 仍为非空;
- tenant、Agent、状态、retire 和 collect 语义不变。
不新增表、状态、外键或默认 snapshot 配置。
## 6. 生命周期与产品行为
### 6.1 Agent 创建、导入和复制
以下路径不再调用 `AgentHomeSnapshotService.create_initial()`
- roster Agent 创建;
- Agent App backing Agent 创建;
- Composer 中补建 backing Agent
- Workflow-only inline Agent 创建;
- DSL 导入或复制时创建目标 Agent/config snapshot。
这些路径仍然创建首个 `AgentConfigSnapshot` 和 revision,但
`home_snapshot_id=None`。创建 Agent 因而不再依赖 Dify Agent、E2B、Local shellctl 或
Enterprise Gateway。
从已有配置创建 Draft 时,原样继承其可空 `home_snapshot_id`。复制到另一个 Agent 时不共享
来源 Agent 所拥有的 Home Snapshot;目标配置使用 `None`,直到目标 Agent 自己产生 checkpoint。
### 6.2 Publish
Publish 允许 Draft 的 `home_snapshot_id``None`
- `None` 可以原样写入新的 `AgentConfigSnapshot`
- 非空时仍校验 snapshot 归属和 `ACTIVE` 状态;
- `validate_home_snapshot_binding()``None` 表示“无需校验逻辑 snapshot”,对非空值保持严格
校验;
- draft 与 active version 的一致性判断直接比较可空值。
Publish 不创建新的 Home Snapshot,因为 Home 内容只在 Build 期间发生变化。
### 6.3 首次运行
Agent App conversation、preview、Build Draft 和 Workflow Agent node 在首次需要参与者时调用
`AgentWorkspaceService.create_binding()`
1. 配置中的 `home_snapshot_id` 原样作为可空 `base_home_snapshot_id`
2. 非空时解析 ledger row,空值时跳过 ledger 查询;
3. 发送可空 `home_snapshot_ref` 给 Dify Agent
4. backend 创建 Materialized Home 和 Workspace
5. Dify API 持久化 Binding,其中 `base_home_snapshot_id` 可以为 `NULL`
后续运行只 acquire 已有 Binding,不重新解释默认 Home,也不隐式创建替代 Binding。
### 6.4 Build Draft Apply
Build Draft 可以从 `home_snapshot_id=None` 开始。其 Binding 的 Materialized Home 由 backend
default 初始化,Build 期间的所有 Home 修改都发生在该 Binding 中。
Apply 保持现有流程:
1. 通过 `agent_workspace_binding_id` 解析确切 Build participant
2. generation 校验允许 `base_home_snapshot_id=None` 与 Draft 的
`home_snapshot_id=None` 匹配;
3. `create_home_snapshot_from_binding` 从该 RuntimeLease 的当前 Home 创建物理 checkpoint
4. Dify API 新建 `AgentHomeSnapshot` ledger row
5. normal draft 的 `home_snapshot_id` 更新为新 ledger id
6. Build Binding 按现有流程 retire,异步 collect。
因此,第一次真正需要不可变 Home 的动作仍然是 Build Draft Apply,而不是 Agent 创建或首次运行。
### 6.5 Retirement 和 collection
nullable 不引入新的资源:
- `NULL` 不参与 snapshot retirement
- snapshot reference 查询自然忽略 `NULL`
- 非空 snapshot 的 retire/collect 逻辑不变;
- Binding 和 Workspace 的 retire/collect 逻辑不变。
## 7. Dify API 服务与调用链
### 7.1 `AgentHomeSnapshotService`
删除 baseline initialization 能力:
- 删除 `create_initial()`
- 删除仅为该方法服务的 `InitializeHomeSnapshotRequest` 调用;
- 保留 `create_for_build_apply()`、retire、collect 和 delete
- `validate_home_snapshot_binding()` 接受 `str | None`,仅对非空值查询并校验 ledger。
不保留未使用的 initialization fallback。
### 7.2 `AgentWorkspaceService`
`create_binding()``base_home_snapshot_id` 改为 `str | None`
- 非空时解析 `AgentHomeSnapshot.snapshot_ref`
- `None` 时设置 `home_snapshot_ref=None`
- 非空但不存在、归属错误或已退休时继续抛出
`AgentWorkspaceNotFoundError`
- 创建的 `AgentWorkspaceBinding.base_home_snapshot_id` 原样保存可空值。
`validate_binding_generation()` 的对应参数改为 `str | None`,继续做直接相等比较。
### 7.3 配置与 runtime caller
以下类型和 helper 接受并传递 `str | None`
- Composer `_create_config_version()``_update_current_version()`
`_agent_soul_matches_active_config()`
- Agent App `AgentAppRunner``AgentAppSessionScope`
`AgentAppWorkspaceStore`
- Workflow Agent `load_or_create_node_execution_session()`
- 其他以 config draft/snapshot 为来源的 resolver 和 request builder。
调用链不得在中途把 `None` 替换成临时 UUID、空字符串或新建 snapshot。
## 8. Dify Agent 协议
### 8.1 Execution Binding request
`CreateExecutionBindingRequest.home_snapshot_ref`
`ExecutionBindingCreateSpec.home_snapshot_ref` 改为 `str | None`,默认 `None`
- 请求字段缺失或 JSON `null` 都表示 backend default Home
- 非空字符串表示必须精确物化该 backend snapshot ref
- 空字符串继续作为非法输入拒绝。
`ExecutionBindingBackend.create_binding()` 的 protocol 文档明确要求:
1. `None` 时创建独立、可变、可供当前 Binding 使用的默认 Home
2. 非空时必须从指定 snapshot 初始化;
3. 指定 snapshot 不可用时失败,不得回退;
4. backend default Home 不得隐式创建不可变 snapshot;
5. 失败返回前清理本次创建的局部物理资源,不损坏已有 Workspace。
### 8.2 删除 initialization 协议
默认 Home 已不需要先制作一个空 snapshot,因此删除:
- `InitializeHomeSnapshotSpec`
- `HomeSnapshotBackend.initialize()`
- `InitializeHomeSnapshotRequest`
- `/home-snapshots/initialize` route/service
- client 的同步和异步 `initialize_home_snapshot` 方法;
- Local/E2B/Enterprise 对应 initialization 实现和测试。
`HomeSnapshotBackend` 只保留从 RuntimeLease 创建 checkpoint 和删除 snapshot 两项职责。
## 9. 物理后端实现
### 9.1 Local
`LocalExecutionBindingBackend.create_binding()`
-`home_snapshot_ref`:保持当前行为,验证 snapshot 目录并复制到新的 per-Binding Home
-`home_snapshot_ref`:不访问 snapshot root,直接创建空的 per-Binding Home 目录;
- 两种路径都设置 Home/Workspace 权限,并保持 shared Workspace 语义;
- 创建失败时继续只清理本次新建的 Home,以及由本次请求新建的 Workspace。
删除 `LocalHomeSnapshotBackend.initialize()`。Build Apply 的
`create_from_runtime()` 和 snapshot delete 保持不变。
### 9.2 E2B
`E2BExecutionBindingBackend` 显式接收部署配置的 E2B template
-`home_snapshot_ref`:以该 snapshot ref 创建 Sandbox
-`home_snapshot_ref`:以配置的 `DIFY_AGENT_E2B_TEMPLATE` 创建 Sandbox
- 不能找到显式 snapshot 时把 SDK 错误映射为 `BindingCreateError`,不尝试 template fallback
- 两种路径都清理并重新创建 canonical Workspace,然后 pause 并返回 Sandbox id 作为 Binding 和
Workspace ref。
template 从 `E2BHomeSnapshotBackend` 移到 Execution Binding backend,因为它现在只负责默认
Binding 的初始环境。删除 E2B 的临时 initialization Sandbox 流程;
`create_from_runtime()` 仍从 Build Binding 的 E2B Sandbox 创建 immutable snapshot。
### 9.3 Enterprise
Enterprise 必须实现无 snapshot 的 Binding 创建:
- `existing_workspace_ref` 非空时,在任何外部写入前继续 fail fast,因为当前 Gateway 不能附加共享
Workspace
- `home_snapshot_ref=None` 时调用既有 Gateway `POST /v1/sandboxes`,请求包含
`tenantId`,不传 `template`,由 Gateway 选择部署默认 sandbox template
- 创建后通过 shellctl proxy 确保 canonical Home 存在,并清理/创建空 Workspace
- 关闭 operation-local transport,保留 Sandbox,返回 Sandbox id 同时作为 Binding 和
Workspace ref
- 初始化失败时 best-effort 删除本次新建的 Sandbox。
Enterprise Gateway 当前没有 immutable snapshot API。因此
`home_snapshot_ref` 非空时明确抛出 `BindingCreateError`,而不是忽略该值或回退到默认 Home;
`create_from_runtime()` 和 snapshot delete 仍保持显式不支持。这样 Enterprise 完整支持本方案要求的
无 snapshot 首次运行,同时不虚构 snapshot 能力。
## 10. Migration 方案
### 10.1 为什么不能只在 head 后加 nullable migration
如果保留 `2f39536b3feb` 中的 `nullable=False`,一个已有 config draft、config snapshot 或
runtime session 的数据库会在执行该 revision 时失败。Alembic 不可能跳过失败 revision 到达后面的
nullable migration。
因此必须修正前面的 migration 本身;仅新增后续 migration 不能解决旧数据库升级失败。
### 10.2 修正历史 revision
修正 `2f39536b3feb`
- `agent_config_drafts.home_snapshot_id``nullable=True` 添加;
- `agent_config_snapshots.home_snapshot_id``nullable=True` 添加;
- 中间态 `agent_runtime_sessions.home_snapshot_id``nullable=True` 添加。
历史行保持 `NULL`,不伪造 ledger id,也不尝试在数据库 migration 中调用外部 backend 创建物理
snapshot。
修正 `f6e4c5686857`
- 创建 `agent_workspace_bindings.base_home_snapshot_id` 时使用 `nullable=True`
- downgrade 重建中间态 `agent_runtime_sessions.home_snapshot_id` 时也使用
`nullable=True`
这样从旧 schema 首次升级到 head 的路径不会在添加列时失败。
### 10.3 已执行旧 revision 的数据库
只修改 migration 文件不会重新执行已经记录在 `alembic_version` 中的 revision。为使已经运行到
`f6e4c5686857` 的数据库也获得 nullable schema,在当前 head 后生成一个新的线性 convergence
revision
```text
... -> 6f5a9c2d8e1b -> 2f39536b3feb -> f6e4c5686857
-> <new make_home_snapshot_references_nullable revision>
```
该 revision 将当前仍存在的三个字段 alter 为 nullable
- `agent_config_drafts.home_snapshot_id`
- `agent_config_snapshots.home_snapshot_id`
- `agent_workspace_bindings.base_home_snapshot_id`
它不处理已被 `f6e4c5686857` 删除的 `agent_runtime_sessions`
对于使用已修正历史链的新数据库,这三个 alter 是 schema convergence;对于已经执行旧文件的
数据库,它们是真正的修复。该 convergence revision 的 downgrade 不重新引入 `NOT NULL`,因为
修正后的 down revision 本身已经定义 nullable schema,而且含有 `NULL` 的真实数据不能安全回退
`NOT NULL`
migration 链保持一个 head,不创建 Alembic merge revision。
### 10.4 数据处理
不做 Home Snapshot backfill
- 历史配置行的 `NULL` 直接表示 backend default Home
- 已有非空值原样保留;
- `agent_home_snapshots` 现有记录不变;
- 不在 migration 中访问 Dify Agent 或任何外部 backend。
## 11. 错误语义
- `home_snapshot_id=None`:合法,选择 backend default Home。
- 非空逻辑 id 无 ledger row、tenant/Agent 不匹配或状态非 `ACTIVE`Dify API fail fast。
- 非空 backend ref 不存在:backend 返回 Binding 创建失败,不回退。
- backend 无法创建默认 Home:返回 `BindingCreateError`Dify API 不创建 Binding ledger row。
- Enterprise 收到非空 snapshot ref:明确不支持并 fail fast。
- 已有 Binding 丢失:保持现有 `BindingLostError`,不创建默认替代 Binding。
## 12. 测试方案
### 12.1 Migration
- 构造包含历史 `agent_config_drafts``agent_config_snapshots`
`agent_runtime_sessions` 数据的 pre-`2f39536b3feb` schema,执行到 head 成功;
- 断言三个当前字段均 nullable,历史配置行的值为 `NULL`
- 覆盖从旧 `f6e4c5686857` 形态执行 convergence revision
- 在 PostgreSQL 和 MySQL migration CI 中确认唯一线性 head
- 更新 migration unit test 中旧的 `nullable=False` fixture 和断言。
### 12.2 Dify API
- 各 Agent 创建/导入/复制路径不调用 Dify Agent,首个 config snapshot 的
`home_snapshot_id is None`
- Draft 和新 config version 正确继承 `None` 或具体 id
- Publish 接受 `None`,非空 snapshot 仍严格校验;
- `AgentWorkspaceService.create_binding()``None` 时发送
`home_snapshot_ref=None` 并持久化 nullable base id
- 非空无效 snapshot 不回退;
- Agent App、Build Draft 和 Workflow node 调用链能携带 `None`
- 从 default Home 建立的 Build Binding 在 Apply 后创建真实 snapshot,并把新 id 写入 normal
draft。
### 12.3 Dify Agent
- protocol/client/server 正确接受字段缺失和 JSON `null`,拒绝空字符串;
- Local 的 default 路径创建空私有 Home,不读取或复制 snapshot
- E2B 的 default 路径使用 configured template,显式路径使用 snapshot ref
- E2B 显式 snapshot 失败不回退 template
- Enterprise default 路径调用 Gateway 创建、初始化 layout 并返回 refs
- Enterprise 显式 snapshot 和 shared Workspace 在外部写入前 fail fast
- 各 backend 创建中途失败时清理本次局部资源;
- 删除 initialization endpoint、client 和 backend method 后不存在残留调用或测试。
## 13. 代码实现步骤
1. **调整模型与 Alembic**
- 将三个当前 ORM 字段改为可空;
- 修正 `2f39536b3feb``f6e4c5686857` 的 nullable 定义;
- 使用 Alembic 生成 `f6e4c5686857` 之后的单一 convergence revision
- 补充带历史数据的 migration 测试并验证唯一 head。
2. **收敛 Home Snapshot 生命周期**
- 删除 `AgentHomeSnapshotService.create_initial()`
- 让 snapshot validation 接受 `None`
- 保留并验证 Build Apply checkpoint、retire、collect 和 delete。
3. **修改 Agent 配置创建与复制**
- 更新 roster、backing Agent、Composer、Workflow-only Agent 和 DSL import/clone
- 所有初始 config snapshot 使用 `home_snapshot_id=None`
- 将 config version/draft helper 及序列化相关类型改为可空;
- Publish 和版本复制原样传播可空值。
4. **修改 Binding 与 product callers**
- 更新 `AgentWorkspaceService.create_binding()` 和 generation validation
- 更新 Agent App runner/store、Workflow node store/resolver 及其测试;
- 确保 `None` 只在 Binding 创建边界转换为 `home_snapshot_ref=None`
5. **修改 Dify Agent 私有协议**
- 将 Execution Binding request/spec 字段改为可空;
- 更新 protocol implementer guidance
- 删除 initialize-home-snapshot DTO、client、route、service 和 backend protocol method
- 更新 exports 和调用测试。
6. **实现三个 backend 的 default Home**
- Local 增加无 snapshot 的空 Home 分支;
- E2B 将 template 注入 Execution Binding backend,并区分 template/default 与 snapshot
- Enterprise 基于既有 Gateway create API 实现 default Binding,显式 snapshot 保持 fail fast
- 更新 backend profile wiring 和集成测试。
7. **更新文档与生成产物**
- 更新 Runtime Resources、Operations Guide、Get Started 和 Shell layer 文档;
- 删除“Agent 创建先初始化 Home Snapshot”和“Enterprise Binding 全部未实现”的陈述;
- 使用仓库生成工具刷新 API docs/前端契约等生成产物,不手工修改生成文件。
8. **验证**
- 运行 API 相关 unit tests、strict type checks 和 migration SQL/heads 检查;
- 运行 dify-agent format、lint、typecheck、Local/E2B/Enterprise tests
- 运行 PostgreSQL/MySQL migration CI
- 确认 Alembic 仍只有一个线性 head。
## 14. 完成标准
- 创建任何类型的 Agent 都不创建物理或逻辑 Home Snapshot。
- 无 Home Snapshot 的 Agent 可以在 Local、E2B 和 Enterprise 上创建并运行 Binding。
- 非空 snapshot 始终精确解析和物化,任何失败都不会回退。
- Build Draft Apply 可以从 default Home checkpoint 出真实 immutable Home Snapshot。
- 带历史 Agent 数据的数据库可以从 `2f39536b3feb` 之前升级到最新 head。
- 已执行旧 `f6e4c5686857` 的数据库也会被 convergence revision 修正。
- 当前三个 ORM 字段与数据库 schema 均为 nullable。
- migration 链单一线性,无 merge revision。
-1
View File
@@ -63,7 +63,6 @@ jobs:
E2E_ADMIN_PASSWORD: E2eAdmin12345
E2E_FORCE_WEB_BUILD: "1"
E2E_INIT_PASSWORD: E2eInit12345
E2E_START_AGENT_BACKEND: "1"
run: vp run e2e:full
- name: Preserve Chromium E2E report and logs
+2 -2
View File
@@ -634,7 +634,7 @@ class AgentAppRunner:
agent_config_snapshot_id: str,
agent_config_version_kind: Literal["snapshot", "draft", "build_draft"] = "snapshot",
agent_soul: AgentSoulConfig,
home_snapshot_id: str,
home_snapshot_id: str | None,
conversation_id: str,
query: str,
message_id: str,
@@ -736,7 +736,7 @@ class AgentAppRunner:
dify_context: DifyRunContext,
agent_id: str,
agent_config_snapshot_id: str,
home_snapshot_id: str,
home_snapshot_id: str | None,
conversation_id: str,
session_scope_snapshot_id: str | None | _DefaultSessionScopeSnapshotId,
agent_config_version_kind: AgentConfigVersionKind,
+1 -1
View File
@@ -31,7 +31,7 @@ class AgentAppSessionScope:
conversation_id: str
agent_id: str
agent_config_snapshot_id: str
home_snapshot_id: str
home_snapshot_id: str | None
agent_config_version_kind: AgentConfigVersionKind = AgentConfigVersionKind.SNAPSHOT
build_draft_id: str | None = None
@@ -125,7 +125,7 @@ class WorkflowAgentWorkspaceStore:
)
def load_or_create_node_execution_session(
self, scope: WorkflowAgentSessionScope, *, home_snapshot_id: str
self, scope: WorkflowAgentSessionScope, *, home_snapshot_id: str | None
) -> StoredWorkflowAgentSession:
with session_factory.create_session() as session:
execution = self._load_execution(session=session, scope=scope)
@@ -31,13 +31,13 @@ def upgrade():
batch_op.create_index('agent_home_snapshot_tenant_agent_idx', ['tenant_id', 'agent_id'], unique=False)
with op.batch_alter_table('agent_config_drafts', schema=None) as batch_op:
batch_op.add_column(sa.Column('home_snapshot_id', models.types.StringUUID(), nullable=False))
batch_op.add_column(sa.Column('home_snapshot_id', models.types.StringUUID(), nullable=True))
with op.batch_alter_table('agent_config_snapshots', schema=None) as batch_op:
batch_op.add_column(sa.Column('home_snapshot_id', models.types.StringUUID(), nullable=False))
batch_op.add_column(sa.Column('home_snapshot_id', models.types.StringUUID(), nullable=True))
with op.batch_alter_table('agent_runtime_sessions', schema=None) as batch_op:
batch_op.add_column(sa.Column('home_snapshot_id', models.types.StringUUID(), nullable=False))
batch_op.add_column(sa.Column('home_snapshot_id', models.types.StringUUID(), nullable=True))
batch_op.drop_index(batch_op.f('agent_runtime_session_conversation_scope_unique'), postgresql_where='(conversation_id IS NOT NULL)')
batch_op.create_index('agent_runtime_session_conversation_scope_unique', ['tenant_id', 'conversation_id', 'agent_id', 'agent_config_snapshot_id', 'home_snapshot_id'], unique=True, postgresql_where=sa.text('conversation_id IS NOT NULL'))
@@ -24,7 +24,7 @@ def upgrade():
sa.Column('app_id', models.types.StringUUID(), nullable=False),
sa.Column('workspace_id', models.types.StringUUID(), nullable=False),
sa.Column('agent_id', models.types.StringUUID(), nullable=False),
sa.Column('base_home_snapshot_id', models.types.StringUUID(), nullable=False),
sa.Column('base_home_snapshot_id', models.types.StringUUID(), nullable=True),
sa.Column('agent_config_version_id', models.types.StringUUID(), nullable=False),
sa.Column('agent_config_version_kind', sa.String(length=32), nullable=False),
sa.Column('backend_binding_ref', sa.String(length=255), nullable=False),
@@ -127,7 +127,7 @@ def downgrade():
sa.Column('updated_at', postgresql.TIMESTAMP(), server_default=sa.text('CURRENT_TIMESTAMP'), autoincrement=False, nullable=False),
sa.Column('pending_form_id', sa.UUID(), autoincrement=False, nullable=True),
sa.Column('pending_tool_call_id', sa.VARCHAR(length=255), autoincrement=False, nullable=True),
sa.Column('home_snapshot_id', sa.UUID(), autoincrement=False, nullable=False),
sa.Column('home_snapshot_id', sa.UUID(), autoincrement=False, nullable=True),
sa.PrimaryKeyConstraint('id', name=op.f('agent_runtime_session_pkey'))
)
with op.batch_alter_table('agent_runtime_sessions', schema=None) as batch_op:
@@ -0,0 +1,49 @@
"""make home snapshot references nullable
Revision ID: e4708db55c1d
Revises: f6e4c5686857
Create Date: 2026-07-28 23:31:26.284147
The corrected preceding revisions make fresh upgrades nullable. This linear
convergence revision intentionally repeats those alters so databases that
already ran the old NOT NULL revisions are repaired too.
"""
from alembic import op
import models as models
# revision identifiers, used by Alembic.
revision = 'e4708db55c1d'
down_revision = 'f6e4c5686857'
branch_labels = None
depends_on = None
def upgrade():
with op.batch_alter_table("agent_config_drafts", schema=None) as batch_op:
batch_op.alter_column(
"home_snapshot_id",
existing_type=models.types.StringUUID(),
nullable=True,
)
with op.batch_alter_table("agent_config_snapshots", schema=None) as batch_op:
batch_op.alter_column(
"home_snapshot_id",
existing_type=models.types.StringUUID(),
nullable=True,
)
with op.batch_alter_table("agent_workspace_bindings", schema=None) as batch_op:
batch_op.alter_column(
"base_home_snapshot_id",
existing_type=models.types.StringUUID(),
nullable=True,
)
def downgrade():
# The corrected down revision already defines nullable columns, and rows
# created under this revision may contain NULL values.
pass
+3 -3
View File
@@ -307,7 +307,7 @@ class AgentConfigDraft(DefaultFieldsMixin, Base):
account_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
draft_owner_key: Mapped[str] = mapped_column(String(255), nullable=False, default="")
base_snapshot_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
home_snapshot_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
home_snapshot_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
agent_workspace_binding_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
config_snapshot: Mapped[Any] = mapped_column(JSONModelColumn(AgentSoulConfig), nullable=False)
created_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
@@ -344,7 +344,7 @@ class AgentConfigSnapshot(DefaultFieldsMixin, Base):
agent_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
version: Mapped[int] = mapped_column(sa.Integer, nullable=False)
config_snapshot: Mapped[Any] = mapped_column(JSONModelColumn(AgentSoulConfig), nullable=False)
home_snapshot_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
home_snapshot_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
summary: Mapped[str | None] = mapped_column(LongText, nullable=True)
version_note: Mapped[str | None] = mapped_column(LongText, nullable=True)
created_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
@@ -520,7 +520,7 @@ class AgentWorkspaceBinding(DefaultFieldsMixin, Base):
app_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
workspace_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
agent_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
base_home_snapshot_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
base_home_snapshot_id: Mapped[str | None] = mapped_column(StringUUID, nullable=True)
agent_config_version_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
agent_config_version_kind: Mapped[AgentConfigVersionKind] = mapped_column(
EnumText(AgentConfigVersionKind, length=32), nullable=False
+4 -14
View File
@@ -510,11 +510,6 @@ class AgentComposerService:
)
session.add(agent)
session.flush()
home_snapshot = AgentHomeSnapshotService.create_initial(
session=session,
tenant_id=tenant_id,
agent_id=agent.id,
)
initial_version = cls._create_config_version(
session=session,
tenant_id=tenant_id,
@@ -523,7 +518,7 @@ class AgentComposerService:
agent_soul=AgentSoulConfig(),
operation=AgentConfigRevisionOperation.CREATE_VERSION,
version_note=None,
home_snapshot_id=home_snapshot.id,
home_snapshot_id=None,
)
agent.active_config_snapshot_id = initial_version.id
agent.active_config_has_model = False
@@ -601,7 +596,7 @@ class AgentComposerService:
tenant_id: str,
agent: Agent,
agent_soul: AgentSoulConfig,
home_snapshot_id: str,
home_snapshot_id: str | None,
) -> bool:
if not agent.active_config_snapshot_id:
return False
@@ -1767,11 +1762,6 @@ class AgentComposerService:
)
session.add(agent)
session.flush()
home_snapshot = AgentHomeSnapshotService.create_initial(
session=session,
tenant_id=tenant_id,
agent_id=agent.id,
)
version = cls._create_config_version(
session=session,
tenant_id=tenant_id,
@@ -1780,7 +1770,7 @@ class AgentComposerService:
agent_soul=agent_soul,
operation=AgentConfigRevisionOperation.CREATE_VERSION,
version_note=None,
home_snapshot_id=home_snapshot.id,
home_snapshot_id=None,
)
agent.active_config_snapshot_id = version.id
agent.active_config_has_model = agent_soul_has_model(agent_soul)
@@ -1951,7 +1941,7 @@ class AgentComposerService:
agent_soul: AgentSoulConfig,
operation: AgentConfigRevisionOperation,
version_note: str | None,
home_snapshot_id: str,
home_snapshot_id: str | None,
previous_snapshot_id: str | None = None,
) -> AgentConfigSnapshot:
next_version = (
+1 -7
View File
@@ -49,7 +49,6 @@ from services.agent.dsl_entities import (
make_portable_agent_package,
portable_ref,
)
from services.agent.home_snapshot_service import AgentHomeSnapshotService
from services.agent.knowledge_datasets import get_tenant_knowledge_dataset_rows
from services.agent.roster_service import AgentRosterService
from services.entities.dsl_entities import DslImportWarning
@@ -569,17 +568,12 @@ class AgentDslService:
)
or 0
) + 1
home_snapshot = AgentHomeSnapshotService.create_initial(
session=self.session,
tenant_id=tenant_id,
agent_id=agent.id,
)
snapshot = AgentConfigSnapshot(
tenant_id=tenant_id,
agent_id=agent.id,
version=next_version,
config_snapshot=soul,
home_snapshot_id=home_snapshot.id,
home_snapshot_id=None,
created_by=account_id,
)
self.session.add(snapshot)
+4 -30
View File
@@ -5,7 +5,7 @@ from __future__ import annotations
import logging
from dify_agent.client import Client, DifyAgentNotFoundError
from dify_agent.protocol import CreateHomeSnapshotFromBindingRequest, InitializeHomeSnapshotRequest
from dify_agent.protocol import CreateHomeSnapshotFromBindingRequest
from sqlalchemy import select
from sqlalchemy.orm import Session
@@ -36,34 +36,6 @@ class AgentHomeSnapshotUnavailableError(RuntimeError):
class AgentHomeSnapshotService:
"""Create, retire, and collect Agent-owned immutable Home Snapshots."""
@classmethod
def create_initial(
cls,
*,
session: Session,
tenant_id: str,
agent_id: str,
) -> AgentHomeSnapshot:
home_snapshot_id = str(uuidv7())
with cls._client() as client:
response = client.initialize_home_snapshot_sync(
InitializeHomeSnapshotRequest(
tenant_id=tenant_id,
agent_id=agent_id,
home_snapshot_id=home_snapshot_id,
)
)
home_snapshot = AgentHomeSnapshot(
id=home_snapshot_id,
tenant_id=tenant_id,
agent_id=agent_id,
snapshot_ref=response.snapshot_ref,
status=AgentWorkingResourceStatus.ACTIVE,
)
session.add(home_snapshot)
session.flush()
return home_snapshot
@classmethod
def create_for_build_apply(
cls,
@@ -211,7 +183,9 @@ class AgentHomeSnapshotService:
return Client(base_url=base_url)
def validate_home_snapshot_binding(*, session: Session, agent: Agent, home_snapshot_id: str) -> None:
def validate_home_snapshot_binding(*, session: Session, agent: Agent, home_snapshot_id: str | None) -> None:
if home_snapshot_id is None:
return
_require_owned_home_snapshot(session=session, agent=agent, home_snapshot_id=home_snapshot_id)
+2 -12
View File
@@ -350,17 +350,12 @@ class AgentRosterService:
self._session.add(agent)
self._session.flush()
home_snapshot = AgentHomeSnapshotService.create_initial(
session=self._session,
tenant_id=tenant_id,
agent_id=agent.id,
)
version = AgentConfigSnapshot(
tenant_id=tenant_id,
agent_id=agent.id,
version=1,
config_snapshot=payload.agent_soul,
home_snapshot_id=home_snapshot.id,
home_snapshot_id=None,
version_note=payload.version_note,
created_by=account_id,
)
@@ -429,17 +424,12 @@ class AgentRosterService:
self._session.add(agent)
self._session.flush()
home_snapshot = AgentHomeSnapshotService.create_initial(
session=self._session,
tenant_id=tenant_id,
agent_id=agent.id,
)
version = AgentConfigSnapshot(
tenant_id=tenant_id,
agent_id=agent.id,
version=1,
config_snapshot=soul,
home_snapshot_id=home_snapshot.id,
home_snapshot_id=None,
created_by=account_id,
)
self._session.add(version)
+15 -12
View File
@@ -110,7 +110,7 @@ class AgentWorkspaceService:
session: Session,
scope: WorkspaceOwnerScope,
agent_id: str,
base_home_snapshot_id: str,
base_home_snapshot_id: str | None,
agent_config_version_id: str,
agent_config_version_kind: AgentConfigVersionKind,
) -> AgentWorkspaceBinding:
@@ -123,16 +123,19 @@ class AgentWorkspaceService:
fails before the backend returns success.
"""
home_snapshot = session.scalar(
select(AgentHomeSnapshot).where(
AgentHomeSnapshot.id == base_home_snapshot_id,
AgentHomeSnapshot.tenant_id == scope.tenant_id,
AgentHomeSnapshot.agent_id == agent_id,
AgentHomeSnapshot.status == AgentWorkingResourceStatus.ACTIVE,
home_snapshot_ref: str | None = None
if base_home_snapshot_id is not None:
home_snapshot = session.scalar(
select(AgentHomeSnapshot).where(
AgentHomeSnapshot.id == base_home_snapshot_id,
AgentHomeSnapshot.tenant_id == scope.tenant_id,
AgentHomeSnapshot.agent_id == agent_id,
AgentHomeSnapshot.status == AgentWorkingResourceStatus.ACTIVE,
)
)
)
if home_snapshot is None:
raise AgentWorkspaceNotFoundError("base Home Snapshot is unavailable")
if home_snapshot is None:
raise AgentWorkspaceNotFoundError("base Home Snapshot is unavailable")
home_snapshot_ref = home_snapshot.snapshot_ref
workspace = cls.resolve_active_workspace(session=session, scope=scope)
workspace_id = workspace.id if workspace is not None else str(uuidv7())
binding_id = str(uuidv7())
@@ -144,7 +147,7 @@ class AgentWorkspaceService:
binding_id=binding_id,
workspace_id=workspace_id,
existing_workspace_ref=workspace.backend_workspace_ref if workspace is not None else None,
home_snapshot_ref=home_snapshot.snapshot_ref,
home_snapshot_ref=home_snapshot_ref,
)
)
if workspace is not None and allocation.workspace_ref != workspace.backend_workspace_ref:
@@ -439,7 +442,7 @@ class AgentWorkspaceService:
def validate_binding_generation(
binding: AgentWorkspaceBinding,
*,
base_home_snapshot_id: str,
base_home_snapshot_id: str | None,
agent_config_version_id: str,
agent_config_version_kind: AgentConfigVersionKind,
) -> None:
@@ -124,7 +124,6 @@ class TestAppDslService:
patch("services.app_service.ModelManager.for_tenant") as mock_model_manager,
patch("services.app_service.FeatureService") as mock_feature_service,
patch("services.app_service.EnterpriseService") as mock_enterprise_service,
patch("services.agent.home_snapshot_service.AgentHomeSnapshotService._client") as mock_home_snapshot_client,
):
mock_workflow_service.return_value.get_draft_workflow.return_value = None
mock_workflow_service.return_value.sync_draft_workflow.return_value = MagicMock()
@@ -143,10 +142,6 @@ class TestAppDslService:
mock_feature_service.get_system_features.return_value.webapp_auth.enabled = False
mock_enterprise_service.WebAppAuth.update_app_access_mode.return_value = None
mock_enterprise_service.WebAppAuth.cleanup_webapp.return_value = None
mock_home_snapshot_client.return_value.__enter__.return_value.initialize_home_snapshot_sync.side_effect = (
lambda request: SimpleNamespace(snapshot_ref=f"test:{request.home_snapshot_id}")
)
yield {
"workflow_service": mock_workflow_service,
"dependencies_service": mock_dependencies_service,
@@ -1,5 +1,4 @@
from datetime import datetime
from types import SimpleNamespace
from unittest.mock import create_autospec, patch
import pytest
@@ -30,7 +29,6 @@ class TestAppService:
patch("services.app_service.EnterpriseService") as mock_enterprise_service,
patch("services.app_service.ModelManager.for_tenant") as mock_model_manager,
patch("services.account_service.FeatureService") as mock_account_feature_service,
patch("services.agent.home_snapshot_service.AgentHomeSnapshotService._client") as mock_home_snapshot_client,
):
# Setup default mock returns for app service
mock_feature_service.get_system_features.return_value.webapp_auth.enabled = False
@@ -44,10 +42,6 @@ class TestAppService:
mock_model_instance = mock_model_manager.return_value
mock_model_instance.get_default_model_instance.return_value = None
mock_model_instance.get_default_provider_model_name.return_value = ("openai", "gpt-3.5-turbo")
mock_home_snapshot_client.return_value.__enter__.return_value.initialize_home_snapshot_sync.side_effect = (
lambda request: SimpleNamespace(snapshot_ref=f"test:{request.home_snapshot_id}")
)
yield {
"feature_service": mock_feature_service,
"enterprise_service": mock_enterprise_service,
@@ -16,6 +16,7 @@ def _scope(
*,
kind: AgentConfigVersionKind = AgentConfigVersionKind.SNAPSHOT,
build_draft_id: str | None = None,
home_snapshot_id: str | None = "home-1",
) -> AgentAppSessionScope:
return AgentAppSessionScope(
tenant_id="tenant-1",
@@ -23,7 +24,7 @@ def _scope(
conversation_id="conversation-1",
agent_id="agent-1",
agent_config_snapshot_id="config-1",
home_snapshot_id="home-1",
home_snapshot_id=home_snapshot_id,
agent_config_version_kind=kind,
build_draft_id=build_draft_id,
)
@@ -54,7 +55,8 @@ def test_scope_selects_conversation_or_build_draft_workspace_owner() -> None:
assert build_owner.owner_id == "build-draft-1"
def test_load_or_create_persists_new_binding_on_caller(monkeypatch) -> None:
@pytest.mark.parametrize("home_snapshot_id", ["home-1", None])
def test_load_or_create_persists_new_binding_on_caller(monkeypatch, home_snapshot_id: str | None) -> None:
caller = SimpleNamespace(agent_workspace_binding_id=None)
context = MagicMock()
session = context.__enter__.return_value
@@ -64,13 +66,14 @@ def test_load_or_create_persists_new_binding_on_caller(monkeypatch) -> None:
monkeypatch.setattr(store, "_load_caller", MagicMock(return_value=caller))
monkeypatch.setattr(AgentWorkspaceService, "create_binding", create)
stored = store.load_or_create(_scope())
stored = store.load_or_create(_scope(home_snapshot_id=home_snapshot_id))
assert stored.binding_id == "binding-1"
assert stored.workspace_id == "workspace-1"
assert stored.backend_binding_ref == "backend-binding-1"
assert caller.agent_workspace_binding_id == "binding-1"
assert create.call_args.kwargs["session"] is session
assert create.call_args.kwargs["base_home_snapshot_id"] == home_snapshot_id
session.commit.assert_called_once_with()
@@ -155,7 +155,8 @@ def test_load_existing_scope_rejects_unavailable_persisted_binding(monkeypatch:
)
def test_load_or_create_persists_binding_on_node_execution(monkeypatch) -> None:
@pytest.mark.parametrize("home_snapshot_id", ["home-1", None])
def test_load_or_create_persists_binding_on_node_execution(monkeypatch, home_snapshot_id: str | None) -> None:
execution = WorkflowNodeExecutionModel(
agent_workspace_binding_id=None,
process_data=json.dumps({"existing": "value"}),
@@ -171,7 +172,7 @@ def test_load_or_create_persists_binding_on_node_execution(monkeypatch) -> None:
monkeypatch.setattr(store, "_load_execution", MagicMock(return_value=execution))
monkeypatch.setattr(AgentWorkspaceService, "create_binding", create)
stored = store.load_or_create_node_execution_session(_scope(), home_snapshot_id="home-1")
stored = store.load_or_create_node_execution_session(_scope(), home_snapshot_id=home_snapshot_id)
assert stored.binding_id == "binding-1"
assert stored.workspace_id == "workspace-1"
@@ -183,6 +184,7 @@ def test_load_or_create_persists_binding_on_node_execution(monkeypatch) -> None:
}
assert "agent_workspace_binding_id" not in execution.process_data_dict
assert create.call_args.kwargs["session"] is session
assert create.call_args.kwargs["base_home_snapshot_id"] == home_snapshot_id
session.commit.assert_called_once_with()
get_active = MagicMock(return_value=_binding())
@@ -42,7 +42,7 @@ def _create_pre_upgrade_schema(engine: sa.Engine) -> None:
sa.Column("binding_id", sa.String(36)),
sa.Column("agent_id", sa.String(36), nullable=False),
sa.Column("agent_config_snapshot_id", sa.String(36)),
sa.Column("home_snapshot_id", sa.String(36), nullable=False),
sa.Column("home_snapshot_id", sa.String(36), nullable=True),
sa.Column("backend_run_id", sa.String(255)),
sa.Column("status", sa.String(32), nullable=False),
)
@@ -110,7 +110,10 @@ def test_upgrade_replaces_runtime_sessions_with_workspace_schema() -> None:
inspector = sa.inspect(engine)
assert "agent_runtime_sessions" not in inspector.get_table_names()
binding_columns = {column["name"] for column in inspector.get_columns("agent_workspace_bindings")}
binding_column_definitions = {
column["name"]: column for column in inspector.get_columns("agent_workspace_bindings")
}
binding_columns = set(binding_column_definitions)
assert {
"workspace_id",
"agent_id",
@@ -125,6 +128,7 @@ def test_upgrade_replaces_runtime_sessions_with_workspace_schema() -> None:
"pending_tool_call_id",
}.issubset(binding_columns)
assert "active_guard" not in binding_columns
assert binding_column_definitions["base_home_snapshot_id"]["nullable"] is True
binding_indexes = {index["name"] for index in inspector.get_indexes("agent_workspace_bindings")}
assert "agent_workspace_binding_agent_active_unique" not in binding_indexes
workspace_columns = {column["name"] for column in inspector.get_columns("agent_workspaces")}
@@ -0,0 +1,182 @@
from __future__ import annotations
import importlib.util
from collections.abc import Callable
from pathlib import Path
from typing import Protocol, cast
import sqlalchemy as sa
from alembic.migration import MigrationContext
from alembic.operations import Operations
_VERSIONS_DIR = Path(__file__).resolve().parents[3] / "migrations/versions"
_MIGRATION_PATHS = (
_VERSIONS_DIR / "2026_07_21_2251-2f39536b3feb_add_agent_home_snapshot_ledger.py",
_VERSIONS_DIR / "2026_07_23_0203-f6e4c5686857_replace_agent_runtime_sessions_with_.py",
_VERSIONS_DIR / "2026_07_28_2331-e4708db55c1d_make_home_snapshot_references_nullable.py",
)
class _MigrationModule(Protocol):
op: Operations
upgrade: Callable[[], None]
def _load_migration(path: Path) -> _MigrationModule:
spec = importlib.util.spec_from_file_location(path.stem, path)
if spec is None or spec.loader is None:
raise RuntimeError(f"failed to load migration {path.name}")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
return cast(_MigrationModule, module)
def _run_upgrade(module: _MigrationModule, engine: sa.Engine) -> None:
with engine.begin() as connection:
original_op = module.op
module.op = Operations(MigrationContext.configure(connection))
try:
module.upgrade()
finally:
module.op = original_op
def _create_pre_home_snapshot_schema(engine: sa.Engine) -> None:
metadata = sa.MetaData()
drafts = sa.Table("agent_config_drafts", metadata, sa.Column("id", sa.String(36), primary_key=True))
snapshots = sa.Table("agent_config_snapshots", metadata, sa.Column("id", sa.String(36), primary_key=True))
conversations = sa.Table("conversations", metadata, sa.Column("id", sa.String(36), primary_key=True))
executions = sa.Table("workflow_node_executions", metadata, sa.Column("id", sa.String(36), primary_key=True))
runtime_sessions = sa.Table(
"agent_runtime_sessions",
metadata,
sa.Column("id", sa.String(36), primary_key=True),
sa.Column("tenant_id", sa.String(36), nullable=False),
sa.Column("conversation_id", sa.String(36)),
sa.Column("workflow_run_id", sa.String(36)),
sa.Column("node_id", sa.String(255)),
sa.Column("binding_id", sa.String(36)),
sa.Column("agent_id", sa.String(36), nullable=False),
sa.Column("agent_config_snapshot_id", sa.String(36)),
sa.Column("backend_run_id", sa.String(255)),
sa.Column("status", sa.String(32), nullable=False),
)
sa.Index("agent_runtime_session_backend_run_idx", runtime_sessions.c.backend_run_id)
sa.Index(
"agent_runtime_session_conversation_lookup_idx",
runtime_sessions.c.tenant_id,
runtime_sessions.c.conversation_id,
runtime_sessions.c.status,
)
sa.Index(
"agent_runtime_session_conversation_scope_unique",
runtime_sessions.c.tenant_id,
runtime_sessions.c.conversation_id,
runtime_sessions.c.agent_id,
runtime_sessions.c.agent_config_snapshot_id,
unique=True,
)
sa.Index(
"agent_runtime_session_workflow_lookup_idx",
runtime_sessions.c.tenant_id,
runtime_sessions.c.workflow_run_id,
runtime_sessions.c.node_id,
runtime_sessions.c.status,
)
sa.Index(
"agent_runtime_session_workflow_scope_unique",
runtime_sessions.c.tenant_id,
runtime_sessions.c.workflow_run_id,
runtime_sessions.c.node_id,
runtime_sessions.c.binding_id,
runtime_sessions.c.agent_id,
unique=True,
)
metadata.create_all(engine)
with engine.begin() as connection:
connection.execute(drafts.insert().values(id="draft-1"))
connection.execute(snapshots.insert().values(id="snapshot-1"))
connection.execute(conversations.insert().values(id="conversation-1"))
connection.execute(executions.insert().values(id="execution-1"))
connection.execute(
runtime_sessions.insert().values(
id="runtime-1",
tenant_id="tenant-1",
conversation_id="conversation-1",
agent_id="agent-1",
agent_config_snapshot_id="snapshot-1",
status="active",
)
)
def _create_old_f6_schema(engine: sa.Engine) -> None:
metadata = sa.MetaData()
drafts = sa.Table(
"agent_config_drafts",
metadata,
sa.Column("id", sa.String(36), primary_key=True),
sa.Column("home_snapshot_id", sa.String(36), nullable=False),
)
snapshots = sa.Table(
"agent_config_snapshots",
metadata,
sa.Column("id", sa.String(36), primary_key=True),
sa.Column("home_snapshot_id", sa.String(36), nullable=False),
)
bindings = sa.Table(
"agent_workspace_bindings",
metadata,
sa.Column("id", sa.String(36), primary_key=True),
sa.Column("base_home_snapshot_id", sa.String(36), nullable=False),
)
metadata.create_all(engine)
with engine.begin() as connection:
connection.execute(drafts.insert().values(id="draft-1", home_snapshot_id="home-draft"))
connection.execute(snapshots.insert().values(id="snapshot-1", home_snapshot_id="home-snapshot"))
connection.execute(bindings.insert().values(id="binding-1", base_home_snapshot_id="home-binding"))
def test_historical_rows_upgrade_through_nullable_home_snapshot_chain() -> None:
engine = sa.create_engine("sqlite:///:memory:")
_create_pre_home_snapshot_schema(engine)
for path in _MIGRATION_PATHS:
_run_upgrade(_load_migration(path), engine)
inspector = sa.inspect(engine)
for table_name, column_name in (
("agent_config_drafts", "home_snapshot_id"),
("agent_config_snapshots", "home_snapshot_id"),
("agent_workspace_bindings", "base_home_snapshot_id"),
):
columns = {column["name"]: column for column in inspector.get_columns(table_name)}
assert columns[column_name]["nullable"] is True
with engine.connect() as connection:
assert connection.scalar(sa.text("SELECT home_snapshot_id FROM agent_config_drafts")) is None
assert connection.scalar(sa.text("SELECT home_snapshot_id FROM agent_config_snapshots")) is None
def test_old_f6_schema_converges_to_nullable_without_changing_data() -> None:
engine = sa.create_engine("sqlite:///:memory:")
_create_old_f6_schema(engine)
_run_upgrade(_load_migration(_MIGRATION_PATHS[-1]), engine)
inspector = sa.inspect(engine)
for table_name, column_name in (
("agent_config_drafts", "home_snapshot_id"),
("agent_config_snapshots", "home_snapshot_id"),
("agent_workspace_bindings", "base_home_snapshot_id"),
):
columns = {column["name"]: column for column in inspector.get_columns(table_name)}
assert columns[column_name]["nullable"] is True
with engine.connect() as connection:
assert connection.scalar(sa.text("SELECT home_snapshot_id FROM agent_config_drafts")) == "home-draft"
assert connection.scalar(sa.text("SELECT home_snapshot_id FROM agent_config_snapshots")) == "home-snapshot"
assert (
connection.scalar(sa.text("SELECT base_home_snapshot_id FROM agent_workspace_bindings"))
== "home-binding"
)
@@ -684,12 +684,10 @@ def test_resolve_package_soul_preserves_existing_and_marks_missing_knowledge(mon
}
def test_create_snapshot_increments_version_and_records_revision(monkeypatch: pytest.MonkeyPatch) -> None:
def test_create_snapshot_increments_version_and_records_revision() -> None:
session = Mock()
session.scalar.return_value = 2
service = AgentDslService(session)
create_initial = Mock(return_value=SimpleNamespace(id="home-3"))
monkeypatch.setattr("services.agent.dsl_service.AgentHomeSnapshotService.create_initial", create_initial)
snapshot = service._create_snapshot(
tenant_id="tenant-1",
@@ -700,12 +698,7 @@ def test_create_snapshot_increments_version_and_records_revision(monkeypatch: py
)
assert snapshot.version == 3
assert snapshot.home_snapshot_id == "home-3"
create_initial.assert_called_once_with(
session=session,
tenant_id="tenant-1",
agent_id="agent-1",
)
assert snapshot.home_snapshot_id is None
assert isinstance(session.add.call_args_list[0].args[0], AgentConfigSnapshot)
revision = session.add.call_args_list[1].args[0]
assert isinstance(revision, AgentConfigRevision)
@@ -129,15 +129,6 @@ class FakeSession:
self.rollbacks += 1
@pytest.fixture
def _stub_home_snapshot_backend(monkeypatch: pytest.MonkeyPatch):
create_initial = MagicMock(return_value=SimpleNamespace(id="home-initial", snapshot_ref="backend-home-initial"))
create_for_build_apply = MagicMock(return_value=SimpleNamespace(id="home-build", snapshot_ref="backend-home-build"))
monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial)
monkeypatch.setattr(AgentHomeSnapshotService, "create_for_build_apply", create_for_build_apply)
return SimpleNamespace(create_initial=create_initial, create_for_build_apply=create_for_build_apply)
def _agent_soul_with_model() -> AgentSoulConfig:
return AgentSoulConfig.model_validate(
{
@@ -612,7 +603,6 @@ def test_publish_save_strategies_run_publish_validation(strategy: ComposerSaveSt
composer_service._validate_composer_payload_for_strategy(_duplicate_env_secret_payload(strategy))
@pytest.mark.usefixtures("_stub_home_snapshot_backend")
def test_save_agent_app_composer_creates_agent_when_missing(monkeypatch: pytest.MonkeyPatch):
fake_session = FakeSession(scalar=[None])
saved_draft = SimpleNamespace(
@@ -646,7 +636,7 @@ def test_save_agent_app_composer_creates_agent_when_missing(monkeypatch: pytest.
assert result == {"loaded": True}
assert fake_session.added[0].name == "Analyst"
assert fake_session.added[0].active_config_snapshot_id == fake_session.added[1].id
assert fake_session.added[1].home_snapshot_id == "home-initial"
assert fake_session.added[1].home_snapshot_id is None
assert fake_session.added[0].active_config_is_published is False
assert fake_session.flushes >= 1
@@ -869,7 +859,7 @@ def test_publish_agent_app_draft_creates_published_snapshot(monkeypatch: pytest.
draft_type=AgentConfigDraftType.DRAFT,
draft_owner_key="",
base_snapshot_id="version-1",
home_snapshot_id="home-1",
home_snapshot_id=None,
config_snapshot=_agent_soul_with_model(),
)
version = SimpleNamespace(id="version-2")
@@ -905,7 +895,7 @@ def test_publish_agent_app_draft_creates_published_snapshot(monkeypatch: pytest.
assert result["draft"]["base_snapshot_id"] == "version-2"
assert created["operation"] == AgentConfigRevisionOperation.PUBLISH_DRAFT
assert created["previous_snapshot_id"] == "version-1"
assert created["home_snapshot_id"] == "home-1"
assert created["home_snapshot_id"] is None
assert calls == ["validate_home", "create_version"]
assert agent.active_config_snapshot_id == "version-2"
assert agent.active_config_has_model is True
@@ -939,7 +929,6 @@ def test_repeated_publish_reuses_normal_draft_home_without_creating_resources(
session = FakeSession(scalar=[agent, draft, agent, draft])
published_homes: list[str] = []
versions = iter([SimpleNamespace(id="version-2"), SimpleNamespace(id="version-3")])
create_initial = MagicMock()
create_from_build = MagicMock()
monkeypatch.setattr(composer_service.ComposerConfigValidator, "validate_publish_payload", lambda _payload: None)
monkeypatch.setattr(AgentComposerService, "validate_knowledge_datasets", lambda **_kwargs: None)
@@ -951,7 +940,6 @@ def test_repeated_publish_reuses_normal_draft_home_without_creating_resources(
)
monkeypatch.setattr(AgentComposerService, "_serialize_version", lambda version: {"id": version.id})
monkeypatch.setattr(AgentComposerService, "_serialize_draft", lambda value: {"id": value.id})
monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial)
monkeypatch.setattr(AgentHomeSnapshotService, "create_for_build_apply", create_from_build)
first = AgentComposerService.publish_agent_app_draft(
@@ -971,7 +959,6 @@ def test_repeated_publish_reuses_normal_draft_home_without_creating_resources(
assert second["active_config_snapshot_id"] == "version-3"
assert published_homes == ["home-1", "home-1"]
assert draft.home_snapshot_id == "home-1"
create_initial.assert_not_called()
create_from_build.assert_not_called()
@@ -1233,7 +1220,7 @@ def test_build_apply_checkpoints_binding_updates_normal_draft_then_collects(monk
draft_type=AgentConfigDraftType.DEBUG_BUILD,
account_id="account-1",
draft_owner_key="account-1",
home_snapshot_id="home-old",
home_snapshot_id=None,
agent_workspace_binding_id="binding-1",
config_snapshot=AgentSoulConfig(),
)
@@ -1243,14 +1230,14 @@ def test_build_apply_checkpoints_binding_updates_normal_draft_then_collects(monk
agent_id="agent-1",
draft_type=AgentConfigDraftType.DRAFT,
draft_owner_key="",
home_snapshot_id="home-old",
home_snapshot_id=None,
config_snapshot=AgentSoulConfig(),
)
source_binding = SimpleNamespace(
id="binding-1",
backend_binding_ref="backend-binding-1",
agent_id="agent-1",
base_home_snapshot_id="home-old",
base_home_snapshot_id=None,
agent_config_version_id="build-1",
agent_config_version_kind=AgentConfigVersionKind.BUILD_DRAFT,
)
@@ -1762,10 +1749,8 @@ def test_build_draft_save_and_discard_do_not_manage_home_resources(monkeypatch:
home_snapshot_id="home-existing",
config_snapshot=AgentSoulConfig(),
)
create_initial = MagicMock()
create_from_build = MagicMock()
delete_home = MagicMock()
monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial)
monkeypatch.setattr(AgentHomeSnapshotService, "create_for_build_apply", create_from_build)
monkeypatch.setattr(AgentHomeSnapshotService, "delete", delete_home)
monkeypatch.setattr(AgentComposerService, "_require_agent", lambda **_kwargs: agent)
@@ -1791,7 +1776,6 @@ def test_build_draft_save_and_discard_do_not_manage_home_resources(monkeypatch:
account_id="account-1",
)
create_initial.assert_not_called()
create_from_build.assert_not_called()
delete_home.assert_not_called()
@@ -3168,7 +3152,6 @@ def test_composer_create_agents_syncs_active_config_has_model(
session = fake_session
created_apps = []
hidden_backing_apps = []
create_initial = MagicMock(return_value=SimpleNamespace(id="home-initial", snapshot_ref="backend-home-initial"))
backing_agent = Agent(
id="roster-agent-1",
tenant_id="tenant-1",
@@ -3199,7 +3182,6 @@ def test_composer_create_agents_syncs_active_config_has_model(
monkeypatch.setattr(composer_service, "AppService", FakeAppService)
monkeypatch.setattr(composer_service, "AgentRosterService", FakeAgentRosterService)
monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial)
monkeypatch.setattr(AgentComposerService, "_require_account", lambda **kwargs: SimpleNamespace(id="account-1"))
monkeypatch.setattr(
AgentComposerService,
@@ -3208,14 +3190,11 @@ def test_composer_create_agents_syncs_active_config_has_model(
id="empty-version-1",
tenant_id="tenant-1",
agent_id="roster-agent-1",
home_snapshot_id="home-roster-initial",
home_snapshot_id=None,
),
)
monkeypatch.setattr(
AgentComposerService,
"_create_config_version",
lambda **kwargs: SimpleNamespace(id="version-with-model"),
)
create_config_version = MagicMock(return_value=SimpleNamespace(id="version-with-model"))
monkeypatch.setattr(AgentComposerService, "_create_config_version", create_config_version)
workflow_agent = AgentComposerService._create_workflow_only_agent(
session=session,
@@ -3239,11 +3218,8 @@ def test_composer_create_agents_syncs_active_config_has_model(
assert workflow_agent.active_config_snapshot_id == "version-with-model"
assert workflow_agent.active_config_has_model is True
assert workflow_agent.backing_app_id == "hidden-app-1"
create_initial.assert_called_once_with(
session=session,
tenant_id="tenant-1",
agent_id=workflow_agent.id,
)
assert create_config_version.call_count == 2
assert all(call.kwargs["home_snapshot_id"] is None for call in create_config_version.call_args_list)
assert hidden_backing_apps[0]["name"] == "Workflow Agent node-1"
assert roster_agent.active_config_snapshot_id == "version-with-model"
assert roster_agent.active_config_has_model is True
@@ -3962,13 +3938,6 @@ def test_roster_create_detail_and_lookup_helpers(monkeypatch: pytest.MonkeyPatch
scalars=[[AgentConfigSnapshot(id="version-1", agent_id="agent-1", version=1)]],
)
service = AgentRosterService(fake_session)
create_initial = MagicMock(
side_effect=[
SimpleNamespace(id="home-roster", snapshot_ref="backend-home-roster"),
SimpleNamespace(id="home-backing", snapshot_ref="backend-home-backing"),
]
)
monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial)
monkeypatch.setattr(
AgentRosterService,
"_get_or_create_agent_app_debug_conversation",
@@ -4006,10 +3975,9 @@ def test_roster_create_detail_and_lookup_helpers(monkeypatch: pytest.MonkeyPatch
assert created.role == "Research assistant"
assert created.source == AgentSource.ROSTER
assert created.active_config_snapshot_id is not None
assert create_initial.call_count == 2
assert [
snapshot.home_snapshot_id for snapshot in fake_session.added if isinstance(snapshot, AgentConfigSnapshot)
] == ["home-roster", "home-backing"]
] == [None, None]
assert created.active_config_has_model is False
assert backing_agent.role == "Support agent"
assert backing_agent.active_config_snapshot_id is not None
@@ -4553,9 +4521,6 @@ class TestAgentAppBackingAgent:
def test_create_backing_agent_for_app_links_app_and_seeds_default_soul(self, monkeypatch: pytest.MonkeyPatch):
session = FakeSession()
service = AgentRosterService(session)
create_initial = MagicMock(return_value=SimpleNamespace(id="home-1", snapshot_ref="backend-home-1"))
monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial)
agent = service.create_backing_agent_for_app(
tenant_id="tenant-1",
account_id="account-1",
@@ -4577,12 +4542,7 @@ class TestAgentAppBackingAgent:
snapshots = [a for a in session.added if isinstance(a, AgentConfigSnapshot)]
assert len(snapshots) == 1
assert snapshots[0].version == 1
assert snapshots[0].home_snapshot_id == "home-1"
create_initial.assert_called_once_with(
session=session,
tenant_id="tenant-1",
agent_id=agent.id,
)
assert snapshots[0].home_snapshot_id is None
assert agent.active_config_snapshot_id == snapshots[0].id
revisions = [
a for a in session.added if getattr(a, "operation", None) == AgentConfigRevisionOperation.CREATE_VERSION
@@ -4966,7 +4926,6 @@ class TestAgentAppBackingAgent:
scalars=[[]],
)
captured: dict[str, object] = {}
create_initial = MagicMock(return_value=SimpleNamespace(id="home-target", snapshot_ref="backend-home-target"))
class FakeAppService:
def create_app(self, tenant_id: str, params, account: object, *, session) -> object:
@@ -4992,7 +4951,6 @@ class TestAgentAppBackingAgent:
return target_app
monkeypatch.setattr(roster_service, "AppService", FakeAppService)
monkeypatch.setattr(AgentHomeSnapshotService, "create_initial", create_initial)
monkeypatch.setattr(
AgentRosterService,
"_get_or_create_agent_app_debug_conversation",
@@ -5027,16 +4985,11 @@ class TestAgentAppBackingAgent:
target_version = captured["target_version"]
assert target_version.config_snapshot.model.model == "gpt-4o"
assert source_version.home_snapshot_id == "home-source"
assert target_version.home_snapshot_id == "home-target"
assert target_version.home_snapshot_id is None
assert target_version.summary == "configured"
assert target_version.version_note == "v1"
assert target_agent.active_config_has_model is True
assert target_agent.updated_by == "account-1"
create_initial.assert_called_once_with(
session=session,
tenant_id="tenant-1",
agent_id=target_agent.id,
)
assert session.commits == 1
def test_duplicate_agent_app_inherits_webapp_access_mode(self, monkeypatch: pytest.MonkeyPatch):
@@ -14,11 +14,11 @@ from models.agent import (
)
from models.agent_config_entities import AgentSoulConfig
from services.agent.errors import AgentBuildSandboxNotFoundError
from services.agent.home_snapshot_service import AgentHomeSnapshotService
from services.agent.home_snapshot_service import AgentHomeSnapshotService, validate_home_snapshot_binding
from services.agent.workspace_service import AgentWorkspaceService
def _build_draft() -> AgentConfigDraft:
def _build_draft(*, home_snapshot_id: str | None = "home-old") -> AgentConfigDraft:
return AgentConfigDraft(
id="build-1",
tenant_id="tenant-1",
@@ -26,7 +26,7 @@ def _build_draft() -> AgentConfigDraft:
draft_type=AgentConfigDraftType.DEBUG_BUILD,
account_id="account-1",
draft_owner_key="account-1",
home_snapshot_id="home-old",
home_snapshot_id=home_snapshot_id,
agent_workspace_binding_id="binding-1",
config_snapshot=AgentSoulConfig(),
)
@@ -34,42 +34,19 @@ def _build_draft() -> AgentConfigDraft:
def _client(*, snapshot_ref: str = "snapshot-ref-1") -> MagicMock:
client = MagicMock()
client.initialize_home_snapshot_sync.return_value = SimpleNamespace(snapshot_ref=snapshot_ref)
client.create_home_snapshot_from_binding_sync.return_value = SimpleNamespace(snapshot_ref=snapshot_ref)
return client
def test_create_initial_persists_backend_snapshot_ref(monkeypatch: pytest.MonkeyPatch) -> None:
def test_validate_home_snapshot_binding_accepts_default_home_without_ledger_lookup() -> None:
session = MagicMock()
client = _client()
monkeypatch.setattr(AgentHomeSnapshotService, "_client", lambda: nullcontext(client))
snapshot = AgentHomeSnapshotService.create_initial(
validate_home_snapshot_binding(
session=session,
tenant_id="tenant-1",
agent_id="agent-1",
agent=SimpleNamespace(id="agent-1"),
home_snapshot_id=None,
)
assert snapshot.snapshot_ref == "snapshot-ref-1"
assert snapshot.status is AgentWorkingResourceStatus.ACTIVE
session.add.assert_called_once_with(snapshot)
session.flush.assert_called_once_with()
def test_create_initial_flush_failure_does_not_delete_backend_snapshot(monkeypatch: pytest.MonkeyPatch) -> None:
session = MagicMock()
session.flush.side_effect = RuntimeError("flush failed")
client = _client()
monkeypatch.setattr(AgentHomeSnapshotService, "_client", lambda: nullcontext(client))
with pytest.raises(RuntimeError, match="flush failed"):
AgentHomeSnapshotService.create_initial(
session=session,
tenant_id="tenant-1",
agent_id="agent-1",
)
client.delete_home_snapshot_sync.assert_not_called()
session.scalar.assert_not_called()
@pytest.mark.parametrize(
@@ -98,7 +75,8 @@ def test_build_apply_checkpoints_exact_active_binding(
client = _client(snapshot_ref="snapshot-ref-2")
monkeypatch.setattr(AgentHomeSnapshotService, "_client", lambda: nullcontext(client))
monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", get_binding)
monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", MagicMock())
validate_generation = MagicMock()
monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", validate_generation)
snapshot = AgentHomeSnapshotService.create_for_build_apply(
session=session,
@@ -110,6 +88,32 @@ def test_build_apply_checkpoints_exact_active_binding(
request = client.create_home_snapshot_from_binding_sync.call_args.args[0]
assert request.backend_binding_ref == "binding-ref-1"
assert snapshot.snapshot_ref == "snapshot-ref-2"
assert validate_generation.call_args.kwargs["base_home_snapshot_id"] == "home-old"
def test_build_apply_forwards_default_home_generation(monkeypatch: pytest.MonkeyPatch) -> None:
session = MagicMock()
session.scalar.return_value = SimpleNamespace(app_id="app-1", backing_app_id=None)
binding = SimpleNamespace(
backend_binding_ref="binding-ref-1",
agent_id="agent-1",
base_home_snapshot_id=None,
agent_config_version_id="build-1",
agent_config_version_kind="build_draft",
)
client = _client(snapshot_ref="snapshot-ref-2")
validate_generation = MagicMock()
monkeypatch.setattr(AgentHomeSnapshotService, "_client", lambda: nullcontext(client))
monkeypatch.setattr(AgentWorkspaceService, "get_active_binding", MagicMock(return_value=binding))
monkeypatch.setattr(AgentWorkspaceService, "validate_binding_generation", validate_generation)
snapshot = AgentHomeSnapshotService.create_for_build_apply(
session=session,
build_draft=_build_draft(home_snapshot_id=None),
)
assert snapshot.snapshot_ref == "snapshot-ref-2"
assert validate_generation.call_args.kwargs["base_home_snapshot_id"] is None
def test_build_apply_fails_fast_without_source_binding() -> None:
@@ -15,7 +15,7 @@ from models.agent import (
AgentWorkspaceBinding,
AgentWorkspaceOwnerType,
)
from services.agent.workspace_service import AgentWorkspaceService, WorkspaceOwnerScope
from services.agent.workspace_service import AgentWorkspaceNotFoundError, AgentWorkspaceService, WorkspaceOwnerScope
def _scope() -> WorkspaceOwnerScope:
@@ -127,6 +127,59 @@ def test_create_binding_success_persists_new_workspace_and_binding(
request = client.create_execution_binding_sync.call_args.args[0]
assert request.existing_workspace_ref is None
assert request.workspace_id == workspace.id
assert request.home_snapshot_ref == "home-ref"
@pytest.mark.parametrize(
"sqlite_session",
[(AgentHomeSnapshot, AgentWorkspace, AgentWorkspaceBinding)],
indirect=True,
)
def test_create_binding_without_home_snapshot_uses_backend_default(
monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
) -> None:
client = _backend_client()
monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client))
binding = AgentWorkspaceService.create_binding(
session=sqlite_session,
scope=_scope(),
agent_id="agent-1",
base_home_snapshot_id=None,
agent_config_version_id="config-1",
agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT,
)
sqlite_session.commit()
stored_binding = sqlite_session.get(AgentWorkspaceBinding, binding.id)
assert stored_binding is not None
assert stored_binding.base_home_snapshot_id is None
request = client.create_execution_binding_sync.call_args.args[0]
assert request.home_snapshot_ref is None
@pytest.mark.parametrize(
"sqlite_session",
[(AgentHomeSnapshot, AgentWorkspace, AgentWorkspaceBinding)],
indirect=True,
)
def test_create_binding_rejects_missing_explicit_home_snapshot_before_backend_call(
monkeypatch: pytest.MonkeyPatch, sqlite_session: Session
) -> None:
client = _backend_client()
monkeypatch.setattr(AgentWorkspaceService, "_client", lambda: nullcontext(client))
with pytest.raises(AgentWorkspaceNotFoundError, match="base Home Snapshot is unavailable"):
AgentWorkspaceService.create_binding(
session=sqlite_session,
scope=_scope(),
agent_id="agent-1",
base_home_snapshot_id="missing-home",
agent_config_version_id="config-1",
agent_config_version_kind=AgentConfigVersionKind.SNAPSHOT,
)
client.create_execution_binding_sync.assert_not_called()
@pytest.mark.parametrize(
+29
View File
@@ -112,6 +112,31 @@ integration-up:
fi; \
sleep 2; \
done
$(eval CONTAINER_NAME_BROKEN_ISO := sandbox-rt-broken-iso-$(TEST_ID))
$(eval HOST_PORT_BROKEN_ISO := $(shell python3 -c 'import socket; s=socket.socket(); s.bind(("",0)); print(s.getsockname()[1]); s.close()'))
$(eval AUTH_TOKEN_BROKEN_ISO := test-token-broken-iso-$(TEST_ID))
@echo "Starting broken-isolation container $(CONTAINER_NAME_BROKEN_ISO) on port $(HOST_PORT_BROKEN_ISO)..."
docker run -d --name $(CONTAINER_NAME_BROKEN_ISO) \
-p $(HOST_PORT_BROKEN_ISO):5004 \
-e SHELLCTL_AUTH_TOKEN=$(AUTH_TOKEN_BROKEN_ISO) \
-e SHELLCTL_LANDLOCK_RO_PATHS=/etc/hosts \
$(IMAGE_NAME)
@echo 'CONTAINER_NAME_BROKEN_ISO=$(CONTAINER_NAME_BROKEN_ISO)' >> $(STATE_FILE)
@echo 'HOST_PORT_BROKEN_ISO=$(HOST_PORT_BROKEN_ISO)' >> $(STATE_FILE)
@echo 'AUTH_TOKEN_BROKEN_ISO=$(AUTH_TOKEN_BROKEN_ISO)' >> $(STATE_FILE)
@for i in $$(seq 1 30); do \
if curl -sf http://localhost:$(HOST_PORT_BROKEN_ISO)/healthz > /dev/null 2>&1; then \
echo "Broken-isolation runtime is ready on port $(HOST_PORT_BROKEN_ISO)"; \
break; \
fi; \
if [ "$$i" -eq 30 ]; then \
echo "ERROR: broken-isolation runtime not ready after 60s" >&2; \
docker logs $(CONTAINER_NAME_BROKEN_ISO); \
docker rm -f $(CONTAINER_NAME_BROKEN_ISO) 2>/dev/null; \
exit 1; \
fi; \
sleep 2; \
done
integration-test:
@test -f $(STATE_FILE) || { echo "ERROR: run 'make integration-up' first" >&2; exit 1; }
@@ -120,6 +145,9 @@ integration-test:
SHELLCTL_TEST_TOKEN=$$AUTH_TOKEN \
SHELLCTL_GO_URL_NO_ISOLATION=http://localhost:$$HOST_PORT_NOISO \
SHELLCTL_TEST_TOKEN_NO_ISOLATION=$$AUTH_TOKEN_NOISO \
SHELLCTL_GO_URL_BROKEN_ISOLATION=http://localhost:$$HOST_PORT_BROKEN_ISO \
SHELLCTL_TEST_TOKEN_BROKEN_ISOLATION=$$AUTH_TOKEN_BROKEN_ISO \
SHELLCTL_TEST_CONTAINER_BROKEN_ISOLATION=$$CONTAINER_NAME_BROKEN_ISO \
go test -tags=integration -v -count=1 -timeout=300s ./tests/...
integration-logs:
@@ -131,6 +159,7 @@ integration-down:
. ./$(STATE_FILE); \
docker rm -f $$CONTAINER_NAME 2>/dev/null || true; \
docker rm -f $$CONTAINER_NAME_NOISO 2>/dev/null || true; \
docker rm -f $$CONTAINER_NAME_BROKEN_ISO 2>/dev/null || true; \
rm -f $(STATE_FILE); \
fi
+15 -3
View File
@@ -59,18 +59,30 @@ Each agent job runs inside a Landlock sandbox that restricts filesystem access:
| Access | Paths (defaults) |
| -------------------- | ------------------------------------------------------------------------------------------------------------- |
| **Read-Write** | `$HOME` (always, includes `$CWD/.tmp` as `TMPDIR`) |
| **Read-Write** | `$HOME` and `$CWD` |
| **Read-Write (dev)** | `/dev/null`, `/dev/zero`, `/dev/urandom`, `/dev/random`, `/dev/tty` |
| **Read-Only + Exec** | `/usr`, `/bin`, `/sbin`, `/lib`, `/lib64`, `/etc`, `/proc`, `/opt/dify-agent-tools`, `/opt/homebrew`, `/snap` |
| **Denied** | Everything else (`/tmp`, other agents' homes, `/var`, `/srv`, etc.) |
The runner automatically creates `$CWD/.tmp` and sets `TMPDIR`, `TMP`, `TEMP` to it, so temp files stay isolated per workspace.
The runner creates no temp directory. It always sets `TMPDIR`, `TMP`, and
`TEMP` to `$CWD`, overriding command-provided values so temporary files remain
inside the existing writable layout.
When `SHELLCTL_ENABLE_PATH_ISOLATION=true`, Landlock setup is fail-closed: a
restriction error prevents the user script from executing.
Landlock denies opening, listing, or mutating paths outside the allow-list, but
Linux does not currently let Landlock restrict `stat(2)`. A job may therefore
observe sibling path metadata even though it cannot read or write sibling Home
content.
### Environment Variables
See [here](./internal/envvar/envvar.go)
Requires Linux ≥ 5.13. On unsupported kernels, a warning is printed to stderr.
Path isolation requires Landlock ABI V3 (upstream Linux 6.2+; actual availability
depends on the running kernel's reported ABI). Unsupported kernels or other
Landlock setup failures stop the job before its script runs.
## Dependencies
+50 -6
View File
@@ -4,21 +4,25 @@
package main
import (
"context"
"fmt"
"os"
"os/signal"
"strings"
"syscall"
"github.com/langgenius/dify/dify-agent-runtime/internal/agentcli"
"github.com/spf13/cobra"
)
// knownRootCommands mirrors the Python CLI's _KNOWN_ROOT_COMMANDS. Anything else
// on the command line is treated as argv forwarded to an implicit connect.
// knownRootCommands identifies names dispatched through the Cobra command tree.
// Any other bare first argument is treated as argv for an implicit connect.
var knownRootCommands = map[string]struct{}{
"config": {},
"connect": {},
"drive": {},
"file": {},
"config": {},
"connect": {},
"drive": {},
"file": {},
"home-snapshot": {},
}
func main() {
@@ -31,6 +35,9 @@ func main() {
//go:generate sh -c "go run . __dump-cli-help > ../../../dify-agent/src/dify_agent/layers/_agent_cli_help.json"
func runCLI(args []string) error {
ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
defer stop()
// Hidden developer intercept: emit the JSON help snapshot from the cobra
// tree. Not registered as a cobra command, so the visible command surface
// stays identical to the Python CLI. See dumphelp.go.
@@ -56,6 +63,7 @@ func runCLI(args []string) error {
}
root := newRootCommand()
root.SetContext(ctx)
root.SetArgs(args)
return root.Execute()
}
@@ -78,10 +86,46 @@ func newRootCommand() *cobra.Command {
newFileCommand(),
newDriveCommand(),
newConfigCommand(),
newHomeSnapshotCommand(),
)
return root
}
func newHomeSnapshotCommand() *cobra.Command {
command := &cobra.Command{
Use: "home-snapshot",
Short: "Transfer a backend-controlled Home Snapshot archive.",
Hidden: true,
}
var excludes []string
upload := &cobra.Command{
Use: "upload",
Short: "Stream the current HOME to the Home Snapshot gateway.",
Hidden: true,
Args: cobra.NoArgs,
RunE: withEnv(func(env *agentcli.Environment, _ []string, command *cobra.Command) error {
return agentcli.RunHomeSnapshotUpload(command.Context(), env, excludes)
}),
}
upload.Flags().StringArrayVar(&excludes, "exclude", nil, "Exclude one normalized path relative to HOME.")
if err := upload.Flags().MarkHidden("exclude"); err != nil {
panic(fmt.Sprintf("hide home-snapshot upload --exclude flag: %v", err))
}
download := &cobra.Command{
Use: "download",
Short: "Restore the Home Snapshot gateway archive into the current HOME.",
Hidden: true,
Args: cobra.NoArgs,
RunE: withEnv(func(env *agentcli.Environment, _ []string, command *cobra.Command) error {
return agentcli.RunHomeSnapshotDownload(command.Context(), env)
}),
}
command.AddCommand(upload, download)
return command
}
func newConnectCommand() *cobra.Command {
var jsonOutput bool
// connect is normally intercepted in runCLI; this definition exists so
@@ -226,3 +226,26 @@ func TestHiddenAliasesRemainHidden(t *testing.T) {
}
}
}
func TestHomeSnapshotControlCommandIsHiddenButDispatched(t *testing.T) {
root := newRootCommand()
command := findCommand(root, "home-snapshot")
if command == nil || !command.Hidden {
t.Fatal("home-snapshot must be registered as a hidden root command")
}
for _, name := range []string{"upload", "download"} {
child := findCommand(command, name)
if child == nil || !child.Hidden {
t.Fatalf("home-snapshot %s must be registered and hidden", name)
}
}
if strings.Contains(executeHelp(t, "--help"), "home-snapshot") {
t.Fatal("root help must not expose home-snapshot")
}
if uploadHelp := executeHelp(t, "home-snapshot", "upload", "--help"); strings.Contains(uploadHelp, "--exclude") {
t.Fatalf("home-snapshot upload help must not expose --exclude\n--- help output ---\n%s", uploadHelp)
}
if isUnknownBareCommand([]string{"home-snapshot"}) {
t.Fatal("home-snapshot must not fall through to implicit connect")
}
}
+32 -14
View File
@@ -76,6 +76,7 @@ func parentMode() {
})
envOverlay := loadEnvJSON(envPath)
removeIsolationOverrides(envOverlay)
env = mergeEnv(env, envOverlay)
// Ensure HOME exists.
@@ -84,13 +85,12 @@ func parentMode() {
cmdutil.HandleError(os.MkdirAll(home, 0755), 125, "mkdir HOME %s", home)
}
// Create a per-workspace temp directory under cwd and inject TMPDIR.
// This avoids granting RW access to the shared /tmp.
agentTmp := filepath.Join(cwd, ".tmp")
cmdutil.HandleError(os.MkdirAll(agentTmp, 0755), 125, "mkdir TMPDIR %s", agentTmp)
env = setEnvIfEmpty(env, "TMPDIR", agentTmp)
env = setEnvIfEmpty(env, "TMP", agentTmp)
env = setEnvIfEmpty(env, "TEMP", agentTmp)
// Temp files stay inside the already-authorized cwd. The runner does not
// create a separate directory and command-provided temp paths cannot widen
// the Landlock write surface.
env = setEnv(env, "TMPDIR", cwd)
env = setEnv(env, "TMP", cwd)
env = setEnv(env, "TEMP", cwd)
// Determine if path isolation is enabled.
enableIsolation := envvar.PathIsolationEnabled()
@@ -173,8 +173,7 @@ func childMode() {
jobDir := filepath.Dir(scriptPath)
cfg := landlock.ConfigFromEnv(home, cwd, jobDir)
if err := landlock.Restrict(cfg); err != nil {
// the landlock is best-effort, so we just log the error whatever it is
fmt.Fprintf(os.Stderr, "shellctl-runner: WARNING: %v — running without filesystem isolation\n", err)
cmdutil.HandleError(err, 126, "apply Landlock filesystem isolation")
}
}
@@ -234,6 +233,20 @@ func filterEnv(env []string, remove []string) []string {
return result
}
// removeIsolationOverrides keeps Landlock policy under deployment control.
// Job env may carry application variables, but it must not widen or replace
// the runner's inherited filesystem allow-list.
func removeIsolationOverrides(env map[string]string) {
for _, key := range []string{
envvar.EnvEnablePathIsolation,
envvar.EnvRWPaths,
envvar.EnvROPaths,
envvar.EnvRWDevPaths,
} {
delete(env, key)
}
}
// mergeEnv applies overlay key=value pairs onto the env slice.
func mergeEnv(env []string, overlay map[string]string) []string {
for k, v := range overlay {
@@ -264,12 +277,17 @@ func envGet(env []string, key string) string {
return ""
}
// setEnvIfEmpty sets key=value in the env slice only if the key is not already present.
func setEnvIfEmpty(env []string, key, value string) []string {
if envGet(env, key) != "" {
return env
// setEnv sets key=value in the env slice even when the command supplied a
// different value.
func setEnv(env []string, key, value string) []string {
prefix := key + "="
for index, entry := range env {
if strings.HasPrefix(entry, prefix) {
env[index] = prefix + value
return env
}
}
return append(env, key+"="+value)
return append(env, prefix+value)
}
// writeAtomic writes value to dest via a temp file + rename.
@@ -0,0 +1,33 @@
package main
import (
"testing"
"github.com/langgenius/dify/dify-agent-runtime/internal/envvar"
)
func TestRemoveIsolationOverridesPreservesApplicationEnvironment(t *testing.T) {
environment := map[string]string{
envvar.EnvEnablePathIsolation: "false",
envvar.EnvRWPaths: "/tmp",
envvar.EnvROPaths: "",
envvar.EnvRWDevPaths: "/dev",
"APPLICATION_VALUE": "kept",
}
removeIsolationOverrides(environment)
if got := environment["APPLICATION_VALUE"]; got != "kept" {
t.Fatalf("application environment = %q, want kept", got)
}
for _, key := range []string{
envvar.EnvEnablePathIsolation,
envvar.EnvRWPaths,
envvar.EnvROPaths,
envvar.EnvRWDevPaths,
} {
if _, exists := environment[key]; exists {
t.Fatalf("runner-controlled environment %q was not removed", key)
}
}
}
+2 -1
View File
@@ -26,7 +26,8 @@ ARG UV_VERSION=0.8.9
ENV PYTHONDONTWRITEBYTECODE=1 \
PYTHONUNBUFFERED=1 \
PIP_NO_CACHE_DIR=1
PIP_NO_CACHE_DIR=1 \
SHELLCTL_ENABLE_PATH_ISOLATION=true
RUN apt-get update \
&& apt-get install -y --no-install-recommends \
+1
View File
@@ -3,6 +3,7 @@ module github.com/langgenius/dify/dify-agent-runtime
go 1.26
require (
github.com/klauspost/compress v1.18.4
github.com/landlock-lsm/go-landlock v0.9.0
github.com/spf13/cobra v1.10.2
google.golang.org/grpc v1.82.1
+2
View File
@@ -17,6 +17,8 @@ github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
github.com/klauspost/compress v1.18.4 h1:RPhnKRAQ4Fh8zU2FY/6ZFDwTVTxgJ/EMydqSTzE9a2c=
github.com/klauspost/compress v1.18.4/go.mod h1:R0h/fSBs8DE4ENlcrlib3PsXS61voFxhIs2DeRhCvJ4=
github.com/landlock-lsm/go-landlock v0.9.0 h1:2q8G8yx9Hsd5bV+R6PJfgQl0zszNxC8KO+SIqGwfxlw=
github.com/landlock-lsm/go-landlock v0.9.0/go.mod h1:mn5GSi81Jf7yMs5WSi+SUi4sUeNLUGVdbT4Id6wXNQw=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
@@ -1,6 +1,9 @@
package agentcli
import "context"
import (
"context"
"io"
)
// StubClient abstracts Agent Stub control-plane and data-plane operations.
// Business logic depends only on this interface, never on HTTP/gRPC details.
@@ -22,6 +25,10 @@ type StubClient interface {
PatchConfigEnv(ctx context.Context, envText string) ([]byte, error)
PutConfigNote(ctx context.Context, note string) ([]byte, error)
// Home Snapshot transfer (HTTP-only, unbounded total duration).
UploadHomeSnapshot(ctx context.Context, archive io.Reader) error
DownloadHomeSnapshot(ctx context.Context) (io.ReadCloser, error)
// Data-plane (always HTTP, signed URLs)
UploadFileToURL(uploadURL, filePath, filename, mimetype string) ([]byte, error)
DownloadFromURL(downloadURL string) ([]byte, error)
@@ -4,6 +4,7 @@ import (
"context"
"encoding/json"
"fmt"
"io"
)
// httpStubClient implements StubClient using pure HTTP transport.
@@ -209,6 +210,14 @@ func (c *httpStubClient) PutConfigNote(_ context.Context, note string) ([]byte,
return body, nil
}
func (c *httpStubClient) UploadHomeSnapshot(ctx context.Context, archive io.Reader) error {
return c.http.uploadHomeSnapshot(ctx, archive)
}
func (c *httpStubClient) DownloadHomeSnapshot(ctx context.Context) (io.ReadCloser, error) {
return c.http.downloadHomeSnapshot(ctx)
}
func (c *httpStubClient) UploadFileToURL(uploadURL, filePath, filename, mimetype string) ([]byte, error) {
return c.http.uploadFile(uploadURL, filePath, filename, mimetype)
}
@@ -0,0 +1,380 @@
package agentcli
import (
"archive/tar"
"context"
"errors"
"fmt"
"io"
"io/fs"
"os"
pathpkg "path"
"path/filepath"
"strings"
"github.com/klauspost/compress/zstd"
)
const homeSnapshotCopyBufferSize = 1024 * 1024
// RunHomeSnapshotUpload streams a tar.zst representation of $HOME to the
// purpose-scoped Agent Stub gateway.
func RunHomeSnapshotUpload(ctx context.Context, env *Environment, excludes []string) error {
home, err := snapshotHome()
if err != nil {
return err
}
normalizedExcludes, err := normalizeSnapshotExcludes(home, excludes)
if err != nil {
return err
}
client, err := NewStubClient(env)
if err != nil {
return err
}
defer func() { _ = client.Close() }()
reader, writer := io.Pipe()
producerDone := make(chan error, 1)
go func() {
producerErr := writeHomeSnapshotArchive(ctx, writer, home, normalizedExcludes)
_ = writer.CloseWithError(producerErr)
producerDone <- producerErr
}()
uploadErr := client.UploadHomeSnapshot(ctx, reader)
if uploadErr != nil {
_ = reader.CloseWithError(uploadErr)
} else {
_ = reader.Close()
}
producerErr := <-producerDone
if uploadErr != nil {
return uploadErr
}
if producerErr != nil {
return producerErr
}
return nil
}
// RunHomeSnapshotDownload restores one gateway archive into an existing empty
// $HOME. The backend owns cleanup if any validation or write fails.
func RunHomeSnapshotDownload(ctx context.Context, env *Environment) error {
home, err := snapshotHome()
if err != nil {
return err
}
entries, err := os.ReadDir(home)
if err != nil {
return fmt.Errorf("read HOME: %w", err)
}
if len(entries) != 0 {
return errors.New("HOME must be empty before Home Snapshot download")
}
client, err := NewStubClient(env)
if err != nil {
return err
}
defer func() { _ = client.Close() }()
body, err := client.DownloadHomeSnapshot(ctx)
if err != nil {
return err
}
defer func() { _ = body.Close() }()
return extractHomeSnapshotArchive(body, home)
}
func snapshotHome() (string, error) {
home := strings.TrimSpace(os.Getenv("HOME"))
if home == "" {
return "", errors.New("HOME must be set")
}
if !filepath.IsAbs(home) {
return "", errors.New("HOME must be an absolute path")
}
home = filepath.Clean(home)
info, err := os.Stat(home)
if err != nil {
return "", fmt.Errorf("stat HOME: %w", err)
}
if !info.IsDir() {
return "", errors.New("HOME must be a directory")
}
return home, nil
}
func normalizeSnapshotExcludes(home string, excludes []string) (map[string]struct{}, error) {
normalized := make(map[string]struct{}, len(excludes))
for _, excluded := range excludes {
if excluded == "" || filepath.IsAbs(excluded) || strings.Contains(excluded, "\\") {
return nil, fmt.Errorf("invalid Home Snapshot exclude %q", excluded)
}
if strings.ContainsAny(excluded, "*?[") {
return nil, fmt.Errorf("Home Snapshot exclude must not be a glob: %q", excluded)
}
cleaned := filepath.Clean(excluded)
if cleaned == "." || cleaned == ".." || strings.HasPrefix(cleaned, ".."+string(filepath.Separator)) {
return nil, fmt.Errorf("invalid Home Snapshot exclude %q", excluded)
}
if cleaned != excluded {
return nil, fmt.Errorf("Home Snapshot exclude must be normalized: %q", excluded)
}
candidate := filepath.Join(home, cleaned)
relative, err := filepath.Rel(home, candidate)
if err != nil || relative != cleaned {
return nil, fmt.Errorf("Home Snapshot exclude escapes HOME: %q", excluded)
}
normalized[filepath.ToSlash(cleaned)] = struct{}{}
}
return normalized, nil
}
func writeHomeSnapshotArchive(
ctx context.Context,
destination io.Writer,
home string,
excludes map[string]struct{},
) (retErr error) {
encoder, err := zstd.NewWriter(
destination,
zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(1)),
)
if err != nil {
return fmt.Errorf("create zstd encoder: %w", err)
}
tarWriter := tar.NewWriter(encoder)
defer func() {
if err := tarWriter.Close(); retErr == nil && err != nil {
retErr = fmt.Errorf("close tar archive: %w", err)
}
if err := encoder.Close(); retErr == nil && err != nil {
retErr = fmt.Errorf("close zstd archive: %w", err)
}
}()
copyBuffer := make([]byte, homeSnapshotCopyBufferSize)
err = filepath.WalkDir(home, func(filePath string, entry fs.DirEntry, walkErr error) error {
if walkErr != nil {
return walkErr
}
if err := ctx.Err(); err != nil {
return err
}
if filePath == home {
return nil
}
relative, err := filepath.Rel(home, filePath)
if err != nil {
return err
}
entryName := filepath.ToSlash(relative)
if isSnapshotExcluded(entryName, excludes) {
if entry.IsDir() {
return filepath.SkipDir
}
return nil
}
info, err := entry.Info()
if err != nil {
return err
}
mode := info.Mode()
if !mode.IsRegular() && !mode.IsDir() && mode&os.ModeSymlink == 0 {
return nil
}
linkTarget := ""
if mode&os.ModeSymlink != 0 {
linkTarget, err = os.Readlink(filePath)
if err != nil {
return err
}
linkTarget = relocatableSymlinkTarget(home, filePath, linkTarget, excludes)
}
header, err := tar.FileInfoHeader(info, linkTarget)
if err != nil {
return err
}
header.Name = entryName
header.Uid = 0
header.Gid = 0
header.Uname = ""
header.Gname = ""
if err := tarWriter.WriteHeader(header); err != nil {
return err
}
if !mode.IsRegular() {
return nil
}
file, err := os.Open(filePath)
if err != nil {
return err
}
_, copyErr := io.CopyBuffer(tarWriter, &contextReader{ctx: ctx, reader: file}, copyBuffer)
closeErr := file.Close()
if copyErr != nil {
return copyErr
}
return closeErr
})
if err != nil {
return fmt.Errorf("pack HOME: %w", err)
}
return nil
}
func relocatableSymlinkTarget(
home string,
linkPath string,
target string,
excludes map[string]struct{},
) string {
if !filepath.IsAbs(target) {
return target
}
originalTarget := target
cleanedTarget := filepath.Clean(target)
relative, err := filepath.Rel(home, cleanedTarget)
if err != nil || relative == ".." || strings.HasPrefix(relative, ".."+string(filepath.Separator)) {
return originalTarget
}
if isSnapshotExcluded(filepath.ToSlash(relative), excludes) {
return originalTarget
}
relocated, err := filepath.Rel(filepath.Dir(linkPath), cleanedTarget)
if err != nil {
return originalTarget
}
return filepath.ToSlash(relocated)
}
func isSnapshotExcluded(entryName string, excludes map[string]struct{}) bool {
for excluded := range excludes {
if entryName == excluded || strings.HasPrefix(entryName, excluded+"/") {
return true
}
}
return false
}
func extractHomeSnapshotArchive(source io.Reader, home string) error {
decoder, err := zstd.NewReader(source)
if err != nil {
return fmt.Errorf("create zstd decoder: %w", err)
}
defer decoder.Close()
tarReader := tar.NewReader(decoder)
copyBuffer := make([]byte, homeSnapshotCopyBufferSize)
seen := map[string]struct{}{}
type directoryMode struct {
path string
mode os.FileMode
}
var directoryModes []directoryMode
for {
header, err := tarReader.Next()
if errors.Is(err, io.EOF) {
break
}
if err != nil {
return fmt.Errorf("read tar archive: %w", err)
}
entryName, err := normalizedArchiveEntryName(header)
if err != nil {
return err
}
if _, duplicate := seen[entryName]; duplicate {
return fmt.Errorf("duplicate Home Snapshot entry %q", header.Name)
}
seen[entryName] = struct{}{}
destination := filepath.Join(home, filepath.FromSlash(entryName))
relative, err := filepath.Rel(home, destination)
if err != nil || filepath.ToSlash(relative) != entryName {
return fmt.Errorf("Home Snapshot entry escapes HOME: %q", header.Name)
}
switch header.Typeflag {
case tar.TypeDir:
if err := os.MkdirAll(destination, 0o700); err != nil {
return err
}
directoryModes = append(directoryModes, directoryMode{
path: destination,
mode: os.FileMode(header.Mode) & os.ModePerm,
})
case tar.TypeReg, tar.TypeRegA:
if err := os.MkdirAll(filepath.Dir(destination), 0o700); err != nil {
return err
}
file, err := os.OpenFile(destination, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0o600)
if err != nil {
return err
}
_, copyErr := io.CopyBuffer(file, tarReader, copyBuffer)
chmodErr := file.Chmod(os.FileMode(header.Mode) & os.ModePerm)
closeErr := file.Close()
if copyErr != nil {
return copyErr
}
if chmodErr != nil {
return chmodErr
}
if closeErr != nil {
return closeErr
}
case tar.TypeSymlink:
if err := os.MkdirAll(filepath.Dir(destination), 0o700); err != nil {
return err
}
if err := os.Symlink(header.Linkname, destination); err != nil {
return err
}
case tar.TypeLink:
return fmt.Errorf("Home Snapshot hardlink entry is not allowed: %q", header.Name)
default:
return fmt.Errorf("Home Snapshot special entry is not allowed: %q", header.Name)
}
}
if _, err := io.CopyBuffer(io.Discard, decoder, copyBuffer); err != nil {
return fmt.Errorf("finish zstd archive: %w", err)
}
for index := len(directoryModes) - 1; index >= 0; index-- {
if err := os.Chmod(directoryModes[index].path, directoryModes[index].mode); err != nil {
return err
}
}
return nil
}
func normalizedArchiveEntryName(header *tar.Header) (string, error) {
name := header.Name
if header.Typeflag == tar.TypeDir {
name = strings.TrimSuffix(name, "/")
}
if name == "" || strings.Contains(name, "\\") || pathpkg.IsAbs(name) {
return "", fmt.Errorf("invalid Home Snapshot entry path %q", header.Name)
}
cleaned := pathpkg.Clean(name)
if cleaned == "." || cleaned == ".." || strings.HasPrefix(cleaned, "../") || cleaned != name {
return "", fmt.Errorf("invalid Home Snapshot entry path %q", header.Name)
}
return cleaned, nil
}
type contextReader struct {
ctx context.Context
reader io.Reader
}
func (r *contextReader) Read(buffer []byte) (int, error) {
select {
case <-r.ctx.Done():
return 0, r.ctx.Err()
default:
return r.reader.Read(buffer)
}
}
@@ -0,0 +1,551 @@
package agentcli
import (
"archive/tar"
"bytes"
"context"
"crypto/rand"
"errors"
"fmt"
"io"
"net/http"
"net/http/httptest"
"os"
"path/filepath"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/klauspost/compress/zstd"
)
func TestHomeSnapshotArchiveRoundTripAndSymlinkRelocation(t *testing.T) {
source := t.TempDir()
target := t.TempDir()
mustMkdirAll(t, filepath.Join(source, "data"))
mustWriteFile(t, filepath.Join(source, "data", "value.txt"), []byte("checkpoint"))
mustMkdirAll(t, filepath.Join(source, "workspace"))
mustWriteFile(t, filepath.Join(source, "workspace", "shared.txt"), []byte("excluded"))
mustSymlink(t, "data/value.txt", filepath.Join(source, "relative"))
mustSymlink(t, filepath.Join(source, "data", "value.txt"), filepath.Join(source, "absolute-internal"))
mustSymlink(t, "/usr/bin/env", filepath.Join(source, "absolute-external"))
absoluteExternalWithDotDot := source + string(filepath.Separator) + ".." + string(filepath.Separator) + "outside"
mustSymlink(t, absoluteExternalWithDotDot, filepath.Join(source, "absolute-external-dotdot"))
mustSymlink(t, filepath.Join(source, "workspace", "shared.txt"), filepath.Join(source, "absolute-excluded"))
excludes, err := normalizeSnapshotExcludes(source, []string{"workspace"})
if err != nil {
t.Fatalf("normalize excludes: %v", err)
}
var archive bytes.Buffer
if err := writeHomeSnapshotArchive(context.Background(), &archive, source, excludes); err != nil {
t.Fatalf("write archive: %v", err)
}
if err := extractHomeSnapshotArchive(bytes.NewReader(archive.Bytes()), target); err != nil {
t.Fatalf("extract archive: %v", err)
}
content, err := os.ReadFile(filepath.Join(target, "data", "value.txt"))
if err != nil || string(content) != "checkpoint" {
t.Fatalf("restored content = %q, %v", content, err)
}
assertSymlinkTarget(t, filepath.Join(target, "relative"), "data/value.txt")
assertSymlinkTarget(t, filepath.Join(target, "absolute-internal"), "data/value.txt")
assertSymlinkTarget(t, filepath.Join(target, "absolute-external"), "/usr/bin/env")
assertSymlinkTarget(t, filepath.Join(target, "absolute-external-dotdot"), absoluteExternalWithDotDot)
assertSymlinkTarget(t, filepath.Join(target, "absolute-excluded"), filepath.Join(source, "workspace", "shared.txt"))
if _, err := os.Lstat(filepath.Join(target, "workspace")); !os.IsNotExist(err) {
t.Fatalf("excluded workspace was restored: %v", err)
}
}
func TestNormalizeSnapshotExcludesRejectsUnsafeValues(t *testing.T) {
home := t.TempDir()
for _, value := range []string{"", ".", "..", "../outside", "/absolute", "foo/../bar", "foo*", `foo\bar`} {
t.Run(value, func(t *testing.T) {
if _, err := normalizeSnapshotExcludes(home, []string{value}); err == nil {
t.Fatalf("expected exclude %q to fail", value)
}
})
}
}
func TestExtractHomeSnapshotArchiveRejectsUnsafeEntries(t *testing.T) {
cases := []tar.Header{
{Name: "/absolute", Typeflag: tar.TypeReg, Mode: 0o600},
{Name: "../outside", Typeflag: tar.TypeReg, Mode: 0o600},
{Name: "not/../normalized", Typeflag: tar.TypeReg, Mode: 0o600},
{Name: "hardlink", Typeflag: tar.TypeLink, Linkname: "target"},
{Name: "fifo", Typeflag: tar.TypeFifo},
}
for _, header := range cases {
t.Run(header.Name, func(t *testing.T) {
archive := compressedTar(t, header)
if err := extractHomeSnapshotArchive(bytes.NewReader(archive), t.TempDir()); err == nil {
t.Fatalf("expected entry %#v to be rejected", header)
}
})
}
}
func TestExtractHomeSnapshotArchiveRejectsCorruptStream(t *testing.T) {
if err := extractHomeSnapshotArchive(bytes.NewReader([]byte("not-zstd")), t.TempDir()); err == nil {
t.Fatal("expected corrupt zstd stream to fail")
}
}
func TestExtractHomeSnapshotArchiveRejectsTruncatedZstdTail(t *testing.T) {
archive := compressedTar(t, tar.Header{Name: "file", Typeflag: tar.TypeReg, Mode: 0o600})
truncated := archive[:len(archive)-1]
if err := extractHomeSnapshotArchive(bytes.NewReader(truncated), t.TempDir()); err == nil {
t.Fatal("expected truncated zstd tail to fail")
}
}
func TestExtractHomeSnapshotArchiveDoesNotRestoreSetIDBits(t *testing.T) {
target := t.TempDir()
archive := compressedTar(t, tar.Header{Name: "file", Typeflag: tar.TypeReg, Mode: 0o6755})
if err := extractHomeSnapshotArchive(bytes.NewReader(archive), target); err != nil {
t.Fatalf("extract archive: %v", err)
}
info, err := os.Stat(filepath.Join(target, "file"))
if err != nil {
t.Fatal(err)
}
if info.Mode()&(os.ModeSetuid|os.ModeSetgid) != 0 {
t.Fatalf("restored privileged mode bits: %v", info.Mode())
}
if info.Mode().Perm() != 0o755 {
t.Fatalf("restored permissions = %o, want 755", info.Mode().Perm())
}
}
func TestRunHomeSnapshotUploadRejectsInvalidHome(t *testing.T) {
homeFile := filepath.Join(t.TempDir(), "not-a-directory")
mustWriteFile(t, homeFile, []byte("file"))
for name, home := range map[string]string{
"empty": "",
"relative": "relative/home",
"missing": filepath.Join(t.TempDir(), "missing"),
"regular-file": homeFile,
} {
t.Run(name, func(t *testing.T) {
t.Setenv("HOME", home)
if err := RunHomeSnapshotUpload(context.Background(), &Environment{}, nil); err == nil {
t.Fatalf("expected HOME %q to be rejected", home)
}
})
}
}
func TestRunHomeSnapshotDownloadRejectsNonEmptyHome(t *testing.T) {
home := t.TempDir()
mustWriteFile(t, filepath.Join(home, "existing"), []byte("keep"))
t.Setenv("HOME", home)
if err := RunHomeSnapshotDownload(context.Background(), &Environment{}); err == nil {
t.Fatal("expected non-empty HOME to be rejected")
}
content, err := os.ReadFile(filepath.Join(home, "existing"))
if err != nil || string(content) != "keep" {
t.Fatalf("existing HOME content changed: %q, %v", content, err)
}
}
func TestRunHomeSnapshotUploadAndDownloadStreamThroughExistingStubEnvironment(t *testing.T) {
source := t.TempDir()
mustWriteFile(t, filepath.Join(source, "checkpoint.txt"), []byte("streamed"))
var archive []byte
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
if request.URL.Path != "/agent-stub/home-snapshots/archive" {
http.NotFound(response, request)
return
}
if request.Header.Get("Authorization") != "Bearer purpose-token" {
http.Error(response, "unauthorized", http.StatusUnauthorized)
return
}
switch request.Method {
case http.MethodPut:
var err error
archive, err = io.ReadAll(request.Body)
if err != nil {
http.Error(response, err.Error(), http.StatusBadRequest)
return
}
response.WriteHeader(http.StatusNoContent)
case http.MethodGet:
_, _ = response.Write(archive)
default:
response.WriteHeader(http.StatusMethodNotAllowed)
}
}))
defer server.Close()
env := &Environment{URL: server.URL + "/agent-stub", AuthJWE: "purpose-token"}
t.Setenv("HOME", source)
if err := RunHomeSnapshotUpload(context.Background(), env, nil); err != nil {
t.Fatalf("upload: %v", err)
}
if len(archive) == 0 {
t.Fatal("gateway received no archive bytes")
}
sourceContent, err := os.ReadFile(filepath.Join(source, "checkpoint.txt"))
if err != nil || string(sourceContent) != "streamed" {
t.Fatalf("upload changed source HOME: %q, %v", sourceContent, err)
}
sourceEntries, err := os.ReadDir(source)
if err != nil || len(sourceEntries) != 1 {
t.Fatalf("upload left source artifacts: %v, %v", sourceEntries, err)
}
target := t.TempDir()
t.Setenv("HOME", target)
if err := RunHomeSnapshotDownload(context.Background(), env); err != nil {
t.Fatalf("download: %v", err)
}
content, err := os.ReadFile(filepath.Join(target, "checkpoint.txt"))
if err != nil || string(content) != "streamed" {
t.Fatalf("restored content = %q, %v", content, err)
}
entries, err := os.ReadDir(target)
if err != nil || len(entries) != 1 {
t.Fatalf("target contains unexpected temp artifacts: %v, %v", entries, err)
}
}
func TestRunHomeSnapshotUploadPropagatesGatewayFailure(t *testing.T) {
home := t.TempDir()
mustWriteFile(t, filepath.Join(home, "checkpoint.txt"), []byte(strings.Repeat("x", 1024)))
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
http.Error(response, "store unavailable", http.StatusBadGateway)
}))
defer server.Close()
t.Setenv("HOME", home)
err := RunHomeSnapshotUpload(
context.Background(),
&Environment{URL: server.URL + "/agent-stub", AuthJWE: "purpose-token"},
nil,
)
if err == nil || !strings.Contains(err.Error(), "502") {
t.Fatalf("upload error = %v", err)
}
}
func TestRunHomeSnapshotUploadEarlyTerminationJoinsProducer(t *testing.T) {
for _, test := range []struct {
name string
cancelOnStart bool
}{
{name: "http-failure"},
{name: "cancellation", cancelOnStart: true},
} {
t.Run(test.name, func(t *testing.T) {
home := t.TempDir()
content := make([]byte, 16*1024*1024)
if _, err := io.ReadFull(rand.Reader, content); err != nil {
t.Fatal(err)
}
mustWriteFile(t, filepath.Join(home, "large.bin"), content)
t.Setenv("HOME", home)
requestStarted := make(chan struct{})
handlerFinished := make(chan struct{})
releaseHandler := make(chan struct{})
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
defer close(handlerFinished)
close(requestStarted)
if test.cancelOnStart {
select {
case <-releaseHandler:
case <-time.After(5 * time.Second):
http.Error(response, "test handler timed out", http.StatusGatewayTimeout)
}
return
}
http.Error(response, "early failure", http.StatusBadGateway)
}))
defer server.Close()
handlerReleased := false
defer func() {
if !handlerReleased {
close(releaseHandler)
}
}()
ctx, cancel := context.WithCancel(context.Background())
uploadAndProducerFinished := make(chan error, 1)
uploadFinished := false
go func() {
uploadAndProducerFinished <- RunHomeSnapshotUpload(
ctx,
&Environment{URL: server.URL + "/agent-stub", AuthJWE: "purpose-token"},
nil,
)
}()
defer func() {
cancel()
if !uploadFinished {
select {
case <-uploadAndProducerFinished:
case <-time.After(5 * time.Second):
t.Errorf("owned upload goroutine did not finish during cleanup")
}
}
}()
select {
case <-requestStarted:
case <-time.After(5 * time.Second):
t.Fatal("upload request never reached the HTTP handler")
}
if test.cancelOnStart {
cancel()
}
var uploadErr error
select {
case uploadErr = <-uploadAndProducerFinished:
uploadFinished = true
case <-time.After(5 * time.Second):
t.Fatal("upload and its archive producer did not finish after early termination")
}
if uploadErr == nil {
t.Fatal("expected upload to fail")
}
if test.cancelOnStart {
if !errors.Is(uploadErr, context.Canceled) && !errors.Is(uploadErr, io.ErrClosedPipe) {
t.Fatalf("canceled upload error = %v", uploadErr)
}
close(releaseHandler)
handlerReleased = true
} else if !strings.Contains(uploadErr.Error(), "502") {
t.Fatalf("HTTP failure upload error = %v", uploadErr)
}
select {
case <-handlerFinished:
case <-time.After(5 * time.Second):
t.Fatal("HTTP handler did not finish after upload termination")
}
})
}
}
func TestHomeSnapshotTransfersRejectNonSuccessStatuses(t *testing.T) {
for _, status := range []int{http.StatusBadRequest, http.StatusInternalServerError} {
for _, operation := range []string{"upload", "download"} {
t.Run(fmt.Sprintf("%s-%d", operation, status), func(t *testing.T) {
client := NewHTTPClient(&Environment{URL: "http://agent-stub.invalid", AuthJWE: "purpose-token"})
client.homeSnapshotClient = httpDoerFunc(func(*http.Request) (*http.Response, error) {
return &http.Response{
StatusCode: status,
Body: io.NopCloser(strings.NewReader(`{"detail":"transfer failed"}`)),
}, nil
})
var err error
switch operation {
case "upload":
err = client.uploadHomeSnapshot(context.Background(), strings.NewReader("archive"))
case "download":
var body io.ReadCloser
body, err = client.downloadHomeSnapshot(context.Background())
if body != nil {
_ = body.Close()
t.Fatal("download returned a body for a non-success response")
}
default:
t.Fatalf("unknown operation %q", operation)
}
if err == nil || !strings.Contains(err.Error(), fmt.Sprintf("HTTP %d", status)) {
t.Fatalf("%s error = %v", operation, err)
}
})
}
}
}
func TestHomeSnapshotTransfersRejectRedirects(t *testing.T) {
var followedRedirects atomic.Int32
redirectTarget := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
followedRedirects.Add(1)
response.WriteHeader(http.StatusNoContent)
}))
defer redirectTarget.Close()
redirectSource := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
http.Redirect(response, request, redirectTarget.URL, http.StatusTemporaryRedirect)
}))
defer redirectSource.Close()
client := NewHTTPClient(&Environment{URL: redirectSource.URL, AuthJWE: "purpose-token"})
tests := []struct {
name string
transfer func() error
}{
{
name: "upload",
transfer: func() error {
return client.uploadHomeSnapshot(context.Background(), strings.NewReader("archive"))
},
},
{
name: "download",
transfer: func() error {
body, err := client.downloadHomeSnapshot(context.Background())
if body != nil {
_ = body.Close()
}
return err
},
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
err := test.transfer()
if err == nil || !strings.Contains(err.Error(), "307") {
t.Fatalf("%s redirect error = %v", test.name, err)
}
})
}
if got := followedRedirects.Load(); got != 0 {
t.Fatalf("Home Snapshot client followed %d redirects", got)
}
}
func TestRunHomeSnapshotDownloadHonorsContextCancellation(t *testing.T) {
handlerStarted := make(chan struct{})
handlerFinished := make(chan struct{})
server := httptest.NewServer(http.HandlerFunc(func(response http.ResponseWriter, request *http.Request) {
defer close(handlerFinished)
response.WriteHeader(http.StatusOK)
response.(http.Flusher).Flush()
close(handlerStarted)
select {
case <-request.Context().Done():
case <-time.After(5 * time.Second):
http.Error(response, "test handler timed out", http.StatusGatewayTimeout)
}
}))
defer server.Close()
t.Setenv("HOME", t.TempDir())
ctx, cancel := context.WithCancel(context.Background())
result := make(chan error, 1)
downloadFinished := false
handlerDone := false
go func() {
result <- RunHomeSnapshotDownload(
ctx,
&Environment{URL: server.URL + "/agent-stub", AuthJWE: "purpose-token"},
)
}()
defer func() {
cancel()
if !downloadFinished {
select {
case <-result:
case <-time.After(5 * time.Second):
t.Errorf("owned download goroutine did not finish during cleanup")
}
}
if !handlerDone {
select {
case <-handlerFinished:
case <-time.After(5 * time.Second):
t.Errorf("HTTP handler did not finish during cleanup")
}
}
}()
select {
case <-handlerStarted:
case <-time.After(5 * time.Second):
t.Fatal("download request never reached the HTTP handler")
}
cancel()
var downloadErr error
select {
case downloadErr = <-result:
downloadFinished = true
case <-time.After(5 * time.Second):
t.Fatal("download did not return promptly after cancellation")
}
if downloadErr == nil {
t.Fatal("expected canceled download to fail")
}
if !errors.Is(downloadErr, context.Canceled) && !strings.Contains(downloadErr.Error(), context.Canceled.Error()) {
t.Fatalf("canceled download error = %v", downloadErr)
}
select {
case <-handlerFinished:
handlerDone = true
case <-time.After(5 * time.Second):
t.Fatal("HTTP handler did not finish after download cancellation")
}
}
func compressedTar(t *testing.T, header tar.Header) []byte {
t.Helper()
var output bytes.Buffer
encoder, err := zstd.NewWriter(&output, zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(1)))
if err != nil {
t.Fatal(err)
}
writer := tar.NewWriter(encoder)
if err := writer.WriteHeader(&header); err != nil {
t.Fatal(err)
}
if err := writer.Close(); err != nil {
t.Fatal(err)
}
if err := encoder.Close(); err != nil {
t.Fatal(err)
}
return output.Bytes()
}
func mustMkdirAll(t *testing.T, path string) {
t.Helper()
if err := os.MkdirAll(path, 0o700); err != nil {
t.Fatal(err)
}
}
func mustWriteFile(t *testing.T, path string, content []byte) {
t.Helper()
if err := os.WriteFile(path, content, 0o600); err != nil {
t.Fatal(err)
}
}
func mustSymlink(t *testing.T, target, path string) {
t.Helper()
if err := os.Symlink(target, path); err != nil {
t.Fatal(err)
}
}
func assertSymlinkTarget(t *testing.T, path, want string) {
t.Helper()
got, err := os.Readlink(path)
if err != nil {
t.Fatal(err)
}
if got != want {
t.Fatalf("symlink %s target = %q, want %q", path, got, want)
}
}
type httpDoerFunc func(request *http.Request) (*http.Response, error)
func (function httpDoerFunc) Do(request *http.Request) (*http.Response, error) {
return function(request)
}
@@ -2,7 +2,9 @@ package agentcli
import (
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"mime/multipart"
@@ -12,28 +14,45 @@ import (
"time"
)
const homeSnapshotArchivePath = "/home-snapshots/archive"
// HTTPClient wraps HTTP interactions with the Agent Stub server.
type HTTPClient struct {
baseURL string
authJWE string
client *http.Client
baseURL string
authJWE string
client *http.Client
homeSnapshotClient httpDoer
}
type httpDoer interface {
Do(request *http.Request) (*http.Response, error)
}
func newHomeSnapshotHTTPClient() *http.Client {
return &http.Client{
CheckRedirect: func(_ *http.Request, _ []*http.Request) error {
return http.ErrUseLastResponse
},
}
}
// NewHTTPClient creates a new HTTP client for the Agent Stub API.
func NewHTTPClient(env *Environment) *HTTPClient {
return &HTTPClient{
baseURL: env.URL,
authJWE: env.AuthJWE,
client: &http.Client{Timeout: 30 * time.Second},
baseURL: env.URL,
authJWE: env.AuthJWE,
client: &http.Client{Timeout: 30 * time.Second},
homeSnapshotClient: newHomeSnapshotHTTPClient(),
}
}
// NewHTTPClientWithTimeout creates a client with a custom timeout.
func NewHTTPClientWithTimeout(env *Environment, timeout time.Duration) *HTTPClient {
return &HTTPClient{
baseURL: env.URL,
authJWE: env.AuthJWE,
client: &http.Client{Timeout: timeout},
baseURL: env.URL,
authJWE: env.AuthJWE,
client: &http.Client{Timeout: timeout},
homeSnapshotClient: newHomeSnapshotHTTPClient(),
}
}
@@ -215,9 +234,58 @@ func (c *HTTPClient) downloadFromURL(downloadURL string) ([]byte, error) {
return io.ReadAll(resp.Body)
}
// checkHTTPError returns a formatted error if status >= 400.
func (c *HTTPClient) uploadHomeSnapshot(ctx context.Context, archive io.Reader) error {
req, err := http.NewRequestWithContext(ctx, http.MethodPut, c.baseURL+homeSnapshotArchivePath, archive)
if err != nil {
return fmt.Errorf("create Home Snapshot upload request: %w", err)
}
req.Header.Set("Authorization", "Bearer "+c.authJWE)
req.Header.Set("Content-Type", "application/zstd")
resp, err := c.homeSnapshotClient.Do(req)
if err != nil {
return fmt.Errorf("Home Snapshot upload request failed: %w", err)
}
if resp == nil || resp.Body == nil {
return errors.New("Home Snapshot upload response is missing a body")
}
defer func() { _ = resp.Body.Close() }()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
body, _ := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
return checkHTTPError(body, resp.StatusCode, "Home Snapshot upload")
}
_, err = io.Copy(io.Discard, resp.Body)
if err != nil {
return fmt.Errorf("read Home Snapshot upload response: %w", err)
}
return nil
}
func (c *HTTPClient) downloadHomeSnapshot(ctx context.Context) (io.ReadCloser, error) {
req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.baseURL+homeSnapshotArchivePath, nil)
if err != nil {
return nil, fmt.Errorf("create Home Snapshot download request: %w", err)
}
req.Header.Set("Authorization", "Bearer "+c.authJWE)
resp, err := c.homeSnapshotClient.Do(req)
if err != nil {
return nil, fmt.Errorf("Home Snapshot download request failed: %w", err)
}
if resp == nil || resp.Body == nil {
return nil, errors.New("Home Snapshot download response is missing a body")
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
defer func() { _ = resp.Body.Close() }()
body, _ := io.ReadAll(io.LimitReader(resp.Body, 64*1024))
return nil, checkHTTPError(body, resp.StatusCode, "Home Snapshot download")
}
return resp.Body, nil
}
// checkHTTPError returns a formatted error for any non-2xx status.
func checkHTTPError(body []byte, statusCode int, operation string) error {
if statusCode < 400 {
if statusCode >= 200 && statusCode < 300 {
return nil
}
var detail struct {
@@ -62,7 +62,9 @@ func DefaultConfig(home, cwd, jobDir string) *Config {
// falling back to defaults for any unset variable.
//
// If a variable is set (even to empty string), its value replaces the default.
// Set to empty to grant no additional paths beyond $HOME.
// An empty value removes only that category's additional configured paths;
// buildRules still grants mandatory read-write access to HOME and Cwd and
// read-only access to JobDir.
func ConfigFromEnv(home, cwd, jobDir string) *Config {
cfg := DefaultConfig(home, cwd, jobDir)
@@ -8,11 +8,11 @@ import (
golandlock "github.com/landlock-lsm/go-landlock/landlock"
)
// Restrict applies Landlock V1 filesystem restrictions for the current process.
// Restrict applies Landlock V3 filesystem restrictions for the current process.
func Restrict(cfg *Config) error {
rules := buildRules(cfg)
if err := golandlock.V1.RestrictPaths(rules...); err != nil {
if err := golandlock.V3.RestrictPaths(rules...); err != nil {
return err
}
return nil
+285 -17
View File
@@ -11,15 +11,20 @@
package tests
import (
"archive/tar"
"bytes"
"encoding/base64"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"os/exec"
"strings"
"testing"
"time"
"github.com/klauspost/compress/zstd"
)
var (
@@ -29,6 +34,10 @@ var (
goURLNoIsolation = os.Getenv("SHELLCTL_GO_URL_NO_ISOLATION")
authTokenNoIsolation = os.Getenv("SHELLCTL_TEST_TOKEN_NO_ISOLATION")
goURLBrokenIsolation = os.Getenv("SHELLCTL_GO_URL_BROKEN_ISOLATION")
authTokenBrokenIsolation = os.Getenv("SHELLCTL_TEST_TOKEN_BROKEN_ISOLATION")
containerBrokenIsolation = os.Getenv("SHELLCTL_TEST_CONTAINER_BROKEN_ISOLATION")
httpClient = &http.Client{Timeout: 120 * time.Second}
)
@@ -51,6 +60,13 @@ func noIsolationTarget() (target, bool) {
return target{name: "go-no-isolation", baseURL: goURLNoIsolation}, true
}
func brokenIsolationTarget() (target, bool) {
if goURLBrokenIsolation == "" {
return target{}, false
}
return target{name: "go-broken-isolation", baseURL: goURLBrokenIsolation}, true
}
func TestMain(m *testing.M) {
// Warmup: wait for both servers to be ready before running tests
for _, tgt := range targets() {
@@ -65,6 +81,12 @@ func TestMain(m *testing.M) {
os.Exit(1)
}
}
if tgt, ok := brokenIsolationTarget(); ok {
if !waitForServer(tgt) {
fmt.Fprintf(os.Stderr, "ERROR: %s server not ready at %s\n", tgt.name, tgt.baseURL)
os.Exit(1)
}
}
for _, tgt := range targets() {
warmupJob(tgt)
@@ -554,19 +576,38 @@ func TestLandlockCanReadSystemBinaries(t *testing.T) {
func TestLandlockCanWriteTmpdir(t *testing.T) {
for _, tgt := range targets() {
t.Run(tgt.name, func(t *testing.T) {
// TMPDIR ($CWD/.tmp) should be writable; /tmp should be denied.
// All temp variables are forced to CWD, caller values are ignored,
// no .tmp directory is created, and shared /tmp remains denied.
result := runJob(t, tgt, map[string]any{
"script": "echo TMPDIR=$TMPDIR && touch $TMPDIR/landlock-tmp-test && echo tmpdir_ok && touch /tmp/landlock-denied 2>&1; echo tmp_exit=$?",
"env": map[string]string{"HOME": "/home/dify"},
"script": `set -eu
printf 'CWD=%s TMPDIR=%s TMP=%s TEMP=%s\n' "$PWD" "$TMPDIR" "$TMP" "$TEMP"
test "$TMPDIR" = "$PWD"
test "$TMP" = "$PWD"
test "$TEMP" = "$PWD"
test ! -e "$PWD/.tmp"
touch "$TMPDIR/landlock-tmp-test"
test -f "$TMPDIR/landlock-tmp-test"
printf 'marker=%s\n' "$TMPDIR/landlock-tmp-test"
if touch /tmp/landlock-denied 2>/dev/null; then
echo outside_tmp_unexpectedly_writable
exit 1
fi
echo outside_tmp_denied`,
"cwd": "/home/dify",
"env": map[string]string{"HOME": "/home/dify", "TMPDIR": "/tmp", "TMP": "/tmp", "TEMP": "/tmp"},
"timeout": 10,
})
assertJobDone(t, result)
assertExitCode(t, result, 0)
output := result["output"].(string)
if !strings.Contains(output, "tmpdir_ok") {
t.Errorf("expected write to $TMPDIR to succeed, got %q", output)
}
if !strings.Contains(output, "tmp_exit=1") && !strings.Contains(output, "Permission denied") {
t.Errorf("expected write to /tmp to be denied, got %q", output)
for _, marker := range []string{
"CWD=/home/dify TMPDIR=/home/dify TMP=/home/dify TEMP=/home/dify",
"marker=/home/dify/landlock-tmp-test",
"outside_tmp_denied",
} {
if !strings.Contains(output, marker) {
t.Errorf("expected output marker %q, got %q", marker, output)
}
}
})
}
@@ -592,28 +633,168 @@ func TestLandlockCannotWriteOutsideHome(t *testing.T) {
}
}
func TestLandlockCannotReadOtherAgentHome(t *testing.T) {
func TestLandlockCannotReadOrWriteOtherAgentHome(t *testing.T) {
for _, tgt := range targets() {
t.Run(tgt.name, func(t *testing.T) {
// First, create a file in one agent's home.
// First, create a file with a control layout whose HOME and cwd are
// the same participant directory.
setup := runJob(t, tgt, map[string]any{
"script": "mkdir -p /home/agent-a && touch /home/agent-a/secret",
"env": map[string]string{"HOME": "/home/agent-a"},
"script": "mkdir -p /home/dify/agent-a /home/dify/agent-b && printf secret > /home/dify/agent-a/secret",
"cwd": "/home/dify",
"env": map[string]string{"HOME": "/home/dify"},
"timeout": 10,
})
assertJobDone(t, setup)
assertExitCode(t, setup, 0)
// Now run as a different agent and try to read the other's file.
// A sibling control layout can use its own Home, but content reads
// and Landlock-controlled writes to the first participant must fail.
result := runJob(t, tgt, map[string]any{
"script": "cat /home/agent-a/secret 2>&1; echo exit=$?",
"env": map[string]string{"HOME": "/home/agent-b"},
"script": "touch \"$HOME/own\"; cat /home/dify/agent-a/secret >/dev/null 2>&1; echo read_exit=$?; touch /home/dify/agent-a/sibling-write >/dev/null 2>&1; echo write_exit=$?; python -c 'import os; os.truncate(\"/home/dify/agent-a/secret\", 0)' >/dev/null 2>&1; echo truncate_exit=$?",
"cwd": "/home/dify/agent-b",
"env": map[string]string{"HOME": "/home/dify/agent-b"},
"timeout": 10,
})
assertJobDone(t, result)
assertExitCode(t, result, 0)
output := result["output"].(string)
if strings.Contains(output, "exit=0") {
t.Errorf("expected read of other agent's file to be denied, but it succeeded: %q", output)
for _, forbidden := range []string{"read_exit=0", "write_exit=0", "truncate_exit=0"} {
if strings.Contains(output, forbidden) {
t.Errorf("expected sibling Home access to be denied, got %q", output)
}
}
})
}
}
func TestLandlockNormalLayoutWritesHomeAndWorkspaceButLifecycleLayoutRejectsWorkspace(t *testing.T) {
for _, tgt := range targets() {
t.Run(tgt.name, func(t *testing.T) {
const (
home = "/home/dify/layout-home"
workspace = "/home/dify/layout-workspace"
)
setup := runJob(t, tgt, map[string]any{
"script": fmt.Sprintf("rm -rf %s %s && mkdir -p %s %s", home, workspace, home, workspace),
"cwd": "/home/dify",
"env": map[string]string{"HOME": "/home/dify"},
"timeout": 10,
})
assertJobDone(t, setup)
assertExitCode(t, setup, 0)
normal := runJob(t, tgt, map[string]any{
"script": "printf home > \"$HOME/private\"; printf workspace > \"$PWD/shared\"; cat \"$HOME/private\" \"$PWD/shared\"",
"cwd": workspace,
"env": map[string]string{"HOME": home},
"timeout": 10,
})
assertJobDone(t, normal)
assertExitCode(t, normal, 0)
if output := normal["output"].(string); !strings.Contains(output, "homeworkspace") {
t.Fatalf("normal layout could not use Home and Workspace: %q", output)
}
lifecycle := runJob(t, tgt, map[string]any{
"script": fmt.Sprintf("touch %s/lifecycle-denied >/dev/null 2>&1; echo workspace_exit=$?", workspace),
"cwd": home,
"env": map[string]string{"HOME": home},
"timeout": 10,
})
assertJobDone(t, lifecycle)
assertExitCode(t, lifecycle, 0)
if output := lifecycle["output"].(string); strings.Contains(output, "workspace_exit=0") {
t.Fatalf("lifecycle layout wrote shared Workspace: %q", output)
}
})
}
}
func TestLifecycleLayoutRejectsArchiveSymlinkWriteOutsideHome(t *testing.T) {
for _, tgt := range targets() {
t.Run(tgt.name, func(t *testing.T) {
const (
home = "/home/dify/archive-home"
workspace = "/home/dify/archive-workspace"
outsideFile = workspace + "/symlink-escape"
)
setup := runJob(t, tgt, map[string]any{
"script": fmt.Sprintf("rm -rf %s %s && mkdir -p %s %s", home, workspace, home, workspace),
"cwd": "/home/dify",
"env": map[string]string{"HOME": "/home/dify"},
"timeout": 10,
})
assertJobDone(t, setup)
assertExitCode(t, setup, 0)
archive := base64.StdEncoding.EncodeToString(workspaceSymlinkArchive(t, workspace))
script := fmt.Sprintf(`
printf '%%s' '%s' | base64 -d > "$HOME/archive.tar.zst"
python - "$HOME/archive.tar.zst" "$HOME/server-ready" <<'PY' &
import http.server
from pathlib import Path
import sys
archive = Path(sys.argv[1]).read_bytes()
ready = Path(sys.argv[2])
class Handler(http.server.BaseHTTPRequestHandler):
def do_GET(self):
self.send_response(200)
self.send_header("Content-Type", "application/zstd")
self.send_header("Content-Length", str(len(archive)))
self.end_headers()
self.wfile.write(archive)
def log_message(self, *_args):
pass
server = http.server.ThreadingHTTPServer(("127.0.0.1", 18081), Handler)
ready.touch()
server.serve_forever()
PY
server_pid=$!
for _ in $(seq 1 100); do
[ -e "$HOME/server-ready" ] && break
sleep 0.01
done
if [ ! -e "$HOME/server-ready" ]; then
kill "$server_pid" >/dev/null 2>&1 || true
exit 1
fi
rm -f "$HOME/archive.tar.zst" "$HOME/server-ready"
dify-agent home-snapshot download >/dev/null 2>&1
download_status=$?
kill "$server_pid" >/dev/null 2>&1 || true
wait "$server_pid" >/dev/null 2>&1 || true
echo download_exit=$download_status
`, archive)
result := runJob(t, tgt, map[string]any{
"script": script,
"cwd": home,
"env": map[string]string{
"HOME": home,
"DIFY_AGENT_STUB_API_BASE_URL": "http://127.0.0.1:18081/agent-stub",
"DIFY_AGENT_STUB_AUTH_JWE": "purpose-token",
},
"timeout": 10,
})
assertJobDone(t, result)
assertExitCode(t, result, 0)
if output := result["output"].(string); strings.Contains(output, "download_exit=0") {
t.Fatalf("malicious archive unexpectedly restored outside Home: %q", output)
}
verify := runJob(t, tgt, map[string]any{
"script": fmt.Sprintf("test ! -e %s && echo outside_clean", outsideFile),
"cwd": "/home/dify",
"env": map[string]string{"HOME": "/home/dify"},
"timeout": 10,
})
assertJobDone(t, verify)
assertExitCode(t, verify, 0)
if output := verify["output"].(string); !strings.Contains(output, "outside_clean") {
t.Fatalf("archive wrote outside Home: %q", output)
}
})
}
@@ -668,6 +849,56 @@ func TestLandlockEnvBypassBlocked(t *testing.T) {
}
}
func TestLandlockAllowListEnvBypassBlocked(t *testing.T) {
for _, tgt := range targets() {
t.Run(tgt.name, func(t *testing.T) {
result := runJob(t, tgt, map[string]any{
"script": "touch /tmp/landlock-allow-list-bypass 2>&1; echo exit=$?",
"env": map[string]string{
"HOME": "/home/dify",
"SHELLCTL_LANDLOCK_RW_PATHS": "/tmp",
},
"timeout": 10,
})
assertJobDone(t, result)
output := result["output"].(string)
if strings.Contains(output, "exit=0") {
t.Errorf("expected command env to be unable to widen the Landlock allow-list, got %q", output)
}
})
}
}
func TestLandlockInitializationFailureIsFailClosed(t *testing.T) {
tgt, ok := brokenIsolationTarget()
if !ok || authTokenBrokenIsolation == "" || containerBrokenIsolation == "" {
t.Skip("broken-isolation acceptance container is not available")
}
marker := fmt.Sprintf("/home/dify/landlock-user-script-marker-%d", time.Now().UnixNano())
result := runJobWithToken(t, tgt, authTokenBrokenIsolation, map[string]any{
"script": fmt.Sprintf("touch %s", marker),
"env": map[string]string{"HOME": "/home/dify"},
"timeout": 10,
})
assertJobDone(t, result)
assertExitCode(t, result, 126)
if output, ok := result["output"].(string); !ok || !strings.Contains(output, "apply Landlock filesystem isolation") {
t.Fatalf("expected Landlock initialization failure, got output %q", output)
}
if output, err := exec.Command(
"docker",
"exec",
containerBrokenIsolation,
"test",
"!",
"-e",
marker,
).CombinedOutput(); err != nil {
t.Fatalf("user script marker was created or could not be checked: %v: %s", err, output)
}
}
// --- Helpers ---
func runJob(t *testing.T, tgt target, payload map[string]any) map[string]any {
@@ -682,6 +913,43 @@ func runJob(t *testing.T, tgt target, payload map[string]any) map[string]any {
return result
}
func workspaceSymlinkArchive(t *testing.T, workspace string) []byte {
t.Helper()
var output bytes.Buffer
encoder, err := zstd.NewWriter(&output, zstd.WithEncoderLevel(zstd.EncoderLevelFromZstd(1)))
if err != nil {
t.Fatal(err)
}
writer := tar.NewWriter(encoder)
if err := writer.WriteHeader(&tar.Header{
Name: "escape",
Typeflag: tar.TypeSymlink,
Linkname: workspace,
Mode: 0o777,
}); err != nil {
t.Fatal(err)
}
content := []byte("must stay inside Home")
if err := writer.WriteHeader(&tar.Header{
Name: "escape/symlink-escape",
Typeflag: tar.TypeReg,
Mode: 0o600,
Size: int64(len(content)),
}); err != nil {
t.Fatal(err)
}
if _, err := writer.Write(content); err != nil {
t.Fatal(err)
}
if err := writer.Close(); err != nil {
t.Fatal(err)
}
if err := encoder.Close(); err != nil {
t.Fatal(err)
}
return output.Bytes()
}
func waitJob(t *testing.T, tgt target, jobID string, payload map[string]any) map[string]any {
t.Helper()
resp := doPost(t, tgt, fmt.Sprintf("/v1/jobs/%s/wait", jobID), payload, true)
+8 -3
View File
@@ -27,14 +27,14 @@ DIFY_AGENT_INNER_API_URL=http://localhost:5001
DIFY_AGENT_INNER_API_KEY=
# Runtime resources
# Select one coherent Home Snapshot + Execution Binding backend: local, enterprise, or e2b.
# Select one coherent Home Snapshot + Execution Binding backend: local, enterprise, e2b, or e2b_s3.
DIFY_AGENT_RUNTIME_BACKEND=local
# Local backend: shellctl data-plane URL and optional bearer token.
# Leave the endpoint empty when this server will not provide dify.runtime or resource endpoints.
DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT=
DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN=
# Enterprise resource operations currently fail fast with NotImplementedError.
# These names are retained for the configured Enterprise Gateway boundary.
# Enterprise supports default-Home Binding creation, acquisition, and coupled destroy.
# Immutable Home Snapshot operations and explicit snapshot materialization remain unsupported.
DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_ENDPOINT=
DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_AUTH_TOKEN=
DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_TIMEOUT=30
@@ -47,6 +47,11 @@ DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox
DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS=3600
DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN=
DIFY_AGENT_E2B_SHELLCTL_PORT=5004
# e2b_s3 reuses the E2B settings above and stores immutable Home archives in OpenDAL storage.
# The URI selects and configures the OpenDAL service. It must support read, write,
# multi-write, conditional create, and delete.
# It also requires an HTTP(S) DIFY_AGENT_STUB_API_BASE_URL and DIFY_AGENT_SERVER_SECRET_KEY below.
DIFY_AGENT_E2B_S3_URI=
# Maximum whole-file size read through /workspace/files for an Agent Stub ToolFile upload (50 MiB).
# The environment variable keeps its existing SANDBOX name for deployment compatibility.
DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES=52428800
@@ -15,9 +15,9 @@ Agent can therefore have multiple active Bindings in one Workspace:
each has an independent Home and session, while all may share Workspace files.
Home and Workspace are logically independent. A backend may still couple their
physical representation. For example, current E2B maps one Binding and its
Workspace to one E2B resource, while Local can attach multiple materialized
Homes to one shared Workspace.
physical representation. For example, native `e2b` maps one Binding and its
Workspace to one E2B resource, while Local and `e2b_s3` can attach multiple
materialized Homes to one shared Workspace.
## Runtime layer graph
@@ -67,12 +67,17 @@ streams are observability state, not the Home/Workspace/Binding ledger.
## Creation and execution flow
Home Snapshot initialization uses `POST /home-snapshots/initialize`. Build Draft
Apply uses `POST /home-snapshots/from-binding`: Dify Agent acquires the exact
source Binding, snapshots its materialized Home through the backend-native
operation, releases the lease, and returns a new opaque snapshot ref. Dify API
then stores a new immutable `agent_home_snapshots` row and records its logical id
on the resulting config version. There is no replay or initialization fallback
Agent creation does not create a Home Snapshot. A config with no logical Home
Snapshot asks the selected backend to materialize its deployment-default Home
when the Binding is created. This default Home is mutable and private to the
Binding; it does not produce an `agent_home_snapshots` row or an implicit
snapshot ref.
Build Draft Apply uses `POST /home-snapshots/from-binding`: Dify Agent acquires
the exact source Binding, snapshots its materialized Home through the
backend-native operation, releases the lease, and returns a new opaque snapshot
ref. Dify API then stores a new immutable `agent_home_snapshots` row and records
its logical id on the resulting config version. There is no replay or fallback
when the source Binding is unavailable.
Before an Agent request, Dify API loads the specific product context. If it has
@@ -82,10 +87,12 @@ its owner and config/Home generation. Missing, retired, or mismatched Bindings
fail fast; Dify API does not search by Agent, Workspace, candidate count, or
recency, and it does not create a replacement implicitly.
`POST /execution-bindings` materializes the selected Home Snapshot and returns
opaque Binding and Workspace refs. Every create request represents a new
participant, even when the Agent, Snapshot, config generation, and Workspace
match another Binding. The request composition contains:
`POST /execution-bindings` accepts either an exact `home_snapshot_ref` or
`null`. An exact ref must be materialized without fallback; `null` selects the
backend's deployment-default Home. It returns opaque Binding and Workspace
refs. Every create request represents a new participant, even when the Agent,
Snapshot, config generation, and Workspace match another Binding. The request
composition contains:
```json
{
@@ -96,11 +103,13 @@ match another Binding. The request composition contains:
```
Each Agent request acquires that ref for the duration of the run and releases it
afterward. Local release closes the operation's shellctl connection. E2B release
also pauses the underlying E2B resource with memory preserved. A later request
or Workspace file operation acquires a new lease for the same Binding ref. If a
backend confirms the resource is gone, acquisition fails; it does not create an
empty replacement Workspace.
afterward. Local release closes the operation's shellctl connection. Native
`e2b` release also pauses the underlying E2B resource with memory preserved;
`e2b_s3` release closes only operation-local clients because other Bindings may
still use the shared Sandbox. A later request or Workspace file operation
acquires a new lease for the same Binding ref. If a backend confirms the
resource is gone, acquisition fails; it does not create an empty replacement
Workspace.
## Retirement and collection
@@ -168,15 +177,18 @@ inside the backend execution namespace. They are not host paths, product ids,
or request configuration. Shell commands start in `workspace_dir`, and `HOME`
is forced to `home_dir`. On Local, sibling materialized Homes may exist in the
same shellctl namespace, while path isolation restricts the active lease to its
own Home plus the shared Workspace.
own Home plus the shared Workspace for content access. Linux Landlock does not
currently restrict `stat(2)`, so sibling path metadata can remain visible even
though sibling Home content cannot be read or written.
## Backend support
| Backend | Home Snapshot operations | Binding operations | Physical relationship |
| --- | --- | --- | --- |
| Local | Supported | Supported, including attaching multiple Bindings to one Workspace | Snapshot directory, per-Binding materialized Home, and Workspace directory are separate. |
| E2B | Supported | Supported without shared-Workspace attachment | Binding and Workspace refs map to the same E2B resource; Home initialization/checkpoint uses E2B snapshots. |
| Enterprise | Not implemented | Not implemented | Configuration is accepted, but every resource operation fails fast with `NotImplementedError`. |
| Local | Supported | Supported, including default empty Homes and attaching multiple Bindings to one Workspace | Snapshot directory, per-Binding materialized Home, and Workspace directory are separate. |
| E2B | Supported | Supported with template-backed default Homes, without shared-Workspace attachment | Binding and Workspace refs map to the same E2B resource; checkpoints use E2B snapshots. |
| E2B + S3 (`e2b_s3`) | Supported as immutable OpenDAL tar.zst streams | Multiple private Binding Homes attach to one shared E2B Workspace Sandbox | Workspace ref is the Sandbox id, Binding ref is `sandbox-id:binding-id`, and Home Snapshot ref is an OpenDAL path relative to the operator root configured by `DIFY_AGENT_E2B_S3_URI`. |
| Enterprise | Not implemented | Default-Home Binding creation, acquire, and coupled destroy are supported | Binding and Workspace refs map to one Gateway sandbox. Explicit Home Snapshot materialization fails fast. |
Local creates a new Home for every Binding id. Destroying one Binding without
the Workspace leaves sibling Homes and the shared Workspace intact. Current E2B
@@ -184,10 +196,19 @@ rejects `existing_workspace_ref` with `shared_workspace_unsupported`, because
its Binding and Workspace are one Sandbox. It also rejects binding-only destroy.
Neither path creates a fallback Workspace or switches backends.
`e2b_s3` keeps `/home/dify/workspace` shared while placing each private Home
under `/home/dify/.dify-agent-materialized-homes/<binding-id>`. Runtime leases
force `$HOME` to that participant directory and cwd to the shared Workspace.
Home save/restore is streamed by the hidden Sandbox CLI through a 10-minute,
method-scoped JWE HTTP gateway to an immutable S3 object; credentials and object
selection never enter the Sandbox. Releasing one Binding lease closes only its
operation-local clients and does not pause the shared Sandbox. Binding-only
collection removes one Home, Workspace collection kills the Sandbox, and Home
Snapshot collection deletes only the object.
`DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS` limits continuous active time for an E2B
resource. Runtime resources pause on timeout; temporary Home initialization
resources are killed. It is not a retention TTL and does not delete paused
resources or immutable snapshots.
resource. Runtime resources pause on timeout. It is not a retention TTL and
does not delete paused resources or immutable snapshots.
See the [Shell layer](../../user-manual/shell-layer/index.md) for request
composition and the [Operations Guide](../../guide/index.md) for Local and E2B
@@ -85,10 +85,12 @@ DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES=52428800
E2B requires `DIFY_AGENT_E2B_API_KEY` and defaults to the prepared
`difys-default-team/dify-agent-local-sandbox` template. The E2B active timeout
pauses the physical resource behind a Binding or kills temporary Home
initialization resources; it is not a retention TTL. Enterprise settings are
accepted, but current Home Snapshot and Binding operations fail fast with
`NotImplementedError`.
pauses the physical resource behind a Binding; it is not a retention TTL.
The `e2b_s3` profile additionally stores immutable Home Snapshots in an
S3-compatible bucket and lets multiple Binding Homes share one E2B Workspace
Sandbox; see the Operations Guide for its gateway and credential settings.
Enterprise supports Bindings created from its deployment-default Home.
Immutable Home Snapshot creation and materialization remain unsupported there.
A shell-enabled request includes Execution Context, `dify.runtime`, and
`dify.shell`. Dify API creates or resolves the specific persistent Binding for
+35 -11
View File
@@ -39,18 +39,19 @@ also reads `.env` and `dify-agent/.env` when present.
| `DIFY_AGENT_PLUGIN_DAEMON_API_KEY` | empty | API key sent to the Dify plugin daemon. |
| `DIFY_AGENT_INNER_API_URL` | `http://localhost:5001` | Dify API service root used when dify-agent calls `/inner/api/...` endpoints. |
| `DIFY_AGENT_INNER_API_KEY` | empty | API key sent to Dify API inner plugin endpoints. Set this to Dify API `INNER_API_KEY_FOR_PLUGIN` (Docker: `PLUGIN_DIFY_INNER_API_KEY`). |
| `DIFY_AGENT_RUNTIME_BACKEND` | `local` | Selects one coherent `local`, `enterprise`, or `e2b` Home Snapshot + Execution Binding backend profile. |
| `DIFY_AGENT_RUNTIME_BACKEND` | `local` | Selects one coherent `local`, `enterprise`, `e2b`, or `e2b_s3` Home Snapshot + Execution Binding backend profile. |
| `DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT` | empty | Local shellctl data-plane URL. With the default Local selection, leaving it empty disables `dify.runtime` and resource endpoints. |
| `DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN` | empty | Optional bearer token sent to Local shellctl. |
| `DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_ENDPOINT` | empty | Enterprise Gateway endpoint required by configuration. Current Home Snapshot and Binding operations fail fast with `NotImplementedError`. |
| `DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_ENDPOINT` | empty | Enterprise Gateway endpoint required by configuration. Default-Home Bindings are supported; immutable Home Snapshot operations remain unsupported. |
| `DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_AUTH_TOKEN` | empty | Optional `X-Inner-Api-Key` sent to the Enterprise Gateway. |
| `DIFY_AGENT_ENTERPRISE_SANDBOX_GATEWAY_TIMEOUT` | `30` | Enterprise control-plane timeout in seconds. |
| `DIFY_AGENT_ENTERPRISE_SANDBOX_PROXY_TIMEOUT` | `60` | Enterprise shellctl-proxy timeout in seconds. |
| `DIFY_AGENT_E2B_API_KEY` | empty | E2B API key; required for E2B. |
| `DIFY_AGENT_E2B_TEMPLATE` | `difys-default-team/dify-agent-local-sandbox` | Prepared E2B template containing shellctl and the initial Home environment. |
| `DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS` | `3600` | Maximum continuous active time, up to 3600 seconds. Binding resources pause on timeout; temporary Home initialization resources are killed. This is not a retention TTL. |
| `DIFY_AGENT_E2B_TEMPLATE` | `difys-default-team/dify-agent-local-sandbox` | Prepared E2B template containing shellctl and the deployment-default Home environment. |
| `DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS` | `3600` | Maximum continuous active time, up to 3600 seconds. Binding resources pause on timeout. This is not a retention TTL. |
| `DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN` | empty | Optional bearer token expected by shellctl inside the E2B template. |
| `DIFY_AGENT_E2B_SHELLCTL_PORT` | `5004` | shellctl port exposed by the E2B template. |
| `DIFY_AGENT_E2B_S3_URI` | empty | OpenDAL scheme URI required by `e2b_s3`. The selected service must support `read`, `write`, `write_can_multi`, `write_with_if_not_exists`, and `delete`; startup fails and lists missing capabilities otherwise. |
| `DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES` | `52428800` | Standalone Dify Agent maximum for whole-file Workspace upload capture; 50 MiB by default. Docker Compose derives it from `PLUGIN_MAX_FILE_SIZE`. |
| `DIFY_AGENT_SHELL_REDACT_PATTERNS` | empty | JSON array of additional regex patterns redacted from Shell output. |
| `DIFY_AGENT_STUB_API_BASE_URL` | empty | Public Agent Stub API base URL reachable from shellctl-managed remote machines. HTTP may be the service root or `/agent-stub`; gRPC must be `grpc://host:port`. Enables `DIFY_AGENT_STUB_*` env injection for user `shell.run` jobs. |
@@ -177,12 +178,19 @@ docker compose \
```
`DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS` controls continuous active E2B time.
The physical resource behind a Binding pauses when that timeout fires; a
temporary Home initialization resource is killed. Pausing preserves the current
Workspace. The setting does not delete an aged paused resource or immutable
snapshot, and Dify Agent currently has no resource-age TTL, reconciler, or
eventual cleanup guarantee. Dify API retirement followed by Binding collection
kills the coupled E2B resource.
The physical resource behind a Binding pauses when that timeout fires, preserving
the current Workspace. The setting is not a resource-age TTL and does not delete
paused resources or immutable snapshots.
For `e2b_s3`, export `DIFY_AGENT_RUNTIME_BACKEND=e2b_s3` plus
`DIFY_AGENT_E2B_S3_URI` before starting the same overlay. The URI selects and
configures an OpenDAL service; its scheme is not restricted to `s3`.
`DIFY_AGENT_STUB_API_BASE_URL` must be an HTTP(S) URL reachable from the E2B Sandbox and
`DIFY_AGENT_SERVER_SECRET_KEY` must be configured. The template must enable
`SHELLCTL_ENABLE_PATH_ISOLATION=true`. Compose does not start MinIO or another
object store. A Home Snapshot ref such as
`home-snapshots/<tenant>/<agent>/<snapshot>.tar.zst` is an OpenDAL path relative
to the operator root configured by the URI.
## Run runtime-backend integration contracts
@@ -216,7 +224,7 @@ DIFY_AGENT_TEST_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox \
DIFY_AGENT_TEST_E2B_ACTIVE_TIMEOUT_SECONDS=900 \
pdm run pytest --import-mode=importlib \
tests/integration/dify_agent/runtime_backend/test_working_environment.py \
-k e2b -q -rs
-k 'e2b_binding' -q -rs
```
The Local auth token is optional when shellctl has authentication disabled.
@@ -224,6 +232,22 @@ The E2B test timeout value `900` means up to 15 minutes of continuous active
test time; it is not a post-test retention TTL. Both contracts create unique
resources and perform explicit cleanup in `finally` blocks.
The `e2b_s3` integration path additionally requires a public HTTP Agent Stub
server configured with the same OpenDAL storage URI and server secret.
`DIFY_AGENT_TEST_E2B_S3_URI` must address the same storage and root as the
server's `DIFY_AGENT_E2B_S3_URI`:
```bash
cd dify-agent
DIFY_AGENT_TEST_E2B_API_KEY="$E2B_API_TOKEN" \
DIFY_AGENT_TEST_E2B_S3_URI='s3://integration-bucket/dify-agent-integration?region=us-east-1' \
DIFY_AGENT_TEST_E2B_S3_STUB_API_BASE_URL=https://agent.example.com/agent-stub \
DIFY_AGENT_TEST_SERVER_SECRET_KEY="$DIFY_AGENT_SERVER_SECRET_KEY" \
pdm run pytest --import-mode=importlib \
tests/integration/dify_agent/runtime_backend/test_working_environment.py \
-k e2b_s3 -q -rs
```
## Scheduling and shutdown semantics
`POST /runs` persists a `running` run record and starts an `asyncio` task in the
@@ -40,11 +40,11 @@ the opaque Binding ref belongs to `DifyRuntimeLayerConfig`.
## Runtime requirements
The server constructs one coherent runtime backend profile. Local and E2B
implement Home Snapshot and Execution Binding operations. Enterprise settings
can be selected, but resource operations currently fail fast with
`NotImplementedError`; there is no compatibility fallback to the retired
Sandbox protocol.
The server constructs one coherent runtime backend profile. Local, native E2B,
and `e2b_s3` implement Home Snapshot and Execution Binding operations.
Enterprise implements default-Home Binding creation, acquisition, and coupled
destruction, while immutable Home Snapshot operations fail fast; there is no
compatibility fallback to the retired Sandbox protocol.
```python
from dify_agent.runtime.compositor_factory import create_default_layer_providers
@@ -73,8 +73,11 @@ DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT=http://127.0.0.1:5004
DIFY_AGENT_LOCAL_SANDBOX_AUTH_TOKEN=replace-with-shellctl-token
```
The auth token may be empty when shellctl authentication is disabled. E2B uses
`DIFY_AGENT_E2B_API_KEY`, the prepared template, and its shellctl settings.
The auth token may be empty when shellctl authentication is disabled. Both E2B
profiles use `DIFY_AGENT_E2B_API_KEY`, the prepared template, and its shellctl
settings. `e2b_s3` additionally requires its S3 bucket settings, an HTTP(S)
Agent Stub URL reachable from the Sandbox, and a server secret; its template
must enable fail-closed path isolation.
To let shell jobs call the Agent Stub with `dify-agent ...`, configure a public
Agent Stub URL and a unique production secret:
@@ -205,12 +208,15 @@ browse the current Workspace through Dify Agent's private
`/workspace/files/list`, `/workspace/files/read`, and
`/workspace/files/upload` routes, each of which acquires a fresh lease.
On Local, multiple Bindings may share a Workspace while each receives a
separate materialized Home. Those directories may be siblings in one shellctl
namespace; path isolation restricts a lease to its Home and Workspace. On E2B,
one physical E2B resource currently represents both Binding and Workspace, so
shared Workspace attachment is unsupported.
On Local and `e2b_s3`, multiple Bindings may share a Workspace while each
receives a separate materialized Home. Those directories may be siblings in one
shellctl namespace; path isolation restricts content access to the lease's Home
and Workspace. Linux Landlock does not currently restrict `stat(2)`, so sibling
path metadata can remain visible even though sibling Home content cannot be
read or written.
On native `e2b`, one physical E2B resource represents both Binding and
Workspace, so shared Workspace attachment is unsupported.
See [Runtime resources](../../concepts/runtime-resources/index.md) for the
ledger and lifecycle contract. The [Operations Guide](../../guide/index.md)
covers Local and E2B validation.
covers Local and both E2B profiles.
+9
View File
@@ -25,6 +25,7 @@ server = [
"jsonschema>=4.23.0,<5.0.0",
"jwcrypto>=1.5.6,<2",
"logfire[fastapi,httpx,redis]>=4.37.0,<5.0.0",
"opendal==0.47.5",
"pydantic-ai-slim[anthropic,google,openai]>=1.85.1,<2.0.0",
"pydantic-settings>=2.12.0,<3.0.0",
"redis>=7.4.0,<8.0.0",
@@ -50,6 +51,14 @@ testpaths = ["tests"]
python_files = ["test_*.py", "*_test.py"]
markers = ["integration: requires a real external service or exercises multiple concrete adapters"]
[[tool.uv.index]]
name = "testpypi"
url = "https://test.pypi.org/simple"
explicit = true
[tool.uv.sources]
opendal = { index = "testpypi" }
[tool.ruff]
line-length = 120
target-version = "py312"
@@ -25,6 +25,7 @@ def create_agent_stub_app(settings: ServerSettings | None = None) -> FastAPI:
file_request_handler=resolved_settings.create_agent_stub_file_request_handler(),
config_request_handler=resolved_settings.create_agent_stub_config_request_handler(),
drive_request_handler=resolved_settings.create_agent_stub_drive_request_handler(),
home_snapshot_gateway=resolved_settings.build_home_snapshot_gateway(),
)
)
return app
@@ -0,0 +1,168 @@
"""Authenticated streaming gateway between Sandbox CLI and Home archive store."""
from __future__ import annotations
import asyncio
from collections.abc import AsyncIterable, AsyncIterator, Awaitable
from dataclasses import dataclass
import logging
from typing import Protocol
from dify_agent.agent_stub.server.tokens.home_snapshot import (
HOME_SNAPSHOT_SCOPE_READ,
HOME_SNAPSHOT_SCOPE_WRITE,
HomeSnapshotTransferScope,
HomeSnapshotTransferTokenCodec,
HomeSnapshotTransferTokenError,
)
from dify_agent.runtime_backend.errors import (
HomeArchiveConflictError,
HomeArchiveStoreError,
HomeSnapshotNotFoundError,
)
_ARCHIVE_READ_SIZE = 1024 * 1024
logger = logging.getLogger(__name__)
class AsyncArchiveFile(Protocol):
"""Opened archive stream that its caller must always close or finalize."""
def read(self, size: int | None = None, /) -> Awaitable[bytes]: ...
def write(self, data: bytes, /) -> Awaitable[int]: ...
def close(self) -> Awaitable[None]: ...
class HomeArchiveStore(Protocol):
"""Open immutable archive streams with normalized storage failures.
Writers must use conditional create-only semantics. An existing object must
raise ``HomeArchiveConflictError``; providers may surface that conflict from
``open_writer()``, ``AsyncArchiveFile.write()``, or
``AsyncArchiveFile.close()``. A missing reader object must raise
``HomeSnapshotNotFoundError``. All other storage I/O failures must raise
``HomeArchiveStoreError``.
A successful open transfers stream ownership to the caller. The caller must
close readers and close/finalize writers on success, failure, or cancellation.
"""
def open_writer(self, snapshot_ref: str, /) -> Awaitable[AsyncArchiveFile]: ...
def open_reader(self, snapshot_ref: str, /) -> Awaitable[AsyncArchiveFile]: ...
class HomeSnapshotGatewayError(RuntimeError):
status_code: int
detail: str
def __init__(self, status_code: int, detail: str) -> None:
super().__init__(detail)
self.status_code = status_code
self.detail = detail
@dataclass(slots=True)
class HomeSnapshotGatewayService:
"""Authorize one immutable object and relay bytes without archive parsing."""
token_codec: HomeSnapshotTransferTokenCodec
archive_store: HomeArchiveStore
async def upload(self, *, authorization: str | None, chunks: AsyncIterable[bytes]) -> None:
principal = self._authorize(authorization, required_scope=HOME_SNAPSHOT_SCOPE_WRITE)
try:
writer = await self.archive_store.open_writer(principal.snapshot_ref)
except HomeArchiveConflictError as exc:
raise HomeSnapshotGatewayError(409, str(exc)) from exc
except HomeArchiveStoreError as exc:
raise HomeSnapshotGatewayError(502, str(exc)) from exc
primary_error: BaseException | None = None
try:
async for chunk in chunks:
if not chunk:
continue
written = await writer.write(chunk)
if written != len(chunk):
raise HomeArchiveStoreError("Home Snapshot store performed a short write")
except HomeArchiveConflictError as exc:
primary_error = exc
raise HomeSnapshotGatewayError(409, str(exc)) from exc
except HomeArchiveStoreError as exc:
primary_error = exc
raise HomeSnapshotGatewayError(502, str(exc)) from exc
except BaseException as exc:
primary_error = exc
raise
finally:
try:
await writer.close()
except HomeArchiveConflictError as exc:
if primary_error is None:
raise HomeSnapshotGatewayError(409, str(exc)) from exc
logger.warning("Home Snapshot archive writer conflicted after upload failed", exc_info=True)
except HomeArchiveStoreError as exc:
if primary_error is None:
raise HomeSnapshotGatewayError(502, str(exc)) from exc
logger.warning("failed to close Home Snapshot archive writer after upload failed", exc_info=True)
except asyncio.CancelledError:
raise
except Exception as exc:
if primary_error is None:
raise HomeSnapshotGatewayError(502, str(exc)) from exc
logger.warning("failed to close Home Snapshot archive writer after upload failed", exc_info=True)
async def download(
self,
*,
authorization: str | None,
) -> AsyncIterator[bytes]:
principal = self._authorize(authorization, required_scope=HOME_SNAPSHOT_SCOPE_READ)
try:
reader = await self.archive_store.open_reader(principal.snapshot_ref)
except HomeSnapshotNotFoundError as exc:
raise HomeSnapshotGatewayError(404, str(exc)) from exc
except HomeArchiveStoreError as exc:
raise HomeSnapshotGatewayError(502, str(exc)) from exc
return self._stream_reader(reader)
async def _stream_reader(self, reader: AsyncArchiveFile) -> AsyncIterator[bytes]:
primary_error: BaseException | None = None
try:
while chunk := await reader.read(_ARCHIVE_READ_SIZE):
yield chunk
except BaseException as exc:
primary_error = exc
raise
finally:
try:
await reader.close()
except BaseException:
if primary_error is None:
raise
logger.warning("failed to close Home Snapshot archive reader after download failed", exc_info=True)
def _authorize(
self,
authorization: str | None,
*,
required_scope: HomeSnapshotTransferScope,
):
try:
return self.token_codec.decode_authorization_header(
authorization,
required_scope=required_scope,
)
except HomeSnapshotTransferTokenError as exc:
raise HomeSnapshotGatewayError(401, str(exc)) from exc
__all__ = [
"AsyncArchiveFile",
"HomeArchiveStore",
"HomeSnapshotGatewayError",
"HomeSnapshotGatewayService",
]
@@ -15,6 +15,7 @@ from fastapi import APIRouter
from dify_agent.agent_stub.server.agent_stub_config import AgentStubConfigRequestHandler
from dify_agent.agent_stub.server.agent_stub_drive import AgentStubDriveRequestHandler
from dify_agent.agent_stub.server.agent_stub_files import AgentStubFileRequestHandler
from dify_agent.agent_stub.server.home_snapshots import HomeSnapshotGatewayService
from dify_agent.agent_stub.server.routes.agent_stub import create_agent_stub_http_router
from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubTokenCodec
@@ -25,6 +26,7 @@ def create_agent_stub_router(
file_request_handler: AgentStubFileRequestHandler | None = None,
drive_request_handler: AgentStubDriveRequestHandler | None = None,
config_request_handler: AgentStubConfigRequestHandler | None = None,
home_snapshot_gateway: HomeSnapshotGatewayService | None = None,
) -> APIRouter:
"""Build the embeddable stub router from pre-built server dependencies."""
return create_agent_stub_http_router(
@@ -32,6 +34,7 @@ def create_agent_stub_router(
file_request_handler,
drive_request_handler,
config_request_handler,
home_snapshot_gateway,
)
@@ -8,7 +8,8 @@ delegation with the gRPC transport.
from __future__ import annotations
from fastapi import APIRouter, Header, HTTPException, Response
from fastapi import APIRouter, Header, HTTPException, Request, Response, status
from fastapi.responses import StreamingResponse
from dify_agent.agent_stub.protocol.agent_stub import (
AgentStubConnectRequest,
@@ -30,6 +31,7 @@ from dify_agent.agent_stub.server.agent_stub_config import AgentStubConfigReques
from dify_agent.agent_stub.server.agent_stub_drive import AgentStubDriveRequestHandler
from dify_agent.agent_stub.server.agent_stub_files import AgentStubFileRequestHandler
from dify_agent.agent_stub.server.control_plane import AgentStubControlPlaneError, AgentStubControlPlaneService
from dify_agent.agent_stub.server.home_snapshots import HomeSnapshotGatewayError, HomeSnapshotGatewayService
from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubTokenCodec
@@ -38,6 +40,7 @@ def create_agent_stub_http_router(
file_request_handler: AgentStubFileRequestHandler | None = None,
drive_request_handler: AgentStubDriveRequestHandler | None = None,
config_request_handler: AgentStubConfigRequestHandler | None = None,
home_snapshot_gateway: HomeSnapshotGatewayService | None = None,
) -> APIRouter:
"""Create HTTP routes bound to the application's Agent Stub dependencies."""
router = APIRouter(prefix="/agent-stub", tags=["agent-stub"])
@@ -45,6 +48,35 @@ def create_agent_stub_http_router(
token_codec, file_request_handler, config_request_handler, drive_request_handler
)
def require_home_snapshot_gateway() -> HomeSnapshotGatewayService:
if home_snapshot_gateway is None:
raise HTTPException(status_code=503, detail="Home Snapshot gateway is not configured")
return home_snapshot_gateway
@router.put("/home-snapshots/archive", status_code=status.HTTP_204_NO_CONTENT)
async def upload_home_snapshot_archive(
request: Request,
authorization: str | None = Header(default=None, alias="Authorization"),
) -> Response:
try:
await require_home_snapshot_gateway().upload(
authorization=authorization,
chunks=request.stream(),
)
except HomeSnapshotGatewayError as exc:
raise HTTPException(status_code=exc.status_code, detail=exc.detail) from exc
return Response(status_code=status.HTTP_204_NO_CONTENT)
@router.get("/home-snapshots/archive")
async def download_home_snapshot_archive(
authorization: str | None = Header(default=None, alias="Authorization"),
) -> StreamingResponse:
try:
stream = await require_home_snapshot_gateway().download(authorization=authorization)
except HomeSnapshotGatewayError as exc:
raise HTTPException(status_code=exc.status_code, detail=exc.detail) from exc
return StreamingResponse(stream, media_type="application/zstd")
@router.post("/connections", response_model=AgentStubConnectResponse)
async def create_connection(
request: AgentStubConnectRequest,
@@ -0,0 +1,129 @@
"""Private compact-JWE primitives shared by strict server token families."""
from __future__ import annotations
import base64
import binascii
import hashlib
import hmac
import json
import re
from typing import TypeVar
from jwcrypto import jwe, jwk
from jwcrypto.common import JWException
from pydantic import BaseModel, ValidationError
_REQUIRED_SERVER_SECRET_BYTES = 32
_BASE64URL_TEXT_PATTERN = re.compile(r"^[A-Za-z0-9_-]+$")
_DEFAULT_SERVER_SECRET_ENV_VAR = "DIFY_AGENT_SERVER_SECRET_KEY"
ClaimsT = TypeVar("ClaimsT", bound=BaseModel)
def decode_server_secret_key(
server_secret_key: str,
*,
env_var_name: str = _DEFAULT_SERVER_SECRET_ENV_VAR,
) -> bytes:
"""Decode one strict unpadded base64url 32-byte server root secret."""
normalized = server_secret_key.strip()
if not normalized or not _BASE64URL_TEXT_PATTERN.fullmatch(normalized):
raise ValueError(f"{env_var_name} must be valid unpadded base64url text")
try:
decoded = _base64url_decode(normalized)
except ValueError as exc:
raise ValueError(f"{env_var_name} must be valid unpadded base64url text") from exc
if len(decoded) != _REQUIRED_SERVER_SECRET_BYTES:
raise ValueError(f"{env_var_name} must decode to exactly {_REQUIRED_SERVER_SECRET_BYTES} decoded bytes")
return decoded
def derive_server_jwe_key(server_secret_key: str, *, purpose: bytes) -> bytes:
"""Derive a purpose-isolated 32-byte content-encryption key."""
return _hkdf_sha256(decode_server_secret_key(server_secret_key), info=purpose, length=32)
def build_symmetric_jwe_key(content_encryption_key: bytes) -> jwk.JWK:
return jwk.JWK(kty="oct", k=_base64url_encode(content_encryption_key))
def encode_compact_jwe(claims: BaseModel, *, key: jwk.JWK) -> str:
token = jwe.JWE(
plaintext=json.dumps(claims.model_dump(mode="json", exclude_none=True), separators=(",", ":")).encode("utf-8"),
protected=json.dumps({"alg": "dir", "enc": "A256GCM"}),
)
token.add_recipient(key)
return token.serialize(compact=True)
def decode_compact_jwe(
token: str,
*,
key: jwk.JWK,
claims_type: type[ClaimsT],
token_name: str,
error_type: type[RuntimeError],
) -> ClaimsT:
decrypted = jwe.JWE()
try:
decrypted.deserialize(token, key=key)
except JWException as exc:
raise error_type(f"failed to decrypt {token_name}") from exc
try:
return claims_type.model_validate_json(decrypted.payload)
except ValidationError as exc:
raise error_type(f"{token_name} payload is invalid") from exc
def extract_bearer_token(
authorization: str | None,
*,
token_name: str,
error_type: type[RuntimeError],
) -> str:
if authorization is None or not authorization.startswith("Bearer "):
raise error_type(f"Authorization must be a Bearer {token_name}")
token = authorization.removeprefix("Bearer ").strip()
if not token:
raise error_type("Authorization bearer token must not be empty")
return token
def _hkdf_sha256(input_key_material: bytes, *, info: bytes, length: int) -> bytes:
hash_len = hashlib.sha256().digest_size
salt = b"\x00" * hash_len
pseudorandom_key = hmac.new(salt, input_key_material, hashlib.sha256).digest()
output = bytearray()
previous_block = b""
counter = 1
while len(output) < length:
previous_block = hmac.new(
pseudorandom_key,
previous_block + info + bytes([counter]),
hashlib.sha256,
).digest()
output.extend(previous_block)
counter += 1
return bytes(output[:length])
def _base64url_decode(value: str) -> bytes:
padding = "=" * (-len(value) % 4)
try:
return base64.b64decode(f"{value}{padding}", altchars=b"-_", validate=True)
except binascii.Error as exc:
raise ValueError("invalid base64url") from exc
def _base64url_encode(value: bytes) -> str:
return base64.urlsafe_b64encode(value).rstrip(b"=").decode("ascii")
__all__ = [
"build_symmetric_jwe_key",
"decode_compact_jwe",
"decode_server_secret_key",
"derive_server_jwe_key",
"encode_compact_jwe",
"extract_bearer_token",
]
@@ -9,21 +9,22 @@ families that may reuse the same root secret.
from __future__ import annotations
import base64
import binascii
from dataclasses import dataclass
import hashlib
import hmac
import json
import re
import time
from typing import ClassVar
from uuid import uuid4
from jwcrypto import jwe, jwk
from jwcrypto.common import JWException
from pydantic import BaseModel, ConfigDict, ValidationError
from jwcrypto import jwk
from pydantic import BaseModel, ConfigDict
from dify_agent.agent_stub.server.tokens._compact_jwe import (
build_symmetric_jwe_key,
decode_compact_jwe,
decode_server_secret_key,
derive_server_jwe_key,
encode_compact_jwe,
extract_bearer_token,
)
from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig
@@ -32,9 +33,6 @@ AGENT_STUB_TOKEN_AUDIENCE = "dify-agent-agent-stub"
AGENT_STUB_TOKEN_SCOPE_CONNECT = "agent_stub:connect"
AGENT_STUB_TOKEN_TTL_SECONDS = 24 * 60 * 60
_AGENT_STUB_JWE_PURPOSE = b"dify-agent:agent-stub:jwe:v1"
_REQUIRED_SERVER_SECRET_BYTES = 32
_BASE64URL_TEXT_PATTERN = re.compile(r"^[A-Za-z0-9_-]+$")
_DEFAULT_SERVER_SECRET_ENV_VAR = "DIFY_AGENT_SERVER_SECRET_KEY"
class AgentStubTokenError(RuntimeError):
@@ -83,10 +81,7 @@ class AgentStubTokenCodec:
def __init__(self, content_encryption_key: bytes) -> None:
self._content_encryption_key = content_encryption_key
self._jwe_key = jwk.JWK(
kty="oct",
k=_base64url_encode(content_encryption_key),
)
self._jwe_key = build_symmetric_jwe_key(content_encryption_key)
@classmethod
def from_server_secret(cls, server_secret_key: str) -> AgentStubTokenCodec:
@@ -127,36 +122,26 @@ class AgentStubTokenCodec:
def encode_claims(self, claims: AgentStubTokenClaims) -> str:
"""Encrypt one validated Agent Stub claim set as compact JWE."""
token = jwe.JWE(
plaintext=json.dumps(claims.model_dump(mode="json", exclude_none=True), separators=(",", ":")).encode(
"utf-8"
),
protected=json.dumps({"alg": "dir", "enc": "A256GCM"}),
)
token.add_recipient(self._jwe_key)
return token.serialize(compact=True)
return encode_compact_jwe(claims, key=self._jwe_key)
def decode_authorization_header(self, authorization: str | None, *, now: int | None = None) -> AgentStubPrincipal:
"""Decode a ``Bearer <compact-jwe>`` header into a request principal."""
if authorization is None or not authorization.startswith("Bearer "):
raise AgentStubTokenError("Authorization must be a Bearer compact JWE token")
token = authorization.removeprefix("Bearer ").strip()
if not token:
raise AgentStubTokenError("Authorization bearer token must not be empty")
token = extract_bearer_token(
authorization,
token_name="compact JWE token",
error_type=AgentStubTokenError,
)
return self.decode_token(token, now=now)
def decode_token(self, token: str, *, now: int | None = None) -> AgentStubPrincipal:
"""Decrypt and validate one compact JWE token string."""
decrypted = jwe.JWE()
try:
decrypted.deserialize(token, key=self._jwe_key)
except JWException as exc:
raise AgentStubTokenError("failed to decrypt Agent Stub bearer token") from exc
try:
claims = AgentStubTokenClaims.model_validate_json(decrypted.payload)
except ValidationError as exc:
raise AgentStubTokenError("Agent Stub bearer token payload is invalid") from exc
claims = decode_compact_jwe(
token,
key=self._jwe_key,
claims_type=AgentStubTokenClaims,
token_name="Agent Stub bearer token",
error_type=AgentStubTokenError,
)
current_time = _timestamp(now)
_validate_claims(claims, now=current_time)
@@ -168,29 +153,9 @@ class AgentStubTokenCodec:
)
def decode_server_secret_key(server_secret_key: str, *, env_var_name: str = _DEFAULT_SERVER_SECRET_ENV_VAR) -> bytes:
"""Decode and validate the configured server root secret.
The secret must be strict unpadded base64url text and must decode to
exactly 32 bytes. Settings validation uses this helper so operator
misconfiguration fails fast before the server starts issuing or accepting
Agent Stub tokens.
"""
normalized = server_secret_key.strip()
if not normalized or not _BASE64URL_TEXT_PATTERN.fullmatch(normalized):
raise ValueError(f"{env_var_name} must be valid unpadded base64url text")
try:
decoded = _base64url_decode(normalized)
except ValueError as exc:
raise ValueError(f"{env_var_name} must be valid unpadded base64url text") from exc
if len(decoded) != _REQUIRED_SERVER_SECRET_BYTES:
raise ValueError(f"{env_var_name} must decode to exactly {_REQUIRED_SERVER_SECRET_BYTES} decoded bytes")
return decoded
def derive_agent_stub_jwe_key(server_secret_key: str) -> bytes:
"""Derive the purpose-scoped 32-byte JWE content-encryption key."""
return _hkdf_sha256(decode_server_secret_key(server_secret_key), info=_AGENT_STUB_JWE_PURPOSE, length=32)
return derive_server_jwe_key(server_secret_key, purpose=_AGENT_STUB_JWE_PURPOSE)
def _validate_claims(claims: AgentStubTokenClaims, *, now: int) -> None:
@@ -206,40 +171,10 @@ def _validate_claims(claims: AgentStubTokenClaims, *, now: int) -> None:
raise AgentStubTokenError(f"Agent Stub bearer token scope must include {AGENT_STUB_TOKEN_SCOPE_CONNECT!r}")
def _hkdf_sha256(input_key_material: bytes, *, info: bytes, length: int) -> bytes:
hash_len = hashlib.sha256().digest_size
salt = b"\x00" * hash_len
pseudorandom_key = hmac.new(salt, input_key_material, hashlib.sha256).digest()
output = bytearray()
previous_block = b""
counter = 1
while len(output) < length:
previous_block = hmac.new(
pseudorandom_key,
previous_block + info + bytes([counter]),
hashlib.sha256,
).digest()
output.extend(previous_block)
counter += 1
return bytes(output[:length])
def _timestamp(value: int | None) -> int:
return int(time.time() if value is None else value)
def _base64url_decode(value: str) -> bytes:
padding = "=" * (-len(value) % 4)
try:
return base64.b64decode(f"{value}{padding}", altchars=b"-_", validate=True)
except binascii.Error as exc:
raise ValueError("invalid base64url") from exc
def _base64url_encode(value: bytes) -> str:
return base64.urlsafe_b64encode(value).rstrip(b"=").decode("ascii")
__all__ = [
"AGENT_STUB_TOKEN_AUDIENCE",
"AGENT_STUB_TOKEN_ISSUER",
@@ -0,0 +1,126 @@
"""Purpose-scoped compact-JWE tokens for Home Snapshot byte transfer."""
from __future__ import annotations
from dataclasses import dataclass
import time
from typing import ClassVar, Literal
from jwcrypto import jwk
from pydantic import BaseModel, ConfigDict
from dify_agent.agent_stub.server.tokens._compact_jwe import (
build_symmetric_jwe_key,
decode_compact_jwe,
derive_server_jwe_key,
encode_compact_jwe,
extract_bearer_token,
)
from dify_agent.runtime_backend.home_snapshot_refs import validate_home_snapshot_ref
HOME_SNAPSHOT_SCOPE_READ = "home_snapshot:read"
HOME_SNAPSHOT_SCOPE_WRITE = "home_snapshot:write"
HOME_SNAPSHOT_TOKEN_TTL_SECONDS = 10 * 60
_HOME_SNAPSHOT_JWE_PURPOSE = b"dify-agent:home-snapshot-transfer:jwe:v1"
type HomeSnapshotTransferScope = Literal["home_snapshot:read", "home_snapshot:write"]
class HomeSnapshotTransferTokenError(RuntimeError):
"""Raised when a Home Snapshot transfer bearer token is invalid."""
class HomeSnapshotTransferTokenClaims(BaseModel):
exp: int
scope: HomeSnapshotTransferScope
snapshot_ref: str
model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")
@dataclass(frozen=True, slots=True)
class HomeSnapshotTransferPrincipal:
scope: HomeSnapshotTransferScope
snapshot_ref: str
class HomeSnapshotTransferTokenCodec:
"""Encode and validate compact JWE tokens for one archive and direction."""
_jwe_key: jwk.JWK
def __init__(self, content_encryption_key: bytes) -> None:
self._jwe_key = build_symmetric_jwe_key(content_encryption_key)
@classmethod
def from_server_secret(cls, server_secret_key: str) -> "HomeSnapshotTransferTokenCodec":
return cls(derive_server_jwe_key(server_secret_key, purpose=_HOME_SNAPSHOT_JWE_PURPOSE))
def encode_token(
self,
*,
scope: HomeSnapshotTransferScope,
snapshot_ref: str,
now: int | None = None,
) -> str:
issued_at = _timestamp(now)
claims = HomeSnapshotTransferTokenClaims(
exp=issued_at + HOME_SNAPSHOT_TOKEN_TTL_SECONDS,
scope=scope,
snapshot_ref=validate_home_snapshot_ref(snapshot_ref),
)
return encode_compact_jwe(claims, key=self._jwe_key)
def decode_authorization_header(
self,
authorization: str | None,
*,
required_scope: HomeSnapshotTransferScope,
now: int | None = None,
) -> HomeSnapshotTransferPrincipal:
token = extract_bearer_token(
authorization,
token_name="compact JWE token",
error_type=HomeSnapshotTransferTokenError,
)
return self.decode_token(token, required_scope=required_scope, now=now)
def decode_token(
self,
token: str,
*,
required_scope: HomeSnapshotTransferScope,
now: int | None = None,
) -> HomeSnapshotTransferPrincipal:
claims = decode_compact_jwe(
token,
key=self._jwe_key,
claims_type=HomeSnapshotTransferTokenClaims,
token_name="Home Snapshot bearer token",
error_type=HomeSnapshotTransferTokenError,
)
if _timestamp(now) >= claims.exp:
raise HomeSnapshotTransferTokenError("Home Snapshot bearer token is expired")
if claims.scope != required_scope:
raise HomeSnapshotTransferTokenError(f"Home Snapshot bearer token scope must be {required_scope!r}")
try:
snapshot_ref = validate_home_snapshot_ref(claims.snapshot_ref)
except ValueError as exc:
raise HomeSnapshotTransferTokenError("Home Snapshot bearer token ref is invalid") from exc
return HomeSnapshotTransferPrincipal(scope=claims.scope, snapshot_ref=snapshot_ref)
def _timestamp(value: int | None) -> int:
return int(time.time() if value is None else value)
__all__ = [
"HOME_SNAPSHOT_SCOPE_READ",
"HOME_SNAPSHOT_SCOPE_WRITE",
"HOME_SNAPSHOT_TOKEN_TTL_SECONDS",
"HomeSnapshotTransferPrincipal",
"HomeSnapshotTransferScope",
"HomeSnapshotTransferTokenClaims",
"HomeSnapshotTransferTokenCodec",
"HomeSnapshotTransferTokenError",
]
@@ -36,7 +36,6 @@ from dify_agent.protocol import (
DeleteHomeSnapshotRequest,
DestroyExecutionBindingRequest,
HomeSnapshotResponse,
InitializeHomeSnapshotRequest,
RUN_EVENT_ADAPTER,
RunEvent,
RunEventsResponse,
@@ -529,24 +528,6 @@ class Client:
response = self._post_sync_json("destroy_execution_binding_sync", "/execution-bindings/destroy", request)
_raise_for_status(response)
async def initialize_home_snapshot(self, request: InitializeHomeSnapshotRequest) -> HomeSnapshotResponse:
"""Create a backend-native initial Home Snapshot."""
response = await self._post_async_json(
"initialize_home_snapshot",
"/home-snapshots/initialize",
request,
)
return _parse_model_response(response, HomeSnapshotResponse)
def initialize_home_snapshot_sync(self, request: InitializeHomeSnapshotRequest) -> HomeSnapshotResponse:
"""Synchronous variant of ``initialize_home_snapshot``."""
response = self._post_sync_json(
"initialize_home_snapshot_sync",
"/home-snapshots/initialize",
request,
)
return _parse_model_response(response, HomeSnapshotResponse)
async def create_home_snapshot_from_binding(
self,
request: CreateHomeSnapshotFromBindingRequest,
@@ -6,9 +6,8 @@ from collections.abc import Sequence
import json
import logging
import re
import time
from dataclasses import dataclass, field
from typing import ClassVar, Literal, NotRequired, Protocol, TypedDict, runtime_checkable
from typing import ClassVar, NotRequired, Protocol, TypedDict, runtime_checkable
from pydantic import BaseModel, ConfigDict, Field, NonNegativeInt, field_validator, model_validator
from pydantic_ai import Tool
@@ -36,6 +35,7 @@ from dify_agent.layers.runtime.layer import DifyRuntimeLayer
from dify_agent.layers.shell.configs import DIFY_SHELL_LAYER_TYPE_ID, DifyShellLayerConfig
from dify_agent.layers.shell.output_text import normalized_output_text, utf8_prefix, utf8_suffix
from dify_agent.runtime_backend import RuntimeLease
from dify_agent.runtime_backend.shellctl import execute_complete_with_commands
logger = logging.getLogger(__name__)
@@ -545,77 +545,6 @@ class DifyShellLayer(PydanticAILayer[DifyShellLayerDeps, object, DifyShellLayerC
return text
async def execute_complete_with_commands(
commands: ShellCommandProtocol,
script: str,
*,
cwd: str | None,
env: dict[str, str] | None,
timeout: float,
max_output_bytes: int,
) -> CompleteShellCommandResult:
deadline = time.monotonic() + timeout
job_id: str | None = None
result: ShellCommandResult | None = None
output_parts: list[str] = []
captured_bytes = 0
incomplete_reason: Literal["output_limit", "timeout"] | None = None
try:
result = await commands.run(script, cwd=cwd, env=env, timeout=_remaining_time(deadline))
job_id = result.job_id
while True:
remaining_bytes = max(max_output_bytes - captured_bytes, 0)
limited_output = utf8_prefix(result.output, remaining_bytes)
output_parts.append(limited_output)
captured_bytes += len(limited_output.encode("utf-8"))
if limited_output != result.output:
incomplete_reason = "output_limit"
break
if captured_bytes >= max_output_bytes and (result.truncated or not result.done):
incomplete_reason = "output_limit"
break
if result.truncated:
result = await commands.read_output(result.job_id, offset=result.offset)
continue
if result.done:
break
remaining_time = _remaining_time(deadline)
if remaining_time <= 0.0:
incomplete_reason = "timeout"
break
result = await commands.wait(result.job_id, offset=result.offset, timeout=remaining_time)
assert result is not None
final_status = result.status
final_done = result.done
final_exit_code = result.exit_code
final_offset = result.offset
final_output_path = result.output_path
if incomplete_reason is not None and not result.done:
terminal_status = await commands.interrupt(result.job_id, grace_seconds=DEFAULT_TERMINATE_GRACE_SECONDS)
final_status = terminal_status.status
final_done = terminal_status.done
final_exit_code = terminal_status.exit_code
final_offset = terminal_status.offset
return CompleteShellCommandResult(
job_id=result.job_id,
status=final_status,
done=final_done,
exit_code=final_exit_code,
output="".join(output_parts),
output_complete=incomplete_reason is None,
incomplete_reason=incomplete_reason,
offset=final_offset,
output_path=final_output_path,
)
finally:
if job_id is not None:
try:
await commands.delete(job_id, force=True)
except RuntimeError as exc:
logger.warning("Failed to delete transient shell job %s: %s", job_id, exc)
async def render_prompt_observation_from_result(
commands: ShellCommandProtocol,
result: ShellCommandResult,
@@ -765,10 +694,6 @@ def _tagged_shell_observation(metadata: dict[str, object], output: str) -> str:
return f"<metadata>\n{compact_metadata}\n</metadata>\n\n<output>\n{output}\n</output>"
def _remaining_time(deadline: float) -> float:
return max(0.0, deadline - time.monotonic())
__all__ = [
"CompleteRemoteCommandResult",
"DifyShellLayer",
@@ -776,6 +701,5 @@ __all__ = [
"DifyShellRuntimeState",
"DEFAULT_TERMINATE_GRACE_SECONDS",
"DEFAULT_TIMEOUT_SECONDS",
"execute_complete_with_commands",
"render_prompt_observation_from_result",
]
@@ -46,7 +46,6 @@ from .home_snapshot import (
CreateHomeSnapshotFromBindingRequest,
DeleteHomeSnapshotRequest,
HomeSnapshotResponse,
InitializeHomeSnapshotRequest,
)
from .workspace import (
WorkspaceFileEntry,
@@ -79,7 +78,6 @@ __all__ = [
"EmptyRunEventData",
"LayerExitSignals",
"HomeSnapshotResponse",
"InitializeHomeSnapshotRequest",
"PydanticAIStreamRunEvent",
"RUN_EVENT_ADAPTER",
"RunCancelledEvent",
@@ -11,7 +11,7 @@ class CreateExecutionBindingRequest(BaseModel):
binding_id: str = Field(min_length=1)
workspace_id: str = Field(min_length=1)
existing_workspace_ref: str | None = None
home_snapshot_ref: str = Field(min_length=1)
home_snapshot_ref: str | None = Field(default=None, min_length=1)
model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")
@@ -5,14 +5,6 @@ from typing import ClassVar
from pydantic import BaseModel, ConfigDict, Field
class InitializeHomeSnapshotRequest(BaseModel):
tenant_id: str = Field(min_length=1)
agent_id: str = Field(min_length=1)
home_snapshot_id: str = Field(min_length=1)
model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid")
class CreateHomeSnapshotFromBindingRequest(BaseModel):
tenant_id: str = Field(min_length=1)
agent_id: str = Field(min_length=1)
@@ -38,5 +30,4 @@ __all__ = [
"CreateHomeSnapshotFromBindingRequest",
"DeleteHomeSnapshotRequest",
"HomeSnapshotResponse",
"InitializeHomeSnapshotRequest",
]
@@ -22,7 +22,6 @@ from .protocols import (
FileSystem,
HomeSnapshotBackend,
HomeSnapshotCreateSpec,
InitializeHomeSnapshotSpec,
RuntimeBackendProfile,
RuntimeLayout,
RuntimeLease,
@@ -46,7 +45,6 @@ __all__ = [
"HomeSnapshotCreateError",
"HomeSnapshotCreateSpec",
"HomeSnapshotNotFoundError",
"InitializeHomeSnapshotSpec",
"RuntimeBackendError",
"RuntimeBackendProfile",
"RuntimeLayout",
@@ -30,7 +30,6 @@ from dify_agent.runtime_backend.protocols import (
ExecutionBindingDestroySpec,
FileSystem,
HomeSnapshotCreateSpec,
InitializeHomeSnapshotSpec,
RuntimeLayout,
RuntimeLease,
)
@@ -42,27 +41,48 @@ if TYPE_CHECKING:
E2B_MAX_ACTIVE_TIMEOUT_SECONDS = 60 * 60
class _E2BControlPlaneNotFoundError(RuntimeError):
class E2BControlPlaneNotFoundError(RuntimeError):
"""Typed boundary error for SDK resources that no longer exist."""
class _E2BFileSystem(Protocol):
class E2BFileType(Protocol):
@property
def value(self) -> str: ...
class E2BFileInfo(Protocol):
@property
def type(self) -> E2BFileType | None: ...
@property
def symlink_target(self) -> str | None: ...
class E2BSandboxFileSystem(Protocol):
async def make_dir(self, path: str) -> bool: ...
async def exists(self, path: str) -> bool: ...
async def get_info(self, path: str) -> E2BFileInfo: ...
async def remove(self, path: str) -> None: ...
async def read(self, path: str) -> str: ...
async def write(self, path: str, data: str | bytes) -> object: ...
class _E2BSnapshotInfo(Protocol):
snapshot_id: str
names: list[str]
class _E2BSandbox(Protocol):
class E2BSandbox(Protocol):
sandbox_id: str
traffic_access_token: str | None
files: _E2BFileSystem
@property
def files(self) -> E2BSandboxFileSystem: ...
def get_host(self, port: int) -> str: ...
@@ -81,9 +101,9 @@ class E2BControlPlane(Protocol):
timeout: int,
metadata: dict[str, str],
on_timeout: Literal["kill", "pause"],
) -> _E2BSandbox: ...
) -> E2BSandbox: ...
async def connect(self, handle: str, *, timeout: int) -> _E2BSandbox: ...
async def connect(self, handle: str, *, timeout: int) -> E2BSandbox: ...
async def kill(self, handle: str) -> bool: ...
@@ -111,12 +131,12 @@ class E2BSDKControlPlane:
timeout: int,
metadata: dict[str, str],
on_timeout: Literal["kill", "pause"],
) -> _E2BSandbox:
) -> E2BSandbox:
from e2b import AsyncSandbox, NotFoundException, SandboxNotFoundException
try:
return cast(
_E2BSandbox,
E2BSandbox,
cast(
object,
await AsyncSandbox.create(
@@ -129,18 +149,18 @@ class E2BSDKControlPlane:
),
)
except (SandboxNotFoundException, NotFoundException) as exc:
raise _E2BControlPlaneNotFoundError(str(exc)) from exc
raise E2BControlPlaneNotFoundError(str(exc)) from exc
async def connect(self, handle: str, *, timeout: int) -> _E2BSandbox:
async def connect(self, handle: str, *, timeout: int) -> E2BSandbox:
from e2b import AsyncSandbox, NotFoundException, SandboxNotFoundException
try:
return cast(
_E2BSandbox,
E2BSandbox,
cast(object, await AsyncSandbox.connect(handle, timeout=timeout, **self._options())),
)
except (SandboxNotFoundException, NotFoundException) as exc:
raise _E2BControlPlaneNotFoundError(str(exc)) from exc
raise E2BControlPlaneNotFoundError(str(exc)) from exc
async def kill(self, handle: str) -> bool:
from e2b import AsyncSandbox, NotFoundException, SandboxNotFoundException
@@ -148,7 +168,7 @@ class E2BSDKControlPlane:
try:
return await AsyncSandbox.kill(handle, **self._options())
except (SandboxNotFoundException, NotFoundException) as exc:
raise _E2BControlPlaneNotFoundError(str(exc)) from exc
raise E2BControlPlaneNotFoundError(str(exc)) from exc
async def delete_snapshot(self, snapshot_ref: str) -> bool:
from e2b import AsyncSandbox, NotFoundException, SandboxNotFoundException
@@ -156,51 +176,19 @@ class E2BSDKControlPlane:
try:
return await AsyncSandbox.delete_snapshot(snapshot_ref, **self._options())
except (SandboxNotFoundException, NotFoundException) as exc:
raise _E2BControlPlaneNotFoundError(str(exc)) from exc
raise E2BControlPlaneNotFoundError(str(exc)) from exc
@dataclass(slots=True)
class E2BHomeSnapshotBackend:
"""Implement immutable Home Snapshot operations with E2B snapshots.
Initialization snapshots the prepared deployment template and releases its
temporary E2B resource. Build Apply snapshots the E2B resource behind the
supplied ``RuntimeLease``. Dify API stores the returned value as an opaque
backend ref; this adapter keeps no cross-request state.
Build Apply snapshots the E2B resource behind the supplied ``RuntimeLease``.
Dify API stores the returned value as an opaque backend ref; this adapter
keeps no cross-request state.
"""
control_plane: E2BControlPlane
template: str
active_timeout_seconds: int
home_dir: str = "/home/dify"
async def initialize(self, spec: InitializeHomeSnapshotSpec) -> str:
sandbox: _E2BSandbox | None = None
try:
sandbox = await self.control_plane.create(
self.template,
timeout=self.active_timeout_seconds,
metadata={
"dify.resource": "home-snapshot-initialize",
"dify.tenant_id": spec.tenant_id,
"dify.agent_id": spec.agent_id,
"dify.home_snapshot_id": spec.home_snapshot_id,
},
on_timeout="kill",
)
_ = await sandbox.files.make_dir(self.home_dir)
snapshot = await sandbox.create_snapshot()
return snapshot.snapshot_id
except BaseException as exc:
if isinstance(exc, Exception):
raise HomeSnapshotCreateError(str(exc)) from exc
raise
finally:
if sandbox is not None:
try:
_ = await sandbox.kill()
except BaseException:
pass
async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str:
"""Create an immutable E2B snapshot from the source Binding's active lease."""
@@ -218,7 +206,7 @@ class E2BHomeSnapshotBackend:
async def delete(self, snapshot_ref: str) -> None:
try:
_ = await self.control_plane.delete_snapshot(snapshot_ref)
except _E2BControlPlaneNotFoundError:
except E2BControlPlaneNotFoundError:
return
except Exception as exc:
raise BindingDestroyError(str(exc)) from exc
@@ -236,6 +224,7 @@ class E2BExecutionBindingBackend:
"""
control_plane: E2BControlPlane
template: str
active_timeout_seconds: int
shellctl_auth_token: str = ""
shellctl_port: int = 5004
@@ -244,13 +233,13 @@ class E2BExecutionBindingBackend:
)
async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation:
"""Create one paused E2B resource from an immutable Home Snapshot ref."""
"""Create one paused E2B resource from a snapshot or deployment template."""
if spec.existing_workspace_ref is not None:
raise SharedWorkspaceUnsupportedError("current E2B backend cannot attach to an existing Workspace")
sandbox: _E2BSandbox | None = None
sandbox: E2BSandbox | None = None
try:
sandbox = await self.control_plane.create(
spec.home_snapshot_ref,
self.template if spec.home_snapshot_ref is None else spec.home_snapshot_ref,
timeout=self.active_timeout_seconds,
metadata={
"dify.resource": "runtime-sandbox",
@@ -274,20 +263,20 @@ class E2BExecutionBindingBackend:
except BaseException:
pass
if isinstance(exc, Exception):
if isinstance(exc, (BindingCreateError, SharedWorkspaceUnsupportedError)):
if isinstance(exc, BindingCreateError):
raise
raise BindingCreateError(str(exc)) from exc
raise
async def acquire(self, binding_ref: str) -> RuntimeLease:
"""Acquire operation-scoped shellctl access for an opaque Binding ref."""
sandbox: _E2BSandbox | None = None
sandbox: E2BSandbox | None = None
try:
sandbox = await self.control_plane.connect(binding_ref, timeout=self.active_timeout_seconds)
if not await sandbox.files.exists(self.layout.workspace_dir):
raise BindingLostError(f"E2B Binding {binding_ref!r} no longer contains its Workspace")
return await self._lease(sandbox)
except _E2BControlPlaneNotFoundError as exc:
except E2BControlPlaneNotFoundError as exc:
raise BindingLostError(f"E2B Binding {binding_ref!r} no longer exists") from exc
except BindingLostError:
await _best_effort_pause(sandbox)
@@ -324,40 +313,17 @@ class E2BExecutionBindingBackend:
raise BindingDestroyError("E2B Workspace ref must equal its Binding ref")
try:
_ = await self.control_plane.kill(spec.binding_ref)
except _E2BControlPlaneNotFoundError:
except E2BControlPlaneNotFoundError:
return
except Exception as exc:
raise BindingDestroyError(str(exc)) from exc
async def _lease(self, sandbox: _E2BSandbox) -> "E2BRuntimeLease":
entrypoint = f"https://{sandbox.get_host(self.shellctl_port)}"
traffic_token = sandbox.traffic_access_token
headers = {"X-Access-Token": traffic_token} if isinstance(traffic_token, str) and traffic_token else {}
http_client = httpx.AsyncClient(
base_url=entrypoint,
headers=headers,
follow_redirects=True,
timeout=httpx.Timeout(60.0),
)
def client_factory() -> ShellctlClientProtocol:
from shellctl.client import ShellctlClient
return cast(
ShellctlClientProtocol,
cast(
object,
ShellctlClient(entrypoint, token=self.shellctl_auth_token, client=http_client),
),
)
data_plane = await create_owned_shellctl_lease(
handle=sandbox.sandbox_id,
async def _lease(self, sandbox: E2BSandbox) -> "E2BRuntimeLease":
data_plane = await create_e2b_shellctl_lease(
sandbox=sandbox,
layout=self.layout,
entrypoint=entrypoint,
token=self.shellctl_auth_token,
client_factory=client_factory,
owned_transport=http_client,
shellctl_auth_token=self.shellctl_auth_token,
shellctl_port=self.shellctl_port,
)
return E2BRuntimeLease(sandbox=sandbox, data_plane=data_plane)
@@ -366,7 +332,7 @@ class E2BExecutionBindingBackend:
class E2BRuntimeLease:
"""Invocation-local E2B SDK object plus the owned shellctl data-plane lease."""
sandbox: _E2BSandbox
sandbox: E2BSandbox
data_plane: ShellctlRuntimeLease
@property
@@ -386,7 +352,46 @@ class E2BRuntimeLease:
return self.data_plane.files
async def _best_effort_pause(sandbox: _E2BSandbox | None) -> None:
async def create_e2b_shellctl_lease(
*,
sandbox: E2BSandbox,
layout: RuntimeLayout,
shellctl_auth_token: str,
shellctl_port: int,
) -> ShellctlRuntimeLease:
"""Create one operation-local shellctl data plane for a dynamic E2B layout."""
entrypoint = f"https://{sandbox.get_host(shellctl_port)}"
traffic_token = sandbox.traffic_access_token
headers = {"X-Access-Token": traffic_token} if isinstance(traffic_token, str) and traffic_token else {}
http_client = httpx.AsyncClient(
base_url=entrypoint,
headers=headers,
follow_redirects=True,
timeout=httpx.Timeout(60.0),
)
def client_factory() -> ShellctlClientProtocol:
from shellctl.client import ShellctlClient
return cast(
ShellctlClientProtocol,
cast(
object,
ShellctlClient(entrypoint, token=shellctl_auth_token, client=http_client),
),
)
return await create_owned_shellctl_lease(
handle=sandbox.sandbox_id,
layout=layout,
entrypoint=entrypoint,
token=shellctl_auth_token,
client_factory=client_factory,
owned_transport=http_client,
)
async def _best_effort_pause(sandbox: E2BSandbox | None) -> None:
if sandbox is None:
return
try:
@@ -398,8 +403,14 @@ async def _best_effort_pause(sandbox: _E2BSandbox | None) -> None:
__all__ = [
"E2B_MAX_ACTIVE_TIMEOUT_SECONDS",
"E2BControlPlane",
"E2BControlPlaneNotFoundError",
"E2BExecutionBindingBackend",
"E2BFileInfo",
"E2BFileType",
"E2BHomeSnapshotBackend",
"E2BSDKControlPlane",
"E2BSandbox",
"E2BSandboxFileSystem",
"E2BRuntimeLease",
"create_e2b_shellctl_lease",
]
@@ -0,0 +1,600 @@
"""E2B Workspace backend with independently persisted S3 Home Snapshots."""
from __future__ import annotations
import asyncio
from collections.abc import Awaitable
from dataclasses import dataclass
import logging
import shlex
from typing import Protocol
import opendal
from dify_agent.adapters.shell.protocols import ShellCommandProtocol
from dify_agent.agent_stub.protocol import (
AGENT_STUB_API_BASE_URL_ENV_VAR,
AGENT_STUB_AUTH_JWE_ENV_VAR,
)
from dify_agent.agent_stub.server.tokens.home_snapshot import (
HOME_SNAPSHOT_SCOPE_READ,
HOME_SNAPSHOT_SCOPE_WRITE,
HomeSnapshotTransferScope,
HomeSnapshotTransferTokenCodec,
)
from dify_agent.runtime_backend.e2b import (
E2BControlPlane,
E2BControlPlaneNotFoundError,
E2BSandbox,
E2BSandboxFileSystem,
create_e2b_shellctl_lease,
)
from dify_agent.runtime_backend.errors import (
BindingAcquireError,
BindingCreateError,
BindingDestroyError,
BindingLostError,
HomeArchiveConflictError,
HomeArchiveStoreError,
HomeSnapshotCreateError,
HomeSnapshotNotFoundError,
)
from dify_agent.runtime_backend.home_snapshot_refs import (
build_home_snapshot_ref,
validate_home_snapshot_ref,
validate_ref_segment,
)
from dify_agent.runtime_backend.protocols import (
ExecutionBindingAllocation,
ExecutionBindingCreateSpec,
ExecutionBindingDestroySpec,
FileSystem,
HomeSnapshotCreateSpec,
RuntimeLayout,
RuntimeLease,
)
from dify_agent.runtime_backend.shellctl import (
CONTROL_COMMAND_OUTPUT_LIMIT,
ShellctlRuntimeLease,
execute_complete_with_commands,
)
_HOME_ROOT = "/home/dify/.dify-agent-materialized-homes"
_RUNTIME_ROOT = "/home/dify/.dify-agent-runtime"
_WORKSPACE_DIR = "/home/dify/workspace"
_WORKSPACE_ID_FILE = f"{_RUNTIME_ROOT}/workspace-id"
_BOOTSTRAP_HOME = "/home/dify"
_BINDING_REF_SEPARATOR = ":"
_BASELINE_EXCLUDES = (
"workspace",
".dify-agent-materialized-homes",
".dify-agent-runtime",
".local/share/shellctl",
)
_REQUIRED_HOME_ARCHIVE_CAPABILITIES = (
"read",
"write",
"write_can_multi",
"write_with_if_not_exists",
"delete",
)
logger = logging.getLogger(__name__)
class OpenDALAsyncFile(Protocol):
def read(self, size: int | None = None, /) -> Awaitable[bytes]: ...
def write(self, data: bytes, /) -> Awaitable[int]: ...
def close(self) -> Awaitable[None]: ...
@dataclass(slots=True)
class _OpenDALArchiveFile:
"""Normalize AsyncFile I/O failures at the OpenDAL backend boundary."""
file: OpenDALAsyncFile
async def read(self, size: int | None = None) -> bytes:
try:
return await self.file.read(size)
except Exception as exc:
raise HomeArchiveStoreError(str(exc)) from exc
async def write(self, data: bytes) -> int:
try:
return await self.file.write(data)
except (opendal.exceptions.AlreadyExists, opendal.exceptions.ConditionNotMatch) as exc:
raise HomeArchiveConflictError("Home Snapshot object already exists") from exc
except Exception as exc:
raise HomeArchiveStoreError(str(exc)) from exc
async def close(self) -> None:
try:
await self.file.close()
except (opendal.exceptions.AlreadyExists, opendal.exceptions.ConditionNotMatch) as exc:
raise HomeArchiveConflictError("Home Snapshot object already exists") from exc
except Exception as exc:
raise HomeArchiveStoreError(str(exc)) from exc
@dataclass(slots=True)
class OpenDALHomeArchiveStore:
"""Small OpenDAL boundary for immutable Home Snapshot archive objects."""
operator: opendal.AsyncOperator
@classmethod
def create_from_uri(cls, uri: str) -> "OpenDALHomeArchiveStore":
operator = opendal.AsyncOperator.from_uri(uri)
capability = operator.capability()
missing = [
name for name in _REQUIRED_HOME_ARCHIVE_CAPABILITIES if not getattr(capability, name, False)
]
if missing:
raise ValueError(
"OpenDAL Home Snapshot storage is missing required capabilities: " + ", ".join(missing)
)
return cls(operator=operator)
async def open_writer(self, snapshot_ref: str) -> OpenDALAsyncFile:
path = validate_home_snapshot_ref(snapshot_ref)
try:
return _OpenDALArchiveFile(
await self.operator.open(path, "wb", if_not_exists=True) # pyright: ignore[reportUnknownMemberType]
)
except (opendal.exceptions.AlreadyExists, opendal.exceptions.ConditionNotMatch) as exc:
raise HomeArchiveConflictError(f"Home Snapshot object {path!r} already exists") from exc
except Exception as exc:
raise HomeArchiveStoreError(str(exc)) from exc
async def open_reader(self, snapshot_ref: str) -> OpenDALAsyncFile:
path = validate_home_snapshot_ref(snapshot_ref)
try:
return _OpenDALArchiveFile(
await self.operator.open(path, "rb") # pyright: ignore[reportUnknownMemberType]
)
except opendal.exceptions.NotFound as exc:
raise HomeSnapshotNotFoundError(f"Home Snapshot object {path!r} does not exist") from exc
except Exception as exc:
raise HomeArchiveStoreError(str(exc)) from exc
async def delete(self, snapshot_ref: str) -> None:
path = validate_home_snapshot_ref(snapshot_ref)
try:
await self.operator.delete(path)
except opendal.exceptions.NotFound:
return
except Exception as exc:
raise HomeArchiveStoreError(str(exc)) from exc
class HomeSnapshotTransferError(RuntimeError):
"""A lifecycle CLI job could not transfer one immutable archive."""
@dataclass(slots=True)
class E2BHomeSnapshotCLI:
"""Invoke hidden archive commands in one operation-local shellctl lease."""
token_codec: HomeSnapshotTransferTokenCodec
agent_stub_api_base_url: str
shellctl_auth_token: str
shellctl_port: int
async def upload(
self,
*,
sandbox: E2BSandbox,
home_dir: str,
snapshot_ref: str,
excludes: tuple[str, ...] = (),
) -> None:
argv = ["dify-agent", "home-snapshot", "upload"]
for excluded in excludes:
argv.extend(("--exclude", excluded))
await self._invoke(
sandbox=sandbox,
home_dir=home_dir,
snapshot_ref=snapshot_ref,
scope=HOME_SNAPSHOT_SCOPE_WRITE,
argv=argv,
)
async def download(
self,
*,
sandbox: E2BSandbox,
home_dir: str,
snapshot_ref: str,
) -> None:
await self._invoke(
sandbox=sandbox,
home_dir=home_dir,
snapshot_ref=snapshot_ref,
scope=HOME_SNAPSHOT_SCOPE_READ,
argv=["dify-agent", "home-snapshot", "download"],
)
async def _invoke(
self,
*,
sandbox: E2BSandbox,
home_dir: str,
snapshot_ref: str,
scope: HomeSnapshotTransferScope,
argv: list[str],
) -> None:
layout = RuntimeLayout(home_dir=home_dir, workspace_dir=home_dir)
lease = await create_e2b_shellctl_lease(
sandbox=sandbox,
layout=layout,
shellctl_auth_token=self.shellctl_auth_token,
shellctl_port=self.shellctl_port,
)
primary_error: BaseException | None = None
try:
token = self.token_codec.encode_token(scope=scope, snapshot_ref=snapshot_ref)
result = await execute_complete_with_commands(
lease.commands,
shlex.join(argv),
cwd=None,
env={
AGENT_STUB_API_BASE_URL_ENV_VAR: self.agent_stub_api_base_url,
AGENT_STUB_AUTH_JWE_ENV_VAR: token,
},
timeout=None,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
if result.exit_code != 0 or not result.output_complete:
detail = result.output.strip() or result.incomplete_reason or result.status
raise HomeSnapshotTransferError(f"Home Snapshot CLI failed: {detail}")
except BaseException as exc:
primary_error = exc
raise
finally:
try:
await lease.close()
except BaseException:
if primary_error is None:
raise
logger.warning("failed to close Home Snapshot control lease", exc_info=True)
@dataclass(slots=True)
class E2BS3HomeSnapshotBackend:
"""Persist Home archives independently from E2B Sandbox lifecycle."""
control_plane: E2BControlPlane
archive_store: OpenDALHomeArchiveStore
lifecycle_cli: E2BHomeSnapshotCLI
template: str
active_timeout_seconds: int
async def initialize(self, spec: HomeSnapshotCreateSpec) -> str:
"""Capture the deployment template's baseline Home into a new object."""
sandbox: E2BSandbox | None = None
try:
snapshot_ref = build_home_snapshot_ref(spec)
sandbox = await self.control_plane.create(
self.template,
timeout=self.active_timeout_seconds,
metadata={
"dify.resource": "home-snapshot-bootstrap",
"dify.tenant_id": spec.tenant_id,
"dify.agent_id": spec.agent_id,
"dify.home_snapshot_id": spec.home_snapshot_id,
},
on_timeout="kill",
)
await self.lifecycle_cli.upload(
sandbox=sandbox,
home_dir=_BOOTSTRAP_HOME,
snapshot_ref=snapshot_ref,
excludes=_BASELINE_EXCLUDES,
)
return snapshot_ref
except BaseException as exc:
if isinstance(exc, HomeSnapshotCreateError):
raise
if isinstance(exc, Exception):
raise HomeSnapshotCreateError(str(exc)) from exc
raise
finally:
if sandbox is not None:
try:
_ = await sandbox.kill()
except asyncio.CancelledError:
raise
except Exception:
logger.warning("failed to kill temporary Home Snapshot Sandbox", exc_info=True)
async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str:
if not isinstance(source, E2BS3RuntimeLease):
raise HomeSnapshotCreateError("e2b_s3 Home Snapshot requires an E2BS3RuntimeLease")
try:
snapshot_ref = build_home_snapshot_ref(spec)
await self.lifecycle_cli.upload(
sandbox=source.sandbox,
home_dir=source.layout.home_dir,
snapshot_ref=snapshot_ref,
)
return snapshot_ref
except BaseException as exc:
if isinstance(exc, HomeSnapshotCreateError):
raise
if isinstance(exc, Exception):
raise HomeSnapshotCreateError(str(exc)) from exc
raise
async def delete(self, snapshot_ref: str) -> None:
try:
await self.archive_store.delete(snapshot_ref)
except ValueError as exc:
raise HomeArchiveStoreError(str(exc)) from exc
@dataclass(slots=True)
class E2BS3ExecutionBindingBackend:
"""Map one shared E2B Sandbox to a Workspace and private participant Homes."""
control_plane: E2BControlPlane
lifecycle_cli: E2BHomeSnapshotCLI
template: str
active_timeout_seconds: int
shellctl_auth_token: str
shellctl_port: int
async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation:
binding_id = ""
sandbox: E2BSandbox | None = None
creates_workspace = spec.existing_workspace_ref is None
home_created = False
try:
binding_id = validate_ref_segment(spec.binding_id)
workspace_id = validate_ref_segment(spec.workspace_id)
snapshot_ref = (
validate_home_snapshot_ref(spec.home_snapshot_ref) if spec.home_snapshot_ref is not None else None
)
if creates_workspace:
sandbox = await self.control_plane.create(
self.template,
timeout=self.active_timeout_seconds,
metadata={
"dify.resource": "runtime-workspace",
"dify.tenant_id": spec.tenant_id,
"dify.workspace_id": workspace_id,
},
on_timeout="pause",
)
sandbox_id = validate_ref_segment(sandbox.sandbox_id)
await self._prepare_new_workspace(sandbox, workspace_id=workspace_id)
else:
sandbox_id = validate_ref_segment(spec.existing_workspace_ref or "")
sandbox = await self.control_plane.connect(
sandbox_id,
timeout=self.active_timeout_seconds,
)
await self._validate_workspace(sandbox, workspace_id=workspace_id)
home_dir = _home_dir(binding_id)
if await sandbox.files.exists(home_dir):
raise BindingCreateError(f"Materialized Home for Binding {binding_id!r} already exists")
home_created = True
_ = await sandbox.files.make_dir(home_dir)
if not await _is_real_directory(sandbox.files, home_dir):
raise BindingCreateError(f"Materialized Home for Binding {binding_id!r} is not a real directory")
await self._chmod_home(sandbox, home_dir=home_dir)
if snapshot_ref is not None:
await self.lifecycle_cli.download(
sandbox=sandbox,
home_dir=home_dir,
snapshot_ref=snapshot_ref,
)
allocation = ExecutionBindingAllocation(
workspace_ref=sandbox_id,
binding_ref=_binding_ref(sandbox_id=sandbox_id, binding_id=binding_id),
)
if creates_workspace:
_ = await sandbox.pause(keep_memory=True)
return allocation
except BaseException as exc:
if sandbox is not None:
if creates_workspace:
try:
_ = await sandbox.kill()
except BaseException:
logger.warning("failed to kill partial e2b_s3 Workspace Sandbox", exc_info=True)
elif home_created:
await _remove_home_best_effort(sandbox, binding_id=binding_id)
if isinstance(exc, BindingCreateError):
raise
if isinstance(exc, Exception):
raise BindingCreateError(str(exc)) from exc
raise
async def acquire(self, binding_ref: str) -> RuntimeLease:
sandbox: E2BSandbox | None = None
try:
sandbox_id, binding_id = _parse_binding_ref(binding_ref)
sandbox = await self.control_plane.connect(
sandbox_id,
timeout=self.active_timeout_seconds,
)
home_dir = _home_dir(binding_id)
required_directories = (home_dir, _WORKSPACE_DIR, _HOME_ROOT, _RUNTIME_ROOT)
if not all([await _is_real_directory(sandbox.files, path) for path in required_directories]):
raise BindingLostError(f"e2b_s3 Binding {binding_ref!r} no longer contains its Home or Workspace")
data_plane = await create_e2b_shellctl_lease(
sandbox=sandbox,
layout=RuntimeLayout(home_dir=home_dir, workspace_dir=_WORKSPACE_DIR),
shellctl_auth_token=self.shellctl_auth_token,
shellctl_port=self.shellctl_port,
)
return E2BS3RuntimeLease(sandbox=sandbox, data_plane=data_plane)
except E2BControlPlaneNotFoundError as exc:
raise BindingLostError(f"e2b_s3 Binding {binding_ref!r} no longer exists") from exc
except BindingLostError:
raise
except BaseException as exc:
if isinstance(exc, Exception):
raise BindingAcquireError(str(exc)) from exc
raise
async def release(self, lease: RuntimeLease) -> None:
if not isinstance(lease, E2BS3RuntimeLease):
raise TypeError("E2BS3ExecutionBindingBackend can only release its own RuntimeLease")
try:
await lease.data_plane.close()
except Exception as exc:
raise BindingAcquireError(str(exc)) from exc
async def destroy_binding(self, spec: ExecutionBindingDestroySpec) -> None:
try:
sandbox_id, binding_id = _parse_binding_ref(spec.binding_ref)
except ValueError as exc:
raise BindingDestroyError(str(exc)) from exc
if spec.destroy_workspace:
try:
workspace_ref = validate_ref_segment(spec.workspace_ref or "")
except ValueError as exc:
raise BindingDestroyError(str(exc)) from exc
if workspace_ref != sandbox_id:
raise BindingDestroyError("e2b_s3 Workspace ref must match the Binding Sandbox")
try:
_ = await self.control_plane.kill(sandbox_id)
except E2BControlPlaneNotFoundError:
return
except Exception as exc:
raise BindingDestroyError(str(exc)) from exc
return
try:
sandbox = await self.control_plane.connect(
sandbox_id,
timeout=self.active_timeout_seconds,
)
home_dir = _home_dir(binding_id)
if await sandbox.files.exists(home_dir):
await sandbox.files.remove(home_dir)
except E2BControlPlaneNotFoundError:
return
except Exception as exc:
raise BindingDestroyError(str(exc)) from exc
async def _prepare_new_workspace(self, sandbox: E2BSandbox, *, workspace_id: str) -> None:
for path in (_WORKSPACE_DIR, _HOME_ROOT, _RUNTIME_ROOT):
if await sandbox.files.exists(path):
await sandbox.files.remove(path)
_ = await sandbox.files.make_dir(_WORKSPACE_DIR)
_ = await sandbox.files.make_dir(_HOME_ROOT)
_ = await sandbox.files.make_dir(_RUNTIME_ROOT)
for path in (_WORKSPACE_DIR, _HOME_ROOT, _RUNTIME_ROOT):
if not await _is_real_directory(sandbox.files, path):
raise BindingCreateError(f"new e2b_s3 layout path {path!r} is not a real directory")
_ = await sandbox.files.write(_WORKSPACE_ID_FILE, f"{workspace_id}\n")
async def _validate_workspace(self, sandbox: E2BSandbox, *, workspace_id: str) -> None:
for path in (_WORKSPACE_DIR, _HOME_ROOT, _RUNTIME_ROOT):
if not await _is_real_directory(sandbox.files, path):
raise BindingCreateError(f"existing e2b_s3 layout path {path!r} is not a real directory")
if not await sandbox.files.exists(_WORKSPACE_ID_FILE):
raise BindingCreateError("existing e2b_s3 Workspace metadata is missing")
stored_workspace_id = (await sandbox.files.read(_WORKSPACE_ID_FILE)).strip()
if stored_workspace_id != workspace_id:
raise BindingCreateError("existing e2b_s3 Workspace ref does not match workspace_id")
async def _chmod_home(self, sandbox: E2BSandbox, *, home_dir: str) -> None:
lease = await create_e2b_shellctl_lease(
sandbox=sandbox,
layout=RuntimeLayout(home_dir=home_dir, workspace_dir=home_dir),
shellctl_auth_token=self.shellctl_auth_token,
shellctl_port=self.shellctl_port,
)
primary_error: BaseException | None = None
try:
result = await execute_complete_with_commands(
lease.commands,
'chmod 700 -- "$HOME"',
cwd=None,
env=None,
timeout=30.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
if result.exit_code != 0 or not result.output_complete:
raise BindingCreateError(result.output or "failed to set Materialized Home permissions")
except BaseException as exc:
primary_error = exc
raise
finally:
try:
await lease.close()
except BaseException:
if primary_error is None:
raise
logger.warning("failed to close Materialized Home control lease", exc_info=True)
@dataclass(slots=True)
class E2BS3RuntimeLease:
"""Operation-local access to one private Home in a shared E2B Workspace."""
sandbox: E2BSandbox
data_plane: ShellctlRuntimeLease
@property
def handle(self) -> str:
return self.data_plane.handle
@property
def layout(self) -> RuntimeLayout:
return self.data_plane.layout
@property
def commands(self) -> ShellCommandProtocol:
return self.data_plane.commands
@property
def files(self) -> FileSystem:
return self.data_plane.files
def _home_dir(binding_id: str) -> str:
return f"{_HOME_ROOT}/{validate_ref_segment(binding_id)}"
def _binding_ref(*, sandbox_id: str, binding_id: str) -> str:
return f"{validate_ref_segment(sandbox_id)}{_BINDING_REF_SEPARATOR}{validate_ref_segment(binding_id)}"
def _parse_binding_ref(binding_ref: str) -> tuple[str, str]:
parts = binding_ref.split(_BINDING_REF_SEPARATOR)
if len(parts) != 2:
raise ValueError("e2b_s3 Binding ref is invalid")
return validate_ref_segment(parts[0]), validate_ref_segment(parts[1])
async def _is_real_directory(files: E2BSandboxFileSystem, path: str) -> bool:
if not await files.exists(path):
return False
info = await files.get_info(path)
return info.symlink_target is None and info.type is not None and info.type.value == "dir"
async def _remove_home_best_effort(sandbox: E2BSandbox, *, binding_id: str) -> None:
try:
home_dir = _home_dir(binding_id)
if await sandbox.files.exists(home_dir):
await sandbox.files.remove(home_dir)
except BaseException:
logger.warning("failed to remove partial e2b_s3 Materialized Home", exc_info=True)
__all__ = [
"E2BHomeSnapshotCLI",
"E2BS3ExecutionBindingBackend",
"E2BS3HomeSnapshotBackend",
"E2BS3RuntimeLease",
"HomeSnapshotTransferError",
"OpenDALHomeArchiveStore",
]
@@ -1,10 +1,10 @@
"""Enterprise Gateway adapter for the working-environment protocol.
The existing Gateway can reconnect to and delete an already allocated sandbox,
but it cannot materialize a Home Snapshot or create a protocol-compliant
Binding. One physical sandbox owns both the materialized Home and Workspace, so
their cleanup is coupled. Runtime access remains operation-local and is routed
through the Gateway's shellctl proxy.
The existing Gateway can allocate, reconnect to, and delete a sandbox, but it
does not expose immutable Home Snapshot operations. One physical sandbox owns
both the materialized Home and Workspace, so their cleanup is coupled. Runtime
access remains operation-local and is routed through the Gateway's shellctl
proxy.
"""
from __future__ import annotations
@@ -21,8 +21,10 @@ from dify_agent.adapters.shell.protocols import ShellCommandProtocol, ShellProvi
from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol, ShellctlCommands
from dify_agent.runtime_backend.errors import (
BindingAcquireError,
BindingCreateError,
BindingDestroyError,
BindingLostError,
SharedWorkspaceUnsupportedError,
WorkspacePreservationUnsupportedError,
)
from dify_agent.runtime_backend.protocols import (
@@ -31,35 +33,27 @@ from dify_agent.runtime_backend.protocols import (
ExecutionBindingDestroySpec,
FileSystem,
HomeSnapshotCreateSpec,
InitializeHomeSnapshotSpec,
RuntimeLayout,
RuntimeLease,
)
from dify_agent.runtime_backend.shellctl import (
CONTROL_COMMAND_OUTPUT_LIMIT,
ShellctlRuntimeLease,
create_owned_shellctl_lease,
run_shellctl_control_command,
execute_complete_with_commands,
)
logger = logging.getLogger(__name__)
def _not_implemented() -> NotImplementedError:
return NotImplementedError("Enterprise Gateway does not implement the Execution Binding protocol")
return NotImplementedError("Enterprise Gateway does not implement immutable Home Snapshot operations")
@dataclass(slots=True)
class EnterpriseHomeSnapshotBackend:
"""Reject Home Snapshot operations until the Gateway exposes immutable snapshots."""
gateway_endpoint: str
auth_token: str
gateway_timeout: float = 30.0
async def initialize(self, spec: InitializeHomeSnapshotSpec) -> str:
del spec
raise _not_implemented()
async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str:
del spec, source
raise _not_implemented()
@@ -71,7 +65,7 @@ class EnterpriseHomeSnapshotBackend:
@dataclass(slots=True)
class EnterpriseExecutionBindingBackend:
"""Access and destroy legacy Gateway sandboxes as coupled physical Bindings."""
"""Manage Gateway sandboxes as coupled physical Bindings and Workspaces."""
gateway_endpoint: str
auth_token: str
@@ -82,8 +76,60 @@ class EnterpriseExecutionBindingBackend:
)
async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation:
del spec
raise _not_implemented()
"""Create a default Gateway sandbox and initialize its canonical layout."""
if spec.existing_workspace_ref is not None:
raise SharedWorkspaceUnsupportedError("current Enterprise backend cannot attach to an existing Workspace")
if spec.home_snapshot_ref is not None:
raise BindingCreateError("current Enterprise backend cannot materialize an immutable Home Snapshot")
sandbox_id: str | None = None
data_plane: ShellctlRuntimeLease | None = None
headers = {"X-Inner-Api-Key": self.auth_token} if self.auth_token else {}
try:
async with httpx.AsyncClient(
base_url=self.gateway_endpoint.rstrip("/"),
headers=headers,
timeout=httpx.Timeout(self.gateway_timeout),
) as client:
response = await client.post("/v1/sandboxes", json={"tenantId": spec.tenant_id})
_ = response.raise_for_status()
payload = response.json()
sandbox_id_value = payload.get("sandboxId") if isinstance(payload, dict) else None
if not isinstance(sandbox_id_value, str) or not sandbox_id_value:
raise BindingCreateError("Enterprise Gateway returned an invalid sandbox id")
sandbox_id = sandbox_id_value
data_plane = await self._create_data_plane(sandbox_id)
result = await execute_complete_with_commands(
ShellctlCommands(client=data_plane.client),
"\n".join(
[
"set -eu",
f"mkdir -p {shlex.quote(self.layout.home_dir)}",
f"rm -rf -- {shlex.quote(self.layout.workspace_dir)}",
f"mkdir -p {shlex.quote(self.layout.workspace_dir)}",
f"chmod 700 {shlex.quote(self.layout.home_dir)} {shlex.quote(self.layout.workspace_dir)}",
]
),
cwd=None,
env=None,
timeout=30.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
if result.exit_code != 0 or not result.output_complete:
raise BindingCreateError(result.output)
await data_plane.close()
data_plane = None
return ExecutionBindingAllocation(binding_ref=sandbox_id, workspace_ref=sandbox_id)
except BaseException as exc:
await _close_best_effort(data_plane, binding_ref=sandbox_id or spec.binding_id)
if sandbox_id is not None:
await self._delete_sandbox_best_effort(sandbox_id)
if isinstance(exc, BindingCreateError):
raise
if isinstance(exc, Exception):
raise BindingCreateError(str(exc)) from exc
raise
async def acquire(self, binding_ref: str) -> RuntimeLease:
"""Reconnect to one existing Gateway sandbox without creating a replacement."""
@@ -91,7 +137,7 @@ class EnterpriseExecutionBindingBackend:
try:
data_plane = await self._create_data_plane(binding_ref)
validation_commands = ShellctlCommands(client=data_plane.client)
result = await run_shellctl_control_command(
result = await execute_complete_with_commands(
validation_commands,
"\n".join(
[
@@ -100,8 +146,15 @@ class EnterpriseExecutionBindingBackend:
f"test -d {shlex.quote(self.layout.workspace_dir)}",
]
),
cwd=None,
env=None,
timeout=5.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
if not result.output_complete:
raise BindingAcquireError(
f"Enterprise Binding {binding_ref!r} validation did not complete: {result.incomplete_reason}"
)
if result.exit_code != 0:
raise BindingLostError(f"Enterprise Binding {binding_ref!r} no longer contains its Home or Workspace")
return EnterpriseRuntimeLease(data_plane=data_plane)
@@ -137,21 +190,34 @@ class EnterpriseExecutionBindingBackend:
if spec.workspace_ref != spec.binding_ref:
raise BindingDestroyError("Enterprise Workspace ref must equal its Binding ref")
headers = {"X-Inner-Api-Key": self.auth_token} if self.auth_token else {}
encoded_binding_ref = quote(spec.binding_ref, safe="")
try:
async with httpx.AsyncClient(
base_url=self.gateway_endpoint.rstrip("/"),
headers=headers,
timeout=httpx.Timeout(self.gateway_timeout),
) as client:
response = await client.delete(f"/v1/sandboxes/{encoded_binding_ref}")
if response.status_code == 404:
return
_ = response.raise_for_status()
await self._delete_sandbox(spec.binding_ref)
except (httpx.TimeoutException, httpx.RequestError, httpx.HTTPStatusError) as exc:
raise BindingDestroyError(str(exc)) from exc
async def _delete_sandbox(self, sandbox_id: str) -> None:
headers = {"X-Inner-Api-Key": self.auth_token} if self.auth_token else {}
encoded_sandbox_id = quote(sandbox_id, safe="")
async with httpx.AsyncClient(
base_url=self.gateway_endpoint.rstrip("/"),
headers=headers,
timeout=httpx.Timeout(self.gateway_timeout),
) as client:
response = await client.delete(f"/v1/sandboxes/{encoded_sandbox_id}")
if response.status_code == 404:
return
_ = response.raise_for_status()
async def _delete_sandbox_best_effort(self, sandbox_id: str) -> None:
try:
await self._delete_sandbox(sandbox_id)
except BaseException:
logger.warning(
"failed to delete Enterprise sandbox after Binding creation failed",
exc_info=True,
extra={"binding_ref": sandbox_id},
)
async def _create_data_plane(self, binding_ref: str) -> ShellctlRuntimeLease:
proxy_base_url = f"{self.gateway_endpoint.rstrip('/')}/proxy/"
headers = {"X-Sandbox-Id": binding_ref}
@@ -13,6 +13,14 @@ class HomeSnapshotNotFoundError(RuntimeBackendError):
pass
class HomeArchiveStoreError(RuntimeBackendError):
pass
class HomeArchiveConflictError(HomeArchiveStoreError):
pass
class BindingCreateError(RuntimeBackendError):
pass
@@ -62,6 +70,8 @@ __all__ = [
"BindingCreateError",
"BindingDestroyError",
"BindingLostError",
"HomeArchiveConflictError",
"HomeArchiveStoreError",
"HomeSnapshotCreateError",
"HomeSnapshotNotFoundError",
"RuntimeBackendError",
@@ -0,0 +1,46 @@
"""Canonical opaque refs used by the e2b_s3 Home Snapshot backend."""
from __future__ import annotations
import re
from dify_agent.runtime_backend.protocols import HomeSnapshotCreateSpec
_SAFE_REF_SEGMENT = re.compile(r"^[A-Za-z0-9._-]+$")
_SNAPSHOT_PREFIX = "home-snapshots"
_SNAPSHOT_SUFFIX = ".tar.zst"
def validate_ref_segment(value: str) -> str:
if value in {"", ".", ".."} or "\\" in value or _SAFE_REF_SEGMENT.fullmatch(value) is None:
raise ValueError("runtime backend ref must be a safe path segment")
return value
def build_home_snapshot_ref(spec: HomeSnapshotCreateSpec) -> str:
return "/".join(
(
_SNAPSHOT_PREFIX,
validate_ref_segment(spec.tenant_id),
validate_ref_segment(spec.agent_id),
f"{validate_ref_segment(spec.home_snapshot_id)}{_SNAPSHOT_SUFFIX}",
)
)
def validate_home_snapshot_ref(snapshot_ref: str) -> str:
if snapshot_ref.startswith("/") or "\\" in snapshot_ref:
raise ValueError("Home Snapshot ref is invalid")
parts = snapshot_ref.split("/")
if len(parts) != 4 or parts[0] != _SNAPSHOT_PREFIX:
raise ValueError("Home Snapshot ref is invalid")
tenant_id = validate_ref_segment(parts[1])
agent_id = validate_ref_segment(parts[2])
filename = parts[3]
if not filename.endswith(_SNAPSHOT_SUFFIX):
raise ValueError("Home Snapshot ref is invalid")
snapshot_id = validate_ref_segment(filename.removesuffix(_SNAPSHOT_SUFFIX))
return f"{_SNAPSHOT_PREFIX}/{tenant_id}/{agent_id}/{snapshot_id}{_SNAPSHOT_SUFFIX}"
__all__ = ["build_home_snapshot_ref", "validate_home_snapshot_ref", "validate_ref_segment"]
@@ -27,14 +27,14 @@ from dify_agent.runtime_backend.protocols import (
ExecutionBindingCreateSpec,
ExecutionBindingDestroySpec,
HomeSnapshotCreateSpec,
InitializeHomeSnapshotSpec,
RuntimeLayout,
RuntimeLease,
)
from dify_agent.runtime_backend.shellctl import (
CONTROL_COMMAND_OUTPUT_LIMIT,
ShellctlRuntimeLease,
create_shellctl_lease,
run_shellctl_control_command,
execute_complete_with_commands,
)
_SAFE_REF_PART = re.compile(r"^[A-Za-z0-9._-]+$")
@@ -49,28 +49,6 @@ class LocalHomeSnapshotBackend:
snapshot_root: str = "/home/dify/.dify-agent-home-snapshots"
client_factory: ShellctlClientFactory | None = None
async def initialize(self, spec: InitializeHomeSnapshotSpec) -> str:
snapshot_ref = _local_snapshot_ref(spec.home_snapshot_id)
lease = self._control_lease(snapshot_ref)
target = self._snapshot_dir(snapshot_ref)
try:
result = await run_shellctl_control_command(
lease.commands,
f"set -eu\nmkdir -p {shlex.quote(target)}\nchmod 700 {shlex.quote(target)}",
)
if result.exit_code != 0:
raise HomeSnapshotCreateError(result.output)
return snapshot_ref
except BaseException as exc:
await _remove_partial(lease.commands, target=target, resource_ref=snapshot_ref)
if isinstance(exc, HomeSnapshotCreateError):
raise
if isinstance(exc, Exception):
raise HomeSnapshotCreateError(str(exc)) from exc
raise
finally:
await _close_best_effort(lease, resource_ref=snapshot_ref)
async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str:
snapshot_ref = _local_snapshot_ref(spec.home_snapshot_id)
target = self._snapshot_dir(snapshot_ref)
@@ -85,8 +63,15 @@ class LocalHomeSnapshotBackend:
]
)
try:
result = await run_shellctl_control_command(lease.commands, script)
if result.exit_code != 0:
result = await execute_complete_with_commands(
lease.commands,
script,
cwd=None,
env=None,
timeout=30.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
if result.exit_code != 0 or not result.output_complete:
raise HomeSnapshotCreateError(result.output)
return snapshot_ref
except BaseException as exc:
@@ -103,11 +88,15 @@ class LocalHomeSnapshotBackend:
normalized = _validated_ref_part(snapshot_ref)
lease = self._control_lease(normalized)
try:
result = await run_shellctl_control_command(
result = await execute_complete_with_commands(
lease.commands,
f"rm -rf -- {shlex.quote(self._snapshot_dir(normalized))}",
cwd=None,
env=None,
timeout=30.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
if result.exit_code != 0:
if result.exit_code != 0 or not result.output_complete:
raise BindingDestroyError(result.output)
except BaseException:
await _close_best_effort(lease, resource_ref=normalized)
@@ -142,34 +131,41 @@ class LocalExecutionBindingBackend:
async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation:
binding_id = _validated_ref_part(spec.binding_id)
workspace_id = _validated_ref_part(spec.workspace_id)
snapshot_ref = _validated_ref_part(spec.home_snapshot_ref)
workspace_ref = workspace_id
if spec.existing_workspace_ref is not None:
existing_workspace_ref = _validated_ref_part(spec.existing_workspace_ref)
if existing_workspace_ref != workspace_ref:
raise BindingCreateError("existing Workspace ref does not match workspace_id")
snapshot_dir: str | None = None
if spec.home_snapshot_ref is not None:
snapshot_ref = _validated_ref_part(spec.home_snapshot_ref)
snapshot_dir = f"{self.snapshot_root.rstrip('/')}/{snapshot_ref}"
binding_ref = _local_binding_ref(binding_id=binding_id, workspace_id=workspace_id)
lease = self._control_lease(binding_ref)
home_dir = self._home_dir(binding_id)
workspace_dir = self._workspace_dir(workspace_id)
snapshot_dir = f"{self.snapshot_root.rstrip('/')}/{snapshot_ref}"
creates_workspace = spec.existing_workspace_ref is None
workspace_setup = (
f"mkdir -p {shlex.quote(workspace_dir)}" if creates_workspace else f"test -d {shlex.quote(workspace_dir)}"
)
script = "\n".join(
[
"set -eu",
f"test -d {shlex.quote(snapshot_dir)}",
workspace_setup,
f"mkdir -p {shlex.quote(home_dir)}",
f"cp -a {shlex.quote(snapshot_dir)}/. {shlex.quote(home_dir)}/",
f"chmod 700 {shlex.quote(home_dir)} {shlex.quote(workspace_dir)}",
]
)
setup = ["set -eu"]
if snapshot_dir is not None:
setup.append(f"test -d {shlex.quote(snapshot_dir)}")
setup.extend([workspace_setup, f"mkdir -p {shlex.quote(home_dir)}"])
if snapshot_dir is not None:
setup.append(f"cp -a {shlex.quote(snapshot_dir)}/. {shlex.quote(home_dir)}/")
setup.append(f"chmod 700 {shlex.quote(home_dir)} {shlex.quote(workspace_dir)}")
script = "\n".join(setup)
lease = self._control_lease(binding_ref)
try:
result = await run_shellctl_control_command(lease.commands, script)
if result.exit_code != 0:
result = await execute_complete_with_commands(
lease.commands,
script,
cwd=None,
env=None,
timeout=30.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
if result.exit_code != 0 or not result.output_complete:
raise BindingCreateError(result.output)
return ExecutionBindingAllocation(binding_ref=binding_ref, workspace_ref=workspace_ref)
except BaseException as exc:
@@ -193,7 +189,7 @@ class LocalExecutionBindingBackend:
async def acquire(self, binding_ref: str) -> RuntimeLease:
lease = self._lease(binding_ref)
try:
result = await run_shellctl_control_command(
result = await execute_complete_with_commands(
lease.commands,
"\n".join(
[
@@ -202,8 +198,15 @@ class LocalExecutionBindingBackend:
f"test -d {shlex.quote(lease.layout.workspace_dir)}",
]
),
cwd=None,
env=None,
timeout=10.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
if not result.output_complete:
raise BindingAcquireError(
f"Local Binding {binding_ref!r} validation did not complete: {result.incomplete_reason}"
)
if result.exit_code != 0:
raise BindingLostError(f"Local Binding {binding_ref!r} no longer exists")
return lease
@@ -234,11 +237,15 @@ class LocalExecutionBindingBackend:
if spec.destroy_workspace:
targets.append(self._workspace_dir(workspace_id))
try:
result = await run_shellctl_control_command(
result = await execute_complete_with_commands(
lease.commands,
"rm -rf -- " + " ".join(shlex.quote(target) for target in targets),
cwd=None,
env=None,
timeout=30.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
if result.exit_code != 0:
if result.exit_code != 0 or not result.output_complete:
raise BindingDestroyError(result.output)
except BaseException:
await _close_best_effort(lease, resource_ref=spec.binding_ref)
@@ -314,8 +321,15 @@ async def _remove_partial(
) -> None:
try:
target_words = target if target_is_shell_words else shlex.quote(target)
result = await run_shellctl_control_command(commands, f"rm -rf -- {target_words}")
if result.exit_code != 0:
result = await execute_complete_with_commands(
commands,
f"rm -rf -- {target_words}",
cwd=None,
env=None,
timeout=30.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
if result.exit_code != 0 or not result.output_complete:
logger.warning("failed to remove partial local resource", extra={"resource_ref": resource_ref})
except BaseException:
logger.warning("failed to remove partial local resource", exc_info=True, extra={"resource_ref": resource_ref})
@@ -5,7 +5,7 @@ from __future__ import annotations
from typing import ClassVar, Literal, Self
from urllib.parse import urlparse
from pydantic import AliasChoices, Field, model_validator
from pydantic import AliasChoices, Field, field_validator, model_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
from dify_agent.runtime_backend.e2b import (
@@ -14,6 +14,13 @@ from dify_agent.runtime_backend.e2b import (
E2BSDKControlPlane,
E2BExecutionBindingBackend,
)
from dify_agent.runtime_backend.e2b_s3 import (
E2BHomeSnapshotCLI,
E2BS3ExecutionBindingBackend,
E2BS3HomeSnapshotBackend,
OpenDALHomeArchiveStore,
)
from dify_agent.agent_stub.server.tokens.home_snapshot import HomeSnapshotTransferTokenCodec
from dify_agent.runtime_backend.enterprise import EnterpriseExecutionBindingBackend, EnterpriseHomeSnapshotBackend
from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, LocalHomeSnapshotBackend
from dify_agent.runtime_backend.protocols import RuntimeBackendProfile
@@ -24,7 +31,7 @@ DEFAULT_E2B_TEMPLATE = "difys-default-team/dify-agent-local-sandbox"
class RuntimeBackendSettings(BaseSettings):
"""Server-private credentials and endpoints for one coherent backend profile."""
runtime_backend: Literal["local", "enterprise", "e2b"] = "local"
runtime_backend: Literal["local", "enterprise", "e2b", "e2b_s3"] = "local"
local_sandbox_endpoint: str | None = Field(
default=None,
@@ -55,6 +62,9 @@ class RuntimeBackendSettings(BaseSettings):
)
e2b_shellctl_auth_token: str = ""
e2b_shellctl_port: int = Field(default=5004, ge=1, le=65535)
e2b_s3_uri: str | None = None
agent_stub_api_base_url: str | None = None
server_secret_key: str | None = None
model_config: ClassVar[SettingsConfigDict] = SettingsConfigDict(
env_prefix="DIFY_AGENT_",
@@ -63,6 +73,19 @@ class RuntimeBackendSettings(BaseSettings):
populate_by_name=True,
)
@field_validator(
"e2b_api_key",
"e2b_s3_uri",
"agent_stub_api_base_url",
"server_secret_key",
mode="before",
)
@classmethod
def normalize_optional_e2b_s3_value(cls, value: object) -> object:
if isinstance(value, str):
return value.strip() or None
return value
@model_validator(mode="after")
def validate_selected_backend(self) -> Self:
match self.runtime_backend:
@@ -77,13 +100,25 @@ class RuntimeBackendSettings(BaseSettings):
"enterprise_sandbox_gateway_endpoint is required for the enterprise runtime backend"
)
_validate_http_url(endpoint, field_name="enterprise_sandbox_gateway_endpoint")
case "e2b":
case "e2b" | "e2b_s3":
if not self.e2b_api_key or not self.e2b_api_key.strip():
raise ValueError("e2b_api_key is required for the e2b runtime backend")
raise ValueError(f"e2b_api_key is required for the {self.runtime_backend} runtime backend")
if not self.e2b_template.strip():
raise ValueError("e2b_template must not be blank")
if self.runtime_backend == "e2b_s3":
self._validate_e2b_s3()
return self
def _validate_e2b_s3(self) -> None:
if not self.e2b_s3_uri:
raise ValueError("e2b_s3_uri is required for the e2b_s3 runtime backend")
if not self.agent_stub_api_base_url:
raise ValueError("agent_stub_api_base_url is required for the e2b_s3 runtime backend")
_validate_http_url(self.agent_stub_api_base_url, field_name="agent_stub_api_base_url")
if not self.server_secret_key:
raise ValueError("server_secret_key is required for the e2b_s3 runtime backend")
_ = HomeSnapshotTransferTokenCodec.from_server_secret(self.server_secret_key)
def create_runtime_backend_profile(settings: RuntimeBackendSettings) -> RuntimeBackendProfile:
"""Construct one driver pair selected exclusively by server deployment settings."""
@@ -99,11 +134,7 @@ def create_runtime_backend_profile(settings: RuntimeBackendSettings) -> RuntimeB
endpoint = settings.enterprise_sandbox_gateway_endpoint or ""
token = settings.enterprise_sandbox_gateway_auth_token or ""
return RuntimeBackendProfile(
home_snapshots=EnterpriseHomeSnapshotBackend(
gateway_endpoint=endpoint,
auth_token=token,
gateway_timeout=settings.enterprise_sandbox_gateway_timeout,
),
home_snapshots=EnterpriseHomeSnapshotBackend(),
execution_bindings=EnterpriseExecutionBindingBackend(
gateway_endpoint=endpoint,
auth_token=token,
@@ -116,11 +147,37 @@ def create_runtime_backend_profile(settings: RuntimeBackendSettings) -> RuntimeB
return RuntimeBackendProfile(
home_snapshots=E2BHomeSnapshotBackend(
control_plane=control_plane,
template=settings.e2b_template,
active_timeout_seconds=settings.e2b_active_timeout_seconds,
),
execution_bindings=E2BExecutionBindingBackend(
control_plane=control_plane,
template=settings.e2b_template,
active_timeout_seconds=settings.e2b_active_timeout_seconds,
shellctl_auth_token=settings.e2b_shellctl_auth_token,
shellctl_port=settings.e2b_shellctl_port,
),
)
case "e2b_s3":
control_plane = E2BSDKControlPlane(api_key=settings.e2b_api_key or "")
archive_store = OpenDALHomeArchiveStore.create_from_uri(settings.e2b_s3_uri or "")
token_codec = HomeSnapshotTransferTokenCodec.from_server_secret(settings.server_secret_key or "")
lifecycle_cli = E2BHomeSnapshotCLI(
token_codec=token_codec,
agent_stub_api_base_url=settings.agent_stub_api_base_url or "",
shellctl_auth_token=settings.e2b_shellctl_auth_token,
shellctl_port=settings.e2b_shellctl_port,
)
return RuntimeBackendProfile(
home_snapshots=E2BS3HomeSnapshotBackend(
control_plane=control_plane,
archive_store=archive_store,
lifecycle_cli=lifecycle_cli,
template=settings.e2b_template,
active_timeout_seconds=settings.e2b_active_timeout_seconds,
),
execution_bindings=E2BS3ExecutionBindingBackend(
control_plane=control_plane,
lifecycle_cli=lifecycle_cli,
template=settings.e2b_template,
active_timeout_seconds=settings.e2b_active_timeout_seconds,
shellctl_auth_token=settings.e2b_shellctl_auth_token,
shellctl_port=settings.e2b_shellctl_port,
@@ -13,13 +13,6 @@ from typing import Protocol
from dify_agent.adapters.shell.protocols import ShellCommandProtocol
@dataclass(frozen=True, slots=True)
class InitializeHomeSnapshotSpec:
tenant_id: str
agent_id: str
home_snapshot_id: str
@dataclass(frozen=True, slots=True)
class HomeSnapshotCreateSpec:
tenant_id: str
@@ -100,7 +93,7 @@ class ExecutionBindingCreateSpec:
binding_id: str
workspace_id: str
existing_workspace_ref: str | None
home_snapshot_ref: str
home_snapshot_ref: str | None = None
@dataclass(frozen=True, slots=True)
@@ -126,12 +119,18 @@ class ExecutionBindingBackend(Protocol):
async def create_binding(self, spec: ExecutionBindingCreateSpec) -> ExecutionBindingAllocation:
"""Materialize a mutable Home and make the requested Workspace ready.
Implementations must initialize the Home from ``home_snapshot_ref`` and
return stable opaque refs only after both resources are usable. With an
``existing_workspace_ref``, they must attach that Workspace without
clearing or replacing its contents; unsupported sharing must fail before
mutating it. On failure, implementations should clean up newly allocated
partial resources and must not damage a pre-existing Workspace.
With a non-null ``home_snapshot_ref``, implementations must initialize
the Home from that exact immutable snapshot and must fail rather than
fall back when it is unavailable. With ``None``, implementations must
create an independent mutable Home from their deployment default without
implicitly creating an immutable snapshot.
Implementations must return stable opaque refs only after Home and
Workspace are usable. With an ``existing_workspace_ref``, they must
attach that Workspace without clearing or replacing its contents;
unsupported sharing must fail before mutating it. On failure,
implementations should clean up newly allocated partial resources and
must not damage a pre-existing Workspace.
"""
...
@@ -150,9 +149,10 @@ class ExecutionBindingBackend(Protocol):
"""End operation-local access without destroying persistent resources.
Implementations must close clients and owned transports and may suspend
the physical runtime. They must not delete or retire the Binding,
materialized Home, or Workspace, and must reject leases created by a
different backend implementation.
a physical runtime only when it is not shared by other Bindings. They
must not suspend a physical runtime still serving another Binding,
delete or retire the Binding, materialized Home, or Workspace, and must
reject leases created by a different backend implementation.
"""
...
@@ -171,16 +171,6 @@ class ExecutionBindingBackend(Protocol):
class HomeSnapshotBackend(Protocol):
"""Manage immutable backend-native Home resources."""
async def initialize(self, spec: InitializeHomeSnapshotSpec) -> str:
"""Create the deployment-defined baseline Home Snapshot.
Implementations may use any backend-native bootstrap mechanism, but must
return a stable opaque ref only after an immutable snapshot is ready for
future Binding creation. Temporary bootstrap resources must not become
part of the logical snapshot lifecycle.
"""
...
async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str:
"""Capture the source lease's current Home as a new immutable snapshot.
@@ -217,7 +207,6 @@ __all__ = [
"FileSystem",
"HomeSnapshotBackend",
"HomeSnapshotCreateSpec",
"InitializeHomeSnapshotSpec",
"RuntimeBackendProfile",
"RuntimeLayout",
"RuntimeLease",
@@ -4,9 +4,14 @@ from __future__ import annotations
from dataclasses import dataclass, field
import logging
from typing import Protocol
import time
from typing import Literal, Protocol
from dify_agent.adapters.shell.protocols import CompleteShellCommandResult, ShellCommandProtocol
from dify_agent.adapters.shell.protocols import (
CompleteShellCommandResult,
ShellCommandProtocol,
ShellCommandResult,
)
from dify_agent.adapters.shell.shellctl import (
ShellctlClientFactory,
ShellctlClientProtocol,
@@ -16,7 +21,9 @@ from dify_agent.adapters.shell.shellctl import (
)
from dify_agent.runtime_backend.protocols import FileSystem, RuntimeLayout
_CONTROL_COMMAND_OUTPUT_LIMIT = 256 * 1024
_COMPLETE_POLL_TIMEOUT_SECONDS = 30.0
_COMPLETE_TERMINATE_GRACE_SECONDS = 10.0
CONTROL_COMMAND_OUTPUT_LIMIT = 256 * 1024
logger = logging.getLogger(__name__)
@@ -113,44 +120,113 @@ async def create_owned_shellctl_lease(
raise
async def run_shellctl_control_command(
async def execute_complete_with_commands(
commands: ShellCommandProtocol,
script: str,
*,
timeout: float = 30.0,
cwd: str | None,
env: dict[str, str] | None,
timeout: float | None,
max_output_bytes: int,
) -> CompleteShellCommandResult:
"""Run one bounded driver control command and always delete its transient job."""
result = await commands.run(script, cwd=None, env=None, timeout=timeout)
job_id = result.job_id
output_parts = [result.output]
"""Run one command to completion and always delete its transient shellctl job.
``timeout=None`` deliberately imposes no total execution deadline. Each
shellctl call still long-polls for at most 30 seconds so cancellation and
terminal state are observed without turning a poll timeout into a job
timeout.
"""
deadline = None if timeout is None else time.monotonic() + timeout
job_id: str | None = None
result: ShellCommandResult | None = None
output_parts: list[str] = []
captured_bytes = 0
incomplete_reason: Literal["output_limit", "timeout"] | None = None
try:
while result.truncated or not result.done:
result = await commands.wait(result.job_id, offset=result.offset, timeout=timeout)
output_parts.append(result.output)
if sum(len(part.encode("utf-8")) for part in output_parts) > _CONTROL_COMMAND_OUTPUT_LIMIT:
raise RuntimeError("shellctl control command exceeded its output limit")
result = await commands.run(
script,
cwd=cwd,
env=env,
timeout=_poll_timeout(deadline),
)
job_id = result.job_id
while True:
remaining_bytes = max(max_output_bytes - captured_bytes, 0)
limited_output = _utf8_prefix(result.output, remaining_bytes)
output_parts.append(limited_output)
captured_bytes += len(limited_output.encode("utf-8"))
if limited_output != result.output:
incomplete_reason = "output_limit"
break
if captured_bytes >= max_output_bytes and (result.truncated or not result.done):
incomplete_reason = "output_limit"
break
if result.truncated:
result = await commands.read_output(result.job_id, offset=result.offset)
continue
if result.done:
break
if deadline is not None and time.monotonic() >= deadline:
incomplete_reason = "timeout"
break
result = await commands.wait(
result.job_id,
offset=result.offset,
timeout=_poll_timeout(deadline),
)
final_status = result.status
final_done = result.done
final_exit_code = result.exit_code
final_offset = result.offset
if incomplete_reason is not None and not result.done:
terminal = await commands.interrupt(
result.job_id,
grace_seconds=_COMPLETE_TERMINATE_GRACE_SECONDS,
)
final_status = terminal.status
final_done = terminal.done
final_exit_code = terminal.exit_code
final_offset = terminal.offset
return CompleteShellCommandResult(
job_id=result.job_id,
status=result.status,
done=result.done,
exit_code=result.exit_code,
status=final_status,
done=final_done,
exit_code=final_exit_code,
output="".join(output_parts),
output_complete=True,
incomplete_reason=None,
offset=result.offset,
output_complete=incomplete_reason is None,
incomplete_reason=incomplete_reason,
offset=final_offset,
output_path=result.output_path,
)
finally:
try:
await commands.delete(job_id, force=True)
except Exception as exc:
logger.warning("Failed to delete transient shellctl control job %s: %s", job_id, exc)
if job_id is not None:
try:
await commands.delete(job_id, force=True)
except Exception as exc:
logger.warning("Failed to delete transient shellctl job %s: %s", job_id, exc)
def _poll_timeout(deadline: float | None) -> float:
if deadline is None:
return _COMPLETE_POLL_TIMEOUT_SECONDS
return max(deadline - time.monotonic(), 0.0)
def _utf8_prefix(value: str, max_bytes: int) -> str:
if max_bytes <= 0:
return ""
encoded = value.encode("utf-8")
if len(encoded) <= max_bytes:
return value
return encoded[:max_bytes].decode("utf-8", errors="ignore")
__all__ = [
"AsyncCloseable",
"CONTROL_COMMAND_OUTPUT_LIMIT",
"ShellctlRuntimeLease",
"create_owned_shellctl_lease",
"create_shellctl_lease",
"run_shellctl_control_command",
"execute_complete_with_commands",
]
+2
View File
@@ -64,6 +64,7 @@ def create_app(settings: ServerSettings | None = None) -> FastAPI:
agent_stub_file_request_handler = resolved_settings.create_agent_stub_file_request_handler()
agent_stub_config_request_handler = resolved_settings.create_agent_stub_config_request_handler()
agent_stub_drive_request_handler = resolved_settings.create_agent_stub_drive_request_handler()
home_snapshot_gateway = resolved_settings.build_home_snapshot_gateway()
runtime_backend_profile = resolved_settings.build_runtime_backend_profile()
layer_providers = create_default_layer_providers(
plugin_daemon_url=resolved_settings.plugin_daemon_url,
@@ -168,6 +169,7 @@ def create_app(settings: ServerSettings | None = None) -> FastAPI:
file_request_handler=agent_stub_file_request_handler,
config_request_handler=agent_stub_config_request_handler,
drive_request_handler=agent_stub_drive_request_handler,
home_snapshot_gateway=home_snapshot_gateway,
)
)
return app
@@ -6,7 +6,6 @@ from dify_agent.protocol.home_snapshot import (
CreateHomeSnapshotFromBindingRequest,
DeleteHomeSnapshotRequest,
HomeSnapshotResponse,
InitializeHomeSnapshotRequest,
)
from dify_agent.runtime_backend import (
BindingAcquireError,
@@ -16,7 +15,6 @@ from dify_agent.runtime_backend import (
HomeSnapshotCreateError,
HomeSnapshotCreateSpec,
HomeSnapshotNotFoundError,
InitializeHomeSnapshotSpec,
)
from dify_agent.runtime_backend.leases import open_runtime_lease
@@ -38,19 +36,6 @@ class HomeSnapshotService:
home_snapshots: HomeSnapshotBackend
execution_bindings: ExecutionBindingBackend
async def initialize(self, request: InitializeHomeSnapshotRequest) -> HomeSnapshotResponse:
try:
snapshot_ref = await self.home_snapshots.initialize(
InitializeHomeSnapshotSpec(
tenant_id=request.tenant_id,
agent_id=request.agent_id,
home_snapshot_id=request.home_snapshot_id,
)
)
except HomeSnapshotCreateError as exc:
raise HomeSnapshotServiceError("home_snapshot_create_failed", str(exc), status_code=502) from exc
return HomeSnapshotResponse(snapshot_ref=snapshot_ref)
async def create_from_binding(
self,
request: CreateHomeSnapshotFromBindingRequest,
@@ -9,7 +9,6 @@ from dify_agent.protocol.home_snapshot import (
CreateHomeSnapshotFromBindingRequest,
DeleteHomeSnapshotRequest,
HomeSnapshotResponse,
InitializeHomeSnapshotRequest,
)
from dify_agent.server.home_snapshots import HomeSnapshotService, HomeSnapshotServiceError
@@ -26,19 +25,6 @@ def create_home_snapshots_router(get_service: Callable[[], HomeSnapshotService |
)
return service
@router.post("/initialize", response_model=HomeSnapshotResponse, status_code=status.HTTP_201_CREATED)
async def initialize_snapshot(
request: InitializeHomeSnapshotRequest,
service: Annotated[HomeSnapshotService, Depends(service_dep)],
) -> HomeSnapshotResponse:
try:
return await service.initialize(request)
except HomeSnapshotServiceError as exc:
raise HTTPException(
status_code=exc.status_code,
detail={"code": exc.code, "message": exc.message},
) from exc
@router.post("/from-binding", response_model=HomeSnapshotResponse, status_code=status.HTTP_201_CREATED)
async def create_snapshot_from_binding(
request: CreateHomeSnapshotFromBindingRequest,
+42 -3
View File
@@ -10,10 +10,10 @@ optional gRPC bind override, and optional Dify inner API bridge settings all
live here under the ``DIFY_AGENT_...`` environment-variable namespace.
"""
import httpx
from typing import ClassVar, Literal, cast
import httpx
from pydantic import AliasChoices, AnyHttpUrl, Field, TypeAdapter, field_validator, model_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
@@ -23,8 +23,11 @@ from dify_agent.agent_stub.server.agent_stub_drive import DifyApiAgentStubDriveR
from dify_agent.agent_stub.server.agent_stub_files import DifyApiAgentStubFileRequestHandler
from dify_agent.agent_stub.server.grpc_bind import normalize_agent_stub_grpc_bind_address
from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubTokenCodec, decode_server_secret_key
from dify_agent.agent_stub.server.home_snapshots import HomeSnapshotGatewayService
from dify_agent.agent_stub.server.tokens.home_snapshot import HomeSnapshotTransferTokenCodec
from dify_agent.runtime_backend import RuntimeBackendProfile
from dify_agent.runtime_backend.e2b import E2B_MAX_ACTIVE_TIMEOUT_SECONDS
from dify_agent.runtime_backend.e2b_s3 import OpenDALHomeArchiveStore
from dify_agent.runtime_backend.profile import RuntimeBackendSettings, create_runtime_backend_profile
DEFAULT_RUN_RETENTION_SECONDS = 3 * 24 * 60 * 60
@@ -41,7 +44,7 @@ class ServerSettings(BaseSettings):
plugin_daemon_api_key: str = ""
inner_api_url: str = "http://localhost:5001"
inner_api_key: str | None = None
runtime_backend: Literal["local", "enterprise", "e2b"] = "local"
runtime_backend: Literal["local", "enterprise", "e2b", "e2b_s3"] = "local"
local_sandbox_endpoint: str | None = Field(
default=None,
validation_alias=AliasChoices("DIFY_AGENT_LOCAL_SANDBOX_ENDPOINT", "DIFY_AGENT_SHELLCTL_ENTRYPOINT"),
@@ -63,6 +66,7 @@ class ServerSettings(BaseSettings):
)
e2b_shellctl_auth_token: str = ""
e2b_shellctl_port: int = Field(default=5004, ge=1, le=65535)
e2b_s3_uri: str | None = None
sandbox_file_upload_max_bytes: int = Field(default=50 * 1024 * 1024, ge=1)
agent_stub_api_base_url: str | None = Field(default=None, validation_alias="DIFY_AGENT_STUB_API_BASE_URL")
agent_stub_grpc_bind_address: str | None = Field(default=None, validation_alias="DIFY_AGENT_STUB_GRPC_BIND_ADDRESS")
@@ -84,6 +88,17 @@ class ServerSettings(BaseSettings):
populate_by_name=True,
)
@field_validator(
"e2b_api_key",
"e2b_s3_uri",
mode="before",
)
@classmethod
def normalize_optional_e2b_s3_value(cls, value: object) -> object:
if isinstance(value, str):
return value.strip() or None
return value
@field_validator("agent_stub_api_base_url")
@classmethod
def normalize_agent_stub_api_base_url_value(cls, value: str | None) -> str | None:
@@ -167,6 +182,19 @@ class ServerSettings(BaseSettings):
)
if not parse_agent_stub_endpoint(self.agent_stub_api_base_url).is_grpc:
raise ValueError("DIFY_AGENT_STUB_GRPC_BIND_ADDRESS requires a grpc:// DIFY_AGENT_STUB_API_BASE_URL.")
if self.runtime_backend == "e2b_s3":
if not self.e2b_api_key or not self.e2b_api_key.strip():
raise ValueError("DIFY_AGENT_E2B_API_KEY is required for the e2b_s3 runtime backend.")
if not self.e2b_template.strip():
raise ValueError("DIFY_AGENT_E2B_TEMPLATE must not be blank.")
if not self.e2b_s3_uri:
raise ValueError("DIFY_AGENT_E2B_S3_URI is required for the e2b_s3 runtime backend.")
if self.agent_stub_api_base_url is None:
raise ValueError("DIFY_AGENT_STUB_API_BASE_URL is required for the e2b_s3 runtime backend.")
if parse_agent_stub_endpoint(self.agent_stub_api_base_url).is_grpc:
raise ValueError("e2b_s3 requires an HTTP(S) DIFY_AGENT_STUB_API_BASE_URL.")
if self.server_secret_key is None:
raise ValueError("DIFY_AGENT_SERVER_SECRET_KEY is required for the e2b_s3 runtime backend.")
return self
def build_runtime_backend_profile(self) -> RuntimeBackendProfile | None:
@@ -187,9 +215,20 @@ class ServerSettings(BaseSettings):
e2b_active_timeout_seconds=self.e2b_active_timeout_seconds,
e2b_shellctl_auth_token=self.e2b_shellctl_auth_token,
e2b_shellctl_port=self.e2b_shellctl_port,
e2b_s3_uri=self.e2b_s3_uri,
agent_stub_api_base_url=self.agent_stub_api_base_url,
server_secret_key=self.server_secret_key,
)
)
def build_home_snapshot_gateway(self) -> HomeSnapshotGatewayService | None:
"""Build the e2b_s3 HTTP byte-stream gateway from independent resources."""
if self.runtime_backend != "e2b_s3":
return None
store = OpenDALHomeArchiveStore.create_from_uri(self.e2b_s3_uri or "")
codec = HomeSnapshotTransferTokenCodec.from_server_secret(self.server_secret_key or "")
return HomeSnapshotGatewayService(token_codec=codec, archive_store=store)
def create_agent_stub_token_codec(self) -> AgentStubTokenCodec | None:
"""Return the Agent Stub token codec when the server secret is configured."""
if self.server_secret_key is None:
@@ -3,6 +3,7 @@
from __future__ import annotations
import os
import shlex
import sys
import uuid
@@ -12,10 +13,17 @@ from dify_agent.runtime_backend import (
ExecutionBindingCreateSpec,
ExecutionBindingDestroySpec,
HomeSnapshotCreateSpec,
InitializeHomeSnapshotSpec,
)
from dify_agent.runtime_backend.e2b import E2BExecutionBindingBackend, E2BHomeSnapshotBackend, E2BSDKControlPlane
from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, LocalHomeSnapshotBackend
from dify_agent.runtime_backend.e2b_s3 import (
E2BHomeSnapshotCLI,
E2BS3ExecutionBindingBackend,
E2BS3HomeSnapshotBackend,
OpenDALHomeArchiveStore,
)
from dify_agent.agent_stub.server.tokens.home_snapshot import HomeSnapshotTransferTokenCodec
from dify_agent.runtime_backend.local import LocalExecutionBindingBackend
from dify_agent.runtime_backend.shellctl import CONTROL_COMMAND_OUTPUT_LIMIT, execute_complete_with_commands
pytestmark = pytest.mark.integration
@@ -32,19 +40,10 @@ async def test_local_two_agents_share_workspace_but_not_home() -> None:
endpoint = _required_env("DIFY_AGENT_TEST_LOCAL_SHELLCTL_ENDPOINT", "real Local shellctl")
token = os.environ.get("DIFY_AGENT_TEST_LOCAL_SHELLCTL_AUTH_TOKEN", "")
marker = uuid.uuid4().hex
snapshots = LocalHomeSnapshotBackend(endpoint=endpoint, auth_token=token)
bindings = LocalExecutionBindingBackend(endpoint=endpoint, auth_token=token)
snapshot_ref: str | None = None
allocations = []
active_leases = []
try:
snapshot_ref = await snapshots.initialize(
InitializeHomeSnapshotSpec(
tenant_id="integration-tenant",
agent_id="integration-agent",
home_snapshot_id=marker,
)
)
first = await bindings.create_binding(
ExecutionBindingCreateSpec(
tenant_id="integration-tenant",
@@ -52,7 +51,7 @@ async def test_local_two_agents_share_workspace_but_not_home() -> None:
binding_id=f"binding-a-{marker}",
workspace_id=f"workspace-{marker}",
existing_workspace_ref=None,
home_snapshot_ref=snapshot_ref,
home_snapshot_ref=None,
)
)
allocations.append(first)
@@ -71,7 +70,7 @@ async def test_local_two_agents_share_workspace_but_not_home() -> None:
binding_id=f"binding-b-{marker}",
workspace_id=f"workspace-{marker}",
existing_workspace_ref=first.workspace_ref,
home_snapshot_ref=snapshot_ref,
home_snapshot_ref=None,
)
)
allocations.append(second)
@@ -102,11 +101,6 @@ async def test_local_two_agents_share_workspace_but_not_home() -> None:
)
except BaseException as exc:
cleanup_errors.append(exc)
if snapshot_ref is not None:
try:
await snapshots.delete(snapshot_ref)
except BaseException as exc:
cleanup_errors.append(exc)
if cleanup_errors and not primary_error:
raise cleanup_errors[0]
@@ -120,22 +114,14 @@ async def test_e2b_binding_checkpoint_and_collection() -> None:
)
marker = uuid.uuid4().hex
control = E2BSDKControlPlane(api_key=api_key)
snapshots = E2BHomeSnapshotBackend(control_plane=control, template=template, active_timeout_seconds=3600)
bindings = E2BExecutionBindingBackend(control_plane=control, active_timeout_seconds=3600)
snapshot_ref: str | None = None
snapshots = E2BHomeSnapshotBackend(control_plane=control)
bindings = E2BExecutionBindingBackend(control_plane=control, template=template, active_timeout_seconds=3600)
checkpoint_ref: str | None = None
allocation = None
checkpoint_allocation = None
lease = None
checkpoint_lease = None
try:
snapshot_ref = await snapshots.initialize(
InitializeHomeSnapshotSpec(
tenant_id="integration-tenant",
agent_id="integration-agent",
home_snapshot_id=marker,
)
)
allocation = await bindings.create_binding(
ExecutionBindingCreateSpec(
tenant_id="integration-tenant",
@@ -143,7 +129,7 @@ async def test_e2b_binding_checkpoint_and_collection() -> None:
binding_id=marker,
workspace_id=marker,
existing_workspace_ref=None,
home_snapshot_ref=snapshot_ref,
home_snapshot_ref=None,
)
)
lease = await bindings.acquire(allocation.binding_ref)
@@ -211,11 +197,173 @@ async def test_e2b_binding_checkpoint_and_collection() -> None:
)
except BaseException as exc:
cleanup_errors.append(exc)
for ref in (checkpoint_ref, snapshot_ref):
if ref is not None:
try:
await snapshots.delete(ref)
except BaseException as exc:
cleanup_errors.append(exc)
if checkpoint_ref is not None:
try:
await snapshots.delete(checkpoint_ref)
except BaseException as exc:
cleanup_errors.append(exc)
if cleanup_errors and not primary_error:
raise cleanup_errors[0]
@pytest.mark.anyio
async def test_e2b_s3_shared_workspace_checkpoint_and_collection() -> None:
api_key = _required_env("DIFY_AGENT_TEST_E2B_API_KEY", "real E2B + S3")
storage_uri = _required_env("DIFY_AGENT_TEST_E2B_S3_URI", "real E2B + OpenDAL storage")
stub_url = _required_env("DIFY_AGENT_TEST_E2B_S3_STUB_API_BASE_URL", "public Home Snapshot gateway")
server_secret = _required_env("DIFY_AGENT_TEST_SERVER_SECRET_KEY", "Home Snapshot transfer JWE")
template = os.environ.get(
"DIFY_AGENT_TEST_E2B_TEMPLATE",
"difys-default-team/dify-agent-local-sandbox",
)
marker = uuid.uuid4().hex
control = E2BSDKControlPlane(api_key=api_key)
store = OpenDALHomeArchiveStore.create_from_uri(storage_uri)
lifecycle_cli = E2BHomeSnapshotCLI(
token_codec=HomeSnapshotTransferTokenCodec.from_server_secret(server_secret),
agent_stub_api_base_url=stub_url,
shellctl_auth_token=os.environ.get("DIFY_AGENT_TEST_E2B_SHELLCTL_AUTH_TOKEN", ""),
shellctl_port=int(os.environ.get("DIFY_AGENT_TEST_E2B_SHELLCTL_PORT", "5004")),
)
snapshots = E2BS3HomeSnapshotBackend(
control_plane=control,
archive_store=store,
lifecycle_cli=lifecycle_cli,
template=template,
active_timeout_seconds=3600,
)
bindings = E2BS3ExecutionBindingBackend(
control_plane=control,
lifecycle_cli=lifecycle_cli,
template=template,
active_timeout_seconds=3600,
shellctl_auth_token=os.environ.get("DIFY_AGENT_TEST_E2B_SHELLCTL_AUTH_TOKEN", ""),
shellctl_port=int(os.environ.get("DIFY_AGENT_TEST_E2B_SHELLCTL_PORT", "5004")),
)
allocations = []
active_leases = []
snapshot_refs: list[str] = []
try:
baseline_ref = await snapshots.initialize(
HomeSnapshotCreateSpec(
tenant_id="integration-tenant",
agent_id="integration-agent",
home_snapshot_id=f"baseline-{marker}",
)
)
snapshot_refs.append(baseline_ref)
first = await bindings.create_binding(
ExecutionBindingCreateSpec(
tenant_id="integration-tenant",
agent_id="integration-agent",
binding_id=f"binding-a-{marker}",
workspace_id=f"workspace-{marker}",
existing_workspace_ref=None,
home_snapshot_ref=baseline_ref,
)
)
allocations.append(first)
first_lease = await bindings.acquire(first.binding_ref)
active_leases.append(first_lease)
await first_lease.files.upload(
content=b"shared",
remote_path="shared.txt",
cwd=first_lease.layout.workspace_dir,
)
await first_lease.files.upload(
content=b"checkpoint-home",
remote_path=".checkpoint-probe",
cwd=first_lease.layout.home_dir,
)
second = await bindings.create_binding(
ExecutionBindingCreateSpec(
tenant_id="integration-tenant",
agent_id="integration-agent",
binding_id=f"binding-b-{marker}",
workspace_id=f"workspace-{marker}",
existing_workspace_ref=first.workspace_ref,
home_snapshot_ref=baseline_ref,
)
)
allocations.append(second)
second_lease = await bindings.acquire(second.binding_ref)
active_leases.append(second_lease)
assert (await second_lease.files.read_bytes(path="shared.txt", max_bytes=1024)).content == b"shared"
assert second_lease.layout.home_dir != first_lease.layout.home_dir
sibling_read = await execute_complete_with_commands(
second_lease.commands,
f"cat {shlex.quote(f'{first_lease.layout.home_dir}/.checkpoint-probe')}",
cwd=None,
env={"HOME": "/tmp/ignored"},
timeout=10.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
assert sibling_read.done and sibling_read.exit_code not in (None, 0)
await bindings.release(second_lease)
active_leases.remove(second_lease)
checkpoint_ref = await snapshots.create_from_runtime(
spec=HomeSnapshotCreateSpec(
tenant_id="integration-tenant",
agent_id="integration-agent",
home_snapshot_id=f"checkpoint-{marker}",
),
source=first_lease,
)
snapshot_refs.append(checkpoint_ref)
third = await bindings.create_binding(
ExecutionBindingCreateSpec(
tenant_id="integration-tenant",
agent_id="integration-agent",
binding_id=f"binding-c-{marker}",
workspace_id=f"workspace-{marker}",
existing_workspace_ref=first.workspace_ref,
home_snapshot_ref=checkpoint_ref,
)
)
allocations.append(third)
third_lease = await bindings.acquire(third.binding_ref)
active_leases.append(third_lease)
restored = await third_lease.files.read_bytes(path="~/.checkpoint-probe", max_bytes=1024)
assert restored.content == b"checkpoint-home"
assert (await third_lease.files.read_bytes(path="shared.txt", max_bytes=1024)).content == b"shared"
await bindings.destroy_binding(
ExecutionBindingDestroySpec(binding_ref=second.binding_ref, destroy_workspace=False)
)
allocations.remove(second)
finally:
primary_error = sys.exc_info()[0] is not None
cleanup_errors: list[BaseException] = []
for lease in active_leases:
try:
await bindings.release(lease)
except BaseException as exc:
cleanup_errors.append(exc)
if allocations:
workspace_owner = allocations.pop(0)
try:
await bindings.destroy_binding(
ExecutionBindingDestroySpec(
binding_ref=workspace_owner.binding_ref,
workspace_ref=workspace_owner.workspace_ref,
destroy_workspace=True,
)
)
except BaseException as exc:
cleanup_errors.append(exc)
for allocation in allocations:
try:
await bindings.destroy_binding(
ExecutionBindingDestroySpec(binding_ref=allocation.binding_ref, destroy_workspace=False)
)
except BaseException as exc:
cleanup_errors.append(exc)
for snapshot_ref in snapshot_refs:
try:
await snapshots.delete(snapshot_ref)
except BaseException as exc:
cleanup_errors.append(exc)
if cleanup_errors and not primary_error:
raise cleanup_errors[0]
@@ -5,8 +5,10 @@ import importlib
import time
import httpx
from fastapi import APIRouter
from fastapi.testclient import TestClient
import dify_agent.agent_stub.server.app as app_module
from dify_agent.agent_stub.server.app import create_agent_stub_app
from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig
from dify_agent.server.settings import ServerSettings
@@ -58,6 +60,32 @@ def test_create_agent_stub_app_can_serve_requests() -> None:
assert response.json()["detail"] == "Agent Stub is not configured"
def test_create_agent_stub_app_injects_e2b_s3_home_snapshot_gateway(monkeypatch) -> None:
gateway = object()
captured: list[object] = []
def build_gateway(_self: ServerSettings) -> object:
return gateway
def capture_router(**kwargs: object) -> APIRouter:
captured.append(kwargs["home_snapshot_gateway"])
return APIRouter()
monkeypatch.setattr(ServerSettings, "build_home_snapshot_gateway", build_gateway)
monkeypatch.setattr(app_module, "create_agent_stub_router", capture_router)
settings = ServerSettings(
runtime_backend="e2b_s3",
e2b_api_key="e2b-secret",
e2b_s3_uri="s3://snapshots/dify-agent",
agent_stub_api_base_url="https://agent.example.com/agent-stub",
server_secret_key=_base64url_secret(b"1" * 32),
)
_ = create_agent_stub_app(settings)
assert captured == [gateway]
def test_create_agent_stub_app_wires_configured_file_handler_for_upload_requests(monkeypatch) -> None:
settings = ServerSettings(
agent_stub_api_base_url="https://agent.example.com/agent-stub",
@@ -0,0 +1,384 @@
from __future__ import annotations
import asyncio
import base64
from collections.abc import AsyncIterator
from dataclasses import dataclass, field
import time
from typing import cast
from fastapi import FastAPI
from fastapi.testclient import TestClient
import pytest
from dify_agent.agent_stub.server.home_snapshots import (
AsyncArchiveFile,
HomeArchiveStore,
HomeSnapshotGatewayService,
)
from dify_agent.agent_stub.server.router import create_agent_stub_router
from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubTokenCodec
from dify_agent.agent_stub.server.tokens.home_snapshot import (
HOME_SNAPSHOT_SCOPE_READ,
HOME_SNAPSHOT_SCOPE_WRITE,
HomeSnapshotTransferTokenCodec,
)
from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig
from dify_agent.runtime_backend.errors import (
HomeArchiveConflictError,
HomeArchiveStoreError,
HomeSnapshotNotFoundError,
)
_SNAPSHOT_REF = "home-snapshots/tenant-1/agent-1/snapshot-1.tar.zst"
def _secret() -> str:
return base64.urlsafe_b64encode(b"s" * 32).rstrip(b"=").decode("ascii")
@dataclass(slots=True)
class _Archive:
reads: list[bytes] = field(default_factory=list)
writes: list[bytes] = field(default_factory=list)
read_error: BaseException | None = None
read_started: asyncio.Event | None = None
read_release: asyncio.Event | None = None
write_error: Exception | None = None
close_error: BaseException | None = None
close_calls: int = 0
async def read(self, size: int | None = None) -> bytes:
assert size == 1024 * 1024
if self.read_started is not None:
self.read_started.set()
assert self.read_release is not None
await asyncio.wait_for(self.read_release.wait(), timeout=1.0)
if self.reads:
return self.reads.pop(0)
if self.read_error is not None:
raise self.read_error
return b""
async def write(self, data: bytes) -> int:
if self.write_error is not None:
raise self.write_error
self.writes.append(data)
return len(data)
async def close(self) -> None:
self.close_calls += 1
if self.close_error is not None:
raise self.close_error
@dataclass(slots=True)
class _Store:
writer: _Archive = field(default_factory=_Archive)
reader: _Archive = field(default_factory=lambda: _Archive(reads=[b"archive-", b"bytes"]))
open_reader_error: Exception | None = None
opened_writes: list[str] = field(default_factory=list)
opened_reads: list[str] = field(default_factory=list)
async def open_writer(self, snapshot_ref: str) -> AsyncArchiveFile:
self.opened_writes.append(snapshot_ref)
return self.writer
async def open_reader(self, snapshot_ref: str) -> AsyncArchiveFile:
self.opened_reads.append(snapshot_ref)
if self.open_reader_error is not None:
raise self.open_reader_error
return self.reader
def _client(store: _Store) -> tuple[TestClient, HomeSnapshotTransferTokenCodec]:
codec = HomeSnapshotTransferTokenCodec.from_server_secret(_secret())
gateway = HomeSnapshotGatewayService(
token_codec=codec,
archive_store=cast(HomeArchiveStore, store),
)
app = FastAPI()
app.include_router(create_agent_stub_router(token_codec=None, home_snapshot_gateway=gateway))
return TestClient(app), codec
def test_home_snapshot_gateway_streams_upload_and_download_from_token_ref() -> None:
store = _Store()
client, codec = _client(store)
write_token = codec.encode_token(scope=HOME_SNAPSHOT_SCOPE_WRITE, snapshot_ref=_SNAPSHOT_REF)
read_token = codec.encode_token(scope=HOME_SNAPSHOT_SCOPE_READ, snapshot_ref=_SNAPSHOT_REF)
upload = client.put(
"/agent-stub/home-snapshots/archive?snapshot_ref=ignored",
headers={"Authorization": f"Bearer {write_token}"},
content=b"request-stream",
)
download = client.get(
"/agent-stub/home-snapshots/archive?snapshot_ref=ignored",
headers={"Authorization": f"Bearer {read_token}"},
)
assert upload.status_code == 204
assert b"".join(store.writer.writes) == b"request-stream"
assert store.writer.close_calls == 1
assert store.opened_writes == [_SNAPSHOT_REF]
assert download.status_code == 200
assert download.content == b"archive-bytes"
assert store.reader.close_calls == 1
assert store.opened_reads == [_SNAPSHOT_REF]
def test_home_snapshot_gateway_rejects_scope_method_mismatch() -> None:
client, codec = _client(_Store())
read_token = codec.encode_token(scope=HOME_SNAPSHOT_SCOPE_READ, snapshot_ref=_SNAPSHOT_REF)
response = client.put(
"/agent-stub/home-snapshots/archive",
headers={"Authorization": f"Bearer {read_token}"},
content=b"bytes",
)
assert response.status_code == 401
def test_home_snapshot_gateway_maps_missing_archive() -> None:
store = _Store(open_reader_error=HomeSnapshotNotFoundError("snapshot is missing"))
client, codec = _client(store)
token = codec.encode_token(scope=HOME_SNAPSHOT_SCOPE_READ, snapshot_ref=_SNAPSHOT_REF)
response = client.get(
"/agent-stub/home-snapshots/archive",
headers={"Authorization": f"Bearer {token}"},
)
assert response.status_code == 404
assert response.json() == {"detail": "snapshot is missing"}
def test_home_snapshot_gateway_maps_store_write_failure_and_closes_writer() -> None:
store = _Store(writer=_Archive(write_error=HomeArchiveStoreError("store unavailable")))
client, codec = _client(store)
token = codec.encode_token(scope=HOME_SNAPSHOT_SCOPE_WRITE, snapshot_ref=_SNAPSHOT_REF)
response = client.put(
"/agent-stub/home-snapshots/archive",
headers={"Authorization": f"Bearer {token}"},
content=b"archive",
)
assert response.status_code == 502
assert response.json() == {"detail": "store unavailable"}
assert store.writer.close_calls == 1
def test_home_snapshot_gateway_maps_conditional_conflict_reported_during_write() -> None:
store = _Store(writer=_Archive(write_error=HomeArchiveConflictError("object exists")))
client, codec = _client(store)
token = codec.encode_token(scope=HOME_SNAPSHOT_SCOPE_WRITE, snapshot_ref=_SNAPSHOT_REF)
response = client.put(
"/agent-stub/home-snapshots/archive",
headers={"Authorization": f"Bearer {token}"},
content=b"archive",
)
assert response.status_code == 409
assert response.json() == {"detail": "object exists"}
assert store.writer.close_calls == 1
def test_home_snapshot_gateway_maps_conditional_conflict_reported_on_close() -> None:
store = _Store(writer=_Archive(close_error=HomeArchiveConflictError("object exists")))
client, codec = _client(store)
token = codec.encode_token(scope=HOME_SNAPSHOT_SCOPE_WRITE, snapshot_ref=_SNAPSHOT_REF)
response = client.put(
"/agent-stub/home-snapshots/archive",
headers={"Authorization": f"Bearer {token}"},
content=b"archive",
)
assert response.status_code == 409
assert response.json() == {"detail": "object exists"}
assert store.writer.close_calls == 1
@pytest.mark.anyio
async def test_home_snapshot_gateway_propagates_writer_close_cancellation() -> None:
store = _Store(writer=_Archive(close_error=asyncio.CancelledError()))
codec = HomeSnapshotTransferTokenCodec.from_server_secret(_secret())
gateway = HomeSnapshotGatewayService(
token_codec=codec,
archive_store=cast(HomeArchiveStore, store),
)
token = codec.encode_token(scope=HOME_SNAPSHOT_SCOPE_WRITE, snapshot_ref=_SNAPSHOT_REF)
async def chunks() -> AsyncIterator[bytes]:
yield b"archive"
with pytest.raises(asyncio.CancelledError):
await gateway.upload(authorization=f"Bearer {token}", chunks=chunks())
assert store.writer.close_calls == 1
@pytest.mark.parametrize(
"request_error",
(RuntimeError("request stream failed"), asyncio.CancelledError()),
ids=("failure", "cancellation"),
)
@pytest.mark.anyio
async def test_home_snapshot_gateway_closes_writer_when_request_stream_stops(
request_error: BaseException,
) -> None:
store = _Store()
codec = HomeSnapshotTransferTokenCodec.from_server_secret(_secret())
gateway = HomeSnapshotGatewayService(
token_codec=codec,
archive_store=cast(HomeArchiveStore, store),
)
token = codec.encode_token(scope=HOME_SNAPSHOT_SCOPE_WRITE, snapshot_ref=_SNAPSHOT_REF)
async def chunks() -> AsyncIterator[bytes]:
yield b"partial"
raise request_error
with pytest.raises(type(request_error)) as exc_info:
await gateway.upload(authorization=f"Bearer {token}", chunks=chunks())
assert exc_info.value is request_error
assert store.writer.writes == [b"partial"]
assert store.writer.close_calls == 1
@pytest.mark.anyio
async def test_home_snapshot_gateway_closes_reader_after_midstream_failure() -> None:
read_error = RuntimeError("read failed")
store = _Store(reader=_Archive(reads=[b"partial"], read_error=read_error))
codec = HomeSnapshotTransferTokenCodec.from_server_secret(_secret())
gateway = HomeSnapshotGatewayService(
token_codec=codec,
archive_store=cast(HomeArchiveStore, store),
)
token = codec.encode_token(scope=HOME_SNAPSHOT_SCOPE_READ, snapshot_ref=_SNAPSHOT_REF)
stream = await gateway.download(authorization=f"Bearer {token}")
assert await anext(stream) == b"partial"
with pytest.raises(RuntimeError, match="read failed") as exc_info:
_ = await anext(stream)
assert exc_info.value is read_error
assert store.reader.close_calls == 1
@pytest.mark.anyio
async def test_home_snapshot_gateway_closes_reader_when_consumer_is_cancelled() -> None:
read_started = asyncio.Event()
read_release = asyncio.Event()
store = _Store(
reader=_Archive(
reads=[b"unreachable"],
read_started=read_started,
read_release=read_release,
)
)
codec = HomeSnapshotTransferTokenCodec.from_server_secret(_secret())
gateway = HomeSnapshotGatewayService(
token_codec=codec,
archive_store=cast(HomeArchiveStore, store),
)
token = codec.encode_token(scope=HOME_SNAPSHOT_SCOPE_READ, snapshot_ref=_SNAPSHOT_REF)
stream = await gateway.download(authorization=f"Bearer {token}")
async def consume_one() -> bytes:
return await anext(stream)
consume = asyncio.create_task(consume_one())
try:
await asyncio.wait_for(read_started.wait(), timeout=1.0)
consume.cancel()
done, _ = await asyncio.wait({consume}, timeout=1.0)
assert consume in done, "download stream ignored consumer cancellation"
with pytest.raises(asyncio.CancelledError):
_ = consume.result()
finally:
read_release.set()
consume.cancel()
done, _ = await asyncio.wait({consume}, timeout=1.0)
if consume not in done:
pytest.fail("download consumer task did not finish during bounded cleanup")
if not consume.cancelled():
_ = consume.exception()
assert store.reader.close_calls == 1
def test_home_snapshot_gateway_rejects_ordinary_agent_stub_token() -> None:
execution_context = DifyExecutionContextLayerConfig(
tenant_id="tenant-1",
user_id="user-1",
user_from="account",
agent_mode="workflow_run",
invoke_from="service-api",
)
ordinary_token = AgentStubTokenCodec.from_server_secret(_secret()).encode_connection_token(execution_context)
client, _ = _client(_Store())
response = client.get(
"/agent-stub/home-snapshots/archive",
headers={"Authorization": f"Bearer {ordinary_token}"},
)
assert response.status_code == 401
def test_ordinary_agent_stub_route_rejects_home_snapshot_transfer_token() -> None:
transfer_codec = HomeSnapshotTransferTokenCodec.from_server_secret(_secret())
transfer_token = transfer_codec.encode_token(
scope=HOME_SNAPSHOT_SCOPE_READ,
snapshot_ref=_SNAPSHOT_REF,
)
app = FastAPI()
app.include_router(
create_agent_stub_router(
token_codec=AgentStubTokenCodec.from_server_secret(_secret()),
home_snapshot_gateway=HomeSnapshotGatewayService(
token_codec=transfer_codec,
archive_store=cast(HomeArchiveStore, _Store()),
),
)
)
response = TestClient(app).post(
"/agent-stub/connections",
headers={"Authorization": f"Bearer {transfer_token}"},
json={"protocol_version": 1, "argv": []},
)
assert response.status_code == 401
def test_home_snapshot_transfer_token_rejects_expiration() -> None:
codec = HomeSnapshotTransferTokenCodec.from_server_secret(_secret())
token = codec.encode_token(
scope=HOME_SNAPSHOT_SCOPE_READ,
snapshot_ref=_SNAPSHOT_REF,
now=int(time.time()) - 601,
)
client, _ = _client(_Store())
response = client.get(
"/agent-stub/home-snapshots/archive",
headers={"Authorization": f"Bearer {token}"},
)
assert response.status_code == 401
def test_home_snapshot_routes_return_503_without_gateway() -> None:
app = FastAPI()
app.include_router(create_agent_stub_router(token_codec=None))
client = TestClient(app)
assert client.put("/agent-stub/home-snapshots/archive", content=b"bytes").status_code == 503
assert client.get("/agent-stub/home-snapshots/archive").status_code == 503
@@ -38,6 +38,7 @@ def test_create_agent_stub_router_mounts_all_agent_stub_routes() -> None:
assert "/agent-stub/connections" in paths
assert "/agent-stub/files/upload-request" in paths
assert "/agent-stub/files/download-request" in paths
assert "/agent-stub/home-snapshots/archive" in paths
def test_create_agent_stub_router_returns_503_for_unconfigured_services() -> None:
@@ -28,7 +28,6 @@ from dify_agent.protocol import (
CreateHomeSnapshotFromBindingRequest,
CreateRunRequest,
DestroyExecutionBindingRequest,
InitializeHomeSnapshotRequest,
RUN_EVENT_ADAPTER,
RunCancelledEvent,
RunEvent,
@@ -311,14 +310,6 @@ def test_async_workspace_methods_post_dtos_and_parse_responses() -> None:
asyncio.run(scenario())
def _initialize_home_snapshot_request() -> InitializeHomeSnapshotRequest:
return InitializeHomeSnapshotRequest(
tenant_id="tenant-1",
agent_id="agent-1",
home_snapshot_id="home-1",
)
def test_sync_execution_binding_client_uses_private_binding_routes() -> None:
def handler(request: httpx.Request) -> httpx.Response:
payload = cast(dict[str, object], json.loads(request.content))
@@ -368,12 +359,9 @@ def _create_home_snapshot_from_binding_request() -> CreateHomeSnapshotFromBindin
)
def test_sync_home_snapshot_client_parses_initialize_checkpoint_and_delete() -> None:
def test_sync_home_snapshot_client_parses_checkpoint_and_delete() -> None:
def handler(request: httpx.Request) -> httpx.Response:
if request.method == "POST":
if request.url.path == "/home-snapshots/initialize":
assert json.loads(request.content) == _initialize_home_snapshot_request().model_dump(mode="json")
return httpx.Response(201, json={"snapshot_ref": "initial-home"})
if request.url.path == "/home-snapshots/from-binding":
assert json.loads(request.content) == _create_home_snapshot_from_binding_request().model_dump(
mode="json"
@@ -387,20 +375,16 @@ def test_sync_home_snapshot_client_parses_initialize_checkpoint_and_delete() ->
http_client = httpx.Client(transport=httpx.MockTransport(handler))
client = Client(base_url="http://testserver", sync_http_client=http_client)
initialized = client.initialize_home_snapshot_sync(_initialize_home_snapshot_request())
created = client.create_home_snapshot_from_binding_sync(_create_home_snapshot_from_binding_request())
client.delete_home_snapshot_sync(created.snapshot_ref)
assert initialized.snapshot_ref == "initial-home"
assert created.snapshot_ref == "team/home 1"
http_client.close()
def test_async_home_snapshot_client_parses_initialize_checkpoint_and_delete() -> None:
def test_async_home_snapshot_client_parses_checkpoint_and_delete() -> None:
def handler(request: httpx.Request) -> httpx.Response:
if request.method == "POST":
if request.url.path == "/home-snapshots/initialize":
return httpx.Response(201, json={"snapshot_ref": "initial-home"})
if request.url.path == "/home-snapshots/from-binding":
return httpx.Response(201, json={"snapshot_ref": "team/home 1"})
assert request.url.path == "/home-snapshots/delete"
@@ -411,11 +395,9 @@ def test_async_home_snapshot_client_parses_initialize_checkpoint_and_delete() ->
http_client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
client = Client(base_url="http://testserver", async_http_client=http_client)
initialized = await client.initialize_home_snapshot(_initialize_home_snapshot_request())
created = await client.create_home_snapshot_from_binding(_create_home_snapshot_from_binding_request())
await client.delete_home_snapshot(created.snapshot_ref)
assert initialized.snapshot_ref == "initial-home"
assert created.snapshot_ref == "team/home 1"
await http_client.aclose()
@@ -430,7 +412,7 @@ def test_home_snapshot_client_maps_sync_validation_and_async_http_errors() -> No
)
with pytest.raises(DifyAgentValidationError):
_ = sync_client.initialize_home_snapshot_sync(_initialize_home_snapshot_request())
_ = sync_client.create_home_snapshot_from_binding_sync(_create_home_snapshot_from_binding_request())
sync_http_client.close()
async def scenario() -> None:
@@ -1,6 +1,7 @@
import asyncio
import json
from types import SimpleNamespace
from typing import Any, cast
import httpx
import pytest
@@ -824,9 +825,11 @@ def test_dify_plugin_tools_layer_rejects_path_without_shell() -> None:
def test_plugin_tool_file_context_uploads_sandbox_path_and_resolves_signed_url() -> None:
class FakeShell:
script: str | None = None
kwargs: dict[str, object] | None = None
async def run_remote_script_complete(self, script: str, **_kwargs: object) -> object:
async def run_remote_script_complete(self, script: str, **kwargs: object) -> object:
self.script = script
self.kwargs = kwargs
return SimpleNamespace(
exit_code=0,
status="done",
@@ -851,9 +854,9 @@ def test_plugin_tool_file_context_uploads_sandbox_path_and_resolves_signed_url()
async def scenario() -> None:
shell = FakeShell()
context = _PluginToolFileContext(
file_client=FakeFileClient(), # type: ignore[arg-type]
file_client=cast(Any, FakeFileClient()),
execution_context=_execution_context_config(),
shell=shell, # type: ignore[arg-type]
shell=cast(Any, shell),
)
result = await context.to_plugin_file_parameter("outputs/report.pdf")
@@ -868,6 +871,7 @@ def test_plugin_tool_file_context_uploads_sandbox_path_and_resolves_signed_url()
}
assert shell.script is not None
assert "outputs/report.pdf" in shell.script
assert shell.kwargs == {"timeout": 60.0, "inject_agent_stub_env": True}
asyncio.run(scenario())
@@ -93,7 +93,8 @@ def _patch_file_help(monkeypatch: pytest.MonkeyPatch) -> list[str]:
timeout: float = 10.0,
inject_agent_stub_env: bool = False,
) -> CompleteRemoteCommandResult:
del self, timeout, inject_agent_stub_env
del self, inject_agent_stub_env
assert timeout == 10.0
captured_scripts.append(script)
return _remote_result(_file_help_output(script))
@@ -9,6 +9,7 @@ from typing import cast
import pytest
import dify_agent.layers.shell.layer as shell_layer_module
import dify_agent.runtime_backend.shellctl as shellctl_backend_module
from dify_agent.layers.shell import (
DIFY_SHELL_LAYER_TYPE_ID,
DifyShellCliToolConfig,
@@ -322,6 +323,7 @@ def test_shell_layer_create_bootstraps_inside_sandbox_workspace() -> None:
def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult:
assert env == {"HOME": expected_home}
assert cwd == expected_workspace_cwd
assert timeout == pytest.approx(30.0, abs=0.01)
assert "apt-get install -y ripgrep" in script
return _command_result("bootstrap-job", status="exited", done=True, exit_code=0)
@@ -959,7 +961,6 @@ def test_run_remote_script_complete_uses_agent_specific_home_and_workspace_cwd()
expected_workspace_cwd = "/home/agent-1/workspace/abc12ff"
def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult:
del timeout
assert script == "pwd"
assert cwd == expected_workspace_cwd
assert env == {"HOME": expected_home}
@@ -1025,10 +1026,10 @@ def test_run_remote_script_uses_agent_specific_home_and_workspace_cwd() -> None:
expected_workspace_cwd = "/home/agent-1/workspace/abc12ff"
def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult:
del timeout
assert script == "pwd"
assert cwd == expected_workspace_cwd
assert env == {"HOME": expected_home}
assert timeout == pytest.approx(60.0, abs=0.01)
return _command_result(
"remote-job",
status="exited",
@@ -1098,7 +1099,7 @@ def test_run_remote_script_complete_returns_incomplete_reason_for_timeout(
def fake_monotonic() -> float:
return now
monkeypatch.setattr(shell_layer_module.time, "monotonic", fake_monotonic)
monkeypatch.setattr(shellctl_backend_module.time, "monotonic", fake_monotonic)
def run_handler(script: str, cwd: str | None, env: Mapping[str, str] | None, timeout: float) -> ShellCommandResult:
nonlocal now
@@ -29,6 +29,29 @@ def test_execution_binding_request_uses_opaque_backend_refs() -> None:
}
def test_execution_binding_request_accepts_missing_or_null_home_snapshot_ref() -> None:
fields = {
"tenant_id": "tenant-1",
"agent_id": "agent-1",
"binding_id": "binding-1",
"workspace_id": "workspace-1",
}
assert CreateExecutionBindingRequest(**fields).home_snapshot_ref is None
assert CreateExecutionBindingRequest(**fields, home_snapshot_ref=None).home_snapshot_ref is None
def test_execution_binding_request_rejects_empty_home_snapshot_ref() -> None:
with pytest.raises(ValidationError, match="home_snapshot_ref"):
CreateExecutionBindingRequest(
tenant_id="tenant-1",
agent_id="agent-1",
binding_id="binding-1",
workspace_id="workspace-1",
home_snapshot_ref="",
)
def test_destroy_workspace_requires_workspace_ref() -> None:
with pytest.raises(ValidationError, match="workspace_ref"):
DestroyExecutionBindingRequest(binding_ref="binding-1", destroy_workspace=True)
@@ -1,6 +1,7 @@
from __future__ import annotations
from dataclasses import dataclass, field
from enum import Enum
from typing import cast
import pytest
@@ -10,7 +11,6 @@ from dify_agent.runtime_backend import (
ExecutionBindingCreateSpec,
ExecutionBindingDestroySpec,
HomeSnapshotCreateSpec,
InitializeHomeSnapshotSpec,
SharedWorkspaceUnsupportedError,
WorkspacePreservationUnsupportedError,
)
@@ -22,9 +22,21 @@ from dify_agent.runtime_backend.e2b import (
from dify_agent.runtime_backend.shellctl import ShellctlRuntimeLease
class _FileType(Enum):
FILE = "file"
DIR = "dir"
@dataclass(frozen=True, slots=True)
class _FileInfo:
type: _FileType | None
symlink_target: str | None = None
@dataclass(slots=True)
class _Files:
paths: set[str] = field(default_factory=set)
contents: dict[str, str | bytes] = field(default_factory=dict)
async def make_dir(self, path: str) -> bool:
self.paths.add(path)
@@ -33,8 +45,23 @@ class _Files:
async def exists(self, path: str) -> bool:
return path in self.paths
async def get_info(self, path: str) -> _FileInfo:
if path not in self.paths:
raise FileNotFoundError(path)
return _FileInfo(type=_FileType.DIR)
async def remove(self, path: str) -> None:
self.paths.discard(path)
self.contents.pop(path, None)
async def read(self, path: str) -> str:
value = self.contents[path]
return value.decode() if isinstance(value, bytes) else value
async def write(self, path: str, data: str | bytes) -> object:
self.paths.add(path)
self.contents[path] = data
return object()
@dataclass(slots=True)
@@ -86,7 +113,7 @@ class _ControlPlane:
sandbox = _Sandbox(sandbox_id=sandbox_id, pause_error=self.pause_error)
self.sandboxes[sandbox_id] = sandbox
self.created.append((template, on_timeout))
assert metadata["dify.resource"] in {"home-snapshot-initialize", "runtime-sandbox"}
assert metadata["dify.resource"] == "runtime-sandbox"
return sandbox
async def connect(self, handle: str, *, timeout: int) -> _Sandbox:
@@ -103,49 +130,57 @@ class _ControlPlane:
@pytest.mark.anyio
async def test_e2b_profile_uses_snapshot_as_runtime_template_and_couples_refs() -> None:
async def test_e2b_binding_uses_default_template_or_exact_snapshot_and_couples_refs() -> None:
control = _ControlPlane()
snapshots = E2BHomeSnapshotBackend(
control_plane=control, # pyright: ignore[reportArgumentType]
)
bindings = E2BExecutionBindingBackend(
control_plane=control, # pyright: ignore[reportArgumentType]
template="prepared-template",
active_timeout_seconds=3600,
)
bindings = E2BExecutionBindingBackend(
control_plane=control, # pyright: ignore[reportArgumentType]
active_timeout_seconds=3600,
)
snapshot_ref = await snapshots.initialize(
InitializeHomeSnapshotSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-1")
)
allocation = await bindings.create_binding(
default_allocation = await bindings.create_binding(
ExecutionBindingCreateSpec(
tenant_id="tenant-1",
agent_id="agent-1",
binding_id="binding-1",
workspace_id="workspace-1",
existing_workspace_ref=None,
home_snapshot_ref=snapshot_ref,
home_snapshot_ref=None,
)
)
snapshot_allocation = await bindings.create_binding(
ExecutionBindingCreateSpec(
tenant_id="tenant-1",
agent_id="agent-1",
binding_id="binding-2",
workspace_id="workspace-2",
existing_workspace_ref=None,
home_snapshot_ref="snapshot-1",
)
)
assert control.created == [("prepared-template", "kill"), (snapshot_ref, "pause")]
assert allocation.binding_ref == allocation.workspace_ref
runtime = control.sandboxes[allocation.binding_ref]
assert control.created == [("prepared-template", "pause"), ("snapshot-1", "pause")]
assert default_allocation.binding_ref == default_allocation.workspace_ref
assert snapshot_allocation.binding_ref == snapshot_allocation.workspace_ref
runtime = control.sandboxes[default_allocation.binding_ref]
assert runtime.files.paths == {"/home/dify/workspace"}
assert runtime.pauses == [True]
await bindings.destroy_binding(
ExecutionBindingDestroySpec(
binding_ref=allocation.binding_ref,
workspace_ref=allocation.workspace_ref,
destroy_workspace=True,
for allocation in (default_allocation, snapshot_allocation):
await bindings.destroy_binding(
ExecutionBindingDestroySpec(
binding_ref=allocation.binding_ref,
workspace_ref=allocation.workspace_ref,
destroy_workspace=True,
)
)
)
await snapshots.delete(snapshot_ref)
await snapshots.delete("snapshot-1")
assert control.killed == [allocation.binding_ref]
assert control.deleted_snapshots == [snapshot_ref]
assert control.killed == [default_allocation.binding_ref, snapshot_allocation.binding_ref]
assert control.deleted_snapshots == ["snapshot-1"]
@pytest.mark.anyio
@@ -153,6 +188,7 @@ async def test_e2b_rejects_shared_workspace_and_binding_only_destroy() -> None:
control = _ControlPlane()
backend = E2BExecutionBindingBackend(
control_plane=control, # pyright: ignore[reportArgumentType]
template="prepared-template",
active_timeout_seconds=3600,
)
spec = ExecutionBindingCreateSpec(
@@ -177,6 +213,7 @@ async def test_e2b_binding_create_kills_sandbox_when_initialization_fails() -> N
control = _ControlPlane(pause_error=RuntimeError("pause failed"))
backend = E2BExecutionBindingBackend(
control_plane=control, # pyright: ignore[reportArgumentType]
template="prepared-template",
active_timeout_seconds=3600,
)
@@ -196,6 +233,36 @@ async def test_e2b_binding_create_kills_sandbox_when_initialization_fails() -> N
assert sandbox.killed == 1
@pytest.mark.anyio
async def test_e2b_missing_explicit_snapshot_does_not_fall_back_to_template() -> None:
class _FailingControlPlane(_ControlPlane):
async def create(self, template: str, *, timeout: int, metadata: dict[str, str], on_timeout: str) -> _Sandbox:
del timeout, metadata
self.created.append((template, on_timeout))
raise RuntimeError("snapshot unavailable")
control = _FailingControlPlane()
backend = E2BExecutionBindingBackend(
control_plane=control, # pyright: ignore[reportArgumentType]
template="prepared-template",
active_timeout_seconds=3600,
)
with pytest.raises(BindingCreateError, match="snapshot unavailable"):
await backend.create_binding(
ExecutionBindingCreateSpec(
tenant_id="tenant-1",
agent_id="agent-1",
binding_id="binding-1",
workspace_id="workspace-1",
existing_workspace_ref=None,
home_snapshot_ref="missing-snapshot",
)
)
assert control.created == [("missing-snapshot", "pause")]
@pytest.mark.anyio
async def test_e2b_checkpoint_uses_exact_source_runtime() -> None:
control = _ControlPlane()
@@ -206,8 +273,6 @@ async def test_e2b_checkpoint_uses_exact_source_runtime() -> None:
)
backend = E2BHomeSnapshotBackend(
control_plane=control, # pyright: ignore[reportArgumentType]
template="prepared-template",
active_timeout_seconds=3600,
)
snapshot_ref = await backend.create_from_runtime(
File diff suppressed because it is too large Load Diff
@@ -9,13 +9,14 @@ import pytest
from dify_agent.runtime_backend import (
BindingAcquireError,
BindingCreateError,
BindingDestroyError,
BindingLostError,
ExecutionBindingCreateSpec,
ExecutionBindingDestroySpec,
HomeSnapshotCreateSpec,
InitializeHomeSnapshotSpec,
RuntimeLease,
SharedWorkspaceUnsupportedError,
WorkspacePreservationUnsupportedError,
)
from dify_agent.runtime_backend.enterprise import (
@@ -282,23 +283,55 @@ async def test_enterprise_destroy_propagates_gateway_failure(
@pytest.mark.anyio
async def test_enterprise_allocation_and_home_snapshots_remain_explicitly_not_implemented() -> None:
snapshots = EnterpriseHomeSnapshotBackend(gateway_endpoint="https://gateway", auth_token="secret")
bindings = EnterpriseExecutionBindingBackend(gateway_endpoint="https://gateway", auth_token="secret")
async def test_enterprise_default_binding_creates_gateway_sandbox_and_layout(
monkeypatch: pytest.MonkeyPatch,
) -> None:
requests: list[httpx.Request] = []
with pytest.raises(NotImplementedError, match="Execution Binding protocol"):
_ = await snapshots.initialize(
InitializeHomeSnapshotSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-1")
def handler(request: httpx.Request) -> httpx.Response:
requests.append(request)
if request.url.path == "/v1/sandboxes":
assert json.loads(request.content) == {"tenantId": "tenant-1"}
return httpx.Response(201, json={"sandboxId": "sandbox-1", "status": "running"})
if request.url.path == "/proxy/v1/jobs/run":
assert request.headers["X-Sandbox-Id"] == "sandbox-1"
payload = cast(dict[str, object], json.loads(request.content))
script = payload["script"]
assert isinstance(script, str)
assert "mkdir -p /home/dify" in script
assert "rm -rf -- /home/dify/workspace" in script
return _job_response()
return httpx.Response(200, json={"job_id": "job-1"})
clients = _mock_http(monkeypatch, handler)
backend = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret")
allocation = await backend.create_binding(
ExecutionBindingCreateSpec(
tenant_id="tenant-1",
agent_id="agent-1",
binding_id="binding-1",
workspace_id="workspace-1",
existing_workspace_ref=None,
home_snapshot_ref=None,
)
with pytest.raises(NotImplementedError, match="Execution Binding protocol"):
_ = await snapshots.create_from_runtime(
spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-2"),
source=cast(RuntimeLease, object()),
)
with pytest.raises(NotImplementedError, match="Execution Binding protocol"):
await snapshots.delete("snapshot-1")
with pytest.raises(NotImplementedError, match="Execution Binding protocol"):
_ = await bindings.create_binding(
)
assert allocation.binding_ref == allocation.workspace_ref == "sandbox-1"
assert requests[0].headers["X-Inner-Api-Key"] == "secret"
assert all(client.is_closed for client in clients)
@pytest.mark.anyio
async def test_enterprise_binding_rejects_snapshot_and_shared_workspace_before_gateway_call(
monkeypatch: pytest.MonkeyPatch,
) -> None:
requests: list[httpx.Request] = []
_ = _mock_http(monkeypatch, lambda request: requests.append(request) or httpx.Response(500))
backend = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret")
with pytest.raises(BindingCreateError, match="immutable Home Snapshot"):
await backend.create_binding(
ExecutionBindingCreateSpec(
tenant_id="tenant-1",
agent_id="agent-1",
@@ -308,3 +341,63 @@ async def test_enterprise_allocation_and_home_snapshots_remain_explicitly_not_im
home_snapshot_ref="snapshot-1",
)
)
with pytest.raises(SharedWorkspaceUnsupportedError):
await backend.create_binding(
ExecutionBindingCreateSpec(
tenant_id="tenant-1",
agent_id="agent-1",
binding_id="binding-1",
workspace_id="workspace-1",
existing_workspace_ref="workspace-1",
home_snapshot_ref=None,
)
)
assert requests == []
@pytest.mark.anyio
async def test_enterprise_binding_create_deletes_new_sandbox_when_layout_setup_fails(
monkeypatch: pytest.MonkeyPatch,
) -> None:
requests: list[httpx.Request] = []
def handler(request: httpx.Request) -> httpx.Response:
requests.append(request)
if request.url.path == "/v1/sandboxes":
return httpx.Response(201, json={"sandboxId": "sandbox-1"})
if request.url.path == "/proxy/v1/jobs/run":
return _job_response(exit_code=1)
if request.method == "DELETE":
return httpx.Response(204)
return httpx.Response(200, json={"job_id": "job-1"})
_ = _mock_http(monkeypatch, handler)
backend = EnterpriseExecutionBindingBackend(gateway_endpoint="http://gateway.example", auth_token="secret")
with pytest.raises(BindingCreateError):
await backend.create_binding(
ExecutionBindingCreateSpec(
tenant_id="tenant-1",
agent_id="agent-1",
binding_id="binding-1",
workspace_id="workspace-1",
existing_workspace_ref=None,
home_snapshot_ref=None,
)
)
assert any(request.method == "DELETE" and request.url.path == "/v1/sandboxes/sandbox-1" for request in requests)
@pytest.mark.anyio
async def test_enterprise_home_snapshots_remain_explicitly_not_implemented() -> None:
snapshots = EnterpriseHomeSnapshotBackend()
with pytest.raises(NotImplementedError, match="immutable Home Snapshot"):
_ = await snapshots.create_from_runtime(
spec=HomeSnapshotCreateSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-2"),
source=cast(RuntimeLease, object()),
)
with pytest.raises(NotImplementedError, match="immutable Home Snapshot"):
await snapshots.delete("snapshot-1")
@@ -13,8 +13,6 @@ from dify_agent.runtime_backend import (
ExecutionBindingCreateSpec,
ExecutionBindingDestroySpec,
HomeSnapshotCreateSpec,
HomeSnapshotCreateError,
InitializeHomeSnapshotSpec,
)
from dify_agent.runtime_backend.local import LocalExecutionBindingBackend, LocalHomeSnapshotBackend
@@ -137,42 +135,6 @@ class _FailThenSucceedFactory:
return tuple(command for run in self.runs for command in run.commands)
@pytest.mark.anyio
async def test_local_snapshot_initialize_creates_private_snapshot_directory() -> None:
factory = _Factory()
snapshots = LocalHomeSnapshotBackend(
endpoint="http://shellctl",
auth_token="",
snapshot_root="/snapshots",
client_factory=factory, # pyright: ignore[reportArgumentType]
)
snapshot_ref = await snapshots.initialize(
InitializeHomeSnapshotSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-1")
)
assert snapshot_ref == "home-home-1"
assert ("mkdir", "-p", "/snapshots/home-home-1") in factory.commands
assert ("chmod", "700", "/snapshots/home-home-1") in factory.commands
@pytest.mark.anyio
async def test_local_snapshot_create_failure_removes_partial_snapshot() -> None:
factory = _FailThenSucceedFactory()
snapshots = LocalHomeSnapshotBackend(
endpoint="http://shellctl",
auth_token="",
snapshot_root="/snapshots",
client_factory=factory, # pyright: ignore[reportArgumentType]
)
with pytest.raises(HomeSnapshotCreateError, match="primary shellctl failure"):
await snapshots.initialize(
InitializeHomeSnapshotSpec(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-1")
)
assert ("rm", "-rf", "--", "/snapshots/home-home-1") in factory.commands
@pytest.mark.anyio
async def test_local_binding_create_materializes_home_and_new_workspace() -> None:
factory = _Factory()
@@ -198,13 +160,41 @@ async def test_local_binding_create_materializes_home_and_new_workspace() -> Non
assert allocation.binding_ref == "binding-1:workspace-1"
assert allocation.workspace_ref == "workspace-1"
assert ("test", "-d", "/snapshots/home-home-1") in factory.commands
assert factory.commands[0] == ("test", "-d", "/snapshots/home-home-1")
assert ("mkdir", "-p", "/workspaces/workspace-1") in factory.commands
assert ("mkdir", "-p", "/homes/binding-1") in factory.commands
assert ("cp", "-a", "/snapshots/home-home-1/.", "/homes/binding-1/") in factory.commands
assert ("chmod", "700", "/homes/binding-1", "/workspaces/workspace-1") in factory.commands
@pytest.mark.anyio
async def test_local_binding_create_uses_empty_default_home_without_snapshot_access() -> None:
factory = _Factory()
backend = LocalExecutionBindingBackend(
endpoint="http://shellctl",
auth_token="",
materialized_home_root="/homes",
workspace_root="/workspaces",
snapshot_root="/snapshots",
client_factory=factory, # pyright: ignore[reportArgumentType]
)
allocation = await backend.create_binding(
ExecutionBindingCreateSpec(
tenant_id="tenant-1",
agent_id="agent-1",
binding_id="binding-1",
workspace_id="workspace-1",
existing_workspace_ref=None,
home_snapshot_ref=None,
)
)
assert allocation.binding_ref == "binding-1:workspace-1"
assert ("mkdir", "-p", "/homes/binding-1") in factory.commands
assert all("/snapshots" not in part for command in factory.commands for part in command)
@pytest.mark.anyio
async def test_local_binding_create_failure_removes_partial_home_and_workspace() -> None:
factory = _FailThenSucceedFactory()
@@ -1,5 +1,7 @@
from __future__ import annotations
import base64
import pytest
from pydantic import ValidationError
@@ -7,6 +9,10 @@ from dify_agent.runtime_backend.e2b import E2B_MAX_ACTIVE_TIMEOUT_SECONDS
from dify_agent.runtime_backend.profile import DEFAULT_E2B_TEMPLATE, RuntimeBackendSettings
def _secret() -> str:
return base64.urlsafe_b64encode(b"s" * 32).rstrip(b"=").decode("ascii")
def test_e2b_backend_uses_prepared_dify_template_by_default() -> None:
settings = RuntimeBackendSettings(runtime_backend="e2b", e2b_api_key="secret")
@@ -32,3 +38,47 @@ def test_e2b_backend_rejects_active_timeout_above_platform_limit() -> None:
def test_local_backend_requires_shellctl_endpoint() -> None:
with pytest.raises(ValidationError, match="local_sandbox_endpoint"):
_ = RuntimeBackendSettings(runtime_backend="local")
def test_e2b_s3_requires_uri_http_gateway_and_server_secret() -> None:
with pytest.raises(ValidationError, match="e2b_s3_uri"):
_ = RuntimeBackendSettings(
runtime_backend="e2b_s3",
e2b_api_key="secret",
agent_stub_api_base_url="https://agent.example/agent-stub",
server_secret_key=_secret(),
)
with pytest.raises(ValidationError, match="agent_stub_api_base_url"):
_ = RuntimeBackendSettings(
runtime_backend="e2b_s3",
e2b_api_key="secret",
e2b_s3_uri="s3://bucket/dify-agent",
server_secret_key=_secret(),
)
with pytest.raises(ValidationError, match="http\\(s\\)"):
_ = RuntimeBackendSettings(
runtime_backend="e2b_s3",
e2b_api_key="secret",
e2b_s3_uri="s3://bucket/dify-agent",
agent_stub_api_base_url="grpc://agent.example:9091",
server_secret_key=_secret(),
)
with pytest.raises(ValidationError, match="server_secret_key"):
_ = RuntimeBackendSettings(
runtime_backend="e2b_s3",
e2b_api_key="secret",
e2b_s3_uri="s3://bucket/dify-agent",
agent_stub_api_base_url="https://agent.example/agent-stub",
)
def test_e2b_s3_normalizes_uri() -> None:
settings = RuntimeBackendSettings(
runtime_backend="e2b_s3",
e2b_api_key="secret",
e2b_s3_uri=" s3://bucket/dify-agent?region=us-east-1 ",
agent_stub_api_base_url="https://agent.example/agent-stub",
server_secret_key=_secret(),
)
assert settings.e2b_s3_uri == "s3://bucket/dify-agent?region=us-east-1"
@@ -1,17 +1,20 @@
from __future__ import annotations
import asyncio
from dataclasses import dataclass, field
from typing import cast
import pytest
import dify_agent.runtime_backend.shellctl as shellctl_backend_module
from dify_agent.adapters.shell.protocols import ShellCommandResult, ShellCommandStatus
from dify_agent.adapters.shell.shellctl import ShellctlClientProtocol
from dify_agent.runtime_backend.protocols import RuntimeLayout
from dify_agent.runtime_backend.shellctl import (
CONTROL_COMMAND_OUTPUT_LIMIT,
create_owned_shellctl_lease,
create_shellctl_lease,
run_shellctl_control_command,
execute_complete_with_commands,
)
@@ -40,9 +43,14 @@ class _FakeTransport:
@dataclass(slots=True)
class _FakeCommands:
initial: ShellCommandResult
wait_error: Exception | None = None
wait_results: list[ShellCommandResult] = field(default_factory=list)
wait_error: BaseException | None = None
delete_error: Exception | None = None
delete_calls: list[tuple[str, bool]] = field(default_factory=list)
run_timeouts: list[float] = field(default_factory=list)
wait_timeouts: list[float] = field(default_factory=list)
wait_started: asyncio.Event | None = None
wait_release: asyncio.Event | None = None
async def run(
self,
@@ -52,13 +60,21 @@ class _FakeCommands:
env: dict[str, str] | None = None,
timeout: float,
) -> ShellCommandResult:
del script, cwd, env, timeout
del script, cwd, env
self.run_timeouts.append(timeout)
return self.initial
async def wait(self, job_id: str, *, offset: int, timeout: float) -> ShellCommandResult:
del job_id, offset, timeout
del job_id, offset
self.wait_timeouts.append(timeout)
if self.wait_started is not None:
self.wait_started.set()
assert self.wait_release is not None
await asyncio.wait_for(self.wait_release.wait(), timeout=1.0)
if self.wait_error is not None:
raise self.wait_error
if self.wait_results:
return self.wait_results.pop(0)
raise AssertionError("wait was not expected")
async def read_output(self, job_id: str, *, offset: int) -> ShellCommandResult:
@@ -166,7 +182,14 @@ async def test_control_command_success_is_preserved_when_delete_fails(
commands = _FakeCommands(initial=_result(), delete_error=RuntimeError("delete failed"))
with caplog.at_level("WARNING", logger="dify_agent.runtime_backend.shellctl"):
result = await run_shellctl_control_command(commands, "true")
result = await execute_complete_with_commands(
commands,
"true",
cwd=None,
env=None,
timeout=30.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
assert result.output == "ok"
assert commands.delete_calls == [("job-1", True)]
@@ -185,7 +208,100 @@ async def test_control_command_error_is_preserved_when_delete_also_fails(
with caplog.at_level("WARNING", logger="dify_agent.runtime_backend.shellctl"):
with pytest.raises(RuntimeError, match="command failed"):
_ = await run_shellctl_control_command(commands, "false")
_ = await execute_complete_with_commands(
commands,
"false",
cwd=None,
env=None,
timeout=30.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
assert commands.delete_calls == [("job-1", True)]
assert "delete failed" in caplog.text
@pytest.mark.anyio
async def test_control_command_wait_cancellation_is_propagated_and_job_is_deleted() -> None:
wait_started = asyncio.Event()
wait_release = asyncio.Event()
commands = _FakeCommands(
initial=_result(done=False),
wait_started=wait_started,
wait_release=wait_release,
)
task = asyncio.create_task(
execute_complete_with_commands(
commands,
"cancelled-control-job",
cwd=None,
env=None,
timeout=30.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
)
try:
await asyncio.wait_for(wait_started.wait(), timeout=1.0)
task.cancel()
done, _ = await asyncio.wait({task}, timeout=1.0)
assert task in done, "control command ignored cancellation"
with pytest.raises(asyncio.CancelledError):
_ = task.result()
finally:
wait_release.set()
task.cancel()
done, _ = await asyncio.wait({task}, timeout=1.0)
if task not in done:
pytest.fail("control command task did not finish during bounded cleanup")
if not task.cancelled():
_ = task.exception()
assert commands.delete_calls == [("job-1", True)]
@pytest.mark.anyio
async def test_complete_executor_without_deadline_repeats_fixed_long_polls() -> None:
commands = _FakeCommands(
initial=_result(done=False),
wait_results=[_result(done=False), _result(done=True)],
)
result = await execute_complete_with_commands(
commands,
"long-running-control-job",
cwd="/home/dify",
env={"TOKEN": "opaque"},
timeout=None,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
assert result.done is True
assert result.output == "okokok"
assert commands.run_timeouts == [30.0]
assert commands.wait_timeouts == [30.0, 30.0]
assert commands.delete_calls == [("job-1", True)]
@pytest.mark.anyio
async def test_complete_executor_with_deadline_preserves_full_remaining_timeout(
monkeypatch: pytest.MonkeyPatch,
) -> None:
now = 100.0
def monotonic() -> float:
return now
monkeypatch.setattr(shellctl_backend_module.time, "monotonic", monotonic)
commands = _FakeCommands(initial=_result(done=True))
_ = await execute_complete_with_commands(
commands,
"bounded-job",
cwd=None,
env=None,
timeout=60.0,
max_output_bytes=CONTROL_COMMAND_OUTPUT_LIMIT,
)
assert commands.run_timeouts == [60.0]
@@ -6,6 +6,7 @@ from typing import ClassVar
import httpx
import pytest
from fastapi import APIRouter
from fastapi.testclient import TestClient
import dify_agent.server.app as app_module
@@ -230,7 +231,9 @@ def test_create_app_creates_scheduler_and_closes_after_shutdown(monkeypatch: pyt
assert execution_context_layer.daemon_api_key == "daemon-secret"
assert shell_layer.agent_stub_token_factory is not None
token = shell_layer.agent_stub_token_factory(_execution_context(), session_id="abc12ff")
decoded = settings.create_agent_stub_token_codec().decode_token(token)
token_codec = settings.create_agent_stub_token_codec()
assert token_codec is not None
decoded = token_codec.decode_token(token)
assert decoded.execution_context == _execution_context()
assert decoded.session_id == "abc12ff"
knowledge_provider = next(provider for provider in layer_providers if provider.type_id == "dify.knowledge_base")
@@ -310,6 +313,33 @@ def test_create_app_wires_authenticated_agent_stub_connection_route(monkeypatch:
assert fake_redis.closed is True
def test_create_app_injects_e2b_s3_home_snapshot_gateway(monkeypatch: pytest.MonkeyPatch) -> None:
gateway = object()
captured: list[object] = []
def build_gateway(_self: ServerSettings) -> object:
return gateway
def capture_router(**kwargs: object) -> APIRouter:
captured.append(kwargs["home_snapshot_gateway"])
return APIRouter()
monkeypatch.setattr(ServerSettings, "build_home_snapshot_gateway", build_gateway)
monkeypatch.setattr(ServerSettings, "build_runtime_backend_profile", lambda _self: None)
monkeypatch.setattr(app_module, "create_agent_stub_router", capture_router)
settings = ServerSettings(
runtime_backend="e2b_s3",
e2b_api_key="e2b-secret",
e2b_s3_uri="s3://snapshots/dify-agent",
agent_stub_api_base_url="https://agent.example.com/agent-stub",
server_secret_key=_base64url_secret(b"1" * 32),
)
_ = create_app(settings)
assert captured == [gateway]
def test_create_app_wires_authenticated_agent_stub_file_upload_route(monkeypatch: pytest.MonkeyPatch) -> None:
fake_redis, fake_http_client = _patch_app_lifecycle(monkeypatch)
settings = ServerSettings(
@@ -8,22 +8,16 @@ import pytest
from dify_agent.protocol import (
CreateHomeSnapshotFromBindingRequest,
DeleteHomeSnapshotRequest,
InitializeHomeSnapshotRequest,
)
from dify_agent.runtime_backend import HomeSnapshotCreateSpec, InitializeHomeSnapshotSpec, RuntimeLease
from dify_agent.runtime_backend import HomeSnapshotCreateSpec, RuntimeLease
from dify_agent.server.home_snapshots import HomeSnapshotService
@dataclass(slots=True)
class _HomeBackend:
initialized: list[InitializeHomeSnapshotSpec] = field(default_factory=list)
checkpointed: list[tuple[HomeSnapshotCreateSpec, RuntimeLease]] = field(default_factory=list)
deleted: list[str] = field(default_factory=list)
async def initialize(self, spec: InitializeHomeSnapshotSpec) -> str:
self.initialized.append(spec)
return "snapshot-initial"
async def create_from_runtime(self, *, spec: HomeSnapshotCreateSpec, source: RuntimeLease) -> str:
self.checkpointed.append((spec, source))
return "snapshot-build"
@@ -47,7 +41,7 @@ class _BindingBackend:
@pytest.mark.anyio
async def test_home_snapshot_service_initializes_and_checkpoints_exact_binding() -> None:
async def test_home_snapshot_service_checkpoints_exact_binding() -> None:
lease = cast(RuntimeLease, object())
homes = _HomeBackend()
bindings = _BindingBackend(lease=lease)
@@ -56,9 +50,6 @@ async def test_home_snapshot_service_initializes_and_checkpoints_exact_binding()
execution_bindings=bindings, # pyright: ignore[reportArgumentType]
)
initial = await service.initialize(
InitializeHomeSnapshotRequest(tenant_id="tenant-1", agent_id="agent-1", home_snapshot_id="home-1")
)
checkpoint = await service.create_from_binding(
CreateHomeSnapshotFromBindingRequest(
tenant_id="tenant-1",
@@ -68,7 +59,6 @@ async def test_home_snapshot_service_initializes_and_checkpoints_exact_binding()
)
)
assert initial.snapshot_ref == "snapshot-initial"
assert checkpoint.snapshot_ref == "snapshot-build"
assert bindings.acquired == ["binding-ref"]
assert bindings.released == [lease]
@@ -13,6 +13,11 @@ from dify_agent.agent_stub.server.agent_stub_files import DifyApiAgentStubFileRe
from dify_agent.agent_stub.server.tokens.agent_stub import AgentStubTokenCodec
from dify_agent.server.settings import ServerSettings
from dify_agent.runtime_backend.e2b import E2BExecutionBindingBackend
from dify_agent.runtime_backend.e2b_s3 import (
E2BS3ExecutionBindingBackend,
E2BS3HomeSnapshotBackend,
OpenDALHomeArchiveStore,
)
from dify_agent.runtime_backend.enterprise import EnterpriseExecutionBindingBackend
from dify_agent.runtime_backend.local import LocalExecutionBindingBackend
@@ -301,6 +306,68 @@ def test_build_runtime_backend_profile_passes_e2b_active_timeout() -> None:
assert profile is not None
assert isinstance(profile.execution_bindings, E2BExecutionBindingBackend)
assert profile.execution_bindings.active_timeout_seconds == 900
assert profile.execution_bindings.template == "difys-default-team/dify-agent-local-sandbox"
def test_build_runtime_backend_profile_and_gateway_for_e2b_s3(monkeypatch: pytest.MonkeyPatch) -> None:
store_calls: list[str] = []
stores: list[OpenDALHomeArchiveStore] = []
def create_store_from_uri(
cls: type[OpenDALHomeArchiveStore],
uri: str,
) -> OpenDALHomeArchiveStore:
del cls
store_calls.append(uri)
store = cast(OpenDALHomeArchiveStore, object())
stores.append(store)
return store
monkeypatch.setattr(OpenDALHomeArchiveStore, "create_from_uri", classmethod(create_store_from_uri))
storage_uri = "s3://snapshots/tenant-root?region=us-east-1"
settings = ServerSettings(
runtime_backend="e2b_s3",
e2b_api_key="e2b-secret",
e2b_s3_uri=storage_uri,
agent_stub_api_base_url="https://agent.example/agent-stub",
server_secret_key=_base64url_secret(b"s" * 32),
)
profile = settings.build_runtime_backend_profile()
gateway = settings.build_home_snapshot_gateway()
assert profile is not None
assert isinstance(profile.execution_bindings, E2BS3ExecutionBindingBackend)
assert isinstance(profile.home_snapshots, E2BS3HomeSnapshotBackend)
assert profile.home_snapshots.control_plane is profile.execution_bindings.control_plane
assert profile.home_snapshots.lifecycle_cli is profile.execution_bindings.lifecycle_cli
assert gateway is not None
assert store_calls == [storage_uri, storage_uri]
assert profile.home_snapshots.archive_store is stores[0]
assert gateway.archive_store is stores[1]
assert gateway.token_codec is not profile.home_snapshots.lifecycle_cli.token_codec
def test_home_snapshot_gateway_is_not_built_for_other_profiles() -> None:
assert ServerSettings().build_home_snapshot_gateway() is None
def test_server_settings_rejects_e2b_s3_without_required_storage_or_http_settings() -> None:
with pytest.raises(ValidationError, match="DIFY_AGENT_E2B_S3_URI"):
_ = ServerSettings(
runtime_backend="e2b_s3",
e2b_api_key="e2b-secret",
agent_stub_api_base_url="https://agent.example/agent-stub",
server_secret_key=_base64url_secret(b"s" * 32),
)
with pytest.raises(ValidationError, match="HTTP\\(S\\)"):
_ = ServerSettings(
runtime_backend="e2b_s3",
e2b_api_key="e2b-secret",
e2b_s3_uri="s3://snapshots/dify-agent",
agent_stub_api_base_url="grpc://agent.example:9091",
server_secret_key=_base64url_secret(b"s" * 32),
)
def test_sandbox_file_upload_limit_defaults_to_tool_file_limit() -> None:
+1
View File
@@ -21,6 +21,7 @@ SERVER_RUNTIME_DEPENDENCIES = {
"jsonschema>=4.23.0,<5.0.0",
"jwcrypto>=1.5.6,<2",
"logfire[fastapi,httpx,redis]>=4.37.0,<5.0.0",
"opendal==0.47.5",
"pydantic-ai-slim[anthropic,google,openai]>=1.85.1,<2.0.0",
"pydantic-settings>=2.12.0,<3.0.0",
"redis>=7.4.0,<8.0.0",
+28
View File
@@ -612,6 +612,7 @@ server = [
{ name = "jsonschema" },
{ name = "jwcrypto" },
{ name = "logfire", extra = ["fastapi", "httpx", "redis"] },
{ name = "opendal" },
{ name = "pydantic-ai-slim", extra = ["anthropic", "google", "openai"] },
{ name = "pydantic-settings" },
{ name = "redis" },
@@ -646,6 +647,7 @@ requires-dist = [
{ name = "jsonschema", marker = "extra == 'server'", specifier = ">=4.23.0,<5.0.0" },
{ name = "jwcrypto", marker = "extra == 'server'", specifier = ">=1.5.6,<2" },
{ name = "logfire", extras = ["fastapi", "httpx", "redis"], marker = "extra == 'server'", specifier = ">=4.37.0,<5.0.0" },
{ name = "opendal", marker = "extra == 'server'", specifier = "==0.47.5", index = "https://test.pypi.org/simple" },
{ name = "protobuf", marker = "extra == 'grpc'", specifier = ">=6.33.5,<7.0.0" },
{ name = "pydantic", specifier = ">=2.12.5,<2.13" },
{ name = "pydantic-ai-slim", specifier = ">=1.102.0,<2.0.0" },
@@ -2048,6 +2050,32 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/1e/c1/d6e64ccd0536bf616556f0cad2b6d94a8125f508d25cfd814b1d2db4e2f1/openai-2.32.0-py3-none-any.whl", hash = "sha256:4dcc9badeb4bf54ad0d187453742f290226d30150890b7890711bda4f32f192f", size = 1162570, upload-time = "2026-04-15T22:28:17.714Z" },
]
[[package]]
name = "opendal"
version = "0.47.5"
source = { registry = "https://test.pypi.org/simple" }
sdist = { url = "https://test-files.pythonhosted.org/packages/48/65/44f2816c43e01a637ac8e77a2dde99c0f131b7617200af97c87637fd63fb/opendal-0.47.5.tar.gz", hash = "sha256:ca44d9eb5bf27d905e55d480d656dab1955b1528718c75dcb5d279fd19634b57", size = 1766937, upload-time = "2026-07-28T08:10:01.164Z" }
wheels = [
{ url = "https://test-files.pythonhosted.org/packages/17/91/b572904880ee8998031fdc7afbaf292e52bf20e523c6d13968e6b3baa009/opendal-0.47.5-cp311-abi3-macosx_10_12_x86_64.whl", hash = "sha256:57a453e32dee21d4d943a59c997b6a41ac658c94b2b40a656bd9f4613a61f8be", size = 17529259, upload-time = "2026-07-28T08:09:17.218Z" },
{ url = "https://test-files.pythonhosted.org/packages/31/e2/caa9c622d1c4604f5e6ee9a3aaf2cc39bf8e3889b2326528b959db0fc3ff/opendal-0.47.5-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:fb2efabdb98beb8eb8ce9bbdb33662f4430a95df9716320f90dd8608b826022b", size = 16118000, upload-time = "2026-07-28T08:09:19.679Z" },
{ url = "https://test-files.pythonhosted.org/packages/28/e3/463213e2fa8e98a8ca4699f2fe13159290ac80bf242bbe70f119d36473c3/opendal-0.47.5-cp311-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:342a51b8e5fae9f049492af72fa4f16aaf6d16177f31edd0ee5178b57067be64", size = 16803733, upload-time = "2026-07-28T08:09:22.393Z" },
{ url = "https://test-files.pythonhosted.org/packages/2e/4f/a2441d286c2534fb2f00d23f18b684b6bf327e74da8a3a4367ebf766dccf/opendal-0.47.5-cp311-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a5101c9b75ad3f4a24593def41f7ce723962b79cb3e601151e808f5d78f78e87", size = 17955349, upload-time = "2026-07-28T08:09:24.774Z" },
{ url = "https://test-files.pythonhosted.org/packages/d8/0e/a63cdeec6ba72270c10d93361f97dcc60105aa06c91b127328badb962aa7/opendal-0.47.5-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:3a403748c68b9587405ede6236d3618a8f5833ccede623be7386f24d4e5c31a8", size = 17026301, upload-time = "2026-07-28T08:09:27.847Z" },
{ url = "https://test-files.pythonhosted.org/packages/8a/5a/86d009a6bab0c02e636b5087b47bb9a7fd7b7b1785f2362d88c846ef9485/opendal-0.47.5-cp311-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:90703efb7f365639dc843d475f2b2477ec561824c0a68ad6a75c35f5c20e7751", size = 17375804, upload-time = "2026-07-28T08:09:30.801Z" },
{ url = "https://test-files.pythonhosted.org/packages/99/11/ff7365587b241aacba5537c78864790b60dc8bb660796bea1fa201a97ca0/opendal-0.47.5-cp311-abi3-musllinux_1_1_armv7l.whl", hash = "sha256:c28395331c92839f91a162ba2b92edf2e92f023ff13adb1bb46b234da8172eb8", size = 17102402, upload-time = "2026-07-28T08:09:33.085Z" },
{ url = "https://test-files.pythonhosted.org/packages/9b/3f/b24d881b970a41133be418876c79ba5b34c6e9811029b5d4fb9c04ccf203/opendal-0.47.5-cp311-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:ee27f1df07608a4ba30c9a5f09341850f7f6c8695d6559522302c2c410db161f", size = 18187917, upload-time = "2026-07-28T08:09:35.408Z" },
{ url = "https://test-files.pythonhosted.org/packages/90/b8/0b387e569b098e923977e36ec1b243c15dfe387eba9026984068bd87cd05/opendal-0.47.5-cp311-abi3-win_amd64.whl", hash = "sha256:1016ad511d8c61611719c9ffe509571ee216791bfcc07dc6aec47a756c30aa67", size = 19199295, upload-time = "2026-07-28T08:09:37.91Z" },
{ url = "https://test-files.pythonhosted.org/packages/3c/9f/63a687163b2a9ba66af8b4206ae4a7ae226e9c54355298cbef15cc54bfa5/opendal-0.47.5-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:012eeed30870754e7edc4895ee75306e32ff82b1d340d897b504ab8b7d02cfaf", size = 17536831, upload-time = "2026-07-28T08:09:41.003Z" },
{ url = "https://test-files.pythonhosted.org/packages/d5/d9/e1cb2dcb2c4d919f30b93ec6c3d777525e476074f6938b7aa9018f8305f9/opendal-0.47.5-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:99cadd469dcecb17866a20826e080e58c9dc0b4205340870603b58faae1e0f97", size = 16101302, upload-time = "2026-07-28T08:09:43.3Z" },
{ url = "https://test-files.pythonhosted.org/packages/d9/c0/d179754422478d879a732a81d9ea6d89607bbc645d6ec07f7036fe4d528d/opendal-0.47.5-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4968311d0f483311273192b90c79e0731eea32598ef253bb41de9c51bff70af8", size = 16797051, upload-time = "2026-07-28T08:09:45.311Z" },
{ url = "https://test-files.pythonhosted.org/packages/a0/2f/fcb7f079badbc7306d778f8b20a0ab83ff34d74bd4fe7e0c3a118dd5a05c/opendal-0.47.5-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47ecf76aa397cb06329d4e69e6c0afdfc34c5b6ae7c6da87efd33c006b38cde5", size = 17949096, upload-time = "2026-07-28T08:09:47.875Z" },
{ url = "https://test-files.pythonhosted.org/packages/c1/0c/bfdd05726f5c1c15b0b4e26e6337b8d8a1cd2a0aa7f5cdfa617b1041c141/opendal-0.47.5-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:b0a31a17f67ad59b7f3baf0c1cec7e60949ce3f7c7ebf5baba493015d8ac89e3", size = 17016235, upload-time = "2026-07-28T08:09:50.176Z" },
{ url = "https://test-files.pythonhosted.org/packages/2c/ef/210fc93d5173cd7e5634b31ebee87466845b90eb92483c50d8f76d240ba7/opendal-0.47.5-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:307f095d31056162a16ce02a44253f1b3de60cf638f837c5bfca5bbc06c53bdb", size = 17369912, upload-time = "2026-07-28T08:09:52.298Z" },
{ url = "https://test-files.pythonhosted.org/packages/31/8e/0a615f7344d85e07b39ba577ebe15211816611f032770d8c7ace6422338a/opendal-0.47.5-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:c526449e91b8442116acc77ec8f3dc7c87bbdd63f06acefc3e1ff880a06eb3a2", size = 17094062, upload-time = "2026-07-28T08:09:54.96Z" },
{ url = "https://test-files.pythonhosted.org/packages/f1/89/103514cebe8fc760741378772657eadfcf156412d72064344e19f53f10a3/opendal-0.47.5-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:b77948ec4ae48b3e9ce4fc6e6e196f82324916458305a01d81282ed520f303a0", size = 18180436, upload-time = "2026-07-28T08:09:57.267Z" },
{ url = "https://test-files.pythonhosted.org/packages/08/85/9e00b103f5c0bbc58e46c506d68b1179bafc6ab24104b62999dda3842082/opendal-0.47.5-cp314-cp314t-win_amd64.whl", hash = "sha256:d03f8dba373bcb1d080bf0b09165c9b6b1e546c30d0a1718bbe00d01e576688f", size = 19183036, upload-time = "2026-07-28T08:09:59.304Z" },
]
[[package]]
name = "openpyxl"
version = "3.1.5"
+1
View File
@@ -281,6 +281,7 @@ DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox
DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS=3600
DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN=
DIFY_AGENT_E2B_SHELLCTL_PORT=5004
DIFY_AGENT_E2B_S3_URI=
DIFY_AGENT_STUB_API_BASE_URL=http://agent_backend:5050/agent-stub
# This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens.
# Replace this development default in production.
+1
View File
@@ -676,6 +676,7 @@ services:
DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS: ${DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS:-3600}
DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN: ${DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN:-}
DIFY_AGENT_E2B_SHELLCTL_PORT: ${DIFY_AGENT_E2B_SHELLCTL_PORT:-5004}
DIFY_AGENT_E2B_S3_URI: ${DIFY_AGENT_E2B_S3_URI:-}
DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES: ${PLUGIN_MAX_FILE_SIZE:-52428800}
DIFY_AGENT_STUB_API_BASE_URL: ${DIFY_AGENT_STUB_API_BASE_URL:-http://agent_backend:5050/agent-stub}
# This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens.
+2 -1
View File
@@ -25,9 +25,10 @@ services:
context: ..
dockerfile: dify-agent/Dockerfile
environment:
DIFY_AGENT_RUNTIME_BACKEND: e2b
DIFY_AGENT_RUNTIME_BACKEND: ${DIFY_AGENT_RUNTIME_BACKEND:-e2b}
DIFY_AGENT_E2B_API_KEY: ${DIFY_AGENT_E2B_API_KEY:-${E2B_API_KEY:-${E2B_API_TOKEN:-}}}
DIFY_AGENT_E2B_TEMPLATE: ${DIFY_AGENT_E2B_TEMPLATE:-difys-default-team/dify-agent-local-sandbox}
DIFY_AGENT_E2B_S3_URI: ${DIFY_AGENT_E2B_S3_URI:-}
ports:
- "${EXPOSE_AGENT_BACKEND_PORT:-15050}:5050"
+1
View File
@@ -682,6 +682,7 @@ services:
DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS: ${DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS:-3600}
DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN: ${DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN:-}
DIFY_AGENT_E2B_SHELLCTL_PORT: ${DIFY_AGENT_E2B_SHELLCTL_PORT:-5004}
DIFY_AGENT_E2B_S3_URI: ${DIFY_AGENT_E2B_S3_URI:-}
DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES: ${PLUGIN_MAX_FILE_SIZE:-52428800}
DIFY_AGENT_STUB_API_BASE_URL: ${DIFY_AGENT_STUB_API_BASE_URL:-http://agent_backend:5050/agent-stub}
# This is security-sensitive: it derives the JWE encryption key for Agent Stub bearer tokens.
@@ -32,6 +32,7 @@ DIFY_AGENT_E2B_TEMPLATE=difys-default-team/dify-agent-local-sandbox
DIFY_AGENT_E2B_ACTIVE_TIMEOUT_SECONDS=3600
DIFY_AGENT_E2B_SHELLCTL_AUTH_TOKEN=
DIFY_AGENT_E2B_SHELLCTL_PORT=5004
DIFY_AGENT_E2B_S3_URI=
# Standalone/direct-container byte limit. Full Docker Compose derives this from
# PLUGIN_MAX_FILE_SIZE in docker/.env.
DIFY_AGENT_SANDBOX_FILE_UPLOAD_MAX_BYTES=52428800