From a02da08fc854769d86e3135b29726738f0b731c5 Mon Sep 17 00:00:00 2001 From: limityan Date: Thu, 23 Jul 2026 20:50:12 +0800 Subject: [PATCH] feat(sdk): add standalone Agent SDK Host baseline --- .github/workflows/ci.yml | 16 +- Cargo.toml | 2 + .../agent-runtime-services-design.md | 8 +- .../agent-sdk-product-architecture.md | 123 +- docs/architecture/cli-product-line-design.md | 2 + docs/architecture/product-architecture.md | 2 +- .../core-boundaries/rules/crate-layout.mjs | 1 + scripts/core-boundaries/rules/crate-rules.mjs | 31 + .../core-boundaries/rules/feature-rules.mjs | 12 + .../rules/source/required-rules.mjs | 39 +- scripts/core-boundaries/self-test.mjs | 27 +- src/apps/cli/src/agent/agentic_system.rs | 10 +- src/apps/cli/src/main.rs | 41 +- src/apps/cli/src/product_assembly.rs | 24 +- src/apps/cli/src/root_handlers.rs | 22 +- src/apps/cli/src/runtime/events.rs | 116 +- src/apps/cli/src/runtime/mod.rs | 44 +- src/apps/cli/src/runtime/services.rs | 246 -- src/apps/sdk-host/Cargo.toml | 31 + src/apps/sdk-host/src/lib.rs | 26 + src/apps/sdk-host/src/main.rs | 58 + src/apps/sdk-host/src/runtime.rs | 112 + src/apps/sdk-host/src/transport.rs | 364 +++ .../sdk-host/tests/process_initialization.rs | 30 + src/apps/sdk-host/tests/stdio_process.rs | 110 + src/apps/sdk-host/tests/stdio_transport.rs | 710 ++++++ .../agentic/agents/definitions/modes/claw.rs | 13 +- .../definitions/subagents/computer_use.rs | 12 + .../src/agentic/agents/prompts/claw_mode.md | 8 +- .../agents/prompts/computer_use_mode.md | 6 +- .../src/agentic/coordination/coordinator.rs | 688 ++++- .../src/agentic/execution/execution_engine.rs | 2 + src/crates/assembly/core/src/agentic/mod.rs | 2 +- .../src/agentic/session/session_manager.rs | 601 ++++- .../core/src/agentic/skill_agent_snapshot.rs | 58 +- .../assembly/core/src/agentic/system.rs | 23 +- .../agentic/tools/pipeline/tool_pipeline.rs | 27 +- .../core/src/agentic/tools/product_runtime.rs | 10 +- .../agentic/tools/product_runtime/catalog.rs | 55 +- .../core/src/agentic/tools/registry.rs | 49 +- .../src/agentic/tools/tool_context_runtime.rs | 44 +- .../assembly/core/src/product_runtime.rs | 164 +- .../src/product_runtime/runtime_services.rs | 89 +- .../core/src/service_agent_runtime.rs | 138 +- .../assembly/product-capabilities/src/lib.rs | 10 +- .../tests/plugin_product_shape.rs | 3 +- .../tests/product_capabilities.rs | 4 +- .../tests/product_sdk_assembly.rs | 53 +- src/crates/contracts/runtime-ports/src/lib.rs | 59 +- .../execution/agent-runtime/src/permission.rs | 38 +- .../execution/agent-runtime/src/runtime.rs | 157 +- src/crates/execution/agent-runtime/src/sdk.rs | 56 +- .../tests/permission_contracts.rs | 31 + .../agent-runtime/tests/sdk_smoke.rs | 89 +- src/crates/interfaces/AGENTS-CN.md | 4 + src/crates/interfaces/AGENTS.md | 5 + src/crates/interfaces/sdk-host/Cargo.toml | 26 + src/crates/interfaces/sdk-host/src/host.rs | 2209 ++++++++++++++++ src/crates/interfaces/sdk-host/src/lib.rs | 8 + .../interfaces/sdk-host/src/protocol.rs | 469 ++++ .../sdk-host/tests/host_lifecycle.rs | 2240 +++++++++++++++++ .../sdk-host/tests/protocol_contracts.rs | 180 ++ src/crates/services/services-core/Cargo.toml | 3 +- src/crates/services/services-core/src/lib.rs | 2 + .../services-core/src/local_runtime_ports.rs | 147 ++ .../tests/local_runtime_ports.rs | 44 + .../services/terminal/src/pty/process.rs | 136 +- .../services/terminal/src/pty/service.rs | 71 +- .../services/terminal/src/session/binding.rs | 52 +- .../services/terminal/src/session/manager.rs | 10 +- 70 files changed, 9500 insertions(+), 802 deletions(-) delete mode 100644 src/apps/cli/src/runtime/services.rs create mode 100644 src/apps/sdk-host/Cargo.toml create mode 100644 src/apps/sdk-host/src/lib.rs create mode 100644 src/apps/sdk-host/src/main.rs create mode 100644 src/apps/sdk-host/src/runtime.rs create mode 100644 src/apps/sdk-host/src/transport.rs create mode 100644 src/apps/sdk-host/tests/process_initialization.rs create mode 100644 src/apps/sdk-host/tests/stdio_process.rs create mode 100644 src/apps/sdk-host/tests/stdio_transport.rs create mode 100644 src/crates/interfaces/sdk-host/Cargo.toml create mode 100644 src/crates/interfaces/sdk-host/src/host.rs create mode 100644 src/crates/interfaces/sdk-host/src/lib.rs create mode 100644 src/crates/interfaces/sdk-host/src/protocol.rs create mode 100644 src/crates/interfaces/sdk-host/tests/host_lifecycle.rs create mode 100644 src/crates/interfaces/sdk-host/tests/protocol_contracts.rs create mode 100644 src/crates/services/services-core/src/local_runtime_ports.rs create mode 100644 src/crates/services/services-core/tests/local_runtime_ports.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8a5ec82759..b06162d606 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -53,8 +53,8 @@ jobs: cache-bin: false save-if: ${{ github.event_name != 'pull_request' }} - - name: Run CLI, ACP, and agent runtime tests - run: cargo test --locked -p bitfun-cli -p bitfun-acp -p bitfun-agent-runtime + - name: Run CLI, ACP, SDK Host, and agent runtime tests + run: cargo test --locked -p bitfun-cli -p bitfun-acp -p bitfun-sdk-host -p bitfun-sdk-host-app -p bitfun-agent-runtime # ── Rust: build check ───────────────────────────────────────────── rust-build-check: @@ -137,6 +137,18 @@ jobs: - name: Check compilation run: cargo check --locked --workspace + - name: Run SDK Host terminal cleanup regressions + # This PR changes only the PTY/session shutdown path required by + # connection-scoped Session cleanup. Keep the gate scoped to those + # regressions instead of adopting unrelated terminal command tests. + run: | + cargo test --locked -p terminal-core shutdown_returns_only_after_process_exit_is_confirmed -- --test-threads=1 + cargo test --locked -p terminal-core shutdown_evicts_a_process_whose_controller_already_confirmed_exit -- --test-threads=1 + cargo test --locked -p terminal-core background_only_binding_is_owned_by_the_session -- --test-threads=1 + + - name: Run standalone SDK Host process smokes + run: cargo test --locked -p bitfun-sdk-host-app --test process_initialization --test stdio_process + - name: Run Windows ConPTY CLI smoke test if: runner.os == 'Windows' run: cargo test --locked -p bitfun-cli --test terminal_process_contracts diff --git a/Cargo.toml b/Cargo.toml index ec562bf7b6..d6b8eb34a0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,10 +1,12 @@ [workspace] members = [ "src/apps/cli", + "src/apps/sdk-host", "src/apps/desktop", "src/apps/server", "src/apps/relay-server", "src/crates/interfaces/acp", + "src/crates/interfaces/sdk-host", "src/crates/assembly/core", "src/crates/assembly/external-sources", "src/crates/adapters/ai-adapters", diff --git a/docs/architecture/agent-runtime-services-design.md b/docs/architecture/agent-runtime-services-design.md index ff8b805b6b..0afeae2323 100644 --- a/docs/architecture/agent-runtime-services-design.md +++ b/docs/architecture/agent-runtime-services-design.md @@ -711,8 +711,10 @@ pub struct HarnessExecutionContext { `src/crates/assembly/core` 仍承担 `bitfun-core` 兼容组装。现有 `ProductAssembler` 是具体结构体， 通过 `assemble(ProductAssemblyInput)` 产生 `ProductRuntimeParts`，本文件不再为它定义第二套目标接口。 -当前 CLI 与 CLI 托管的 ACP server 已使用类型化 `RuntimeServices`，分别以 `DeliveryProfile::Cli` 和 -`DeliveryProfile::Acp` 构造 `ProductRuntimeParts`。Desktop 主交互直接从现有协调器和调度器端口构造窄口径 +当前 CLI、CLI 托管的 ACP server 与独立 SDK Host 已使用类型化 `RuntimeServices`，分别以 +`DeliveryProfile::Cli`、`DeliveryProfile::Acp` 和 `DeliveryProfile::Sdk` 构造 `ProductRuntimeParts`。 +SDK profile 当前从共享产品事实获得与 Headless CLI 相同的能力集合，但保持独立产品身份和 +`AgentSubmissionSource::SdkHost`；这不建立 CLI crate/协议依赖。Desktop 主交互直接从现有协调器和调度器端口构造窄口径 Rust Runtime SDK，不注册未实现的 `RuntimeServices` 能力，也不宣称完整 Desktop profile 可用。CLI 通过 一个调用级上下文把该 Rust 接口、Harness、能力注册、调用级权限和 Agentic 事件广播交给 TUI、Exec、Session、Usage 与 交互模式下的 Peer Host。Rust Runtime SDK 已承接会话创建/列举/删除/基础恢复、重命名/归档、会话模型更新、thread-goal 查询、类型化转录读取、本地分支、用量生成、 @@ -764,6 +766,8 @@ Desktop 与 CLI Peer Host 还各自注入同一个 Core-backed `LocalWorkspaceSn - 具体运行时服务通过 `RuntimeServicesBuilder` / provider registry 构造。 - CLI 只选择 `DeliveryProfile::Cli` 一次；必需服务缺失时组装失败，不回退到静态计划或另一 profile。 - CLI 的 ACP stdio 入口只选择 `DeliveryProfile::Acp` 一次；组装或 Rust Runtime SDK 构造失败时在接受 stdio 请求前退出。 +- 独立 `bitfun-sdk-host` 只选择 `DeliveryProfile::Sdk` 一次；stdio framing 与进程 bootstrap 留在 app， + `interfaces/sdk-host` 只保留版本化协议和连接用例。Host 不通过 CLI 启动，也不使用 CLI submission source。 - CLI 的 `json` 输出为单结果文档，`stream-json` 直接复用现有 `AgenticEventEnvelope`；协议层不新增 `schema_version`、`sequence` 或平行事件 taxonomy。 - 能力计划选择工具提供方组计划和 Harness 描述符；当前不存在供任意模块注册所有对象的通用组装注册表。 diff --git a/docs/architecture/agent-sdk-product-architecture.md b/docs/architecture/agent-sdk-product-architecture.md index b015403a07..9026adbe29 100644 --- a/docs/architecture/agent-sdk-product-architecture.md +++ b/docs/architecture/agent-sdk-product-architecture.md @@ -7,9 +7,12 @@ [`cli-product-line-design.md`](cli-product-line-design.md) 定义；能力导入/导出和外部宿主适配由 [`capability-runtime-integration-design.md`](extensions/capability-runtime-integration-design.md) 定义。 -本文记录目标架构与公开兼容门槛，不表示 TypeScript/Python SDK、SDK Host 协议或下文全部能力 -已经交付。当前代码中的 `agent-runtime::sdk` 是低层 Rust Runtime SDK（当前 preview），服务现有产品入口和 -Rust 嵌入；它不是本文定义的公开 BitFun Agent SDK，也不能据此宣称已达到 Claude Agent SDK 等价能力。 +本文记录目标架构与公开兼容门槛，不表示 TypeScript/Python SDK 或下文全部能力已经交付。当前代码中的 +`agent-runtime::sdk` 是低层 Rust Runtime SDK（当前 preview），服务现有产品入口和 Rust 嵌入； +`interfaces/sdk-host` 是本地协议与连接用例 adapter，`apps/sdk-host` 是独立的原生 Host 进程及 stdio transport。 +当前代码只形成第一切片的实现候选，尚未通过跨内部调用、Headless CLI 与 SDK Host 的同一 Query golden fixture， +因此还不能称为已交付的 SDK Host preview。它们都不是本文定义的 +公开 BitFun Agent SDK，也不能据此宣称已达到 Claude Agent SDK 等价能力。 ## 1. 最终决策 @@ -57,7 +60,7 @@ BitFun 采用以下组合，而不完整复制任一产品： | 公开心智 | Agent、Session、Message、Tool、Hook、MCP、Permission | Thread、Turn、Item、Approval、Event | Project/Session/Route/API | 采用用户已熟悉的 Agent/Session/Turn/Message/Event；不公开内部协议名 | | Agent Loop | SDK 直接获得 Claude Code 同一工具和 loop | Runtime/App Server 持有 loop，SDK 是精选门面 | Server 持有 loop，SDK 接近 API client | 唯一 BitFun Runtime 持有 loop；SDK 只映射用例和回调 | | 双向交互 | 自定义 Tool、权限、Hook、用户输入 callback | App Server request/response 支持审批和输入 | Server API/Event，客户端控制资源 | SDK Host 必须双向；不能只解析单向 JSONL 输出 | -| Session | resume/fork、外部存储、checkpoint 等 | thread start/resume/fork、turn/item、持久化 | session CRUD、prompt/abort、event | Session/Turn 作为稳定公共事实；存储 owner 仍唯一 | +| Session | resume/fork、外部会话存储 | thread start/resume/fork、turn/item、持久化 | session CRUD、prompt/abort、event | Session/Turn 作为稳定公共事实；存储 owner 仍唯一 | | 协议纪律 | SDK API 主导，底层 binary 对用户透明 | 初始化握手、能力协商、schema 生成、稳定/实验分层、背压 | OpenAPI 生成类型，Server URL 是显式概念 | 采用 Codex 式版本/能力/schema/背压纪律，默认隐藏 transport | | Headless CLI | `claude -p` 与 SDK 能力相通，CLI 更适合一次性/CI | `codex exec` 与 SDK 都调用同一 runtime | Server/CLI/TUI 均围绕同一服务 | `bitfun exec` 保持独立入口，与 SDK 做能力和事件等价 | | 扩展暴露 | Skills、Plugins、Agents、MCP、Hooks | Skills、Plugins、MCP、Hooks 等由 Runtime/App Server 投影 | Plugin、Tool、MCP、Agent 等 Server API | SDK 只暴露已由 BitFun owner 提交的能力，不输出生态原始对象 | @@ -119,6 +122,11 @@ SDK 负责定位、启动、握手和关闭匹配 SDK Host；这些内部名词 | BitFun Agent SDK | 面向 Python/TypeScript 的 `query()`、Session、typed callback 产品 | 尚未交付 | SDK Host 只是 BitFun Agent SDK 到 Agent Runtime 的内部跨进程适配器，不是第四套 Runtime，也不是用户默认需要配置的产品。 +当前本地 stdio Host 只支持 `initialize`、Session create/close、Query start/cancel、Event/Result、typed error 和 +`shutdown`；structured output、usage、自定义 Tool、Permission/Hook callback、SDK MCP 配置和预启动 Host 均通过 +capability 明确报告为不可用。当前没有 Permission callback；遇到 `ask` 时 Host 会拒绝请求、取消对应 Turn，并以 +`action_required` 结束 Query，不会无限等待或自动放行。Host 由独立的 `bitfun-sdk-host` 二进制承载； +`bitfun` 不包含隐藏 Host 子命令，也不依赖 SDK Host 协议 crate。 以下事实不能由当前 Rust Runtime SDK 推导： @@ -145,7 +153,7 @@ flowchart TB GuiAdapter["GUI/Tauri/Web Adapter"] TuiAdapter["TUI Adapter"] CliAdapter["JSON / stream-json Adapter"] - SdkAdapter["SDK Host Adapter"] + SdkAdapter["SDK Host protocol adapter"] AcpAdapter["ACP Adapter"] ServerAdapter["HTTP/WebSocket/Remote Adapter"] end @@ -189,7 +197,50 @@ flowchart TB 3. **共享应用层不等于共享 UI 或协议**。GUI/TUI 保留各自 renderer、交互和平台生命周期；CLI、 SDK、ACP、HTTP/WebSocket 保留各自 wire projection。 -### 4.1 各形态能做什么 +### 4.1 进程与依赖视图 + +```mermaid +flowchart LR + Python["Python SDK"] -->|"manage / stdio"| Host["bitfun-sdk-host\nstandalone app"] + TypeScript["TypeScript SDK"] -->|"manage / stdio"| Host + Host --> AppTransport["apps/sdk-host\nprocess bootstrap + stdio framing"] + AppTransport --> HostProtocol["interfaces/sdk-host\nprotocol + connection use cases"] + HostProtocol --> RuntimeApi["Agent Runtime API"] + + Bitfun["bitfun\nTUI + headless CLI + ACP command"] --> RuntimeApi + Desktop["Desktop / GUI"] --> RuntimeApi + Server["Server / Remote"] --> RuntimeApi + + RuntimeApi --> Owners["single Session / Tool / MCP /\nPermission / Hook owners"] + Owners --> TerminalPort["TerminalPort\nshared internal contract"] + TerminalPort --> TerminalCore["terminal-core\ninternal service implementation"] +``` + +| 依赖事实 | 约束 | +|---|---| +| `bitfun` CLI | 不依赖 `bitfun-sdk-host` app 或协议；只依赖共享 Runtime/Capability owner | +| `bitfun-sdk-host` app | 独立 composition root；只选择 `DeliveryProfile::Sdk` 与 `AgentSubmissionSource::SdkHost`。当前与 Headless CLI 共享 Runtime owner 和 assembly-plan 能力上限，但 Host 的实际 wire 能力是明确的严格子集；它不依赖 `bitfun-cli` crate，也不冒充 CLI | +| `interfaces/sdk-host` | 只拥有 JSON-RPC DTO、初始化/能力协商和连接级 Query/Session 生命周期；不拥有 stdin/stdout、进程入口或 Runtime 业务状态 | +| `TerminalPort` | Runtime 内部的抽象端口，供内置 Bash/命令类 Tool 复用；不是 Python/TypeScript SDK 方法，也不是 SDK Host wire | +| `terminal-core` | `product-full` Host 内部的具体终端/PTY 服务；SDK Host app 可像 GUI/TUI/CLI 一样通过共享 Runtime owner 使用它，但 `interfaces/sdk-host`、公开 SDK 类型和 JSON-RPC DTO 禁止依赖或暴露它 | +| Python/TypeScript SDK | 后续随包携带或定位匹配的 `bitfun-sdk-host`；用户不需要安装 CLI | +| HTTP/WebSocket Server | 独立远程产品；不作为本地 SDK callback Host，也不被 CLI 隐式启动 | + +各 composition root 必须在配置归一化读取全局 Tool owner 前选择自己的 delivery profile；选择是进程级且不可替换， +后续 Agentic 初始化只验证同一选择。CLI、ACP 和 SDK Host 复用这一 Core 初始化契约，但不互相依赖 app 或协议： + +```mermaid +flowchart LR + Entry["CLI / ACP / SDK Host entry"] --> Profile["select delivery profile"] + Profile --> Config["initialize config + AI services"] + Config --> Agentic["initialize shared Agentic owners"] + Agentic --> Surface["build surface-specific runtime / transport"] +``` + +独立 Host 进程在接受协议输入前安装进程级 TLS crypto provider，并在受控的 16 MiB worker stack 上启动 +Tokio Runtime；这些是 Host composition root 的运行条件，不进入 SDK Host 协议，也不由 CLI 代为提供。 + +### 4.2 各形态能做什么 | 形态 | 主要用途 | 必须共享的 Agent 能力 | 形态特有职责 | |---|---|---|---| @@ -208,7 +259,7 @@ flowchart TB sequenceDiagram participant App as "Python / TypeScript application" participant SDK as "BitFun Agent SDK" - participant Host as "Matched local SDK Host" + participant Host as "Standalone matched bitfun-sdk-host" participant API as "Agent Runtime API" participant Owners as "Runtime / MCP / Tool / Permission / Hook owners" @@ -252,6 +303,32 @@ SDK Host 协议只负责跨进程/跨语言边界，不拥有业务状态。目 - Agent/Subagent/Skill/Plugin source 的类型化配置与能力状态。 - typed errors、deadline、cancellation、bounded queue、backpressure 和 late-result fence。 +目标方法族按 capability 逐步开放。当前 PR1 实现候选的实际 wire 边界如下，不能从目标拓扑推导尚未交付的 callback： + +| 范围 | PR1 实现候选行为 | 明确未开放 | +|---|---|---| +| Query | `query/start` 通过共享 Dialog Scheduler 提交；每个 Session 同时只允许一个 SDK Query；若 Scheduler 返回 queued，Host 接受并绑定其精确 `turnId`，event、cancel 与 settlement 均只跟踪该 Turn | steer、resume/fork、structured output | +| Event | 只投影封闭的 `assistant_text_delta { text }`；内部事件和未知 payload 不上 wire | Tool、usage、trace 和原始事件总线 payload | +| Result | `query/result` 在 Turn settlement 后发布；失败携带稳定 code、retryable、correlation 和 recovery | 任意 JSON 错误或通过 message 驱动控制流 | +| Permission | capability 为 `false`；提交时关闭用户输入 Tool；出现属于当前精确 Turn 的 `ask` 时，在有界 deadline 内 reject + cancel + `action_required` | Permission callback 与隐式 auto-approve | +| Session | `session/create` 和不带 `sessionId` 的 `query/start` 只创建 connection-scoped transient Session；前者返回 `lifetime=connection`，后者返回 `sessionLifetime=connection`；`query/start.sessionId` 只能引用同一连接已创建的 Session，未知或已持久化的 ID 返回 `capability_unavailable`；`session/close` 由一个总 deadline 约束 scheduler maintenance 与资源清理；断连/close 后不可 resume/list/fork；该 Session 及其 Subagent 不暴露会读取或操纵独立 Session、持久定时任务或全局运行资源的 `SessionControl`、`SessionMessage`、`SessionHistory`、`Cron`、`ControlHub` | durable create、list、resume/fork、跨 Session 消息、持久调度和全局控制 | +| Transport | 独立 `bitfun-sdk-host` 进程使用本地换行分隔 JSON-RPC；有请求/Query/Session 上限、写超时和总 shutdown deadline；Host 用 exact ID 跟踪 transient create，在 graceful budget 后执行 abort → join → 专用 transient discard，永不通过普通 durable delete 补偿；deadline 耗尽时结束当前连接并记录结构化本地告警，当前协议不向已断开的调用方承诺清理完成结果 | managed SDK supervisor 的 Host 进程树 kill/reap、TCP、预启动 Host、远程复用 | + +`query/start` 若提供 `sessionId`，不得同时提供只在创建 Session 时有效的 `sessionName`、`agent`、`cwd` 或 +`model`；Host 返回 `invalid_request`，不会静默忽略。当前只有 `shutdown` 支持 JSON-RPC notification；会创建、返回或关闭 +Session/Query 资源的生命周期方法必须带 request ID，无 ID 时不执行且不响应。有效 JSON 但不符合 request 或 notification +envelope 的输入返回标准 `-32600`，只有 JSON 语法或行 framing 失败返回 `-32700`。 + +这里的 transient 只表示 Session 状态、turn、transcript 和 cache 不持久化；它不是权限或文件副作用沙箱。Task 创建的 Subagent 继承该边界，并在 +父 Session discard 时由现有 Session owner 按父子关系级联释放。discard 同时回收该 Session 仍持有的 Snapshot、Cron +job、Terminal binding 和后台终端进程；任一已存在资源无法确认关闭时 fail closed。文件编辑、已完成命令、MCP 或外部 +服务等 Tool 已产生的真实副作用不会随 Session discard 自动回滚。PR2 不得静默改变 PR1 方法的 lifetime；durable +create/resume 必须以新的 capability/协议版本开放，并先具备原子 commit、崩溃恢复、请求指纹和跨进程 fencing。 + +同一 `ToolRuntimeRestrictions` 必须同时约束首轮 Skill/Agent 列表、Subagent 的工具摘要、模型 tool manifest、 +`GetToolSpec` 与最终执行 admission。Agent 模板提到可选工具时必须要求先以当前 tool list 为准；不能让 prompt +继续指示模型调用已从 manifest 移除的能力，也不能为 SDK Host 单独维护一套 prompt 或工具过滤规则。 + 注册、查询和调用使用不同身份： | Lease scope | 典型内容 | 释放条件 | @@ -459,9 +536,10 @@ flowchart LR - `correlation_id`、`causation_id`、deadline/attempt。 - 可选且闭合的 recovery action；调用方不得解析 message 控制流程。 -错误类别至少区分 validation、authentication、permission denied、action required、capability unavailable、 -version mismatch、host unavailable、overloaded、timeout、cancelled、callback failed、invalid callback output、 -MCP unavailable、process lost 和 internal error。 +错误类别至少区分 validation、authentication、permission denied、action required、provider quota、provider billing、 +capability unavailable、version mismatch、host unavailable、overloaded、timeout、cancelled、cleanup required、 +callback failed、invalid callback output、MCP unavailable、process lost 和 internal error。未确认 Turn settlement 或隐藏 +Session 清理失败属于 `cleanup_required`，不得建议在同一 Host 上自动重试。 四类观测数据保持分离： @@ -481,6 +559,10 @@ MCP unavailable、process lost 和 internal error。 2. **SDK Host protocol version**：wire schema、方法、事件、错误和 capability。 3. **Runtime capability version**：Hook 点、Tool/MCP/Permission 语义和 product capability。 +当前 Host 握手返回 `stability=not_delivered`；其 `protocolVersion` 只是实现候选内部的 wire 修订号，不构成公开 v1 +兼容承诺。进入 preview 前必须从 Rust DTO 生成或机器验证唯一协议 schema，并以 drift check 保证 Python、TypeScript +和 Host 不各自手写一套合同；达到共同 golden fixture 与 schema 门槛后，才把成熟度切换为 `preview`。 + SDK 包默认锁定匹配 SDK Host runtime。连接预启动的受信本地 SDK Host 时先完成双向认证，再完成初始化与 capability negotiation。稳定 API 不得依赖 experimental capability；实验字段必须显式 opt-in，并允许旧 SDK 忽略未知通知。协议必须定义 @@ -522,7 +604,7 @@ SDK 工作可以与 OpenCode/Claude/Codex Hook 适配并行，但只能在以下 Python/TypeScript wire callback。 - SDK Hook callback 的最终接入必须基于已合并的公共 Hook owner。两条分支同时修改 owner 合同或 `agent-runtime` 权威状态时停止并行，先串行冻结合同。 -- 文档、schema fixture 和独立 `interfaces/sdk-host`/SDK package 可以并行；Cargo workspace、公共合同和 +- 文档、schema fixture、独立 `apps/sdk-host`、`interfaces/sdk-host` 和 SDK package 可以并行；Cargo workspace、公共合同和 owner 文件的机械冲突不应通过复制类型解决。 ## 12. 后续实现切片 @@ -540,14 +622,23 @@ flowchart LR ### SDK-PR1：Agent Runtime 合同与本地 SDK Host 基线 +当前状态：独立本地 Host 进程与协议基线已有实现候选，但尚未满足本节退出条件；公开 Python/TypeScript 包也尚未交付。 + - 把当前 Rust Runtime SDK 收敛为共享 Agent Runtime 用例，不扩大为 service locator。 -- 建立 versioned initialize、Session/Turn、query/event/result、cancel、session close、typed error、capability 和 shutdown 协议。 -- 提供同进程 contract fixture 和本地 stdio Host smoke；ACP/CLI/GUI/TUI 行为不迁移到第二 owner。 +- 建立 versioned initialize、Session/Turn、query/event/result、cancel、session close、typed error、capability 和 shutdown 协议； + Host 新建 Session 明确为 connection-scoped transient container，不冒充公开 SDK 的 durable Session。 +- 提供同进程 contract fixture 和独立本地 stdio Host process smoke；`bitfun` CLI 不承载 Host 子命令， + ACP/CLI/GUI/TUI 行为不迁移到第二 owner。Host 使用独立 SDK profile/source；它与 CLI 共享 Runtime owner 和 + assembly-plan 能力上限，但实际 Host capability matrix 保持严格子集，不形成 app 或协议依赖。 - 仅投影现有稳定能力，并建立通用 request lease；不先发布自定义 Tool/Hook callback。 -退出条件：同一 Session/Turn fixture 能经内部调用、Headless CLI 和 SDK Host protocol 得到等价业务事实； +完整退出条件：同一 Session/Turn fixture 能经内部调用、Headless CLI 和 SDK Host protocol 得到等价业务事实； 过载、取消、进程丢失和版本不匹配 fail closed。`session.close` 能取消并回收待处理 owner interaction、MCP 引用、 -子进程和 Subagent，不删除或归档持久化 Session。 +子进程和 Subagent；Host-owned transient Session 及其 transient Subagent 后代由现有 Session owner 级联 discard； +PR1 不接管、删除或归档任何已有 durable Session。 +当前独立 process smoke 已覆盖 Host 的启动、协议握手、stdout 纯净与关闭；Headless CLI 有独立执行测试。 +跨 Headless CLI、SDK Host 和内部直调用的同一 Session/Turn golden fixture，以及 Tool/Permission 等价结果， +仍是该退出条件的未完成证据，不能用各自 smoke 名称替代。 ### SDK-PR2：Python/TypeScript query 与 Session 产品闭环 @@ -559,6 +650,8 @@ flowchart LR - 完成预启动本地 SDK Host 的认证负向测试：错误/过期/重放 secret、端点抢占、ACL/owner、peer credential、 多连接隔离和敏感数据不泄漏；任何失败都发生在业务 payload 之前。 - 提供 bare/default source 选择，避免 CI 隐式加载用户 Hook/Skill/MCP。 +- durable create/resume 必须先完成可判定的持久化 commit、同 ID 幂等恢复、binding/request fingerprint 校验、 + partial cleanup 和 managed Host kill/reap fencing，不能复用 PR1 transient 语义冒充持久会话。 - 用相同示例和 golden fixture 验证 Python/TypeScript/Headless CLI。 退出条件：两个仓库外 pilot 可完成安装、执行、取消、恢复、升级和错误恢复；managed 与预启动本地 SDK Host diff --git a/docs/architecture/cli-product-line-design.md b/docs/architecture/cli-product-line-design.md index 52c2c7d56c..01d29add65 100644 --- a/docs/architecture/cli-product-line-design.md +++ b/docs/architecture/cli-product-line-design.md @@ -242,6 +242,8 @@ CLI 不提供 `--output-schema v1`。Codex/Claude 同类参数表达的是调用 Headless CLI 和公开 Agent SDK 都调用同一 Agent Runtime API，但交付形态不同。本文件只保留 CLI 约束： - `bitfun exec` 面向 shell、CI 和一次性任务，使用 stdin/stdout/stderr、退出码与 `text/json/stream-json`。 +- `bitfun` 不承载隐藏 SDK Host 子命令，也不依赖 SDK Host 协议；独立 `bitfun-sdk-host` app 与 CLI + 分别选择 SDK/CLI profile 和 submission source，只复用同一 Runtime owner，以及由共享产品事实生成的等价能力集合。 - CLI 不在进程内执行 Python/TypeScript Tool、Permission 或 Hook callback。 - 公开 SDK 不解析 `stream-json` 作为正式双向协议；它通过版本化 SDK Host 获得 callback 与连接生命周期。 - 两者的能力对照、共同 fixture 和等价门槛以 diff --git a/docs/architecture/product-architecture.md b/docs/architecture/product-architecture.md index 7d4bf77842..5044e37d0b 100644 --- a/docs/architecture/product-architecture.md +++ b/docs/architecture/product-architecture.md @@ -363,7 +363,7 @@ flowchart LR | ACP | CLI 托管的服务端仍以 `bitfun-core/product-full` 作为兼容执行层 | 入口已选择 `DeliveryProfile::Acp` 并消费 Runtime Parts；组装层在入队前原子拒绝忙碌会话，不改变其他产品入口的排队行为；活动会话模型与模式写入走 Agent Runtime API。`session/load` 先校验和建立临时 MCP，再恢复 Core，在历史回放成功后才发布活动状态；失败会卸载本次内存状态而不删除历史。同 ID 的重叠打开/关闭被明确拒绝。成功的 `session/close` 阻止新轮次、排空队列和后台子会话，再卸载临时 Core 状态并回收 MCP 与连接；失败保留会话所有权和历史，返回可重试阶段。持久化历史仍可重新加载。完整历史、模型/模式目录与配置读取仍留在现有 Core/ACP 归属，不据此宣称完整解耦 | | Server / Remote | Server 可返回共享 v1 control/catalog 快照，`hostCapabilities` 明确为只读，并在解析 mutation payload 前 fail closed；Peer Host 代理同一读写控制动作与既有能力专属操作，连接旧 Peer 时读路径降级到 legacy catalog、来源启停降级到既有 mutation，Safe Mode 明确要求升级 Host；SSH Remote 工作区的外部来源发现与执行仍未实现 | 控制端用 Host 身份与工作区共同隔离异步结果。Peer Host 只处理 Host 上的真实工作区，不在控制端替远端发现；SSH Remote 未接入时返回明确不支持且不回退本机来源 | | Web / Mobile Web | 依赖现有后端入口，不持有插件执行单元 | 对应 profile 当前为空计划或未接入生产，不能据枚举值宣称独立产品能力 | -| Public Agent SDK | 尚未交付；当前 Rust Runtime SDK 的成熟度为 preview，另有 SDK profile 计划和测试替身 | 目标为 Python/TypeScript `query()`、client/session 和 callback；未建立 SDK Host、仓库外消费者与 Claude 等价矩阵前不得宣称 preview 或可发布 | +| Public Agent SDK | Python/TypeScript 产品尚未交付；Rust Runtime SDK 是内部 preview。本地 SDK Host 协议与独立 `bitfun-sdk-host` 进程只有实现候选，在跨内部调用、Headless CLI 与 Host 的同一 Query golden fixture 通过前不称为已交付 preview | SDK Host 独立选择 SDK profile/source；能力集合从共享产品事实生成，不依赖或冒充 CLI。目标为 Python/TypeScript `query()`、client/session 和 callback；未建立语言包、仓库外消费者与 Claude 等价矩阵前不得宣称公开 SDK preview 或可发布 | 对外一级状态统一使用[外部 AI 工作内容设计](extensions/external-ai-work-sources-design.md#7-状态与提示规则)定义的 已发现、已应用、可用、需确认、更新中、沿用上一版本、部分受限、暂时过期、已移除/已停用和不可用，并附带 diff --git a/scripts/core-boundaries/rules/crate-layout.mjs b/scripts/core-boundaries/rules/crate-layout.mjs index 8c5c9904c6..ce20eebdaa 100644 --- a/scripts/core-boundaries/rules/crate-layout.mjs +++ b/scripts/core-boundaries/rules/crate-layout.mjs @@ -27,6 +27,7 @@ export const crateLayoutRules = [ { crateName: 'terminal', layer: 'services', path: 'src/crates/services/terminal' }, { crateName: 'acp', layer: 'interfaces', path: 'src/crates/interfaces/acp' }, + { crateName: 'sdk-host', layer: 'interfaces', path: 'src/crates/interfaces/sdk-host' }, { crateName: 'ai-adapters', layer: 'adapters', path: 'src/crates/adapters/ai-adapters' }, { crateName: 'claude-code-adapter', layer: 'adapters', path: 'src/crates/adapters/claude-code-adapter' }, { crateName: 'codex-adapter', layer: 'adapters', path: 'src/crates/adapters/codex-adapter' }, diff --git a/scripts/core-boundaries/rules/crate-rules.mjs b/scripts/core-boundaries/rules/crate-rules.mjs index 1705b3e6ef..fcb852b0e2 100644 --- a/scripts/core-boundaries/rules/crate-rules.mjs +++ b/scripts/core-boundaries/rules/crate-rules.mjs @@ -12,6 +12,7 @@ export const noCoreDependencyCrates = [ 'product-capabilities', 'runtime-ports', 'runtime-services', + 'sdk-host', 'services-core', 'services-integrations', 'agent-tools', @@ -183,6 +184,36 @@ export const lightweightBoundaryRules = [ 'syntect-tui', ], }, + { + crateName: 'sdk-host', + reason: + 'SDK Host protocol must stay a portable Runtime adapter without concrete product, terminal, service, or CLI implementations', + forbiddenDeps: [ + 'bitfun-core', + 'bitfun-ai-adapters', + 'bitfun-services-core', + 'bitfun-services-integrations', + 'bitfun-agent-tools', + 'bitfun-tool-packs', + 'bitfun-product-capabilities', + 'bitfun-product-domains', + 'bitfun-transport', + 'terminal-core', + 'tool-runtime', + 'tauri', + 'reqwest', + 'git2', + 'rmcp', + 'image', + 'tokio-tungstenite', + 'bitfun-sdk-host-app', + 'bitfun-cli', + 'ratatui', + 'crossterm', + 'arboard', + 'syntect-tui', + ], + }, { crateName: 'harness', reason: diff --git a/scripts/core-boundaries/rules/feature-rules.mjs b/scripts/core-boundaries/rules/feature-rules.mjs index 3d222bb740..5c589db1fa 100644 --- a/scripts/core-boundaries/rules/feature-rules.mjs +++ b/scripts/core-boundaries/rules/feature-rules.mjs @@ -1,6 +1,12 @@ // Boundary rules for feature assembly and optional dependency ownership. export const optionalDependencyFeatureOwnerRules = [ + { + crateName: 'services-core', + reason: + 'services-core workspace runtime dependencies must stay behind the explicit workspace-runtime feature', + dependencies: [{ depName: 'dunce', ownerFeatures: ['workspace-runtime'] }], + }, { crateName: 'runtime-ports', reason: @@ -135,6 +141,12 @@ export const productCoreFeatureAssemblyRules = [ requiredFeatures: ['product-full'], reason: 'CLI must explicitly assemble the full bitfun-core product runtime', }, + { + manifestPath: 'src/apps/sdk-host/Cargo.toml', + dependencyName: 'bitfun-core', + requiredFeatures: ['product-full'], + reason: 'SDK Host must explicitly assemble the full bitfun-core product runtime', + }, { manifestPath: 'src/apps/server/Cargo.toml', dependencyName: 'bitfun-core', diff --git a/scripts/core-boundaries/rules/source/required-rules.mjs b/scripts/core-boundaries/rules/source/required-rules.mjs index 1e5c7000fd..e0f32c7322 100644 --- a/scripts/core-boundaries/rules/source/required-rules.mjs +++ b/scripts/core-boundaries/rules/source/required-rules.mjs @@ -1451,8 +1451,8 @@ export const requiredContentRules = [ }, { regex: - /\bsdk_delivery_profile_builds_minimal_agent_runtime_without_product_full_capabilities\b/, - message: 'missing SDK delivery profile minimal runtime smoke', + /\bsdk_delivery_profile_builds_shared_runtime_owner_ceiling_without_bitfun_core\b/, + message: 'missing SDK delivery profile identity and shared runtime-owner ceiling smoke', }, { regex: /\bProductAssembler::new\(\)/, @@ -1472,7 +1472,38 @@ export const requiredContentRules = [ }, { regex: /\bDeliveryProfile::Sdk\b/, - message: 'product SDK smoke must cover the no-direct-core SDK delivery profile', + message: 'product SDK smoke must cover the distinct SDK delivery profile', + }, + ], + }, + { + path: 'src/apps/sdk-host/src/runtime.rs', + reason: + 'the standalone SDK Host must inject its selected delivery profile into the Core tool owner before agentic system construction', + patterns: [ + { + regex: /\binit_agentic_system_for_profile\b/, + message: 'SDK Host runtime must initialize Core with its selected delivery profile', + }, + { + regex: /\bselect_agentic_system_profile\b/, + message: + 'SDK Host runtime must select its profile before configuration can read the global tool owner', + }, + { + regex: /\bDeliveryProfile::Sdk\b/, + message: 'SDK Host runtime must retain a distinct SDK delivery profile', + }, + ], + }, + { + path: 'src/apps/sdk-host/src/main.rs', + reason: + 'the SDK Host composition root must select its delivery profile before global configuration canonicalization', + patterns: [ + { + regex: /\bselect_process_profile\b/, + message: 'SDK Host process startup must select its delivery profile first', }, ], }, @@ -6970,7 +7001,7 @@ export const requiredContentRules = [ }, { regex: /\bDeliveryProfile::Sdk\b/, - message: 'product tool runtime no-direct-core regression must cover SDK profile', + message: 'product tool runtime explicit profile regression must cover SDK profile', }, ], }, diff --git a/scripts/core-boundaries/self-test.mjs b/scripts/core-boundaries/self-test.mjs index bd8b63531b..bce09f1091 100644 --- a/scripts/core-boundaries/self-test.mjs +++ b/scripts/core-boundaries/self-test.mjs @@ -657,6 +657,15 @@ export function runManifestParserSelfTest({ const servicesOptionalOwnerRule = optionalDependencyFeatureOwnerRules.find( (rule) => rule.crateName === 'services-integrations', ); + const servicesCoreOptionalOwnerRule = optionalDependencyFeatureOwnerRules.find( + (rule) => rule.crateName === 'services-core', + ); + const servicesCoreDunceOwner = servicesCoreOptionalOwnerRule?.dependencies.find( + (dependency) => dependency.depName === 'dunce', + ); + if (!servicesCoreDunceOwner?.ownerFeatures.includes('workspace-runtime')) { + throw new Error('services-core workspace-runtime must own optional dependency dunce'); + } const servicesOptionalOwnerDeps = new Set( servicesOptionalOwnerRule?.dependencies.map((dependency) => dependency.depName) ?? [], ); @@ -1376,6 +1385,22 @@ export function runManifestParserSelfTest({ if (!agentRuntimeProfile?.forbiddenNonOptionalDeps.includes('tool-runtime')) { throw new Error('agent-runtime dependency profile must forbid concrete tool runtime'); } + if (!noCoreDependencyCrates.includes('sdk-host')) { + throw new Error('SDK Host protocol crate must be covered by the no-core dependency guard'); + } + const sdkHostRule = lightweightBoundaryRules.find((rule) => rule.crateName === 'sdk-host'); + for (const dependency of [ + 'bitfun-core', + 'terminal-core', + 'bitfun-services-core', + 'bitfun-services-integrations', + 'tool-runtime', + 'bitfun-cli', + ]) { + if (!sdkHostRule?.forbiddenDeps.includes(dependency)) { + throw new Error(`SDK Host protocol boundary must forbid concrete dependency: ${dependency}`); + } + } const productCapabilitiesRule = lightweightBoundaryRules.find( (rule) => rule.crateName === 'product-capabilities', ); @@ -2694,7 +2719,7 @@ export function runManifestParserSelfTest({ path: 'src/crates/assembly/product-capabilities/tests/product_sdk_assembly.rs', contracts: [ 'product_runtime_parts_can_build_agent_runtime_sdk_without_core', - 'sdk_delivery_profile_builds_minimal_agent_runtime_without_product_full_capabilities', + 'sdk_delivery_profile_builds_shared_runtime_owner_ceiling_without_bitfun_core', 'DeliveryProfile::Cli', 'DeliveryProfile::Sdk', ], diff --git a/src/apps/cli/src/agent/agentic_system.rs b/src/apps/cli/src/agent/agentic_system.rs index 694c11b391..27713295ce 100644 --- a/src/apps/cli/src/agent/agentic_system.rs +++ b/src/apps/cli/src/agent/agentic_system.rs @@ -1,11 +1,17 @@ use anyhow::{Context, Result}; +use bitfun_core::product_assembly::DeliveryProfile; use bitfun_core::product_runtime::CoreRuntimeServicesProvider; pub(crate) use bitfun_core::agentic::system::AgenticSystem; -pub(crate) async fn init_agentic_system() -> Result { - let system = bitfun_core::agentic::system::init_agentic_system() +pub(crate) fn select_agentic_system_profile(profile: DeliveryProfile) -> Result<()> { + bitfun_core::agentic::system::select_agentic_system_profile(profile) + .context("Failed to select agentic system delivery profile") +} + +pub(crate) async fn init_agentic_system(profile: DeliveryProfile) -> Result { + let system = bitfun_core::agentic::system::init_agentic_system_for_profile(profile) .await .context("Failed to initialize agentic system")?; system diff --git a/src/apps/cli/src/main.rs b/src/apps/cli/src/main.rs index 949848ba83..628a5e991f 100644 --- a/src/apps/cli/src/main.rs +++ b/src/apps/cli/src/main.rs @@ -565,6 +565,9 @@ async fn initialize_core_services( ) -> Result> { use bitfun_core::infrastructure::ai::AIClientFactory; + agent::agentic_system::select_agentic_system_profile( + bitfun_core::product_assembly::DeliveryProfile::Cli, + )?; bitfun_core::service::config::initialize_global_config() .await .map_err(|error| anyhow!("Failed to initialize global config service: {error}"))?; @@ -581,9 +584,11 @@ async fn initialize_core_services( initialize_terminal_service().await; - let agentic_system = agent::agentic_system::init_agentic_system() - .await - .map_err(|error| anyhow!("Failed to initialize agentic system: {error}"))?; + let agentic_system = agent::agentic_system::init_agentic_system( + bitfun_core::product_assembly::DeliveryProfile::Cli, + ) + .await + .map_err(|error| anyhow!("Failed to initialize agentic system: {error}"))?; tracing::info!("Agentic system initialized"); let runtime = std::sync::Arc::new(runtime::CliRuntimeContext::build( @@ -960,17 +965,9 @@ async fn run_cli() -> Result<()> { } Some(Commands::Doctor) => { - use std::sync::Arc; - - use runtime::services::{CliClock, CliRuntimeEventSink, CliRuntimeServicesProvider}; - let workspace = std::env::current_dir()?; - let services = CliRuntimeServicesProvider::new( - &workspace, - Arc::new(CliRuntimeEventSink::new(16)), - Arc::new(CliClock), - )? - .build()?; + let (_, services) = + bitfun_core::product_runtime::build_local_runtime_services(&workspace, 16)?; let product_runtime = product_assembly::assemble_cli_runtime_parts(services)?; if !management::print_doctor(&product_runtime).await? { std::process::exit(1); @@ -1411,3 +1408,21 @@ mod final_change_verification_cli_tests { assert!(!final_change_verification_enabled(verify, disable)); } } + +#[cfg(test)] +mod sdk_host_command_tests { + use super::Cli; + use clap::{CommandFactory, Parser}; + + #[test] + fn sdk_host_is_not_a_cli_command() { + let error = match Cli::try_parse_from(["bitfun", "sdk-host"]) { + Ok(_) => panic!("SDK Host must be a sibling application, not a CLI subcommand"), + Err(error) => error, + }; + assert_eq!(error.kind(), clap::error::ErrorKind::InvalidSubcommand); + let help = Cli::command().render_long_help().to_string(); + assert!(!help.contains("sdk-host")); + assert!(!include_str!("../Cargo.toml").contains("bitfun-sdk-host")); + } +} diff --git a/src/apps/cli/src/product_assembly.rs b/src/apps/cli/src/product_assembly.rs index bd4255fdb1..8076fcfb12 100644 --- a/src/apps/cli/src/product_assembly.rs +++ b/src/apps/cli/src/product_assembly.rs @@ -25,12 +25,10 @@ fn assemble_runtime_parts( #[cfg(test)] mod tests { - use std::sync::Arc; - use super::{assemble_acp_runtime_parts, assemble_cli_runtime_parts}; - use crate::runtime::services::{CliClock, CliRuntimeEventSink, CliRuntimeServicesProvider}; use bitfun_core::product_assembly::ProductServiceCapabilityStatus; use bitfun_core::product_assembly::{product_assembly_plan_for_profile, DeliveryProfile}; + use bitfun_core::product_runtime::build_local_runtime_services; use bitfun_runtime_ports::{ PluginRuntimeAvailability, PluginRuntimeUnavailableReason, RuntimeServiceCapability, }; @@ -70,14 +68,8 @@ mod tests { #[test] fn cli_product_assembly_consumes_production_runtime_services() { let workspace = tempfile::tempdir().expect("workspace"); - let services = CliRuntimeServicesProvider::new( - workspace.path(), - Arc::new(CliRuntimeEventSink::new(8)), - Arc::new(CliClock), - ) - .expect("provider") - .build() - .expect("runtime services"); + let (_, services) = + build_local_runtime_services(workspace.path(), 8).expect("runtime services"); let parts = assemble_cli_runtime_parts(services).expect("CLI product runtime parts"); @@ -99,14 +91,8 @@ mod tests { #[test] fn acp_product_assembly_uses_acp_profile_and_production_services() { let workspace = tempfile::tempdir().expect("workspace"); - let services = CliRuntimeServicesProvider::new( - workspace.path(), - Arc::new(CliRuntimeEventSink::new(8)), - Arc::new(CliClock), - ) - .expect("provider") - .build() - .expect("runtime services"); + let (_, services) = + build_local_runtime_services(workspace.path(), 8).expect("runtime services"); let parts = assemble_acp_runtime_parts(services).expect("ACP product runtime parts"); diff --git a/src/apps/cli/src/root_handlers.rs b/src/apps/cli/src/root_handlers.rs index 8b6d410e3e..48604f7eb1 100644 --- a/src/apps/cli/src/root_handlers.rs +++ b/src/apps/cli/src/root_handlers.rs @@ -774,19 +774,10 @@ async fn handle_external_config_action(action: ExternalConfigAction) -> Result<( } pub(crate) fn handle_health_command() -> Result<()> { - use std::sync::Arc; - use bitfun_core::runtime_ports::PluginRuntimeAvailability; - use crate::runtime::services::{CliClock, CliRuntimeEventSink, CliRuntimeServicesProvider}; - let workspace = std::env::current_dir().context("Failed to resolve current directory")?; - let services = CliRuntimeServicesProvider::new( - &workspace, - Arc::new(CliRuntimeEventSink::new(16)), - Arc::new(CliClock), - )? - .build()?; + let (_, services) = bitfun_core::product_runtime::build_local_runtime_services(&workspace, 16)?; let product_runtime = crate::product_assembly::assemble_cli_runtime_parts(services)?; println!("BitFun CLI health"); @@ -817,6 +808,9 @@ pub(crate) fn handle_health_command() -> Result<()> { pub(crate) async fn serve_acp_stdio() -> Result<()> { crate::setup_workspace(); + crate::agent::agentic_system::select_agentic_system_profile( + bitfun_core::product_assembly::DeliveryProfile::Acp, + )?; bitfun_core::service::config::initialize_global_config() .await .context("Failed to initialize global config service")?; @@ -830,9 +824,11 @@ pub(crate) async fn serve_acp_stdio() -> Result<()> { crate::initialize_terminal_service().await; - let agentic_system = crate::agent::agentic_system::init_agentic_system() - .await - .context("Failed to initialize agentic system")?; + let agentic_system = crate::agent::agentic_system::init_agentic_system( + bitfun_core::product_assembly::DeliveryProfile::Acp, + ) + .await + .context("Failed to initialize agentic system")?; tracing::info!("Agentic system initialized"); let workspace_root = std::env::current_dir().unwrap_or_else(|_| std::path::PathBuf::from(".")); diff --git a/src/apps/cli/src/runtime/events.rs b/src/apps/cli/src/runtime/events.rs index 77e65a396a..5c413a4ae7 100644 --- a/src/apps/cli/src/runtime/events.rs +++ b/src/apps/cli/src/runtime/events.rs @@ -1,115 +1 @@ -use std::sync::Arc; - -use bitfun_agent_runtime::sdk::{AgentEventReceiver, AgentEventSource}; -use bitfun_core::agentic::events::EventQueue; - -struct EventQueueDrain { - task: tokio::task::JoinHandle<()>, -} - -impl EventQueueDrain { - fn start(queue: Arc) -> Self { - let task = tokio::spawn(async move { - loop { - queue.wait_for_events().await; - while !queue.dequeue_configured_batch().await.is_empty() {} - } - }); - Self { task } - } -} - -impl Drop for EventQueueDrain { - fn drop(&mut self) { - self.task.abort(); - } -} - -#[derive(Clone)] -pub(crate) struct CliAgentEventSource { - source: AgentEventSource, - _drain: Arc, -} - -impl CliAgentEventSource { - pub(crate) fn new(queue: Arc) -> Self { - Self { - source: AgentEventSource::new(queue.clone()), - _drain: Arc::new(EventQueueDrain::start(queue)), - } - } - - pub(crate) fn subscribe(&self) -> AgentEventReceiver { - self.source.subscribe() - } - - pub(crate) fn runtime_source(&self) -> AgentEventSource { - self.source.clone() - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use bitfun_core::agentic::events::{EventQueue, EventQueueConfig}; - use bitfun_events::AgenticEvent; - - use super::CliAgentEventSource; - - #[tokio::test] - async fn subscribers_observe_every_event_while_the_legacy_queue_stays_bounded() { - let queue = Arc::new(EventQueue::new(EventQueueConfig { - max_queue_size: 4, - batch_size: 2, - })); - let source = CliAgentEventSource::new(queue.clone()); - let mut first = source.subscribe(); - let mut second = source.subscribe(); - - for index in 0..32 { - queue - .enqueue( - AgenticEvent::SessionStateChanged { - session_id: "session-1".to_string(), - new_state: format!("state-{index}"), - }, - None, - ) - .await - .expect("enqueue event"); - } - - let mut first_event = None; - let mut second_event = None; - for _ in 0..32 { - first_event = Some( - tokio::time::timeout(std::time::Duration::from_secs(1), first.recv()) - .await - .expect("first subscriber must not stall") - .expect("first subscriber event"), - ); - second_event = Some( - tokio::time::timeout(std::time::Duration::from_secs(1), second.recv()) - .await - .expect("second subscriber must not stall") - .expect("second subscriber event"), - ); - } - - let first_event = first_event.expect("last first event"); - let second_event = second_event.expect("last second event"); - assert_eq!(first_event.id, second_event.id); - assert!(matches!( - first_event.event, - AgenticEvent::SessionStateChanged { ref new_state, .. } if new_state == "state-31" - )); - tokio::time::timeout(std::time::Duration::from_secs(1), async { - while !queue.is_empty().await { - tokio::task::yield_now().await; - } - }) - .await - .expect("queue drainer must keep the legacy queue bounded"); - } -} +pub(crate) use bitfun_core::product_runtime::CoreProductAgentEventSource as CliAgentEventSource; diff --git a/src/apps/cli/src/runtime/mod.rs b/src/apps/cli/src/runtime/mod.rs index c4b7e1dac9..a2de5ded21 100644 --- a/src/apps/cli/src/runtime/mod.rs +++ b/src/apps/cli/src/runtime/mod.rs @@ -3,11 +3,11 @@ use std::sync::Arc; use anyhow::{Context, Result}; use bitfun_agent_runtime::sdk::AgentRuntime; -use bitfun_core::agentic::coordination::{self, DialogScheduler}; use bitfun_core::agentic::system::AgenticSystem; use bitfun_core::product_assembly::{ProductAssemblyPlan, ProductServiceCapabilityAvailability}; use bitfun_core::product_runtime::{ - CoreAgentRuntimeCompatibility, CoreLocalWorkspaceSnapshot, CoreProductAgentRuntime, + build_local_runtime_services, ensure_product_dialog_scheduler, CoreAgentRuntimeCompatibility, + CoreLocalWorkspaceSnapshot, CoreProductAgentRuntime, }; use bitfun_core::runtime_ports::PluginRuntimeAvailability; use bitfun_runtime_ports::LocalWorkspaceSnapshotPort; @@ -17,11 +17,9 @@ use crate::product_assembly::{assemble_acp_runtime_parts, assemble_cli_runtime_p pub(crate) mod approval; pub(crate) mod events; -pub(crate) mod services; use approval::CliApprovalPolicy; use events::CliAgentEventSource; -use services::{CliClock, CliRuntimeEventSink, CliRuntimeServicesProvider}; const RUNTIME_EVENT_BUFFER: usize = 256; @@ -69,15 +67,10 @@ impl CliRuntimeContext { workspace_root: impl AsRef, approval_policy: CliApprovalPolicy, ) -> Result { - let scheduler = ensure_dialog_scheduler(&agentic_system); - let runtime_events = Arc::new(CliRuntimeEventSink::new(RUNTIME_EVENT_BUFFER)); - let provider = CliRuntimeServicesProvider::new( - workspace_root, - runtime_events.clone(), - Arc::new(CliClock), - )?; - let workspace_root = provider.workspace_root().to_path_buf(); - let parts = assemble_cli_runtime_parts(provider.build()?) + let scheduler = ensure_product_dialog_scheduler(&agentic_system); + let (workspace_root, services) = + build_local_runtime_services(workspace_root, RUNTIME_EVENT_BUFFER)?; + let parts = assemble_cli_runtime_parts(services) .context("Failed to assemble CLI product runtime")?; let product = CliProductRuntimeState { @@ -172,11 +165,9 @@ impl AcpRuntimeContext { agentic_system: AgenticSystem, workspace_root: impl AsRef, ) -> Result { - let scheduler = ensure_dialog_scheduler(&agentic_system); - let runtime_events = Arc::new(CliRuntimeEventSink::new(RUNTIME_EVENT_BUFFER)); - let provider = - CliRuntimeServicesProvider::new(workspace_root, runtime_events, Arc::new(CliClock))?; - let parts = assemble_acp_runtime_parts(provider.build()?) + let scheduler = ensure_product_dialog_scheduler(&agentic_system); + let (_, services) = build_local_runtime_services(workspace_root, RUNTIME_EVENT_BUFFER)?; + let parts = assemble_acp_runtime_parts(services) .context("Failed to assemble ACP product runtime")?; let (services, harness_registry, _disabled_plugin_runtime) = parts.into_runtime_parts(); let agent_events = CliAgentEventSource::new(agentic_system.event_queue.clone()); @@ -203,20 +194,3 @@ impl AcpRuntimeContext { (self.agent_runtime.clone(), self.compatibility.clone()) } } - -fn ensure_dialog_scheduler(agentic_system: &AgenticSystem) -> Arc { - if let Some(scheduler) = coordination::get_global_scheduler() { - return scheduler; - } - - let session_manager = agentic_system.coordinator.get_session_manager().clone(); - let scheduler = DialogScheduler::new(agentic_system.coordinator.clone(), session_manager); - agentic_system - .coordinator - .set_scheduler_notifier(scheduler.outcome_sender()); - agentic_system - .coordinator - .set_round_injection_source(scheduler.round_injection_monitor()); - coordination::set_global_scheduler(scheduler.clone()); - scheduler -} diff --git a/src/apps/cli/src/runtime/services.rs b/src/apps/cli/src/runtime/services.rs deleted file mode 100644 index df8aa5db5a..0000000000 --- a/src/apps/cli/src/runtime/services.rs +++ /dev/null @@ -1,246 +0,0 @@ -use std::fmt; -use std::path::{Path, PathBuf}; -use std::sync::Arc; -use std::time::{SystemTime, UNIX_EPOCH}; - -use bitfun_core::product_runtime::CoreRuntimeServicesProvider; -use bitfun_runtime_ports::{ - ClockPort, FileSystemPort, PortResult, RuntimeEventEnvelope, RuntimeEventSink, - RuntimeServiceCapability, RuntimeServicePort, WorkspacePort, -}; -use bitfun_runtime_services::{ - RuntimeServices, RuntimeServicesBuilder, RuntimeServicesError, RuntimeServicesProvider, - RuntimeServicesRegistry, -}; -use tokio::sync::broadcast; - -#[derive(Debug)] -pub(crate) struct CliFileSystemService { - workspace_root: PathBuf, -} - -impl CliFileSystemService { - fn workspace_root(&self) -> &Path { - &self.workspace_root - } -} - -impl RuntimeServicePort for CliFileSystemService { - fn capability(&self) -> RuntimeServiceCapability { - RuntimeServiceCapability::FileSystem - } -} - -impl FileSystemPort for CliFileSystemService {} - -#[derive(Debug)] -pub(crate) struct CliWorkspaceService { - workspace_root: PathBuf, -} - -impl CliWorkspaceService { - fn workspace_root(&self) -> &Path { - &self.workspace_root - } -} - -impl RuntimeServicePort for CliWorkspaceService { - fn capability(&self) -> RuntimeServiceCapability { - RuntimeServiceCapability::Workspace - } -} - -impl WorkspacePort for CliWorkspaceService {} - -#[derive(Debug, Clone, Copy, Default)] -pub(crate) struct CliClock; - -impl RuntimeServicePort for CliClock { - fn capability(&self) -> RuntimeServiceCapability { - RuntimeServiceCapability::Clock - } -} - -impl ClockPort for CliClock { - fn now_unix_millis(&self) -> i64 { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .map(|duration| duration.as_millis().min(i64::MAX as u128) as i64) - .unwrap_or_default() - } -} - -#[derive(Debug, Clone)] -pub(crate) struct CliRuntimeEventSink { - tx: broadcast::Sender, -} - -impl CliRuntimeEventSink { - pub(crate) fn new(capacity: usize) -> Self { - let (tx, _) = broadcast::channel(capacity.max(1)); - Self { tx } - } - - #[cfg(test)] - pub(crate) fn subscribe(&self) -> broadcast::Receiver { - self.tx.subscribe() - } -} - -#[async_trait::async_trait] -impl RuntimeEventSink for CliRuntimeEventSink { - async fn publish_runtime_event(&self, event: RuntimeEventEnvelope) -> PortResult<()> { - let _ = self.tx.send(event); - Ok(()) - } -} - -#[derive(Clone)] -pub(crate) struct CliRuntimeServicesProvider { - workspace_root: PathBuf, - filesystem: Arc, - workspace: Arc, - events: Arc, - clock: Arc, -} - -impl fmt::Debug for CliRuntimeServicesProvider { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - formatter - .debug_struct("CliRuntimeServicesProvider") - .field("workspace_root", &self.workspace_root) - .finish_non_exhaustive() - } -} - -impl CliRuntimeServicesProvider { - pub(crate) fn new( - workspace_root: impl AsRef, - events: Arc, - clock: Arc, - ) -> anyhow::Result { - let requested_root = workspace_root.as_ref(); - let canonical_root = dunce::canonicalize(requested_root).map_err(|error| { - anyhow::anyhow!( - "workspace root is not available ({}): {error}", - requested_root.display() - ) - })?; - if !canonical_root.is_dir() { - anyhow::bail!( - "workspace root is not a directory: {}", - canonical_root.display() - ); - } - - Ok(Self { - workspace_root: canonical_root.clone(), - filesystem: Arc::new(CliFileSystemService { - workspace_root: canonical_root.clone(), - }), - workspace: Arc::new(CliWorkspaceService { - workspace_root: canonical_root, - }), - events, - clock, - }) - } - - pub(crate) fn workspace_root(&self) -> &Path { - &self.workspace_root - } - - pub(crate) fn build(&self) -> Result { - RuntimeServicesRegistry::new() - .with_provider(CoreRuntimeServicesProvider::new()) - .with_provider(self.clone()) - .build(RuntimeServicesBuilder::new()) - } -} - -impl RuntimeServicesProvider for CliRuntimeServicesProvider { - fn register(&self, builder: RuntimeServicesBuilder) -> RuntimeServicesBuilder { - debug_assert_eq!(self.filesystem.workspace_root(), self.workspace_root); - debug_assert_eq!(self.workspace.workspace_root(), self.workspace_root); - let filesystem: Arc = self.filesystem.clone(); - let workspace: Arc = self.workspace.clone(); - builder - .with_filesystem(filesystem) - .with_workspace(workspace) - .with_events(self.events.clone()) - .with_clock(self.clock.clone()) - } -} - -#[cfg(test)] -mod tests { - use std::sync::Arc; - - use bitfun_runtime_ports::{ - AgentSubmissionSource, RuntimeEventEnvelope, RuntimeEventType, RuntimeServiceCapability, - }; - - use super::{CliClock, CliRuntimeEventSink, CliRuntimeServicesProvider}; - - #[tokio::test] - async fn provider_registers_required_capability_contracts() { - let workspace = tempfile::tempdir().expect("workspace"); - let events = Arc::new(CliRuntimeEventSink::new(8)); - let provider = - CliRuntimeServicesProvider::new(workspace.path(), events.clone(), Arc::new(CliClock)) - .expect("provider"); - - let services = provider.build().expect("runtime services"); - - assert_eq!( - provider.workspace_root(), - dunce::canonicalize(workspace.path()).expect("canonical workspace") - ); - for capability in [ - RuntimeServiceCapability::FileSystem, - RuntimeServiceCapability::Workspace, - RuntimeServiceCapability::SessionStore, - RuntimeServiceCapability::Events, - RuntimeServiceCapability::Clock, - RuntimeServiceCapability::Terminal, - RuntimeServiceCapability::Network, - RuntimeServiceCapability::Git, - ] { - assert!( - services.has_capability(capability), - "missing runtime capability registration {capability}" - ); - } - assert!(services.clock.now_unix_millis() > 0); - - let mut receiver = events.subscribe(); - let envelope = RuntimeEventEnvelope { - session_id: "session-1".to_string(), - turn_id: Some("turn-1".to_string()), - source: Some(AgentSubmissionSource::Cli), - event_type: RuntimeEventType::TurnStarted, - payload: serde_json::json!({ "ready": true }), - }; - services - .events - .publish_runtime_event(envelope.clone()) - .await - .expect("publish runtime event"); - assert_eq!(receiver.recv().await.expect("runtime event"), envelope); - } - - #[test] - fn provider_rejects_a_missing_workspace_root() { - let temp = tempfile::tempdir().expect("tempdir"); - let missing = temp.path().join("missing"); - - let error = CliRuntimeServicesProvider::new( - &missing, - Arc::new(CliRuntimeEventSink::new(8)), - Arc::new(CliClock), - ) - .expect_err("missing workspace must fail"); - - assert!(error.to_string().contains("workspace"), "{error}"); - } -} diff --git a/src/apps/sdk-host/Cargo.toml b/src/apps/sdk-host/Cargo.toml new file mode 100644 index 0000000000..f72eca53c7 --- /dev/null +++ b/src/apps/sdk-host/Cargo.toml @@ -0,0 +1,31 @@ +[package] +name = "bitfun-sdk-host-app" +version.workspace = true +authors.workspace = true +edition.workspace = true +description = "Standalone local Agent SDK Host implementation candidate" +autobins = false + +[[bin]] +name = "bitfun-sdk-host" +path = "src/main.rs" + +[dependencies] +anyhow = { workspace = true } +async-trait = { workspace = true } +bitfun-agent-runtime = { path = "../../crates/execution/agent-runtime" } +bitfun-core = { path = "../../crates/assembly/core", default-features = false, features = ["product-full"] } +bitfun-sdk-host = { path = "../../crates/interfaces/sdk-host" } +futures-util = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true } +tokio-util = { workspace = true, features = ["codec"] } +tracing = { workspace = true } +tracing-subscriber = { workspace = true } + +[dev-dependencies] +rustls = { workspace = true } +tempfile = "3" + +[lints] +workspace = true diff --git a/src/apps/sdk-host/src/lib.rs b/src/apps/sdk-host/src/lib.rs new file mode 100644 index 0000000000..b8a32c20fa --- /dev/null +++ b/src/apps/sdk-host/src/lib.rs @@ -0,0 +1,26 @@ +//! Process-level bootstrap shared by the standalone SDK Host entrypoint and tests. + +pub mod transport; + +/// Stack size used by the SDK Host worker. +/// +/// The Host initializes the same full Agent Runtime as the CLI and preserves +/// the reviewed Windows stack-overflow protection used by that runtime. +pub const SDK_HOST_WORKER_STACK_BYTES: usize = 16 * 1024 * 1024; + +/// Installs process-global prerequisites before any TLS-capable service starts. +pub fn initialize_process_runtime() { + bitfun_core::service::remote_connect::ensure_rustls_crypto_provider(); +} + +/// Spawns the SDK Host runtime on the reviewed worker-stack boundary. +pub fn spawn_sdk_host_worker(task: F) -> std::io::Result> +where + T: Send + 'static, + F: FnOnce() -> T + Send + 'static, +{ + std::thread::Builder::new() + .name("bitfun-sdk-host".to_string()) + .stack_size(SDK_HOST_WORKER_STACK_BYTES) + .spawn(task) +} diff --git a/src/apps/sdk-host/src/main.rs b/src/apps/sdk-host/src/main.rs new file mode 100644 index 0000000000..4768ad5692 --- /dev/null +++ b/src/apps/sdk-host/src/main.rs @@ -0,0 +1,58 @@ +mod runtime; + +use anyhow::{Context, Result}; + +async fn run_host() -> Result<()> { + tracing_subscriber::fmt() + .with_max_level(tracing::Level::WARN) + .with_writer(std::io::stderr) + .with_ansi(false) + .with_target(false) + .init(); + + let workspace_root = std::env::current_dir().context("Failed to resolve SDK workspace")?; + runtime::SdkHostRuntime::select_process_profile()?; + runtime::initialize_terminal_service().await; + + bitfun_core::service::config::initialize_global_config() + .await + .context("Failed to initialize global config service")?; + bitfun_core::infrastructure::ai::AIClientFactory::initialize_global() + .await + .context("Failed to initialize global AI client factory")?; + + let host = runtime::SdkHostRuntime::build(&workspace_root) + .await + .context("Failed to assemble Agent SDK Host")?; + bitfun_sdk_host_app::transport::serve_stdio( + host.agent_runtime().clone(), + host.workspace_root().to_string_lossy().into_owned(), + ) + .await + .context("Agent SDK Host transport failed") +} + +fn main() { + bitfun_sdk_host_app::initialize_process_runtime(); + + let worker = bitfun_sdk_host_app::spawn_sdk_host_worker(|| { + let runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .expect("failed to build SDK Host Tokio runtime"); + runtime.block_on(run_host()) + }) + .expect("failed to spawn SDK Host worker thread"); + + match worker.join() { + Ok(Ok(())) => {} + Ok(Err(error)) => { + eprintln!("Error: {error:#}"); + std::process::exit(1); + } + Err(_) => { + eprintln!("Error: SDK Host worker thread panicked"); + std::process::exit(1); + } + } +} diff --git a/src/apps/sdk-host/src/runtime.rs b/src/apps/sdk-host/src/runtime.rs new file mode 100644 index 0000000000..29f0d459da --- /dev/null +++ b/src/apps/sdk-host/src/runtime.rs @@ -0,0 +1,112 @@ +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result}; +use bitfun_agent_runtime::sdk::AgentRuntime; +use bitfun_core::agentic::system::{self, AgenticSystem}; +use bitfun_core::product_assembly::{DeliveryProfile, ProductAssembler, ProductAssemblyInput}; +use bitfun_core::product_runtime::{ + build_local_runtime_services, ensure_product_dialog_scheduler, CoreProductAgentEventSource, + CoreProductAgentRuntime, CoreRuntimeServicesProvider, +}; + +const RUNTIME_EVENT_BUFFER: usize = 256; +const DELIVERY_PROFILE: DeliveryProfile = DeliveryProfile::Sdk; + +pub(crate) struct SdkHostRuntime { + workspace_root: PathBuf, + agent_runtime: AgentRuntime, + _agent_events: CoreProductAgentEventSource, +} + +impl SdkHostRuntime { + pub(crate) fn select_process_profile() -> Result<()> { + system::select_agentic_system_profile(DELIVERY_PROFILE) + .context("Failed to select SDK Host delivery profile") + } + + pub(crate) async fn build(workspace_root: impl AsRef) -> Result { + let (workspace_root, services) = + build_local_runtime_services(workspace_root, RUNTIME_EVENT_BUFFER)?; + + // The SDK Host keeps its own product identity. The SDK and CLI profiles + // currently select the same assembly-plan ceiling from shared facts. + // The Host's effective wire capability set remains a strict subset. + let parts = ProductAssembler::new() + .assemble(ProductAssemblyInput::new(DELIVERY_PROFILE, services)) + .context("Failed to assemble SDK Host product runtime")?; + let agentic_system = system::init_agentic_system_for_profile(parts.plan().profile()) + .await + .context("Failed to initialize agentic system")?; + bind_core_execution_ports(&agentic_system); + let scheduler = ensure_product_dialog_scheduler(&agentic_system); + let (services, harness_registry, _disabled_plugin_runtime) = parts.into_runtime_parts(); + let agent_events = CoreProductAgentEventSource::new(agentic_system.event_queue.clone()); + let agent_runtime = CoreProductAgentRuntime::build_sdk_host( + agentic_system.coordinator, + scheduler, + agentic_system.token_usage_service, + agent_events.runtime_source(), + services, + harness_registry, + ) + .map_err(anyhow::Error::msg) + .context("Failed to build Agent SDK runtime")?; + + Ok(Self { + workspace_root, + agent_runtime, + _agent_events: agent_events, + }) + } + + pub(crate) fn workspace_root(&self) -> &Path { + &self.workspace_root + } + + pub(crate) fn agent_runtime(&self) -> &AgentRuntime { + &self.agent_runtime + } +} + +fn bind_core_execution_ports(agentic_system: &AgenticSystem) { + agentic_system + .coordinator + .set_terminal_port(CoreRuntimeServicesProvider::terminal_port()); + agentic_system + .coordinator + .set_remote_exec_port(CoreRuntimeServicesProvider::remote_exec_port()); +} + +pub(crate) async fn initialize_terminal_service() { + use bitfun_core::infrastructure::try_get_path_manager_arc; + use bitfun_core::service::runtime::RuntimeManager; + use bitfun_core::service::terminal::{TerminalApi, TerminalConfig}; + + let mut config = TerminalConfig::default(); + match try_get_path_manager_arc() { + Ok(path_manager) => { + config.shell_integration.scripts_dir = + Some(path_manager.user_data_dir().join("sdk-host/temp/scripts")); + config.transcript.root_dir = Some(path_manager.user_data_dir().join("terminals")); + } + Err(error) => { + tracing::warn!( + "Failed to configure SDK Host terminal storage; recording is disabled: {}", + error + ); + } + } + + if let Ok(runtime_manager) = RuntimeManager::new() { + let current_path = std::env::var("PATH").ok(); + if let Some(merged_path) = runtime_manager.merged_path_env(current_path.as_deref()) { + config.env.insert("PATH".to_string(), merged_path.clone()); + #[cfg(windows)] + config.env.insert("Path".to_string(), merged_path); + } + } else { + tracing::warn!("Failed to initialize SDK Host terminal runtime PATH"); + } + + let _terminal_api = TerminalApi::new(config).await; +} diff --git a/src/apps/sdk-host/src/transport.rs b/src/apps/sdk-host/src/transport.rs new file mode 100644 index 0000000000..d18e5c0eac --- /dev/null +++ b/src/apps/sdk-host/src/transport.rs @@ -0,0 +1,364 @@ +//! Local newline-delimited JSON-RPC transport for the standalone SDK Host candidate. + +use std::sync::Arc; +use std::time::Duration; + +use bitfun_agent_runtime::sdk::AgentRuntime; +use futures_util::StreamExt; +use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; +use tokio::sync::Mutex; +use tokio::task::JoinSet; +use tokio::time::{timeout, Instant}; +use tokio_util::codec::{FramedRead, LinesCodec, LinesCodecError}; +use tokio_util::sync::CancellationToken; + +use bitfun_sdk_host::host::{ConnectionControl, HostOutput, SdkHostConfig, SdkHostConnection}; +use bitfun_sdk_host::protocol::{ + JsonRpcErrorResponse, JsonRpcRequest, RequestId, METHOD_INITIALIZE, METHOD_QUERY_CANCEL, + METHOD_SESSION_CLOSE, METHOD_SHUTDOWN, +}; + +#[derive(Debug, Clone)] +pub struct SdkHostTransportConfig { + pub max_line_bytes: usize, + pub write_timeout_ms: u64, + pub shutdown_total_timeout_ms: u64, + pub host: SdkHostConfig, +} + +impl Default for SdkHostTransportConfig { + fn default() -> Self { + Self { + max_line_bytes: 1024 * 1024, + write_timeout_ms: 5_000, + shutdown_total_timeout_ms: 10_000, + host: SdkHostConfig::default(), + } + } +} + +struct JsonLineOutput { + writer: Mutex, + failed: CancellationToken, + write_timeout: Duration, +} + +impl JsonLineOutput { + fn new(writer: Writer, write_timeout_ms: u64) -> Self { + Self { + writer: Mutex::new(writer), + failed: CancellationToken::new(), + write_timeout: Duration::from_millis(write_timeout_ms.max(1)), + } + } + + fn failure_token(&self) -> CancellationToken { + self.failed.clone() + } +} + +#[async_trait::async_trait] +impl HostOutput for JsonLineOutput +where + Writer: AsyncWrite + Unpin + Send, +{ + async fn send(&self, value: serde_json::Value) -> Result<(), ()> { + let mut line = serde_json::to_vec(&value).map_err(|_| ())?; + line.push(b'\n'); + let result = timeout(self.write_timeout, async { + let mut writer = self.writer.lock().await; + writer.write_all(&line).await.map_err(|_| ())?; + writer.flush().await.map_err(|_| ()) + }) + .await; + match result { + Ok(Ok(())) => Ok(()), + Ok(Err(())) | Err(_) => { + self.failed.cancel(); + Err(()) + } + } + } +} + +pub async fn serve_streams( + runtime: AgentRuntime, + default_cwd: impl Into, + reader: Reader, + writer: Writer, + config: SdkHostTransportConfig, +) -> Result<(), std::io::Error> +where + Reader: AsyncRead + Unpin, + Writer: AsyncWrite + Unpin + Send + 'static, +{ + let max_in_flight_requests = config.host.max_in_flight_requests.max(1); + let max_in_flight_control_requests = config.host.max_in_flight_control_requests.max(1); + let shutdown_total_timeout = Duration::from_millis(config.shutdown_total_timeout_ms.max(1)); + let output = Arc::new(JsonLineOutput::new(writer, config.write_timeout_ms)); + let output_failed = output.failure_token(); + let connection = + SdkHostConnection::with_output(runtime, default_cwd, output.clone(), config.host); + let mut lines = FramedRead::new( + reader, + LinesCodec::new_with_max_length(config.max_line_bytes), + ); + let mut parse_error_index = 0u64; + let mut data_requests = JoinSet::new(); + let mut control_requests = JoinSet::new(); + + loop { + let line = tokio::select! { + _ = output_failed.cancelled() => { + data_requests.abort_all(); + control_requests.abort_all(); + graceful_shutdown(&connection, shutdown_total_timeout).await; + return Err(std::io::Error::new( + std::io::ErrorKind::BrokenPipe, + "SDK Host output is unavailable", + )); + } + completed = data_requests.join_next(), if !data_requests.is_empty() => { + if let Some(Ok(ConnectionControl::Shutdown)) = completed { + data_requests.abort_all(); + control_requests.abort_all(); + graceful_shutdown(&connection, shutdown_total_timeout).await; + return Ok(()); + } + continue; + } + completed = control_requests.join_next(), if !control_requests.is_empty() => { + if let Some(Ok(ConnectionControl::Shutdown)) = completed { + data_requests.abort_all(); + control_requests.abort_all(); + graceful_shutdown(&connection, shutdown_total_timeout).await; + return Ok(()); + } + continue; + } + line = lines.next() => line, + }; + let Some(line) = line else { + data_requests.abort_all(); + control_requests.abort_all(); + graceful_shutdown(&connection, shutdown_total_timeout).await; + return Ok(()); + }; + let line = match line { + Ok(line) => line, + Err(error) => { + parse_error_index += 1; + let response = JsonRpcErrorResponse::parse_error( + "SDK Host request line is invalid or exceeds the size limit", + format!("parse:{parse_error_index}"), + ); + if let Ok(value) = serde_json::to_value(response) { + let _ = output.send(value).await; + } + if let LinesCodecError::Io(error) = error { + data_requests.abort_all(); + control_requests.abort_all(); + graceful_shutdown(&connection, shutdown_total_timeout).await; + return Err(error); + } + continue; + } + }; + let value = match serde_json::from_str::(&line) { + Ok(value) => value, + Err(_) => { + parse_error_index += 1; + let response = JsonRpcErrorResponse::parse_error( + "SDK Host request is not valid JSON", + format!("parse:{parse_error_index}"), + ); + if let Ok(value) = serde_json::to_value(response) { + let _ = output.send(value).await; + } + continue; + } + }; + let is_notification = !value + .as_object() + .is_some_and(|object| object.contains_key("id")); + let request_id = request_id_from_value(&value); + if is_notification && !is_json_rpc_notification(&value) { + parse_error_index += 1; + let response = JsonRpcErrorResponse::invalid_request( + None, + "SDK Host notification is not a valid JSON-RPC notification", + format!("invalid:{parse_error_index}"), + ); + if let Ok(value) = serde_json::to_value(response) { + let _ = output.send(value).await; + } + continue; + } + let request = match serde_json::from_value::(value) { + Ok(request) => request, + Err(_) => { + parse_error_index += 1; + let response = JsonRpcErrorResponse::invalid_request( + request_id, + "SDK Host request is not a valid JSON-RPC request", + format!("invalid:{parse_error_index}"), + ); + if let Ok(value) = serde_json::to_value(response) { + let _ = output.send(value).await; + } + continue; + } + }; + if request.id.is_none() && !is_notification { + parse_error_index += 1; + let response = JsonRpcErrorResponse::invalid_request( + None, + "SDK Host request id must be a string or integer", + format!("invalid:{parse_error_index}"), + ); + if let Ok(value) = serde_json::to_value(response) { + let _ = output.send(value).await; + } + continue; + } + + if request.method == METHOD_INITIALIZE { + if !connection.is_initialized().await { + drain_request_sets( + &mut data_requests, + &mut control_requests, + Duration::from_secs(5), + ) + .await; + } + connection.handle_request(request).await; + continue; + } + if request.method == METHOD_SHUTDOWN { + if connection.handle_request(request).await == ConnectionControl::Shutdown { + graceful_shutdown_with_requests( + &mut data_requests, + &mut control_requests, + &connection, + shutdown_total_timeout, + ) + .await; + return Ok(()); + } + continue; + } + + let is_control_request = matches!( + request.method.as_str(), + METHOD_QUERY_CANCEL | METHOD_SESSION_CLOSE + ); + let request_set = if is_control_request { + &mut control_requests + } else { + &mut data_requests + }; + let capacity = if is_control_request { + max_in_flight_control_requests + } else { + max_in_flight_requests + }; + if request_set.len() >= capacity { + connection.reject_overloaded(request.id.clone()).await; + continue; + } + let connection = connection.clone(); + request_set.spawn(async move { connection.handle_request(request).await }); + } +} + +fn is_json_rpc_notification(value: &serde_json::Value) -> bool { + let Some(object) = value.as_object() else { + return false; + }; + !object.contains_key("id") + && object.get("jsonrpc").and_then(serde_json::Value::as_str) == Some("2.0") + && object + .get("method") + .and_then(serde_json::Value::as_str) + .is_some_and(|method| !method.is_empty()) + && object + .get("params") + .is_none_or(|params| params.is_object() || params.is_array()) + && object + .keys() + .all(|key| matches!(key.as_str(), "jsonrpc" | "method" | "params")) +} + +fn request_id_from_value(value: &serde_json::Value) -> Option { + match value.get("id")? { + serde_json::Value::String(value) => Some(RequestId::String(value.clone())), + serde_json::Value::Number(value) => value.as_i64().map(RequestId::Number), + _ => None, + } +} + +async fn graceful_shutdown(connection: &SdkHostConnection, total_timeout: Duration) { + if !connection.shutdown_connection_bounded(total_timeout).await { + tracing::warn!("SDK Host connection cleanup completed with residual errors"); + } +} + +async fn graceful_shutdown_with_requests( + data_requests: &mut JoinSet, + control_requests: &mut JoinSet, + connection: &SdkHostConnection, + total_timeout: Duration, +) { + let started_at = Instant::now(); + drain_request_sets( + data_requests, + control_requests, + total_timeout.min(Duration::from_secs(5)) / 2, + ) + .await; + graceful_shutdown( + connection, + total_timeout.saturating_sub(started_at.elapsed()), + ) + .await; +} + +async fn drain_request_sets( + data_requests: &mut JoinSet, + control_requests: &mut JoinSet, + drain_timeout: Duration, +) { + let started_at = Instant::now(); + drain_requests(control_requests, drain_timeout).await; + drain_requests( + data_requests, + drain_timeout.saturating_sub(started_at.elapsed()), + ) + .await; +} + +async fn drain_requests(requests: &mut JoinSet, drain_timeout: Duration) { + if timeout(drain_timeout, async { + while requests.join_next().await.is_some() {} + }) + .await + .is_err() + { + requests.abort_all(); + while requests.join_next().await.is_some() {} + } +} + +pub async fn serve_stdio( + runtime: AgentRuntime, + default_cwd: impl Into, +) -> Result<(), std::io::Error> { + serve_streams( + runtime, + default_cwd, + tokio::io::stdin(), + tokio::io::stdout(), + SdkHostTransportConfig::default(), + ) + .await +} diff --git a/src/apps/sdk-host/tests/process_initialization.rs b/src/apps/sdk-host/tests/process_initialization.rs new file mode 100644 index 0000000000..2cc267da9e --- /dev/null +++ b/src/apps/sdk-host/tests/process_initialization.rs @@ -0,0 +1,30 @@ +#[test] +fn sdk_host_process_installs_a_rustls_crypto_provider() { + bitfun_sdk_host_app::initialize_process_runtime(); + + assert!( + rustls::crypto::CryptoProvider::get_default().is_some(), + "SDK Host must select a process-level crypto provider before HTTPS AI requests" + ); +} + +#[test] +fn sdk_host_process_uses_the_reviewed_worker_stack_contract() { + let caller = std::thread::current().id(); + let worker = bitfun_sdk_host_app::spawn_sdk_host_worker(|| std::thread::current().id()) + .expect("spawn SDK Host worker"); + + assert_eq!( + bitfun_sdk_host_app::SDK_HOST_WORKER_STACK_BYTES, + 16 * 1024 * 1024 + ); + assert_ne!(worker.join().expect("join SDK Host worker"), caller); +} + +#[test] +fn sdk_host_process_keeps_cleanup_warnings_on_stderr() { + let entrypoint = include_str!("../src/main.rs"); + + assert!(entrypoint.contains(".with_max_level(tracing::Level::WARN)")); + assert!(entrypoint.contains(".with_writer(std::io::stderr)")); +} diff --git a/src/apps/sdk-host/tests/stdio_process.rs b/src/apps/sdk-host/tests/stdio_process.rs new file mode 100644 index 0000000000..3709a7a6ab --- /dev/null +++ b/src/apps/sdk-host/tests/stdio_process.rs @@ -0,0 +1,110 @@ +use std::process::Stdio; +use std::time::Duration; + +use serde_json::{json, Value}; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; + +#[tokio::test] +async fn standalone_sdk_host_negotiates_and_shuts_down_without_cli() { + let temp = tempfile::tempdir().expect("isolated SDK Host environment"); + let workspace = temp.path().join("workspace"); + let user_root = temp.path().join("user-root"); + let home_root = temp.path().join("home"); + let config_root = temp.path().join("config-root"); + for path in [&workspace, &user_root, &home_root, &config_root] { + std::fs::create_dir_all(path).expect("SDK Host fixture directory"); + } + + let mut child = tokio::process::Command::new(env!("CARGO_BIN_EXE_bitfun-sdk-host")) + .current_dir(&workspace) + .env_remove("BITFUN_USER_ROOT") + .env_remove("BITFUN_HOME") + .env("BITFUN_E2E_STORAGE_GUARD", "1") + .env("BITFUN_E2E_USER_ROOT", &user_root) + .env("BITFUN_E2E_HOME", &home_root) + .env("APPDATA", &config_root) + .env("XDG_CONFIG_HOME", &config_root) + .env("HOME", &home_root) + .env("USERPROFILE", &home_root) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .expect("start standalone Agent SDK Host"); + + let mut stdin = child.stdin.take().expect("SDK Host stdin"); + let mut stdout = BufReader::new(child.stdout.take().expect("SDK Host stdout")); + let mut stderr = child.stderr.take().expect("SDK Host stderr"); + + send_request( + &mut stdin, + 1, + "initialize", + json!({ + "protocolVersion": 1, + "clientInfo": { "name": "standalone-process-fixture", "version": "0.1.0" }, + "capabilities": { "serverNotifications": true } + }), + ) + .await; + let initialized = read_response(&mut stdout, "initialize").await; + assert_eq!(initialized["id"], 1); + assert_eq!(initialized["result"]["protocolVersion"], 1); + + send_request(&mut stdin, 2, "shutdown", json!({})).await; + let shutdown = read_response(&mut stdout, "shutdown").await; + assert_eq!(shutdown["id"], 2); + assert_eq!(shutdown["result"]["accepted"], true); + drop(stdin); + + let status = tokio::time::timeout(Duration::from_secs(15), child.wait()) + .await + .expect("SDK Host must stop after shutdown") + .expect("wait for SDK Host"); + let mut stderr_output = Vec::new(); + stderr + .read_to_end(&mut stderr_output) + .await + .expect("read SDK Host stderr"); + assert!( + status.success(), + "SDK Host failed: {}", + String::from_utf8_lossy(&stderr_output) + ); +} + +async fn send_request( + stdin: &mut tokio::process::ChildStdin, + id: i64, + method: &str, + params: Value, +) { + let mut line = serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "id": id, + "method": method, + "params": params, + })) + .expect("serialize SDK Host request"); + line.push(b'\n'); + stdin + .write_all(&line) + .await + .expect("write SDK Host request"); + stdin.flush().await.expect("flush SDK Host request"); +} + +async fn read_response( + stdout: &mut BufReader, + operation: &str, +) -> Value { + let mut line = String::new(); + let bytes = tokio::time::timeout(Duration::from_secs(60), stdout.read_line(&mut line)) + .await + .unwrap_or_else(|_| panic!("SDK Host {operation} timed out")) + .expect("read SDK Host stdout"); + assert_ne!(bytes, 0, "SDK Host stdout closed during {operation}"); + serde_json::from_str(&line) + .unwrap_or_else(|error| panic!("SDK Host stdout was not JSON: {error}: {line}")) +} diff --git a/src/apps/sdk-host/tests/stdio_transport.rs b/src/apps/sdk-host/tests/stdio_transport.rs new file mode 100644 index 0000000000..f072c107f5 --- /dev/null +++ b/src/apps/sdk-host/tests/stdio_transport.rs @@ -0,0 +1,710 @@ +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +use async_trait::async_trait; +use bitfun_agent_runtime::sdk::{ + AgentRuntimeBuilder, AgentSessionClosePort, AgentSessionCreateRequest, + AgentSessionCreateResult, AgentSessionDeleteRequest, AgentSessionListRequest, + AgentSessionManagementPort, AgentSessionSummary, AgentSessionWorkspaceBinding, + AgentSessionWorkspaceRequest, AgentSubmissionPort, AgentSubmissionRequest, + AgentSubmissionResult, AgentTransientSessionDiscardRequest, PortResult, +}; +use bitfun_sdk_host_app::transport::{serve_streams, SdkHostTransportConfig}; +use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; +use tokio::sync::Notify; +use tokio::time::{timeout, Duration}; + +struct MinimalOwner; + +struct BlockingCreateOwner { + calls: AtomicUsize, + deleted: AtomicUsize, + release: Notify, +} + +impl BlockingCreateOwner { + fn new() -> Self { + Self { + calls: AtomicUsize::new(0), + deleted: AtomicUsize::new(0), + release: Notify::new(), + } + } +} + +#[async_trait] +impl AgentSubmissionPort for MinimalOwner { + async fn create_session( + &self, + request: AgentSessionCreateRequest, + ) -> PortResult { + Ok(AgentSessionCreateResult { + session_id: "unused".to_string(), + session_name: request.session_name, + agent_type: request.agent_type, + }) + } + + async fn create_session_with_id( + &self, + session_id: String, + request: AgentSessionCreateRequest, + ) -> PortResult { + Ok(AgentSessionCreateResult { + session_id, + session_name: request.session_name, + agent_type: request.agent_type, + }) + } + + async fn create_transient_session_with_id( + &self, + session_id: String, + request: AgentSessionCreateRequest, + ) -> PortResult { + self.create_session_with_id(session_id, request).await + } + + async fn submit_message( + &self, + request: AgentSubmissionRequest, + ) -> PortResult { + Ok(AgentSubmissionResult { + turn_id: request.turn_id.unwrap_or_else(|| "unused".to_string()), + accepted: true, + }) + } + + async fn resolve_session_agent_type(&self, _session_id: &str) -> PortResult> { + Ok(Some("agentic".to_string())) + } +} + +#[async_trait] +impl AgentSessionClosePort for MinimalOwner { + async fn discard_transient_session( + &self, + _request: AgentTransientSessionDiscardRequest, + ) -> PortResult { + Ok(false) + } +} + +#[async_trait] +impl AgentSubmissionPort for BlockingCreateOwner { + async fn create_session( + &self, + request: AgentSessionCreateRequest, + ) -> PortResult { + if self.calls.fetch_add(1, Ordering::AcqRel) == 0 { + self.release.notified().await; + } + Ok(AgentSessionCreateResult { + session_id: "session-blocking".to_string(), + session_name: request.session_name, + agent_type: request.agent_type, + }) + } + + async fn create_session_with_id( + &self, + session_id: String, + request: AgentSessionCreateRequest, + ) -> PortResult { + if self.calls.fetch_add(1, Ordering::AcqRel) == 0 { + self.release.notified().await; + } + Ok(AgentSessionCreateResult { + session_id, + session_name: request.session_name, + agent_type: request.agent_type, + }) + } + + async fn create_transient_session_with_id( + &self, + session_id: String, + request: AgentSessionCreateRequest, + ) -> PortResult { + self.create_session_with_id(session_id, request).await + } + + async fn submit_message( + &self, + request: AgentSubmissionRequest, + ) -> PortResult { + Ok(AgentSubmissionResult { + turn_id: request.turn_id.unwrap_or_else(|| "unused".to_string()), + accepted: true, + }) + } + + async fn resolve_session_agent_type(&self, _session_id: &str) -> PortResult> { + Ok(Some("agentic".to_string())) + } +} + +#[async_trait] +impl AgentSessionClosePort for BlockingCreateOwner { + async fn discard_transient_session( + &self, + _request: AgentTransientSessionDiscardRequest, + ) -> PortResult { + self.deleted.fetch_add(1, Ordering::AcqRel); + Ok(true) + } +} + +#[async_trait] +impl AgentSessionManagementPort for BlockingCreateOwner { + async fn list_sessions( + &self, + _request: AgentSessionListRequest, + ) -> PortResult> { + Ok(Vec::new()) + } + + async fn delete_session(&self, _request: AgentSessionDeleteRequest) -> PortResult<()> { + self.deleted.fetch_add(1, Ordering::AcqRel); + Ok(()) + } + + async fn resolve_session_workspace_binding( + &self, + _request: AgentSessionWorkspaceRequest, + ) -> PortResult> { + Ok(None) + } +} + +#[tokio::test] +async fn stdio_transport_serves_initialize_and_shutdown_without_non_protocol_stdout() { + let owner = Arc::new(MinimalOwner); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_session_close_port(owner) + .build() + .unwrap(); + let (client, server) = tokio::io::duplex(16 * 1024); + let (client_read, mut client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + let task = tokio::spawn(serve_streams( + runtime, + "D:/workspace/project", + server_read, + server_write, + SdkHostTransportConfig::default(), + )); + client_write + .write_all( + concat!( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"shutdown\",\"params\":{}}\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + client_write.shutdown().await.unwrap(); + + let mut lines = BufReader::new(client_read).lines(); + let initialized: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + let shutdown: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + + assert_eq!(initialized["id"], 1); + assert_eq!(initialized["result"]["protocolVersion"], 1); + assert_eq!(shutdown["id"], 2); + assert_eq!(shutdown["result"]["accepted"], true); + assert!(lines.next_line().await.unwrap().is_none()); + task.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn stdio_transport_executes_json_rpc_notifications_without_replying() { + let owner = Arc::new(MinimalOwner); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_session_close_port(owner) + .build() + .unwrap(); + let (client, server) = tokio::io::duplex(16 * 1024); + let (client_read, mut client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + let task = tokio::spawn(serve_streams( + runtime, + "D:/workspace/project", + server_read, + server_write, + SdkHostTransportConfig::default(), + )); + client_write + .write_all( + concat!( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"method\":\"shutdown\",\"params\":{}}\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + client_write.shutdown().await.unwrap(); + + let mut lines = BufReader::new(client_read).lines(); + let initialized: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(initialized["id"], 1); + timeout(Duration::from_secs(1), task) + .await + .expect("shutdown notification should execute") + .unwrap() + .unwrap(); + assert!(lines.next_line().await.unwrap().is_none()); +} + +#[tokio::test] +async fn malformed_and_oversized_lines_fail_closed_with_standard_parse_errors() { + let owner = Arc::new(MinimalOwner); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_session_close_port(owner) + .build() + .unwrap(); + let (client, server) = tokio::io::duplex(16 * 1024); + let (client_read, mut client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + let task = tokio::spawn(serve_streams( + runtime, + "D:/workspace/project", + server_read, + server_write, + SdkHostTransportConfig { + max_line_bytes: 128, + ..SdkHostTransportConfig::default() + }, + )); + client_write.write_all(b"not-json\n").await.unwrap(); + client_write + .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":7}\n") + .await + .unwrap(); + client_write + .write_all(format!("{}\n", "x".repeat(256)).as_bytes()) + .await + .unwrap(); + client_write.shutdown().await.unwrap(); + + let mut lines = BufReader::new(client_read).lines(); + let malformed: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + let invalid_request: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + let oversized: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(malformed["id"], serde_json::Value::Null); + assert_eq!(malformed["error"]["code"], -32700); + assert_eq!(malformed["error"]["data"]["code"], "invalid_request"); + assert_eq!(invalid_request["id"], 7); + assert_eq!(invalid_request["error"]["code"], -32600); + assert_eq!(invalid_request["error"]["data"]["code"], "invalid_request"); + assert_eq!(oversized["error"]["code"], -32700); + task.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn transport_accepts_input_while_an_owner_call_is_pending_and_bounds_requests() { + let owner = Arc::new(BlockingCreateOwner::new()); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .build() + .unwrap(); + let (client, server) = tokio::io::duplex(16 * 1024); + let (client_read, mut client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + let task = tokio::spawn(serve_streams( + runtime, + "D:/workspace/project", + server_read, + server_write, + SdkHostTransportConfig { + host: bitfun_sdk_host::host::SdkHostConfig { + max_in_flight_requests: 1, + ..Default::default() + }, + ..Default::default() + }, + )); + client_write + .write_all( + concat!( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"session/create\",\"params\":{}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"session/create\",\"params\":{}}\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + + let mut lines = BufReader::new(client_read).lines(); + let initialized: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(initialized["id"], 1); + let overloaded: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(overloaded["id"], 3); + assert_eq!(overloaded["error"]["data"]["code"], "overloaded"); + + owner.release.notify_one(); + let created: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(created["id"], 2); + client_write + .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":4,\"method\":\"session/create\",\"params\":{}}\n") + .await + .unwrap(); + let recovered: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(recovered["id"], 4); + assert!(recovered["result"]["sessionId"].is_string()); + client_write + .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":5,\"method\":\"shutdown\",\"params\":{}}\n") + .await + .unwrap(); + client_write.shutdown().await.unwrap(); + let shutdown: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(shutdown["id"], 5); + task.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn shutdown_remains_available_when_the_data_request_budget_is_exhausted() { + let owner = Arc::new(BlockingCreateOwner::new()); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .build() + .unwrap(); + let (client, server) = tokio::io::duplex(16 * 1024); + let (client_read, mut client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + let task = tokio::spawn(serve_streams( + runtime, + "D:/workspace/project", + server_read, + server_write, + SdkHostTransportConfig { + shutdown_total_timeout_ms: 500, + host: bitfun_sdk_host::host::SdkHostConfig { + max_in_flight_requests: 1, + ..Default::default() + }, + ..Default::default() + }, + )); + client_write + .write_all( + b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + ) + .await + .unwrap(); + let mut lines = BufReader::new(client_read).lines(); + let initialized: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(initialized["id"], 1); + + client_write + .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"session/create\",\"params\":{}}\n") + .await + .unwrap(); + timeout(Duration::from_secs(1), async { + while owner.calls.load(Ordering::Acquire) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("blocking data request must start"); + + client_write + .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"shutdown\",\"params\":{}}\n") + .await + .unwrap(); + client_write.shutdown().await.unwrap(); + + let shutdown: serde_json::Value = serde_json::from_str( + &timeout(Duration::from_secs(1), lines.next_line()) + .await + .expect("shutdown must bypass exhausted data request capacity") + .unwrap() + .unwrap(), + ) + .unwrap(); + assert_eq!(shutdown["id"], 3); + assert_eq!(shutdown["result"]["accepted"], true); + task.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn duplicate_initialize_does_not_abort_an_in_flight_request() { + let owner = Arc::new(BlockingCreateOwner::new()); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .build() + .unwrap(); + let (client, server) = tokio::io::duplex(16 * 1024); + let (client_read, mut client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + let task = tokio::spawn(serve_streams( + runtime, + "D:/workspace/project", + server_read, + server_write, + SdkHostTransportConfig::default(), + )); + client_write + .write_all( + concat!( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"session/create\",\"params\":{}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + + let mut lines = BufReader::new(client_read).lines(); + let initialized: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(initialized["id"], 1); + let duplicate: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(duplicate["id"], 3); + assert_eq!(duplicate["error"]["data"]["code"], "already_initialized"); + + owner.release.notify_one(); + let created: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(created["id"], 2); + client_write + .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":4,\"method\":\"shutdown\",\"params\":{}}\n") + .await + .unwrap(); + client_write.shutdown().await.unwrap(); + let shutdown: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(shutdown["id"], 4); + task.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn connection_eof_cleans_a_session_created_after_its_request_is_aborted() { + let owner = Arc::new(BlockingCreateOwner::new()); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .build() + .unwrap(); + let (client, server) = tokio::io::duplex(16 * 1024); + let (client_read, mut client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + let mut task = tokio::spawn(serve_streams( + runtime, + "D:/workspace/project", + server_read, + server_write, + SdkHostTransportConfig { + shutdown_total_timeout_ms: 100, + ..SdkHostTransportConfig::default() + }, + )); + client_write + .write_all( + concat!( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"session/create\",\"params\":{}}\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + let mut lines = BufReader::new(client_read).lines(); + let initialized: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(initialized["id"], 1); + timeout(Duration::from_secs(1), async { + while owner.calls.load(Ordering::Acquire) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("session creation should start"); + + client_write.shutdown().await.unwrap(); + timeout(Duration::from_secs(1), &mut task) + .await + .expect("transient Session cleanup must stay within the Host deadline") + .unwrap() + .unwrap(); + assert_eq!(owner.deleted.load(Ordering::Acquire), 1); + assert!(lines.next_line().await.unwrap().is_none()); +} + +#[tokio::test] +async fn explicit_shutdown_bounds_request_drain_and_transient_cleanup_together() { + let owner = Arc::new(BlockingCreateOwner::new()); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .build() + .unwrap(); + let (client, server) = tokio::io::duplex(16 * 1024); + let (client_read, mut client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + let mut task = tokio::spawn(serve_streams( + runtime, + "D:/workspace/project", + server_read, + server_write, + SdkHostTransportConfig { + shutdown_total_timeout_ms: 100, + ..SdkHostTransportConfig::default() + }, + )); + client_write + .write_all( + concat!( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"session/create\",\"params\":{}}\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + let mut lines = BufReader::new(client_read).lines(); + let initialized: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(initialized["id"], 1); + timeout(Duration::from_secs(1), async { + while owner.calls.load(Ordering::Acquire) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("session creation should start"); + + client_write + .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"shutdown\",\"params\":{}}\n") + .await + .unwrap(); + let shutdown: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(shutdown["id"], 3); + timeout(Duration::from_secs(1), &mut task) + .await + .expect("request drain and transient cleanup must share one total deadline") + .unwrap() + .unwrap(); + assert_eq!(owner.deleted.load(Ordering::Acquire), 1); +} + +#[tokio::test] +async fn requests_before_a_successful_initialize_cannot_cross_the_handshake() { + let owner = Arc::new(MinimalOwner); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_session_close_port(owner) + .build() + .unwrap(); + let (client, server) = tokio::io::duplex(16 * 1024); + let (client_read, mut client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + let task = tokio::spawn(serve_streams( + runtime, + "D:/workspace/project", + server_read, + server_write, + SdkHostTransportConfig::default(), + )); + client_write + .write_all( + concat!( + "{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":999,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"session/create\",\"params\":{}}\n", + "{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n" + ) + .as_bytes(), + ) + .await + .unwrap(); + + let mut lines = BufReader::new(client_read).lines(); + let mismatch: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + let pre_initialize: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + let initialized: serde_json::Value = + serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(mismatch["id"], 1); + assert_eq!(mismatch["error"]["data"]["code"], "version_mismatch"); + assert_eq!(pre_initialize["id"], 2); + assert_eq!(pre_initialize["error"]["data"]["code"], "not_initialized"); + assert_eq!(initialized["id"], 3); + assert_eq!(initialized["result"]["protocolVersion"], 1); + + client_write + .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":4,\"method\":\"shutdown\",\"params\":{}}\n") + .await + .unwrap(); + client_write.shutdown().await.unwrap(); + assert_eq!( + serde_json::from_str::(&lines.next_line().await.unwrap().unwrap()) + .unwrap()["id"], + 4 + ); + task.await.unwrap().unwrap(); +} + +#[tokio::test] +async fn blocked_output_times_out_and_ends_the_connection() { + let owner = Arc::new(MinimalOwner); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_session_close_port(owner) + .build() + .unwrap(); + let (client, server) = tokio::io::duplex(64); + let (_client_read, mut client_write) = tokio::io::split(client); + let (server_read, server_write) = tokio::io::split(server); + let task = tokio::spawn(serve_streams( + runtime, + "D:/workspace/project", + server_read, + server_write, + SdkHostTransportConfig { + write_timeout_ms: 20, + ..Default::default() + }, + )); + client_write + .write_all( + b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"initialize\",\"params\":{\"protocolVersion\":1,\"clientInfo\":{\"name\":\"fixture\",\"version\":\"0.1\"},\"capabilities\":{\"serverNotifications\":true}}}\n", + ) + .await + .unwrap(); + + let result = timeout(Duration::from_secs(1), task) + .await + .expect("blocked SDK Host output must have a deadline") + .unwrap(); + assert_eq!(result.unwrap_err().kind(), std::io::ErrorKind::BrokenPipe); +} diff --git a/src/crates/assembly/core/src/agentic/agents/definitions/modes/claw.rs b/src/crates/assembly/core/src/agentic/agents/definitions/modes/claw.rs index 3ea2dd6ff4..5ea6cf9cd5 100644 --- a/src/crates/assembly/core/src/agentic/agents/definitions/modes/claw.rs +++ b/src/crates/assembly/core/src/agentic/agents/definitions/modes/claw.rs @@ -91,7 +91,7 @@ impl Agent for ClawMode { #[cfg(test)] mod tests { use super::ClawMode; - use crate::agentic::agents::Agent; + use crate::agentic::agents::{Agent, PromptBuilderContext}; use bitfun_agent_runtime::prompt::UserContextSection; #[test] @@ -107,4 +107,15 @@ mod tests { .user_context_policy() .includes(UserContextSection::MemorySummary)); } + + #[tokio::test] + async fn claw_prompt_conditions_optional_control_and_session_tools() { + let prompt = ClawMode::new() + .get_system_prompt(Some(&PromptBuilderContext::new("/workspace", None, None))) + .await + .expect("Claw prompt"); + + assert!(prompt.contains("only when it appears in your current tool list")); + assert!(prompt.contains("only when both tools appear in your current tool list")); + } } diff --git a/src/crates/assembly/core/src/agentic/agents/definitions/subagents/computer_use.rs b/src/crates/assembly/core/src/agentic/agents/definitions/subagents/computer_use.rs index 57a919ba30..fa9d2211e6 100644 --- a/src/crates/assembly/core/src/agentic/agents/definitions/subagents/computer_use.rs +++ b/src/crates/assembly/core/src/agentic/agents/definitions/subagents/computer_use.rs @@ -84,6 +84,7 @@ impl Agent for ComputerUseMode { #[cfg(test)] mod tests { use super::{Agent, ComputerUseMode}; + use crate::agentic::agents::PromptBuilderContext; #[test] fn computer_use_mode_basics() { @@ -96,4 +97,15 @@ mod tests { assert!(!agent.default_tools().contains(&"Write".to_string())); assert!(!agent.is_readonly()); } + + #[tokio::test] + async fn computer_use_prompt_conditions_optional_browser_control() { + let prompt = ComputerUseMode::new() + .get_system_prompt(Some(&PromptBuilderContext::new("/workspace", None, None))) + .await + .expect("ComputerUse prompt"); + + assert!(prompt.contains("When `ControlHub` appears in your current tool list")); + assert!(prompt.contains("If `ControlHub` is unavailable")); + } } diff --git a/src/crates/assembly/core/src/agentic/agents/prompts/claw_mode.md b/src/crates/assembly/core/src/agentic/agents/prompts/claw_mode.md index e25f05eda0..5794d5c4d6 100644 --- a/src/crates/assembly/core/src/agentic/agents/prompts/claw_mode.md +++ b/src/crates/assembly/core/src/agentic/agents/prompts/claw_mode.md @@ -14,7 +14,7 @@ When a first-class tool exists for an action, use the tool directly instead of a # Control Boundaries -Use `ControlHub` for browser automation, terminal signalling, and routing/capability introspection: +Use `ControlHub` for browser automation, terminal signalling, and routing/capability introspection only when it appears in your current tool list: - `domain: "browser"` for websites and web apps in the user's real browser through CDP. - `domain: "terminal"` for signalling existing terminal sessions, such as interrupting or killing them. @@ -22,13 +22,13 @@ Use `ControlHub` for browser automation, terminal signalling, and routing/capabi Do not use `ControlHub` for local computer, operating-system, or desktop UI work. Desktop and system actions have moved to the dedicated `ComputerUse` tool/agent. This includes screenshots, OCR, mouse, keyboard, app state, app launching, opening files or URLs through the OS, clipboard access, OS facts, and local scripts. -If the user asks you to operate or inspect the local computer, delegate the task to a `ComputerUse` session via SessionControl/SessionMessage when available. Include the user's goal, target app/window/site, safety constraints, and expected verification in the handoff. If delegation is unavailable, explain that the task needs the Computer Use mode. +If the user asks you to operate or inspect the local computer, delegate the task to a `ComputerUse` session via SessionControl/SessionMessage only when both tools appear in your current tool list. Include the user's goal, target app/window/site, safety constraints, and expected verification in the handoff. If delegation is unavailable, explain that the task needs the Computer Use mode. # Session Coordination -For complex coding tasks or office-style multi-step tasks, prefer multi-session coordination over doing everything in the current session. +For complex coding tasks or office-style multi-step tasks, prefer multi-session coordination when the required session tools are available. Otherwise, keep ownership in the current session and use listed `Task` subagents where useful. -Use `SessionControl` to list, reuse, create, and delete sessions. Use `SessionMessage` to hand off a self-contained subtask to another session. +Use `SessionControl` to list, reuse, create, and delete sessions, and `SessionMessage` to hand off a self-contained subtask, only when both tools appear in your current tool list. Never attempt an unavailable tool just because this template describes it. Use this pattern when: diff --git a/src/crates/assembly/core/src/agentic/agents/prompts/computer_use_mode.md b/src/crates/assembly/core/src/agentic/agents/prompts/computer_use_mode.md index 4a3de401f0..950d9f7ab1 100644 --- a/src/crates/assembly/core/src/agentic/agents/prompts/computer_use_mode.md +++ b/src/crates/assembly/core/src/agentic/agents/prompts/computer_use_mode.md @@ -18,10 +18,10 @@ Work in a tight observe -> act -> verify loop. Before acting on a desktop UI, ob Prefer the smallest reliable control surface: -1. Use `ControlHub` with `domain: "browser"` for websites and web apps in the user's real browser. +1. When `ControlHub` appears in your current tool list, use it with `domain: "browser"` for websites and web apps in the user's real browser. 2. Use `ComputerUse` for third-party desktop apps, OS dialogs, system-wide keyboard and mouse, accessibility, OCR, screenshots, app state, app/file/url opening, clipboard access, OS facts, and local scripts. 3. Use `ExecCommand` for local shell commands when that is the clearest path and does not bypass desktop safety expectations. -4. Use `ControlHub` with `domain: "meta"` to inspect non-desktop control capabilities before long or uncertain automation flows. +4. When available, use `ControlHub` with `domain: "meta"` to inspect non-desktop control capabilities before long or uncertain automation flows. Prefer script or command-line automation when it is clearly safer and reversible, but run it step by step. Do not hide a whole GUI workflow in one large script. For GUI work, prefer keyboard shortcuts and accessibility-backed targets before mouse coordinates. @@ -67,7 +67,7 @@ When Runtime Context indicates the primary model does not support image understa # Browser Work -For websites and web apps, prefer `ControlHub` with `domain: "browser"` so cookies, login state, and extensions are preserved. Do not drive browser content through desktop screenshots when browser-domain controls are available. +For websites and web apps, prefer `ControlHub` with `domain: "browser"` when it is available so cookies, login state, and extensions are preserved. If `ControlHub` is unavailable, do not claim browser-domain automation; use `ComputerUse` only for browser chrome or OS-level interaction that it can actually observe and verify. Use desktop-domain controls only for browser chrome, OS dialogs, permission prompts, file pickers, or when browser-domain capabilities are unavailable. diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index f0a6dfd2f2..eb2d707f13 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -200,6 +200,44 @@ fn metadata_bool(metadata: Option<&serde_json::Value>, key: &str) -> Option ToolRuntimeRestrictions { + if !transient { + return restrictions; + } + + for (tool_name, message) in [ + ( + "SessionControl", + "SessionControl is unavailable in connection-scoped transient Sessions.", + ), + ( + "SessionMessage", + "SessionMessage is unavailable in connection-scoped transient Sessions.", + ), + ( + "SessionHistory", + "SessionHistory is unavailable in connection-scoped transient Sessions.", + ), + ( + "Cron", + "Cron is unavailable in connection-scoped transient Sessions.", + ), + ( + "ControlHub", + "ControlHub is unavailable in connection-scoped transient Sessions.", + ), + ] { + restrictions.denied_tool_names.insert(tool_name.to_string()); + restrictions + .denied_tool_messages + .insert(tool_name.to_string(), message.to_string()); + } + restrictions +} + /// Subagent execution result /// /// Contains the text response after subagent execution @@ -455,8 +493,12 @@ pub(crate) struct HiddenSubagentExecutionRequest { runtime_tool_restrictions: ToolRuntimeRestrictions, prompt_cache_source_session_id: Option, session_kind: SessionKind, + transient: bool, emit_lifecycle_events: bool, prepared_session_created: bool, + /// Keeps scheduler maintenance fenced from the moment a hidden Session is + /// prepared until the final execution/cleanup owner releases every clone. + execution_lease: Option>, external_generation_lease: Option, } @@ -537,6 +579,17 @@ struct CancelTokenGuard { dialog_turn_id: String, } +#[derive(Debug)] +struct SessionExecutionLease { + active_counter: Arc, +} + +impl Drop for SessionExecutionLease { + fn drop(&mut self) { + self.active_counter.fetch_sub(1, Ordering::SeqCst); + } +} + impl Drop for CancelTokenGuard { fn drop(&mut self) { let execution_engine = self.execution_engine.clone(); @@ -1649,6 +1702,27 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet /// /// Delegated subagent sessions use the hidden-subagent creation path instead. pub async fn create_session_with_workspace_and_creator( + &self, + session_id: Option, + session_name: String, + agent_type: String, + config: SessionConfig, + workspace_path: String, + created_by: Option, + ) -> BitFunResult { + self.create_session_with_workspace_and_creator_internal( + session_id, + session_name, + agent_type, + config, + workspace_path, + created_by, + false, + ) + .await + } + + async fn create_session_with_workspace_and_creator_internal( &self, session_id: Option, session_name: String, @@ -1656,6 +1730,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet mut config: SessionConfig, workspace_path: String, created_by: Option, + transient: bool, ) -> BitFunResult { // Persist the workspace binding inside the session config so execution can // consistently restore the correct workspace regardless of the entry point. @@ -1664,19 +1739,33 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet let defaults = Self::agent_model_defaults().await; snapshot_normal_session_model(&mut config, &defaults); let agent_type = Self::normalize_agent_type(&agent_type); - let session = self - .session_manager - .create_session_with_id_and_creator( - session_id, - session_name, - agent_type, - config, - created_by, - ) - .await?; + let session = if transient { + self.session_manager + .create_transient_session_with_id_and_details( + session_id, + session_name, + agent_type, + config, + created_by, + SessionKind::Standard, + ) + .await? + } else { + self.session_manager + .create_session_with_id_and_creator( + session_id, + session_name, + agent_type, + config, + created_by, + ) + .await? + }; - Self::track_session_workspace_activity_best_effort(&session.config, "session_created") - .await; + if !transient { + Self::track_session_workspace_activity_best_effort(&session.config, "session_created") + .await; + } // SessionManager::create_session_with_id_and_creator already persists the // session into the effective workspace session storage path. Avoid writing @@ -2172,16 +2261,51 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet created_by: Option, kind: SessionKind, ) -> BitFunResult { - self.session_manager - .create_session_with_id_and_details( - session_id, - session_name, - agent_type, - config, - created_by, - kind, - ) - .await + self.create_hidden_agent_session_with_durability( + session_id, + session_name, + agent_type, + config, + created_by, + kind, + false, + ) + .await + } + + async fn create_hidden_agent_session_with_durability( + &self, + session_id: Option, + session_name: String, + agent_type: String, + config: SessionConfig, + created_by: Option, + kind: SessionKind, + transient: bool, + ) -> BitFunResult { + if transient { + self.session_manager + .create_transient_session_with_id_and_details( + session_id, + session_name, + agent_type, + config, + created_by, + kind, + ) + .await + } else { + self.session_manager + .create_session_with_id_and_details( + session_id, + session_name, + agent_type, + config, + created_by, + kind, + ) + .await + } } async fn load_session_context_messages(&self, session: &Session) -> BitFunResult> { @@ -2236,6 +2360,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet workspace_services: Option<&WorkspaceServices>, enable_tools: bool, skill_agent_context_vars: &HashMap, + runtime_tool_restrictions: &ToolRuntimeRestrictions, ) -> BitFunResult { let agent_registry = get_agent_registry(); agent_registry @@ -2259,6 +2384,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet workspace_services, enable_tools, skill_agent_context_vars, + runtime_tool_restrictions, ) .await; @@ -3647,6 +3773,18 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet { skill_agent_context_vars.insert("acp_transport".to_string(), "true".to_string()); } + let runtime_tool_restrictions = if is_miniapp_headless_agent_run( + user_message_metadata.as_ref(), + session.created_by.as_deref(), + ) { + miniapp_headless_agent_tool_restrictions() + } else { + ToolRuntimeRestrictions::default() + }; + let runtime_tool_restrictions = runtime_tool_restrictions_for_session_lifetime( + runtime_tool_restrictions, + self.session_manager.is_transient_session(&session_id), + ); // Materialize references only when a queued turn is actually being // dispatched. The agent receives local artifact URIs, never a path to @@ -3673,6 +3811,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet workspace_services.as_ref(), session.config.enable_tools, &skill_agent_context_vars, + &runtime_tool_restrictions, ) .await?; let effective_user_input = wrapped_user_input_payload.content.clone(); @@ -3925,14 +4064,6 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .as_ref() .map(|workspace| workspace.session_storage_dir().to_path_buf()); - let runtime_tool_restrictions = if is_miniapp_headless_agent_run( - user_message_metadata.as_ref(), - session.created_by.as_deref(), - ) { - miniapp_headless_agent_tool_restrictions() - } else { - ToolRuntimeRestrictions::default() - }; let persisted_subagent_context = self .load_persisted_subagent_continuation_context(&session) .await; @@ -4186,6 +4317,16 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet } } + fn register_session_execution(&self, session_id: &str) -> Arc { + let active_counter = self + .active_turns_per_session + .entry(session_id.to_string()) + .or_insert_with(|| Arc::new(AtomicUsize::new(0))) + .clone(); + active_counter.fetch_add(1, Ordering::SeqCst); + Arc::new(SessionExecutionLease { active_counter }) + } + pub(crate) async fn wait_for_turn_settlement( &self, session_id: &str, @@ -4514,6 +4655,41 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet Ok(()) } + /// Releases one connection-scoped Session family through the same + /// coordination owner used by durable Session deletion. Coordination rows + /// and live background outcomes are removed before runtime state so a + /// failed cleanup can be retried without losing the family identity. + pub(crate) async fn discard_transient_session( + &self, + workspace_path: &Path, + remote_connection_id: Option<&str>, + remote_ssh_host: Option<&str>, + session_id: &str, + ) -> BitFunResult { + let family = self.session_manager.transient_session_family_postorder( + workspace_path, + remote_connection_id, + remote_ssh_host, + session_id, + )?; + if family.is_empty() { + return Ok(false); + } + for related_session_id in &family { + self.background_subagent_outcomes + .delete_session_references(related_session_id) + .await?; + } + self.session_manager + .discard_transient_session( + workspace_path, + remote_connection_id, + remote_ssh_host, + session_id, + ) + .await + } + pub async fn delete_hidden_subagent_sessions_for_parent_turns( &self, workspace_path: &Path, @@ -5215,8 +5391,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet runtime_tool_restrictions, prompt_cache_source_session_id, session_kind, + transient, emit_lifecycle_events, prepared_session_created, + execution_lease, external_generation_lease: _external_generation_lease, } = request; let prepared_target_session_id = target_session_id.clone(); @@ -5390,13 +5568,14 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet }, None => { let session = self - .create_hidden_agent_session( + .create_hidden_agent_session_with_durability( None, session_name.clone(), logical_agent_type.clone(), session_config.clone(), created_by.clone(), session_kind, + transient, ) .await?; let session_id = session.session_id.clone(); @@ -5420,6 +5599,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet } }; let session_id = session.session_id.clone(); + let _execution_lease = + execution_lease.unwrap_or_else(|| self.register_session_execution(&session_id)); // Sync context window from AI config so subagents with large-context // models are not prematurely capped at SessionConfig::default()'s 128128. if let Err(error) = self @@ -6677,6 +6858,17 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet "session-{}", request.subagent_parent_info.session_id )); + self.session_manager + .get_session(&request.subagent_parent_info.session_id) + .ok_or_else(|| { + BitFunError::NotFound(format!( + "Parent session not found: {}", + request.subagent_parent_info.session_id + )) + })?; + let parent_transient = self + .session_manager + .is_transient_session(&request.subagent_parent_info.session_id); let approved_model_binding = request .external_generation_lease .as_ref() @@ -6738,6 +6930,9 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .await?; initial_messages.push(Message::user(task_description.clone())); + let transient = self + .session_manager + .is_transient_session(&session.session_id); return Ok(HiddenSubagentExecutionRequest { target_session_id: Some(session.session_id.clone()), dialog_turn_id: None, @@ -6752,13 +6947,18 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet context: request.context, permission_runtime_ceiling: Some(request.permission_runtime_ceiling), delegation_policy: request.delegation_policy, - runtime_tool_restrictions: runtime_tool_restrictions_for_delegation_policy( - request.delegation_policy, + runtime_tool_restrictions: runtime_tool_restrictions_for_session_lifetime( + runtime_tool_restrictions_for_delegation_policy( + request.delegation_policy, + ), + transient, ), prompt_cache_source_session_id: None, session_kind: SessionKind::Subagent, + transient, emit_lifecycle_events: true, prepared_session_created: false, + execution_lease: None, external_generation_lease: request.external_generation_lease, }); } @@ -6831,13 +7031,16 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet context: request.context, permission_runtime_ceiling: Some(request.permission_runtime_ceiling), delegation_policy: request.delegation_policy, - runtime_tool_restrictions: runtime_tool_restrictions_for_delegation_policy( - request.delegation_policy, + runtime_tool_restrictions: runtime_tool_restrictions_for_session_lifetime( + runtime_tool_restrictions_for_delegation_policy(request.delegation_policy), + parent_transient, ), prompt_cache_source_session_id: None, session_kind: SessionKind::Subagent, + transient: parent_transient, emit_lifecycle_events: true, prepared_session_created: false, + execution_lease: None, external_generation_lease: request.external_generation_lease, }) } @@ -6915,13 +7118,16 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet context: request.context, permission_runtime_ceiling: Some(request.permission_runtime_ceiling), delegation_policy: request.delegation_policy, - runtime_tool_restrictions: runtime_tool_restrictions_for_delegation_policy( - request.delegation_policy, + runtime_tool_restrictions: runtime_tool_restrictions_for_session_lifetime( + runtime_tool_restrictions_for_delegation_policy(request.delegation_policy), + parent_transient, ), prompt_cache_source_session_id: Some(snapshot.parent_session_id), session_kind: SessionKind::Subagent, + transient: parent_transient, emit_lifecycle_events: true, prepared_session_created: false, + execution_lease: None, external_generation_lease: request.external_generation_lease, }) } @@ -6948,17 +7154,21 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet target_session_id ))); } + if request.execution_lease.is_none() { + request.execution_lease = Some(self.register_session_execution(target_session_id)); + } return Ok(request); } let session = self - .create_hidden_agent_session( + .create_hidden_agent_session_with_durability( None, request.session_name.clone(), request.logical_agent_type.clone(), request.session_config.clone(), request.created_by.clone(), request.session_kind, + request.transient, ) .await?; let session_id = session.session_id.clone(); @@ -6982,6 +7192,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet request.target_session_id = Some(session_id); request.prepared_session_created = true; + request.execution_lease = request + .target_session_id + .as_deref() + .map(|session_id| self.register_session_execution(session_id)); Ok(request) } @@ -7318,8 +7532,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet runtime_tool_restrictions: request.runtime_tool_restrictions, prompt_cache_source_session_id: None, session_kind: request.session_kind, + transient: false, emit_lifecycle_events: request.emit_lifecycle_events, prepared_session_created: false, + execution_lease: None, external_generation_lease: None, }; @@ -7543,9 +7759,9 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet /// Clean up runtime-only subagent resources. /// - /// Subagent sessions are now persisted so users can reopen them from the UI. - /// This cleanup path must only release ephemeral runtime resources such as - /// snapshot bookkeeping; it must not delete the persisted session itself. + /// Durable and reusable Subagent sessions remain available for follow-up. + /// A transient fresh-only child has no supported continuation path, so its + /// existing lifecycle owner releases the Session after terminal cleanup. async fn cleanup_subagent_resources(&self, session_id: &str) -> BitFunResult<()> { let cleanup_started_at = Instant::now(); debug!( @@ -7554,10 +7770,11 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet ); // Clean up snapshot system resources - if let Some(workspace_path) = self - .session_manager - .get_session(session_id) - .and_then(|session| session.config.workspace_path.map(std::path::PathBuf::from)) + let session = self.session_manager.get_session(session_id); + if let Some(workspace_path) = session + .as_ref() + .and_then(|session| session.config.workspace_path.as_deref()) + .map(std::path::PathBuf::from) { debug!( "Subagent cleanup stage starting: session_id={}, stage=snapshot_cleanup, workspace_path={}", @@ -7589,6 +7806,30 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet ); } + if let Some(session) = session.filter(|session| { + self.session_manager.is_transient_session(session_id) + && session.config.continuation_policy == SessionContinuationPolicy::FreshOnly + }) { + let workspace_path = session + .config + .workspace_path + .as_deref() + .map(Path::new) + .ok_or_else(|| { + BitFunError::Validation(format!( + "Transient subagent workspace binding is missing: {session_id}" + )) + })?; + self.session_manager + .discard_transient_session( + workspace_path, + session.config.remote_connection_id.as_deref(), + session.config.remote_ssh_host.as_deref(), + session_id, + ) + .await?; + } + debug!( "Subagent resource cleanup completed: session_id={}, duration_ms={}", session_id, @@ -7844,6 +8085,7 @@ async fn create_agent_session_from_runtime_request( coordinator: &ConversationCoordinator, session_id: Option, request: bitfun_runtime_ports::AgentSessionCreateRequest, + transient: bool, map_core_error: fn(BitFunError) -> bitfun_runtime_ports::PortError, ) -> bitfun_runtime_ports::PortResult { let workspace_path = request.workspace_path.clone().ok_or_else(|| { @@ -7854,7 +8096,7 @@ async fn create_agent_session_from_runtime_request( })?; let created_by = resolve_agent_session_create_created_by(&request.metadata); let session = coordinator - .create_session_with_workspace_and_creator( + .create_session_with_workspace_and_creator_internal( session_id, request.session_name, request.agent_type, @@ -7868,6 +8110,7 @@ async fn create_agent_session_from_runtime_request( }, workspace_path, created_by, + transient, ) .await .map_err(map_core_error)?; @@ -7885,8 +8128,14 @@ impl bitfun_runtime_ports::AgentSubmissionPort for ConversationCoordinator { &self, request: bitfun_runtime_ports::AgentSessionCreateRequest, ) -> bitfun_runtime_ports::PortResult { - create_agent_session_from_runtime_request(self, None, request, runtime_port_backend_error) - .await + create_agent_session_from_runtime_request( + self, + None, + request, + false, + runtime_port_backend_error, + ) + .await } async fn create_session_with_id( @@ -7901,6 +8150,25 @@ impl bitfun_runtime_ports::AgentSubmissionPort for ConversationCoordinator { self, Some(session_id), request, + false, + runtime_port_error_preserving_message, + ) + .await + } + + async fn create_transient_session_with_id( + &self, + session_id: String, + request: bitfun_runtime_ports::AgentSessionCreateRequest, + ) -> bitfun_runtime_ports::PortResult { + bitfun_core_types::validate_session_id(&session_id).map_err(|message| { + runtime_port_error_preserving_message(BitFunError::Validation(message)) + })?; + create_agent_session_from_runtime_request( + self, + Some(session_id), + request, + true, runtime_port_error_preserving_message, ) .await @@ -8672,8 +8940,9 @@ mod tests { merge_prepended_messages_for_turn, normalize_subagent_max_concurrency, resolve_agent_session_create_created_by, resolve_agent_submission_turn_id, resolve_subagent_model_selection, runtime_port_error_preserving_message, - turn_review_manifest_for_agent, BackgroundSubagentWaitMode, ConversationCoordinator, - SessionReferenceLocator, SubagentExecutionRequest, TEST_AGENT_MODEL_DEFAULTS, + runtime_tool_restrictions_for_session_lifetime, turn_review_manifest_for_agent, + BackgroundSubagentWaitMode, ConversationCoordinator, SessionReferenceLocator, + SubagentExecutionRequest, TEST_AGENT_MODEL_DEFAULTS, }; use crate::agentic::coordination::coordination_store::{ BackgroundTaskRegistration, RegisteredBackgroundTask, @@ -8747,6 +9016,40 @@ mod tests { ); } + #[test] + fn transient_session_runtime_restrictions_deny_out_of_band_session_tools() { + let mut base = crate::agentic::tools::ToolRuntimeRestrictions::default(); + base.denied_tool_names.insert("Bash".to_string()); + + let transient = runtime_tool_restrictions_for_session_lifetime(base.clone(), true); + for tool_name in [ + "SessionControl", + "SessionMessage", + "SessionHistory", + "Cron", + "ControlHub", + ] { + assert!( + !transient.is_tool_allowed(tool_name), + "{tool_name} must not cross a connection-scoped Session boundary" + ); + } + assert!(!transient.is_tool_allowed("Bash")); + assert!(transient.is_tool_allowed("Read")); + + let durable = runtime_tool_restrictions_for_session_lifetime(base, false); + for tool_name in [ + "SessionControl", + "SessionMessage", + "SessionHistory", + "Cron", + "ControlHub", + ] { + assert!(durable.is_tool_allowed(tool_name)); + } + assert!(!durable.is_tool_allowed("Bash")); + } + #[test] fn migrated_runtime_ports_preserve_existing_core_error_messages() { let error = runtime_port_error_preserving_message( @@ -9039,8 +9342,9 @@ mod tests { } use tokio::sync::RwLock as TokioRwLock; - fn test_coordinator_with_max_active_sessions( + fn test_coordinator_with_config( max_active_sessions: usize, + enable_persistence: bool, ) -> (ConversationCoordinator, Arc) { let event_queue = Arc::new(EventQueue::new(EventQueueConfig::default())); let coordination_database_file = std::env::temp_dir() @@ -9056,7 +9360,7 @@ mod tests { max_active_sessions, session_idle_timeout: Duration::from_secs(3600), auto_save_interval: Duration::from_secs(300), - enable_persistence: false, + enable_persistence, prompt_cache_policy: PromptCachePolicy::default(), }, )); @@ -9094,6 +9398,12 @@ mod tests { (coordinator, session_manager) } + fn test_coordinator_with_max_active_sessions( + max_active_sessions: usize, + ) -> (ConversationCoordinator, Arc) { + test_coordinator_with_config(max_active_sessions, false) + } + fn test_coordinator() -> (ConversationCoordinator, Arc) { test_coordinator_with_max_active_sessions(100) } @@ -10237,6 +10547,140 @@ mod tests { let _ = std::fs::remove_dir_all(workspace_path); } + #[tokio::test] + async fn transient_session_port_never_persists_or_discards_a_durable_identity() { + let (coordinator, session_manager) = test_coordinator_with_config(100, true); + let workspace_path = std::env::temp_dir().join(format!( + "bitfun-agent-transient-session-port-test-{}", + uuid::Uuid::new_v4() + )); + std::fs::create_dir_all(&workspace_path).expect("workspace dir should exist"); + let workspace = workspace_path.to_string_lossy().into_owned(); + let request = |name: &str| AgentSessionCreateRequest { + session_name: name.to_string(), + agent_type: "agentic".to_string(), + workspace_path: Some(workspace.clone()), + workspace_id: None, + remote_connection_id: None, + remote_ssh_host: None, + model_id: None, + metadata: serde_json::Map::new(), + }; + + let transient = AgentSubmissionPort::create_transient_session_with_id( + &coordinator, + "transient-session".to_string(), + request("Transient"), + ) + .await + .expect("transient Session creation should succeed"); + let loaded = session_manager + .get_session(&transient.session_id) + .expect("transient Session should be loaded"); + assert!(session_manager.is_transient_session(&loaded.session_id)); + let storage_path = session_manager + .resolve_session_workspace_binding(&transient.session_id) + .await + .expect("transient Session should retain its workspace binding") + .session_storage_dir(); + assert!(!session_manager + .persistence_manager() + .session_storage_exists(&storage_path, &transient.session_id) + .expect("persistence probe should succeed")); + + let transient_child = session_manager + .create_transient_session_with_id_and_details( + None, + "Transient child".to_string(), + "Explore".to_string(), + SessionConfig { + workspace_path: Some(workspace.clone()), + ..Default::default() + }, + Some(format!("session-{}", transient.session_id)), + SessionKind::Subagent, + ) + .await + .expect("transient child Session should be created"); + let transient_grandchild = session_manager + .create_transient_session_with_id_and_details( + None, + "Transient grandchild".to_string(), + "Explore".to_string(), + SessionConfig { + workspace_path: Some(workspace.clone()), + ..Default::default() + }, + Some(format!("session-{}", transient_child.session_id)), + SessionKind::Subagent, + ) + .await + .expect("nested transient child Session should be created"); + let background_task = register_test_background_task( + &coordinator, + &transient.session_id, + "transient-parent-turn", + &transient_child.session_id, + ) + .await; + + let discarded = coordinator + .discard_transient_session(&workspace_path, None, None, &transient.session_id) + .await + .expect("owned transient Session discard should succeed"); + assert!(discarded); + assert!(session_manager.get_session(&transient.session_id).is_none()); + assert!( + session_manager + .get_session(&transient_child.session_id) + .is_none(), + "discarding a transient parent must release its transient descendants" + ); + assert!(session_manager + .get_session(&transient_grandchild.session_id) + .is_none()); + assert!(session_manager + .resolve_session_workspace_binding(&transient.session_id) + .await + .is_none()); + assert!(coordinator + .background_subagent_outcomes + .resolve_agent_id(&transient.session_id, &background_task.agent_id) + .await + .is_err()); + + let durable = AgentSubmissionPort::create_session_with_id( + &coordinator, + "durable-session".to_string(), + request("Durable"), + ) + .await + .expect("durable Session creation should succeed"); + let discard_error = session_manager + .discard_transient_session(&workspace_path, None, None, &durable.session_id) + .await + .expect_err("transient discard must reject durable Session ownership"); + assert!(matches!( + discard_error, + crate::util::errors::BitFunError::Validation(_) + )); + session_manager.evict_loaded_session_for_test(&durable.session_id); + let collision = AgentSubmissionPort::create_transient_session_with_id( + &coordinator, + durable.session_id.clone(), + request("Collision"), + ) + .await + .expect_err("transient Session must not shadow persisted durable identity"); + assert_eq!( + collision.kind, + bitfun_runtime_ports::PortErrorKind::InvalidRequest + ); + + let _ = std::fs::remove_dir_all(storage_path); + let _ = std::fs::remove_dir_all(workspace_path); + } + #[tokio::test] async fn agent_submission_create_session_rejects_invalid_requested_session_id() { let (coordinator, _) = test_coordinator(); @@ -10341,6 +10785,146 @@ mod tests { assert_eq!(model_id, "primary"); } + #[tokio::test] + async fn fresh_subagent_inherits_transient_parent_persistence_boundary() { + let (coordinator, session_manager) = test_coordinator(); + let workspace_path = std::env::temp_dir().join(format!( + "bitfun-fresh-subagent-transient-test-{}", + uuid::Uuid::new_v4() + )); + std::fs::create_dir_all(&workspace_path).expect("workspace dir should exist"); + struct TempWorkspaceGuard(std::path::PathBuf); + impl Drop for TempWorkspaceGuard { + fn drop(&mut self) { + let _ = std::fs::remove_dir_all(&self.0); + } + } + let _workspace_guard = TempWorkspaceGuard(workspace_path.clone()); + let workspace = workspace_path.to_string_lossy().into_owned(); + let parent_session = session_manager + .create_transient_session_with_id_and_details( + None, + "Transient parent".to_string(), + "agentic".to_string(), + SessionConfig { + model_id: Some("primary".to_string()), + workspace_path: Some(workspace.clone()), + ..Default::default() + }, + None, + SessionKind::Standard, + ) + .await + .expect("transient parent should be created"); + + let resolved = coordinator + .resolve_hidden_subagent_execution_request(SubagentExecutionRequest { + task_description: "Inspect the workspace".to_string(), + context_mode: SubagentContextMode::Fresh, + target_session_id: None, + subagent_type: Some("Explore".to_string()), + logical_subagent_type: None, + continuation_policy: SessionContinuationPolicy::Reusable, + model_binding_policy: SessionModelBindingPolicy::Mutable, + workspace_path: Some(workspace.clone()), + model_id: Some("primary".to_string()), + inherit_parent_model: false, + subagent_parent_info: SubagentParentInfo { + session_id: parent_session.session_id.clone(), + dialog_turn_id: "parent-turn".to_string(), + tool_call_id: "task-tool".to_string(), + }, + context: HashMap::new(), + permission_runtime_ceiling: PermissionRuntimeCeiling::default(), + delegation_policy: DelegationPolicy::top_level().spawn_child(), + external_generation_lease: None, + }) + .await + .expect("fresh subagent request should resolve"); + + assert!(resolved.transient); + assert!(!resolved + .runtime_tool_restrictions + .is_tool_allowed("SessionControl")); + assert!(!resolved + .runtime_tool_restrictions + .is_tool_allowed("SessionMessage")); + + let prepared = coordinator + .prepare_hidden_subagent_execution_request(resolved) + .await + .expect("transient child should prepare"); + let child_session_id = prepared + .target_session_id() + .expect("prepared child Session id") + .to_string(); + assert!( + coordinator + .ensure_session_execution_drained(&child_session_id, Duration::from_millis(10)) + .await + .is_err(), + "a prepared hidden execution must fence Session maintenance" + ); + drop(prepared); + coordinator + .ensure_session_execution_drained(&child_session_id, Duration::from_millis(50)) + .await + .expect("dropping the final hidden execution lease should release maintenance"); + + coordinator + .cleanup_subagent_resources(&child_session_id) + .await + .expect("transient child cleanup should succeed"); + assert!( + session_manager.get_session(&child_session_id).is_some(), + "a reusable transient Subagent must remain available for send_input until its parent is discarded" + ); + + let fresh_only = coordinator + .resolve_hidden_subagent_execution_request(SubagentExecutionRequest { + task_description: "Run once".to_string(), + context_mode: SubagentContextMode::Fresh, + target_session_id: None, + subagent_type: Some("Explore".to_string()), + logical_subagent_type: None, + continuation_policy: SessionContinuationPolicy::FreshOnly, + model_binding_policy: SessionModelBindingPolicy::Mutable, + workspace_path: Some(workspace), + model_id: Some("primary".to_string()), + inherit_parent_model: false, + subagent_parent_info: SubagentParentInfo { + session_id: parent_session.session_id, + dialog_turn_id: "parent-turn-2".to_string(), + tool_call_id: "task-tool-2".to_string(), + }, + context: HashMap::new(), + permission_runtime_ceiling: PermissionRuntimeCeiling::default(), + delegation_policy: DelegationPolicy::top_level().spawn_child(), + external_generation_lease: None, + }) + .await + .expect("fresh-only transient child should resolve"); + let fresh_only = coordinator + .prepare_hidden_subagent_execution_request(fresh_only) + .await + .expect("fresh-only transient child should prepare"); + let fresh_only_session_id = fresh_only + .target_session_id() + .expect("fresh-only prepared child Session id") + .to_string(); + drop(fresh_only); + coordinator + .cleanup_subagent_resources(&fresh_only_session_id) + .await + .expect("fresh-only transient child cleanup should succeed"); + assert!( + session_manager + .get_session(&fresh_only_session_id) + .is_none(), + "a fresh-only transient Subagent should be released after terminal cleanup" + ); + } + #[tokio::test] async fn reused_subagent_send_input_updates_requested_and_inherited_model() { let (coordinator, session_manager) = test_coordinator(); diff --git a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs index e84bf1584c..8c78c4d7aa 100644 --- a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs @@ -1900,6 +1900,7 @@ impl ExecutionEngine { context.workspace_services.as_ref(), Some(&primary_model_facts), &tool_manifest_context_vars, + &context.runtime_tool_restrictions, ); let tool_manifest = if enable_tools { Some( @@ -2782,6 +2783,7 @@ impl ExecutionEngine { context.workspace_services.as_ref(), Some(&primary_model_facts), &tool_manifest_context_vars, + &context.runtime_tool_restrictions, ); let tool_manifest = if enable_tools { diff --git a/src/crates/assembly/core/src/agentic/mod.rs b/src/crates/assembly/core/src/agentic/mod.rs index d6d91f5f12..3a8c8aef21 100644 --- a/src/crates/assembly/core/src/agentic/mod.rs +++ b/src/crates/assembly/core/src/agentic/mod.rs @@ -75,5 +75,5 @@ pub use round_preempt::{ pub use session::*; pub use side_question::*; pub use skill_agent_snapshot::*; -pub use system::{init_agentic_system, AgenticSystem}; +pub use system::{init_agentic_system, init_agentic_system_for_profile, AgenticSystem}; pub use workspace::{WorkspaceBackend, WorkspaceBinding}; diff --git a/src/crates/assembly/core/src/agentic/session/session_manager.rs b/src/crates/assembly/core/src/agentic/session/session_manager.rs index 9882204fbb..9cf65de64b 100644 --- a/src/crates/assembly/core/src/agentic/session/session_manager.rs +++ b/src/crates/assembly/core/src/agentic/session/session_manager.rs @@ -158,11 +158,22 @@ fn current_unix_secs() -> i64 { .unwrap_or_default() } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum SessionResourceCleanupPolicy { + BestEffort, + Required, +} + /// Session manager pub struct SessionManager { /// Active sessions in memory sessions: Arc>, + /// Process-local durability classification owned by the Session lifecycle. + /// Entries are installed before a transient Session becomes visible and are + /// removed with that Session; they are never serialized into public config. + transient_session_ids: Arc>, + /// Exact admission accounting for loaded sessions. A permit is acquired /// before create/restore publishes runtime state and released on unload/delete/eviction. active_session_capacity: Arc, @@ -280,8 +291,9 @@ impl SessionManager { } #[cfg(test)] - fn evict_loaded_session_for_test(&self, session_id: &str) { + pub(crate) fn evict_loaded_session_for_test(&self, session_id: &str) { self.sessions.remove(session_id); + self.transient_session_ids.remove(session_id); self.release_active_session_reservation(session_id); } @@ -612,8 +624,16 @@ impl SessionManager { } } - fn should_persist_session(session: &Session) -> bool { - Self::should_persist_session_kind(session.kind) + fn should_persist_session_with_transient_ids( + session: &Session, + transient_session_ids: &DashMap, + ) -> bool { + !transient_session_ids.contains_key(&session.session_id) + && Self::should_persist_session_kind(session.kind) + } + + fn should_persist_session(&self, session: &Session) -> bool { + Self::should_persist_session_with_transient_ids(session, &self.transient_session_ids) } fn same_session_version( @@ -626,12 +646,14 @@ impl SessionManager { fn collect_auto_save_snapshots( sessions: &DashMap, + transient_session_ids: &DashMap, ) -> Vec { sessions .iter() .filter_map(|entry| { let session = entry.value(); - if !Self::should_persist_session(session) { + if !Self::should_persist_session_with_transient_ids(session, transient_session_ids) + { return None; } Some(SessionAutoSaveSnapshot { @@ -668,6 +690,7 @@ impl SessionManager { fn collect_expired_session_candidates( sessions: &DashMap, + transient_session_ids: &DashMap, now: SystemTime, timeout: Duration, ) -> Vec { @@ -675,7 +698,12 @@ impl SessionManager { .iter() .filter_map(|entry| { let session = entry.value(); - if !Self::is_session_expired(session, now, timeout) { + // Idle eviction is a restore optimization for durable Sessions. + // Non-persistent Sessions have an explicit lifecycle owner and + // no on-disk state from which they could be restored. + if !Self::should_persist_session_with_transient_ids(session, transient_session_ids) + || !Self::is_session_expired(session, now, timeout) + { return None; } Some(SessionCleanupCandidate { @@ -711,13 +739,18 @@ impl SessionManager { pub fn should_persist_session_id(&self, session_id: &str) -> bool { self.config.enable_persistence + && !self.transient_session_ids.contains_key(session_id) && self .sessions .get(session_id) - .map(|session| Self::should_persist_session(&session)) + .map(|session| self.should_persist_session(&session)) .unwrap_or(true) } + pub(crate) fn is_transient_session(&self, session_id: &str) -> bool { + self.transient_session_ids.contains_key(session_id) + } + async fn effective_storage_path_for_config_with_persistence( persistence_manager: &PersistenceManager, config: &SessionConfig, @@ -1684,6 +1717,7 @@ impl SessionManager { let manager = Self { sessions: Arc::new(DashMap::new()), + transient_session_ids: Arc::new(DashMap::new()), active_session_capacity: Arc::new(Semaphore::new(config.max_active_sessions)), active_session_permits: Arc::new(DashMap::new()), session_storage_path_index: Arc::new(DashMap::new()), @@ -1890,6 +1924,7 @@ impl SessionManager { fn spawn_model_reconciliation_listener(&self) { let sessions = self.sessions.clone(); + let transient_session_ids = self.transient_session_ids.clone(); let active_session_capacity = self.active_session_capacity.clone(); let active_session_permits = self.active_session_permits.clone(); let session_storage_path_index = self.session_storage_path_index.clone(); @@ -1920,6 +1955,7 @@ impl SessionManager { // surface area we need from the cloned shared fields above. let manager = Self { sessions, + transient_session_ids, active_session_capacity, active_session_permits, session_storage_path_index, @@ -2035,6 +2071,49 @@ impl SessionManager { config: SessionConfig, created_by: Option, kind: SessionKind, + ) -> BitFunResult { + self.create_session_with_id_and_details_internal( + session_id, + session_name, + agent_type, + config, + created_by, + kind, + false, + ) + .await + } + + pub(crate) async fn create_transient_session_with_id_and_details( + &self, + session_id: Option, + session_name: String, + agent_type: String, + config: SessionConfig, + created_by: Option, + kind: SessionKind, + ) -> BitFunResult { + self.create_session_with_id_and_details_internal( + session_id, + session_name, + agent_type, + config, + created_by, + kind, + true, + ) + .await + } + + async fn create_session_with_id_and_details_internal( + &self, + session_id: Option, + session_name: String, + agent_type: String, + config: SessionConfig, + created_by: Option, + kind: SessionKind, + transient: bool, ) -> BitFunResult { let _workspace_path = Self::session_workspace_from_config(&config).ok_or_else(|| { BitFunError::Validation("Session workspace_path is required".to_string()) @@ -2066,7 +2145,6 @@ impl SessionManager { ))); } if self.config.enable_persistence - && Self::should_persist_session(&session) && self .persistence_manager .session_storage_exists(&session_storage_path, &session_id)? @@ -2078,6 +2156,9 @@ impl SessionManager { let active_session_permit = self.reserve_active_session()?; let storage_claim = self.claim_session_storage_path(&session_id, &session_storage_path, true)?; + if transient { + self.transient_session_ids.insert(session_id.clone(), ()); + } // 1. Add to memory match self.sessions.entry(session_id.clone()) { @@ -2086,6 +2167,9 @@ impl SessionManager { } Entry::Occupied(entry) => { drop(entry); + if transient { + self.transient_session_ids.remove(&session_id); + } self.release_failed_session_storage_path_claim( &session_id, &session_storage_path, @@ -2106,7 +2190,7 @@ impl SessionManager { // 3. Persist to local path (handles remote workspaces correctly) // Use the local `session` directly -- no need to re-fetch from DashMap, // which would hold a Ref guard across the async save_session call. - if self.config.enable_persistence && Self::should_persist_session(&session) { + if self.config.enable_persistence && self.should_persist_session(&session) { if let Err(error) = self .persistence_manager .create_session_if_absent(&session_storage_path, &session) @@ -2119,6 +2203,9 @@ impl SessionManager { .delete_session(&session_id); self.file_read_state_store.delete_session(&session_id); self.evidence_ledger.delete_session(&session_id); + if transient { + self.transient_session_ids.remove(&session_id); + } self.release_failed_session_storage_path_claim( &session_id, &session_storage_path, @@ -2984,7 +3071,7 @@ impl SessionManager { session.updated_at = SystemTime::now(); session.last_activity_at = SystemTime::now(); - self.config.enable_persistence && Self::should_persist_session(&session) + self.config.enable_persistence && self.should_persist_session(&session) } else { return Err(BitFunError::NotFound(format!( "Session not found: {}", @@ -3041,7 +3128,7 @@ impl SessionManager { session.updated_at = SystemTime::now(); session.last_activity_at = SystemTime::now(); - self.config.enable_persistence && Self::should_persist_session(&session) + self.config.enable_persistence && self.should_persist_session(&session) } else { return Err(BitFunError::NotFound(format!( "Session not found: {}", @@ -3514,6 +3601,199 @@ impl SessionManager { .await } + /// Discards one loaded non-durable Session without touching persisted + /// Session storage. Missing Sessions are an idempotent success. + pub(crate) async fn discard_transient_session( + &self, + workspace_path: &Path, + remote_connection_id: Option<&str>, + remote_ssh_host: Option<&str>, + session_id: &str, + ) -> BitFunResult { + bitfun_core_types::validate_session_id(session_id).map_err(BitFunError::Validation)?; + let Some(root) = self.get_session(session_id) else { + return Ok(false); + }; + self.validate_transient_session_binding( + &root, + workspace_path, + remote_connection_id, + remote_ssh_host, + )?; + + for descendant in self.transient_descendants_postorder(session_id) { + let workspace_path = descendant + .config + .workspace_path + .as_deref() + .map(Path::new) + .ok_or_else(|| { + BitFunError::Validation(format!( + "Transient session workspace binding is missing: {}", + descendant.session_id + )) + })?; + self.discard_one_transient_session( + workspace_path, + descendant.config.remote_connection_id.as_deref(), + descendant.config.remote_ssh_host.as_deref(), + &descendant.session_id, + ) + .await?; + } + + self.discard_one_transient_session( + workspace_path, + remote_connection_id, + remote_ssh_host, + session_id, + ) + .await + } + + pub(crate) fn transient_session_family_postorder( + &self, + workspace_path: &Path, + remote_connection_id: Option<&str>, + remote_ssh_host: Option<&str>, + session_id: &str, + ) -> BitFunResult> { + bitfun_core_types::validate_session_id(session_id).map_err(BitFunError::Validation)?; + let Some(root) = self.get_session(session_id) else { + return Ok(Vec::new()); + }; + self.validate_transient_session_binding( + &root, + workspace_path, + remote_connection_id, + remote_ssh_host, + )?; + let mut family = self + .transient_descendants_postorder(session_id) + .into_iter() + .map(|session| session.session_id) + .collect::>(); + family.push(session_id.to_string()); + Ok(family) + } + + fn transient_descendants_postorder(&self, root_session_id: &str) -> Vec { + fn visit( + parent_session_id: &str, + sessions: &[Session], + transient_session_ids: &DashMap, + visited: &mut HashSet, + ordered: &mut Vec, + ) { + let marker = format!("session-{parent_session_id}"); + for child in sessions.iter().filter(|session| { + transient_session_ids.contains_key(&session.session_id) + && session.created_by.as_deref() == Some(marker.as_str()) + }) { + if !visited.insert(child.session_id.clone()) { + continue; + } + visit( + &child.session_id, + sessions, + transient_session_ids, + visited, + ordered, + ); + ordered.push(child.clone()); + } + } + + let sessions = self + .sessions + .iter() + .map(|entry| entry.value().clone()) + .collect::>(); + let mut ordered = Vec::new(); + let mut visited = HashSet::from([root_session_id.to_string()]); + visit( + root_session_id, + &sessions, + &self.transient_session_ids, + &mut visited, + &mut ordered, + ); + ordered + } + + fn validate_transient_session_binding( + &self, + session: &Session, + workspace_path: &Path, + remote_connection_id: Option<&str>, + remote_ssh_host: Option<&str>, + ) -> BitFunResult<()> { + if !self.is_transient_session(&session.session_id) { + return Err(BitFunError::Validation(format!( + "Cannot discard a durable session as transient: {}", + session.session_id + ))); + } + let expected_workspace = Self::normalize_session_storage_path(workspace_path); + let actual_workspace = session + .config + .workspace_path + .as_deref() + .map(Path::new) + .map(Self::normalize_session_storage_path) + .ok_or_else(|| { + BitFunError::Validation(format!( + "Transient session workspace binding is missing: {}", + session.session_id + )) + })?; + if actual_workspace != expected_workspace + || session.config.remote_connection_id.as_deref() != remote_connection_id + || session.config.remote_ssh_host.as_deref() != remote_ssh_host + { + return Err(BitFunError::Validation(format!( + "Transient session ownership binding does not match: {}", + session.session_id + ))); + } + Ok(()) + } + + async fn discard_one_transient_session( + &self, + workspace_path: &Path, + remote_connection_id: Option<&str>, + remote_ssh_host: Option<&str>, + session_id: &str, + ) -> BitFunResult { + let _mutation_guard = self.lock_session_mutation(session_id).await; + let Some(session) = self.get_session(session_id) else { + return Ok(false); + }; + self.validate_transient_session_binding( + &session, + workspace_path, + remote_connection_id, + remote_ssh_host, + )?; + if matches!(session.state, SessionState::Processing { .. }) { + return Err(BitFunError::Validation(format!( + "Cannot discard a processing transient session: {session_id}" + ))); + } + self.cleanup_session_owned_resources( + workspace_path, + session_id, + SessionResourceCleanupPolicy::Required, + ) + .await?; + self.sessions.remove(session_id); + self.transient_session_ids.remove(session_id); + self.release_active_session_reservation(session_id); + self.session_storage_path_index.remove(session_id); + Ok(true) + } + /// Release one loaded session and its transient runtime stores while keeping /// persisted history and the storage-path binding available for a later restore. /// @@ -3525,13 +3805,18 @@ impl SessionManager { let Some(session) = self.get_session(session_id) else { return Ok(false); }; + if self.is_transient_session(session_id) { + return Err(BitFunError::Validation(format!( + "Cannot unload a transient session; use the owned discard path: {session_id}" + ))); + } if matches!(session.state, SessionState::Processing { .. }) { return Err(BitFunError::Validation(format!( "Cannot unload a processing session: {session_id}" ))); } - if self.config.enable_persistence && Self::should_persist_session(&session) { + if self.config.enable_persistence && self.should_persist_session(&session) { let storage_path = self .effective_session_storage_path(session_id) .await @@ -3562,70 +3847,34 @@ impl SessionManager { Ok(true) } - async fn delete_session_from_paths_locked( + async fn cleanup_session_owned_resources( &self, cleanup_workspace_path: &Path, - session_storage_path: &Path, session_id: &str, + policy: SessionResourceCleanupPolicy, ) -> BitFunResult<()> { - let delete_started_at = Instant::now(); - debug!( - "Session deletion started: session_id={}, cleanup_workspace_path={}, session_storage_path={}, persistence_enabled={}", - session_id, - cleanup_workspace_path.display(), - session_storage_path.display(), - self.config.enable_persistence - ); - - // Persisted deletion is the only fallible required stage. Complete it - // before mutating loaded runtime state so a storage failure leaves the - // active session usable and retryable. - if self.config.enable_persistence { - let persistence_stage_started_at = Instant::now(); - debug!( - "Session deletion stage starting: session_id={}, stage=persistence_delete", - session_id - ); - self.persistence_manager - .delete_session(session_storage_path, session_id) - .await?; - debug!( - "Session deletion stage completed: session_id={}, stage=persistence_delete, duration_ms={}", - session_id, - elapsed_ms_u64(persistence_stage_started_at) + let mut required_error = None; + let mut record_error = |stage: &'static str, error: String| { + warn!( + "Session resource cleanup failed: session_id={}, stage={}, error={}", + session_id, stage, error ); - } + if policy == SessionResourceCleanupPolicy::Required && required_error.is_none() { + required_error = Some(BitFunError::Session(format!( + "Session resource cleanup is incomplete: session_id={session_id}, stage={stage}, error={error}" + ))); + } + }; - // 1. Clean up snapshot system resources (including physical snapshot files) - let snapshot_stage_started_at = Instant::now(); - debug!( - "Session deletion stage starting: session_id={}, stage=snapshot_cleanup", - session_id - ); if let Ok(snapshot_manager) = ensure_snapshot_manager_for_workspace(cleanup_workspace_path) { let snapshot_service = snapshot_manager.get_snapshot_service(); let snapshot_service = snapshot_service.read().await; - if let Err(e) = snapshot_service.accept_session(session_id).await { - warn!("Failed to cleanup snapshot system resources: {}", e); - } else { - debug!( - "Snapshot system resources cleaned up: session_id={}", - session_id - ); + if let Err(error) = snapshot_service.accept_session(session_id).await { + record_error("snapshot", error.to_string()); } } - debug!( - "Session deletion stage completed: session_id={}, stage=snapshot_cleanup, duration_ms={}", - session_id, - elapsed_ms_u64(snapshot_stage_started_at) - ); - let context_stage_started_at = Instant::now(); - debug!( - "Session deletion stage starting: session_id={}, stage=context_store_delete", - session_id - ); clear_session_runtime_stores( session_id, self.context_store.as_ref(), @@ -3636,67 +3885,73 @@ impl SessionManager { self.file_read_state_store.as_ref(), self.evidence_ledger.as_ref(), ); - debug!( - "Session deletion stage completed: session_id={}, stage=context_store_delete, duration_ms={}", - session_id, - elapsed_ms_u64(context_stage_started_at) - ); if let Some(cron) = crate::service::cron::get_global_cron_service() { - let cron_stage_started_at = Instant::now(); - debug!( - "Session deletion stage starting: session_id={}, stage=cron_cleanup", - session_id - ); match cron.delete_jobs_for_session(session_id).await { - Ok(removed) if removed > 0 => { - info!( - "Removed {} scheduled job(s) for deleted session_id={}", - removed, session_id - ); - } + Ok(removed) if removed > 0 => info!( + "Removed {} scheduled job(s) for session_id={}", + removed, session_id + ), Ok(_) => {} - Err(e) => { - warn!( - "Failed to remove scheduled jobs for session_id={}: {}", - session_id, e - ); - } + Err(error) => record_error("cron", error.to_string()), } - debug!( - "Session deletion stage completed: session_id={}, stage=cron_cleanup, duration_ms={}", - session_id, - elapsed_ms_u64(cron_stage_started_at) - ); } - // 3. Clean up associated Terminal session use crate::service::terminal::TerminalApi; if let Ok(terminal_api) = TerminalApi::from_singleton() { let binding = terminal_api.session_manager().binding(); - let terminal_stage_started_at = Instant::now(); + if let Err(error) = binding.remove(session_id).await { + record_error("terminal", error.to_string()); + } + } + + if let Some(error) = required_error { + return Err(error); + } + Ok(()) + } + + async fn delete_session_from_paths_locked( + &self, + cleanup_workspace_path: &Path, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult<()> { + let delete_started_at = Instant::now(); + debug!( + "Session deletion started: session_id={}, cleanup_workspace_path={}, session_storage_path={}, persistence_enabled={}", + session_id, + cleanup_workspace_path.display(), + session_storage_path.display(), + self.config.enable_persistence + ); + + // Persisted deletion is the only fallible required stage. Complete it + // before mutating loaded runtime state so a storage failure leaves the + // active session usable and retryable. + if self.config.enable_persistence { + let persistence_stage_started_at = Instant::now(); debug!( - "Session deletion stage starting: session_id={}, stage=terminal_binding_cleanup, has_binding={}", - session_id, - binding.has(session_id) + "Session deletion stage starting: session_id={}, stage=persistence_delete", + session_id ); - if binding.has(session_id) { - if let Err(e) = binding.remove(session_id).await { - warn!("Failed to cleanup associated Terminal session: {}", e); - } else { - debug!( - "Associated Terminal session cleaned up: session_id={}", - session_id - ); - } - } + self.persistence_manager + .delete_session(session_storage_path, session_id) + .await?; debug!( - "Session deletion stage completed: session_id={}, stage=terminal_binding_cleanup, duration_ms={}", + "Session deletion stage completed: session_id={}, stage=persistence_delete, duration_ms={}", session_id, - elapsed_ms_u64(terminal_stage_started_at) + elapsed_ms_u64(persistence_stage_started_at) ); } + self.cleanup_session_owned_resources( + cleanup_workspace_path, + session_id, + SessionResourceCleanupPolicy::BestEffort, + ) + .await?; + // 4. Remove from memory let memory_stage_started_at = Instant::now(); debug!( @@ -3704,6 +3959,7 @@ impl SessionManager { session_id ); self.sessions.remove(session_id); + self.transient_session_ids.remove(session_id); self.release_active_session_reservation(session_id); debug!( "Session deletion stage completed: session_id={}, stage=in_memory_remove, duration_ms={}", @@ -4737,7 +4993,7 @@ impl SessionManager { session.last_activity_at = SystemTime::now(); let should_persist = - Self::should_persist_session(&session) && self.config.enable_persistence; + self.should_persist_session(&session) && self.config.enable_persistence; if should_persist { Some(session.clone()) } else { @@ -5362,7 +5618,7 @@ impl SessionManager { turn.duration_ms = Some(0); turn.status = TurnStatus::Completed; - if self.config.enable_persistence && Self::should_persist_session(&session) { + if self.config.enable_persistence && self.should_persist_session(&session) { self.persistence_manager .save_dialog_turn(&workspace_path, &turn) .await?; @@ -5380,7 +5636,7 @@ impl SessionManager { session.updated_at = SystemTime::now(); session.last_activity_at = SystemTime::now(); - if self.config.enable_persistence && Self::should_persist_session(&session) { + if self.config.enable_persistence && self.should_persist_session(&session) { Some(session.clone()) } else { None @@ -6159,7 +6415,7 @@ impl SessionManager { session.compression_state = compression_state; session.updated_at = SystemTime::now(); session.last_activity_at = SystemTime::now(); - if self.config.enable_persistence && Self::should_persist_session(&session) { + if self.config.enable_persistence && self.should_persist_session(&session) { Some(session.clone()) } else { None @@ -6322,6 +6578,7 @@ impl SessionManager { /// Start auto-save task fn spawn_auto_save_task(&self) { let sessions = self.sessions.clone(); + let transient_session_ids = self.transient_session_ids.clone(); let persistence = self.persistence_manager.clone(); let session_mutation_locks = self.session_mutation_locks.clone(); let interval = self.config.auto_save_interval; @@ -6332,7 +6589,8 @@ impl SessionManager { loop { ticker.tick().await; - for snapshot in Self::collect_auto_save_snapshots(&sessions) { + for snapshot in Self::collect_auto_save_snapshots(&sessions, &transient_session_ids) + { let _mutation_guard = session_mutation_locks.lock(&snapshot.session_id).await; if !Self::auto_save_snapshot_is_current(&sessions, &snapshot) { continue; @@ -6367,6 +6625,7 @@ impl SessionManager { /// Start cleanup task for expired sessions fn spawn_cleanup_task(&self) { let sessions = self.sessions.clone(); + let transient_session_ids = self.transient_session_ids.clone(); let active_session_permits = self.active_session_permits.clone(); let timeout = self.config.session_idle_timeout; let persistence = self.persistence_manager.clone(); @@ -6389,7 +6648,12 @@ impl SessionManager { ticker.tick().await; let now = SystemTime::now(); - let candidates = Self::collect_expired_session_candidates(&sessions, now, timeout); + let candidates = Self::collect_expired_session_candidates( + &sessions, + &transient_session_ids, + now, + timeout, + ); for candidate in candidates { let _mutation_guard = session_mutation_locks.lock(&candidate.session_id).await; @@ -6408,7 +6672,12 @@ impl SessionManager { continue; }; - if enable_persistence && Self::should_persist_session(&session) { + if enable_persistence + && Self::should_persist_session_with_transient_ids( + &session, + &transient_session_ids, + ) + { if let Some(workspace_path) = Self::effective_storage_path_for_config_with_persistence( persistence.as_ref(), @@ -6488,12 +6757,12 @@ mod tests { TurnStatus, UserMessageData, }; use bitfun_runtime_ports::SessionStoragePathRequest; - use dashmap::try_result::TryResult; + use dashmap::{try_result::TryResult, DashMap}; use serde_json::json; use std::collections::HashSet; use std::path::{Path, PathBuf}; use std::sync::Arc; - use std::time::Duration; + use std::time::{Duration, SystemTime}; use uuid::Uuid; struct TestWorkspace { @@ -6540,6 +6809,47 @@ mod tests { )); } + #[test] + fn idle_eviction_only_selects_sessions_that_can_be_restored() { + let now = SystemTime::now(); + let expired_at = now - Duration::from_secs(120); + let mut durable = Session::new( + "Durable".to_string(), + "agentic".to_string(), + SessionConfig::default(), + ); + durable.last_activity_at = expired_at; + let mut transient = Session::new( + "Connection scoped".to_string(), + "agentic".to_string(), + SessionConfig::default(), + ); + transient.last_activity_at = expired_at; + let durable_id = durable.session_id.clone(); + let transient_id = transient.session_id.clone(); + let sessions = DashMap::new(); + sessions.insert(durable_id.clone(), durable); + sessions.insert(transient_id.clone(), transient); + let transient_session_ids = DashMap::new(); + transient_session_ids.insert(transient_id.clone(), ()); + + let candidates = SessionManager::collect_expired_session_candidates( + &sessions, + &transient_session_ids, + now, + Duration::from_secs(60), + ); + + assert_eq!( + candidates + .iter() + .map(|candidate| candidate.session_id.as_str()) + .collect::>(), + [durable_id.as_str()] + ); + assert!(sessions.contains_key(&transient_id)); + } + #[test] fn persisted_round_preserves_deferred_wire_call_and_effective_identity() { let assistant = Message::assistant_with_tools( @@ -6712,6 +7022,62 @@ mod tests { assert_ne!(first.session_id, second.session_id); } + #[tokio::test] + async fn transient_session_cannot_bypass_owned_discard_through_unload() { + let workspace = TestWorkspace::new(); + let manager = in_memory_test_manager(); + let session = manager + .create_transient_session_with_id_and_details( + None, + "Connection Session".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + None, + SessionKind::Standard, + ) + .await + .expect("transient session should create"); + + let error = manager + .unload_session_from_memory(&session.session_id) + .await + .expect_err("transient session must use owned discard"); + assert!(error.to_string().contains("transient session")); + assert!(manager.get_session(&session.session_id).is_some()); + assert!(manager.is_transient_session(&session.session_id)); + } + + #[tokio::test] + async fn internal_delete_compensation_clears_transient_identity() { + let workspace = TestWorkspace::new(); + let manager = in_memory_test_manager(); + let session = manager + .create_transient_session_with_id_and_details( + None, + "Prepared Subagent".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + None, + SessionKind::Subagent, + ) + .await + .expect("transient subagent should create"); + + manager + .delete_session_by_id(&session.session_id) + .await + .expect("internal compensation should delete prepared subagent"); + + assert!(manager.get_session(&session.session_id).is_none()); + assert!(!manager.is_transient_session(&session.session_id)); + } + #[tokio::test] async fn restores_share_the_same_exact_active_session_capacity_as_creates() { let workspace = TestWorkspace::new(); @@ -7741,7 +8107,10 @@ mod tests { .await .expect("session should create"); - let snapshots = SessionManager::collect_auto_save_snapshots(&manager.sessions); + let snapshots = SessionManager::collect_auto_save_snapshots( + &manager.sessions, + &manager.transient_session_ids, + ); assert!(snapshots .iter() .any(|snapshot| snapshot.session_id == session.session_id)); diff --git a/src/crates/assembly/core/src/agentic/skill_agent_snapshot.rs b/src/crates/assembly/core/src/agentic/skill_agent_snapshot.rs index 51144990e7..ac5882e25f 100644 --- a/src/crates/assembly/core/src/agentic/skill_agent_snapshot.rs +++ b/src/crates/assembly/core/src/agentic/skill_agent_snapshot.rs @@ -6,6 +6,7 @@ use crate::agentic::tools::implementations::skills::{get_skill_registry, SkillIn use crate::agentic::tools::manifest_resolver::{resolve_tool_manifest, ResolvedToolManifest}; use crate::agentic::tools::product_runtime::GetToolSpecTool; use crate::agentic::tools::tool_context_runtime; +use crate::agentic::tools::ToolRuntimeRestrictions; use crate::agentic::workspace::WorkspaceServices; use crate::agentic::WorkspaceBinding; pub use bitfun_agent_runtime::skill_agent_snapshot::{ @@ -26,6 +27,7 @@ pub async fn resolve_skill_agent_snapshot( workspace_services: Option<&WorkspaceServices>, enable_tools: bool, context_vars: &std::collections::HashMap, + runtime_tool_restrictions: &ToolRuntimeRestrictions, ) -> SkillAgentSnapshotResolution { if !enable_tools { return SkillAgentSnapshotResolution { @@ -53,6 +55,7 @@ pub async fn resolve_skill_agent_snapshot( workspace_services, None, context_vars, + runtime_tool_restrictions, ); let manifest = resolve_tool_manifest( &tool_policy.allowed_tools, @@ -61,8 +64,14 @@ pub async fn resolve_skill_agent_snapshot( ) .await; - let snapshot = - build_skill_agent_snapshot(workspace, workspace_services, agent_type, &manifest).await; + let snapshot = build_skill_agent_snapshot( + workspace, + workspace_services, + agent_type, + &manifest, + runtime_tool_restrictions, + ) + .await; let tool_listing_sections = build_tool_listing_sections(&manifest, &snapshot); SkillAgentSnapshotResolution { @@ -76,6 +85,7 @@ async fn build_skill_agent_snapshot( workspace_services: Option<&WorkspaceServices>, agent_type: &str, manifest: &ResolvedToolManifest, + runtime_tool_restrictions: &ToolRuntimeRestrictions, ) -> TurnSkillAgentSnapshot { let has_tool = |tool_name: &str| { manifest @@ -91,7 +101,8 @@ async fn build_skill_agent_snapshot( } if has_tool("Task") { - snapshot.subagents = load_subagent_entries(workspace, Some(agent_type)).await; + snapshot.subagents = + load_subagent_entries(workspace, Some(agent_type), runtime_tool_restrictions).await; } snapshot @@ -172,6 +183,7 @@ fn skill_snapshot_entry_from_skill_info(skill: SkillInfo) -> SkillSnapshotEntry async fn load_subagent_entries( workspace: Option<&WorkspaceBinding>, agent_type: Option<&str>, + runtime_tool_restrictions: &ToolRuntimeRestrictions, ) -> Vec { let registry = get_agent_registry(); let workspace_root = workspace @@ -189,10 +201,17 @@ async fn load_subagent_entries( agents .into_iter() - .map(|agent| AgentSnapshotEntry { - id: agent.id, - description: agent.description, - default_tools: agent.default_tools, + .map(|agent| { + let default_tools = agent + .default_tools + .into_iter() + .filter(|tool_name| runtime_tool_restrictions.is_tool_allowed(tool_name)) + .collect(); + AgentSnapshotEntry { + id: agent.id, + description: agent.description, + default_tools, + } }) .collect() } @@ -218,3 +237,28 @@ pub async fn build_embedded_user_context_reminder( .build_user_context_reminder(user_context_policy) .await } + +#[cfg(test)] +mod tests { + use super::load_subagent_entries; + use crate::agentic::tools::ToolRuntimeRestrictions; + + #[tokio::test] + async fn subagent_projection_hides_runtime_denied_tools() { + let mut restrictions = ToolRuntimeRestrictions::default(); + restrictions + .denied_tool_names + .insert("ControlHub".to_string()); + + let agents = load_subagent_entries(None, Some("Claw"), &restrictions).await; + let computer_use = agents + .iter() + .find(|agent| agent.id == "ComputerUse") + .expect("Claw should advertise the ComputerUse subagent"); + + assert!(!computer_use + .default_tools + .iter() + .any(|tool_name| !restrictions.is_tool_allowed(tool_name))); + } +} diff --git a/src/crates/assembly/core/src/agentic/system.rs b/src/crates/assembly/core/src/agentic/system.rs index eecfd0923e..f4212a16b5 100644 --- a/src/crates/assembly/core/src/agentic/system.rs +++ b/src/crates/assembly/core/src/agentic/system.rs @@ -15,6 +15,7 @@ use crate::agentic::tools; use crate::infrastructure::ai::AIClientFactory; use crate::infrastructure::try_get_path_manager_arc; use crate::service::token_usage::{TokenUsageService, TokenUsageSubscriber}; +use bitfun_product_capabilities::DeliveryProfile; /// Agentic runtime state shared by host adapters. #[derive(Clone)] @@ -26,7 +27,27 @@ pub struct AgenticSystem { /// Initialize the agentic runtime and register the global coordinator. pub async fn init_agentic_system() -> Result { - info!("Initializing agentic system"); + init_agentic_system_for_profile(DeliveryProfile::ProductFull).await +} + +/// Select the process-wide Agent delivery profile before any service reads the +/// global tool registry. +/// +/// Product composition roots call this before configuration canonicalization; +/// later initialization verifies the same profile and rejects replacement. +pub fn select_agentic_system_profile(delivery_profile: DeliveryProfile) -> Result<()> { + tools::registry::initialize_global_tool_registry_for_profile(delivery_profile) + .map(|_| ()) + .map_err(anyhow::Error::msg) +} + +/// Initialize the single process-wide agentic runtime for one product profile. +pub async fn init_agentic_system_for_profile( + delivery_profile: DeliveryProfile, +) -> Result { + info!("Initializing agentic system for profile {delivery_profile}"); + + select_agentic_system_profile(delivery_profile)?; let _ai_client_factory = AIClientFactory::get_global().await?; diff --git a/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs b/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs index f7443a791f..36eb18ddae 100644 --- a/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs +++ b/src/crates/assembly/core/src/agentic/tools/pipeline/tool_pipeline.rs @@ -748,6 +748,7 @@ impl ToolPipeline { async fn register_permission_requests( &self, requests: Vec, + dialog_turn_id: &str, auto_approve: bool, ) -> BitFunResult> { let manager = self.permission_request_manager.as_ref().ok_or_else(|| { @@ -758,10 +759,15 @@ impl ToolPipeline { let receivers = if auto_approve { manager - .register_batch_non_interactive(requests.clone()) + .register_batch_non_interactive_for_turn( + requests.clone(), + dialog_turn_id.to_string(), + ) .await } else { - manager.register_batch(requests.clone()).await + manager + .register_batch_for_turn(requests.clone(), dialog_turn_id.to_string()) + .await } .map_err(|error| BitFunError::service(error.to_string()))?; @@ -870,8 +876,15 @@ impl ToolPipeline { .first() .and_then(|task_id| self.state_manager.get_task(task_id)) .is_some_and(|task| task.options.auto_approve_ask); + let dialog_turn_id = task_ids + .first() + .and_then(|task_id| self.state_manager.get_task(task_id)) + .map(|task| task.context.dialog_turn_id) + .ok_or_else(|| { + BitFunError::service("Permission batch lost its owning Dialog Turn".to_string()) + })?; let receivers = self - .register_permission_requests(batch_requests, auto_approve) + .register_permission_requests(batch_requests, &dialog_turn_id, auto_approve) .await?; let mut receivers_by_task = HashMap::>::new(); @@ -1058,8 +1071,12 @@ impl ToolPipeline { PermissionExecutionPlan::Rejected { reason } } PermissionPlanDraft::Requests(requests) => PermissionExecutionPlan::Awaiting( - self.register_permission_requests(requests, task.options.auto_approve_ask) - .await?, + self.register_permission_requests( + requests, + &task.context.dialog_turn_id, + task.options.auto_approve_ask, + ) + .await?, ), }; diff --git a/src/crates/assembly/core/src/agentic/tools/product_runtime.rs b/src/crates/assembly/core/src/agentic/tools/product_runtime.rs index 53faae70fe..cb0c0ad498 100644 --- a/src/crates/assembly/core/src/agentic/tools/product_runtime.rs +++ b/src/crates/assembly/core/src/agentic/tools/product_runtime.rs @@ -140,6 +140,15 @@ mod tests { ); } + #[test] + fn sdk_and_cli_profiles_current_tool_plan_ceilings_match_without_sharing_identity() { + let sdk = ProductToolRuntime::for_profile(DeliveryProfile::Sdk).create_registry(); + let cli = ProductToolRuntime::for_profile(DeliveryProfile::Cli).create_registry(); + + assert_eq!(sdk.get_tool_names(), cli.get_tool_names()); + assert_eq!(sdk.get_deferred_tool_names(), cli.get_deferred_tool_names()); + } + #[test] fn product_tool_runtime_keeps_no_direct_core_profiles_empty() { for profile in [ @@ -147,7 +156,6 @@ mod tests { DeliveryProfile::Remote, DeliveryProfile::Web, DeliveryProfile::MobileWeb, - DeliveryProfile::Sdk, ] { let runtime = ProductToolRuntime::for_profile(profile); let registry = runtime.create_registry(); diff --git a/src/crates/assembly/core/src/agentic/tools/product_runtime/catalog.rs b/src/crates/assembly/core/src/agentic/tools/product_runtime/catalog.rs index 6e40dca936..5b3e615fe3 100644 --- a/src/crates/assembly/core/src/agentic/tools/product_runtime/catalog.rs +++ b/src/crates/assembly/core/src/agentic/tools/product_runtime/catalog.rs @@ -148,17 +148,20 @@ impl ProductToolCatalogProvider { exposure_overrides: &AgentToolPolicyOverrides, context: &ToolUseContext, ) -> (Vec, AgentToolPolicyOverrides) { + let allowed_tools = allowed_tools + .iter() + .filter(|tool_name| context.runtime_tool_restrictions.is_tool_allowed(tool_name)) + .cloned() + .collect::>(); if Self::deferred_tool_loading_enabled(context) { - return (allowed_tools.to_vec(), exposure_overrides.clone()); + return (allowed_tools, exposure_overrides.clone()); } let allowed_tools = allowed_tools - .iter() + .into_iter() .filter(|tool_name| { - tool_name.as_str() != GET_TOOL_SPEC_TOOL_NAME - && tool_name.as_str() != CALL_DEFERRED_TOOL_NAME + tool_name != GET_TOOL_SPEC_TOOL_NAME && tool_name != CALL_DEFERRED_TOOL_NAME }) - .cloned() .collect::>(); let exposure_overrides = allowed_tools .iter() @@ -560,6 +563,48 @@ mod tests { ); } + #[tokio::test] + async fn runtime_restrictions_hide_tools_from_manifest_and_get_tool_spec() { + let allowed_tools = vec!["Read".to_string(), "WebFetch".to_string()]; + let mut context = tool_context(Some("agentic")); + context + .runtime_tool_restrictions + .denied_tool_names + .extend(["Read".to_string(), "WebFetch".to_string()]); + + let manifest = resolve_product_tool_manifest( + &allowed_tools, + &AgentToolPolicyOverrides::default(), + &context, + ) + .await; + assert!(!manifest + .allowed_tool_names + .iter() + .any(|name| name == "Read")); + assert!(!manifest + .allowed_tool_names + .iter() + .any(|name| name == "WebFetch")); + assert!(!manifest + .deferred_tool_names + .iter() + .any(|name| name == "WebFetch")); + assert!(!manifest + .tool_definitions + .iter() + .any(|definition| definition.name == "Read")); + + let deferred_names = ProductToolCatalogProvider + .deferred_tools_for_get_tool_spec(Some(&context)) + .await + .expect("contextual GetToolSpec catalog") + .into_iter() + .map(|tool| tool.name().to_string()) + .collect::>(); + assert!(!deferred_names.iter().any(|name| name == "WebFetch")); + } + #[tokio::test] async fn product_resolved_manifest_owner_matches_legacy_shape() { let allowed_tools = vec!["Read".to_string(), "WebFetch".to_string()]; diff --git a/src/crates/assembly/core/src/agentic/tools/registry.rs b/src/crates/assembly/core/src/agentic/tools/registry.rs index b7cde41b28..7c89bbb165 100644 --- a/src/crates/assembly/core/src/agentic/tools/registry.rs +++ b/src/crates/assembly/core/src/agentic/tools/registry.rs @@ -9,6 +9,7 @@ use bitfun_agent_tools::{ DynamicToolDescriptor, DynamicToolProvider, PortResult, ToolDecoratorRef, ToolRegistry as AgentToolRegistry, }; +use bitfun_product_capabilities::DeliveryProfile; use log::{debug, info, trace, warn}; use std::sync::Arc; @@ -34,6 +35,10 @@ impl ToolRegistry { ProductToolRuntime::default().create_registry() } + pub(in crate::agentic) fn for_profile(profile: DeliveryProfile) -> Self { + ProductToolRuntime::for_profile(profile).create_registry() + } + /// Create a registry with an injected decoration boundary. /// /// The default production decorator preserves snapshot-aware wrapping while @@ -238,15 +243,55 @@ pub fn create_tool_registry() -> ToolRegistry { use std::sync::OnceLock; use tokio::sync::RwLock as TokioRwLock; -static GLOBAL_TOOL_REGISTRY: OnceLock>> = OnceLock::new(); +struct GlobalToolRegistry { + profile: DeliveryProfile, + registry: Arc>, +} + +static GLOBAL_TOOL_REGISTRY: OnceLock = OnceLock::new(); + +pub(in crate::agentic) fn initialize_global_tool_registry_for_profile( + profile: DeliveryProfile, +) -> Result>, String> { + if let Some(global) = GLOBAL_TOOL_REGISTRY.get() { + return if global.profile == profile { + Ok(global.registry.clone()) + } else { + Err(format!( + "Global tool registry already uses delivery profile {}; cannot replace it with {}", + global.profile, profile + )) + }; + } + + let candidate = GlobalToolRegistry { + profile, + registry: Arc::new(TokioRwLock::new(ToolRegistry::for_profile(profile))), + }; + let _ = GLOBAL_TOOL_REGISTRY.set(candidate); + let global = GLOBAL_TOOL_REGISTRY + .get() + .expect("global tool registry must be initialized"); + if global.profile != profile { + return Err(format!( + "Global tool registry concurrently selected delivery profile {}; requested {}", + global.profile, profile + )); + } + Ok(global.registry.clone()) +} /// Get global tool registry pub fn get_global_tool_registry() -> Arc> { GLOBAL_TOOL_REGISTRY .get_or_init(|| { info!("Initializing global tool registry"); - Arc::new(TokioRwLock::new(ToolRegistry::new())) + GlobalToolRegistry { + profile: DeliveryProfile::ProductFull, + registry: Arc::new(TokioRwLock::new(ToolRegistry::new())), + } }) + .registry .clone() } diff --git a/src/crates/assembly/core/src/agentic/tools/tool_context_runtime.rs b/src/crates/assembly/core/src/agentic/tools/tool_context_runtime.rs index 6cb879d343..4341b1c374 100644 --- a/src/crates/assembly/core/src/agentic/tools/tool_context_runtime.rs +++ b/src/crates/assembly/core/src/agentic/tools/tool_context_runtime.rs @@ -35,7 +35,9 @@ use bitfun_agent_runtime::checkpoint::{ build_light_checkpoint as build_runtime_light_checkpoint, GitStatusCheckpointFacts, LightCheckpointWorkspaceFacts, }; +use bitfun_agent_runtime::permission::AUTO_APPROVE_ASK_CONTEXT_KEY; use bitfun_agent_runtime::remote_file_delivery::TOOL_CONTEXT_REMOTE_FILE_DELIVERY_KEY; +use bitfun_agent_runtime::user_questions::USER_INPUT_AVAILABLE_CONTEXT_KEY; use bitfun_agent_tools::{ LoadedDeferredToolSpec, PortableToolContextProvider, ToolContextFacts, ToolWorkspaceKind, }; @@ -267,6 +269,7 @@ pub(crate) fn build_tool_description_context( workspace_services: Option<&WorkspaceServices>, primary_model_facts: Option<&PrimaryModelFacts>, context_vars: &HashMap, + runtime_tool_restrictions: &ToolRuntimeRestrictions, ) -> ToolUseContext { let mut custom_data = HashMap::new(); let primary_model_facts = primary_model_facts.cloned().unwrap_or_default(); @@ -284,7 +287,7 @@ pub(crate) fn build_tool_description_context( primary_model_facts, custom_data, computer_use_host: None, - runtime_tool_restrictions: ToolRuntimeRestrictions::default(), + runtime_tool_restrictions: runtime_tool_restrictions.clone(), runtime_handles: core_tool_runtime_handles(workspace_services.cloned(), None, None, None), } } @@ -322,6 +325,19 @@ fn build_tool_context_custom_data(context: &ToolExecutionContext) -> HashMap Some(true), + Some("false") => Some(false), + _ => None, + }; + if let Some(value) = value { + extension_custom_data.insert(key.to_string(), Value::Bool(value)); + } + } build_tool_runtime_custom_data(ToolRuntimeCustomDataInput { context_vars: &context.context_vars, delegation_policy: context.delegation_policy, @@ -1353,12 +1369,17 @@ mod call_runtime_tests { #[cfg(test)] mod context_builder_tests { use super::build_tool_description_context; + use crate::agentic::tools::ToolRuntimeRestrictions; use std::collections::HashMap; use tool_runtime::context::PrimaryModelFacts; #[test] fn tool_description_context_preserves_manifest_custom_data_shape() { let context_vars = HashMap::new(); + let mut runtime_tool_restrictions = ToolRuntimeRestrictions::default(); + runtime_tool_restrictions + .denied_tool_names + .insert("Write".to_string()); let context = build_tool_description_context( "coding", @@ -1371,6 +1392,7 @@ mod context_builder_tests { true, )), &context_vars, + &runtime_tool_restrictions, ); assert_eq!(context.agent_type.as_deref(), Some("coding")); @@ -1381,7 +1403,7 @@ mod context_builder_tests { assert!(context.loaded_deferred_tool_specs.is_empty()); assert!(context.cancellation_token().is_none()); assert!(context.workspace_services().is_none()); - assert!(context.runtime_tool_restrictions.is_tool_allowed("Write")); + assert!(!context.runtime_tool_restrictions.is_tool_allowed("Write")); assert!(context.primary_model_supports_image_understanding()); assert_eq!(context.primary_model_facts().model_id, "model_1"); assert_eq!(context.primary_model_facts().model_name, "vision-model"); @@ -1403,6 +1425,8 @@ mod task_context_tests { SubagentParentInfo, ToolExecutionContext, ToolExecutionOptions, ToolTask, }; use crate::agentic::tools::ToolRuntimeRestrictions; + use bitfun_agent_runtime::permission::AUTO_APPROVE_ASK_CONTEXT_KEY; + use bitfun_agent_runtime::user_questions::USER_INPUT_AVAILABLE_CONTEXT_KEY; use bitfun_agent_tools::LoadedDeferredToolSpec; use bitfun_runtime_ports::DelegationPolicy; use serde_json::json; @@ -1421,6 +1445,14 @@ mod task_context_tests { let mut context_vars = HashMap::new(); context_vars.insert("turn_index".to_string(), "7".to_string()); context_vars.insert("acp_transport".to_string(), "true".to_string()); + context_vars.insert( + USER_INPUT_AVAILABLE_CONTEXT_KEY.to_string(), + "false".to_string(), + ); + context_vars.insert( + AUTO_APPROVE_ASK_CONTEXT_KEY.to_string(), + "false".to_string(), + ); context_vars.insert( "deep_review_run_manifest".to_string(), r#"{"run_id":"run-1"}"#.to_string(), @@ -1513,6 +1545,14 @@ mod task_context_tests { .custom_data .contains_key("primary_model_supports_image_understanding")); assert_eq!(context.custom_data["acp_transport"], json!(true)); + assert_eq!( + context.custom_data[USER_INPUT_AVAILABLE_CONTEXT_KEY], + json!(false) + ); + assert_eq!( + context.custom_data[AUTO_APPROVE_ASK_CONTEXT_KEY], + json!(false) + ); assert_eq!( context.custom_data["deep_review_run_manifest"], json!({ "run_id": "run-1" }) diff --git a/src/crates/assembly/core/src/product_runtime.rs b/src/crates/assembly/core/src/product_runtime.rs index ee29d371ec..49bf56900d 100644 --- a/src/crates/assembly/core/src/product_runtime.rs +++ b/src/crates/assembly/core/src/product_runtime.rs @@ -13,8 +13,8 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; use bitfun_agent_runtime::permission::PermissionRequestManager; use bitfun_agent_runtime::sdk::{ - AgentEventSource, AgentRuntime, AgentSessionForkAtTurnRequest, AgentSessionForkPort, - AgentSessionForkRequest, AgentSessionForkResult, AgentSessionUsagePort, + AgentEventReceiver, AgentEventSource, AgentRuntime, AgentSessionForkAtTurnRequest, + AgentSessionForkPort, AgentSessionForkRequest, AgentSessionForkResult, AgentSessionUsagePort, AgentSessionUsageRequest, AgentTurnSettlementPort, AgentTurnSettlementRequest, }; use bitfun_harness::HarnessRegistry; @@ -31,6 +31,7 @@ use crate::agentic::coordination::{ ConversationCoordinator, DialogScheduler, SessionMaintenancePermit, }; use crate::agentic::core::Session; +use crate::agentic::events::EventQueue; use crate::agentic::keyed_lock::KeyedAsyncLockGuard; use crate::agentic::persistence::session_branch::SessionBranchRequest; use crate::agentic::persistence::{PersistenceManager, SessionMetadataPage}; @@ -46,7 +47,73 @@ use crate::service_agent_runtime::CoreServiceAgentRuntime; use crate::util::errors::{BitFunError, BitFunResult}; pub use bitfun_product_capabilities::ProductRuntimeAssembly as CoreProductRuntimeAssembly; -pub use runtime_services::CoreRuntimeServicesProvider; +pub use runtime_services::{build_local_runtime_services, CoreRuntimeServicesProvider}; + +struct ProductEventQueueDrain { + task: tokio::task::JoinHandle<()>, +} + +impl ProductEventQueueDrain { + fn start(queue: Arc) -> Self { + let task = tokio::spawn(async move { + loop { + queue.wait_for_events().await; + while !queue.dequeue_configured_batch().await.is_empty() {} + } + }); + Self { task } + } +} + +impl Drop for ProductEventQueueDrain { + fn drop(&mut self) { + self.task.abort(); + } +} + +/// Shared product-host event source that keeps the legacy queue bounded. +#[derive(Clone)] +pub struct CoreProductAgentEventSource { + source: AgentEventSource, + _drain: Arc, +} + +impl CoreProductAgentEventSource { + pub fn new(queue: Arc) -> Self { + Self { + source: AgentEventSource::new(queue.clone()), + _drain: Arc::new(ProductEventQueueDrain::start(queue)), + } + } + + pub fn subscribe(&self) -> AgentEventReceiver { + self.source.subscribe() + } + + pub fn runtime_source(&self) -> AgentEventSource { + self.source.clone() + } +} + +/// Returns the process-shared dialog scheduler used by sibling product hosts. +pub fn ensure_product_dialog_scheduler( + agentic_system: &crate::agentic::system::AgenticSystem, +) -> Arc { + if let Some(scheduler) = crate::agentic::coordination::get_global_scheduler() { + return scheduler; + } + + let session_manager = agentic_system.coordinator.get_session_manager().clone(); + let scheduler = DialogScheduler::new(agentic_system.coordinator.clone(), session_manager); + agentic_system + .coordinator + .set_scheduler_notifier(scheduler.outcome_sender()); + agentic_system + .coordinator + .set_round_injection_source(scheduler.round_injection_monitor()); + crate::agentic::coordination::set_global_scheduler(scheduler.clone()); + scheduler +} #[derive(Debug, Clone, Copy, Default)] struct SystemPermissionClock; @@ -316,6 +383,35 @@ impl CoreProductAgentRuntime { harness_registry, ) } + + /// Build the standalone Agent SDK Host implementation candidate. + /// + /// The Host receives the same Core-owned session, query, event, cancellation, + /// and settlement capabilities as the CLI product runtime. It does not own + /// a second agent loop or protocol-specific execution services. + pub fn build_sdk_host( + coordinator: Arc, + scheduler: Arc, + token_usage_service: Arc, + event_source: AgentEventSource, + services: RuntimeServices, + harness_registry: HarnessRegistry, + ) -> Result { + let session_operations = Arc::new(CoreSessionOperationsPort::new( + coordinator.clone(), + token_usage_service, + )); + CoreServiceAgentRuntime::sdk_host_product_agent_runtime( + coordinator, + scheduler, + event_source, + session_operations.clone(), + session_operations.clone(), + session_operations, + services, + harness_registry, + ) + } } /// Core-owned compatibility boundary for product operations not yet exposed by @@ -963,8 +1059,8 @@ mod tests { use super::{ generate_core_session_usage_report, latest_persisted_turn_id, runtime_port_error, validate_latest_turn_fork_scope, validate_persisted_session_id, - CoreAgentRuntimeCompatibility, CoreLocalWorkspaceSnapshot, CoreProductAgentRuntime, - CoreSessionOperationsPort, + CoreAgentRuntimeCompatibility, CoreLocalWorkspaceSnapshot, CoreProductAgentEventSource, + CoreProductAgentRuntime, CoreSessionOperationsPort, }; use crate::agentic::coordination::{ConversationCoordinator, DialogScheduler}; use crate::agentic::events::{EventQueue, EventQueueConfig, EventRouter}; @@ -990,6 +1086,7 @@ mod tests { use bitfun_agent_runtime::sdk::{ AgentSessionForkPort, AgentSessionForkRequest, AgentSessionUsageRequest, PortErrorKind, }; + use bitfun_events::AgenticEvent; use tokio::sync::RwLock as TokioRwLock; struct TestWorkspace { @@ -1024,6 +1121,62 @@ mod tests { } } + #[tokio::test] + async fn product_event_source_broadcasts_while_draining_the_legacy_queue() { + let queue = Arc::new(EventQueue::new(EventQueueConfig { + max_queue_size: 4, + batch_size: 2, + })); + let source = CoreProductAgentEventSource::new(queue.clone()); + let mut first = source.subscribe(); + let mut second = source.subscribe(); + + for index in 0..32 { + queue + .enqueue( + AgenticEvent::SessionStateChanged { + session_id: "session-1".to_string(), + new_state: format!("state-{index}"), + }, + None, + ) + .await + .expect("enqueue event"); + } + + let mut last_first = None; + let mut last_second = None; + for _ in 0..32 { + last_first = Some( + tokio::time::timeout(Duration::from_secs(1), first.recv()) + .await + .expect("first subscriber must not stall") + .expect("first subscriber event"), + ); + last_second = Some( + tokio::time::timeout(Duration::from_secs(1), second.recv()) + .await + .expect("second subscriber must not stall") + .expect("second subscriber event"), + ); + } + + let last_first = last_first.expect("last first event"); + let last_second = last_second.expect("last second event"); + assert_eq!(last_first.id, last_second.id); + assert!(matches!( + last_first.event, + AgenticEvent::SessionStateChanged { ref new_state, .. } if new_state == "state-31" + )); + tokio::time::timeout(Duration::from_secs(1), async { + while !queue.is_empty().await { + tokio::task::yield_now().await; + } + }) + .await + .expect("queue drainer must keep the legacy queue bounded"); + } + #[test] fn product_agent_runtime_exposes_reviewed_full_and_narrow_builders() { fn build( @@ -1045,6 +1198,7 @@ mod tests { let _ = build; let _ = CoreProductAgentRuntime::build_session_surface; let _ = CoreProductAgentRuntime::build_acp; + let _ = CoreProductAgentRuntime::build_sdk_host; } #[test] diff --git a/src/crates/assembly/core/src/product_runtime/runtime_services.rs b/src/crates/assembly/core/src/product_runtime/runtime_services.rs index b6fb081420..2eb192847d 100644 --- a/src/crates/assembly/core/src/product_runtime/runtime_services.rs +++ b/src/crates/assembly/core/src/product_runtime/runtime_services.rs @@ -3,16 +3,19 @@ //! This file registers existing core concrete adapters into typed runtime //! service builders. It does not create new runtime behavior. +use std::path::{Path, PathBuf}; use std::sync::Arc; #[cfg(feature = "ssh-remote")] -use bitfun_runtime_ports::{PortError, PortErrorKind, PortResult, RemoteExecPort}; +use bitfun_runtime_ports::{PortError, PortErrorKind, RemoteExecPort}; use bitfun_runtime_ports::{ - RemoteProjectionPort, RemoteWorkspacePort, SessionStorePort, TerminalPort, + PortResult, RemoteProjectionPort, RemoteWorkspacePort, SessionStorePort, TerminalPort, }; use bitfun_runtime_services::{ - RuntimeServiceMarkerPort, RuntimeServicesBuilder, RuntimeServicesProvider, + RuntimeServiceMarkerPort, RuntimeServices, RuntimeServicesBuilder, RuntimeServicesProvider, + RuntimeServicesRegistry, }; +use bitfun_services_core::local_runtime_ports::LocalRuntimePorts; use terminal_core::TerminalRuntimePort; use crate::agentic::session::CoreSessionStorePort; @@ -105,3 +108,83 @@ impl RuntimeServicesProvider for CoreRuntimeServicesProvider { } } } + +#[derive(Clone)] +struct CoreLocalRuntimeServicesProvider { + ports: LocalRuntimePorts, +} + +impl CoreLocalRuntimeServicesProvider { + fn new(workspace_root: impl AsRef, event_capacity: usize) -> anyhow::Result { + Ok(Self { + ports: LocalRuntimePorts::new(workspace_root, event_capacity)?, + }) + } +} + +impl RuntimeServicesProvider for CoreLocalRuntimeServicesProvider { + fn register(&self, builder: RuntimeServicesBuilder) -> RuntimeServicesBuilder { + builder + .with_filesystem(self.ports.filesystem()) + .with_workspace(self.ports.workspace()) + .with_events(self.ports.events()) + .with_clock(self.ports.clock()) + } +} + +/// Builds the shared local process service set used by sibling product hosts. +/// +/// The caller remains the composition root and selects its delivery profile. +/// This function only binds the existing Core services and the required local +/// workspace, filesystem, event, and clock ports. +pub fn build_local_runtime_services( + workspace_root: impl AsRef, + event_capacity: usize, +) -> anyhow::Result<(PathBuf, RuntimeServices)> { + let local = CoreLocalRuntimeServicesProvider::new(workspace_root, event_capacity)?; + let canonical_root = local.ports.workspace_root().to_path_buf(); + let services = RuntimeServicesRegistry::new() + .with_provider(CoreRuntimeServicesProvider::new()) + .with_provider(local) + .build(RuntimeServicesBuilder::new())?; + Ok((canonical_root, services)) +} + +#[cfg(test)] +mod local_runtime_tests { + use super::build_local_runtime_services; + use bitfun_runtime_ports::RuntimeServiceCapability; + + #[test] + fn local_runtime_services_bind_required_core_and_workspace_ports() { + let workspace = tempfile::tempdir().expect("workspace"); + let (canonical_root, services) = + build_local_runtime_services(workspace.path(), 8).expect("local runtime services"); + + assert_eq!( + canonical_root, + dunce::canonicalize(workspace.path()).unwrap() + ); + for capability in [ + RuntimeServiceCapability::FileSystem, + RuntimeServiceCapability::Workspace, + RuntimeServiceCapability::SessionStore, + RuntimeServiceCapability::Events, + RuntimeServiceCapability::Clock, + RuntimeServiceCapability::Terminal, + RuntimeServiceCapability::Network, + RuntimeServiceCapability::Git, + ] { + assert!(services.has_capability(capability), "missing {capability}"); + } + assert!(services.clock.now_unix_millis() > 0); + } + + #[test] + fn local_runtime_services_reject_a_missing_workspace() { + let temp = tempfile::tempdir().expect("tempdir"); + let error = build_local_runtime_services(temp.path().join("missing"), 8) + .expect_err("missing workspace must fail"); + assert!(error.to_string().contains("workspace"), "{error}"); + } +} diff --git a/src/crates/assembly/core/src/service_agent_runtime.rs b/src/crates/assembly/core/src/service_agent_runtime.rs index 22ee348e98..1ceaa7ada0 100644 --- a/src/crates/assembly/core/src/service_agent_runtime.rs +++ b/src/crates/assembly/core/src/service_agent_runtime.rs @@ -13,11 +13,12 @@ use bitfun_agent_runtime::sdk::{ }; use bitfun_runtime_ports::{ AgentDialogTurnPort, AgentDialogTurnRequest, AgentInputAttachment, AgentLifecycleDeliveryPort, - AgentLocalCommandTurnPort, AgentSessionCreateRequest, AgentSessionManagementPort, - AgentSubmissionPort, AgentSubmissionSource, AgentThreadGoalManagementPort, - AgentTurnCancellationPort, AgentTurnCancellationRequest, RemoteControlStatePort, - RemoteControlStateRequest, RemoteControlStateSnapshot, RemoteSessionWorkspaceIdentity, - RuntimeServiceCapability, RuntimeServicePort, SessionStoragePathRequest, SessionStorePort, + AgentLocalCommandTurnPort, AgentSessionClosePort, AgentSessionCreateRequest, + AgentSessionManagementPort, AgentSubmissionPort, AgentSubmissionSource, + AgentThreadGoalManagementPort, AgentTurnCancellationPort, AgentTurnCancellationRequest, + RemoteControlStatePort, RemoteControlStateRequest, RemoteControlStateSnapshot, + RemoteSessionWorkspaceIdentity, RuntimeServiceCapability, RuntimeServicePort, + SessionStoragePathRequest, SessionStorePort, }; use bitfun_services_integrations::remote_connect::{ agent_input_attachment_from_remote_image_context, build_remote_chat_messages, @@ -497,7 +498,11 @@ impl AgentSessionManagementPort for ScheduledSessionManagementPort { })?; let _maintenance = self .scheduler - .begin_session_deletion(&request.session_id, &storage_path, Duration::from_secs(2)) + .begin_session_deletion( + &request.session_id, + &storage_path, + Duration::from_millis(2_000), + ) .await .map_err(|error| { let kind = match error { @@ -554,6 +559,89 @@ impl AgentSessionManagementPort for ScheduledSessionManagementPort { } } +#[async_trait::async_trait] +impl AgentSessionClosePort for ScheduledSessionManagementPort { + async fn discard_transient_session( + &self, + request: bitfun_runtime_ports::AgentTransientSessionDiscardRequest, + ) -> bitfun_runtime_ports::PortResult { + bitfun_core_types::validate_session_id(&request.session_id).map_err(|message| { + bitfun_runtime_ports::PortError::new( + bitfun_runtime_ports::PortErrorKind::InvalidRequest, + message, + ) + })?; + let storage_path = CoreSessionStorePort::default() + .resolve_session_storage_path(SessionStoragePathRequest { + workspace_path: std::path::PathBuf::from(&request.workspace_path), + remote_connection_id: request.remote_connection_id.clone(), + remote_ssh_host: request.remote_ssh_host.clone(), + }) + .await? + .effective_storage_path; + let session_manager = self.coordinator.get_session_manager(); + session_manager + .validate_session_storage_path_binding(&request.session_id, &storage_path) + .map_err(map_session_close_error)?; + let close_deadline = + tokio::time::Instant::now() + Duration::from_millis(request.wait_timeout_ms.max(1)); + let _maintenance = self + .scheduler + .begin_session_maintenance( + &request.session_id, + &storage_path, + close_deadline.saturating_duration_since(tokio::time::Instant::now()), + ) + .await + .map_err(map_session_close_error)?; + let cleanup_budget = close_deadline.saturating_duration_since(tokio::time::Instant::now()); + if cleanup_budget.is_zero() { + return Err(bitfun_runtime_ports::PortError::new( + bitfun_runtime_ports::PortErrorKind::Timeout, + "Session close deadline was exhausted before transient resource cleanup", + )); + } + tokio::time::timeout( + cleanup_budget, + self.coordinator.discard_transient_session( + std::path::Path::new(&request.workspace_path), + request.remote_connection_id.as_deref(), + request.remote_ssh_host.as_deref(), + &request.session_id, + ), + ) + .await + .map_err(|_| { + bitfun_runtime_ports::PortError::new( + bitfun_runtime_ports::PortErrorKind::Timeout, + "Transient Session resource cleanup exceeded the Session close deadline", + ) + })? + .map_err(map_session_close_error) + } +} + +fn map_session_close_error( + error: crate::util::errors::BitFunError, +) -> bitfun_runtime_ports::PortError { + let kind = match &error { + crate::util::errors::BitFunError::Validation(_) => { + bitfun_runtime_ports::PortErrorKind::InvalidRequest + } + crate::util::errors::BitFunError::NotFound(_) => { + bitfun_runtime_ports::PortErrorKind::NotFound + } + crate::util::errors::BitFunError::Timeout(_) => { + bitfun_runtime_ports::PortErrorKind::Timeout + } + crate::util::errors::BitFunError::Cancelled(_) => { + bitfun_runtime_ports::PortErrorKind::Cancelled + } + _ => bitfun_runtime_ports::PortErrorKind::Backend, + }; + bitfun_runtime_ports::PortError::new(kind, error.to_string()) +} + fn scheduled_session_management_port( coordinator: Arc, scheduler: Arc, @@ -561,6 +649,13 @@ fn scheduled_session_management_port( Arc::new(ScheduledSessionManagementPort::new(coordinator, scheduler)) } +fn scheduled_session_close_port( + coordinator: Arc, + scheduler: Arc, +) -> Arc { + Arc::new(ScheduledSessionManagementPort::new(coordinator, scheduler)) +} + pub(crate) struct CoreServiceAgentRuntime; impl CoreServiceAgentRuntime { @@ -846,6 +941,7 @@ impl CoreServiceAgentRuntime { let submission: Arc = coordinator.clone(); let session_management = scheduled_session_management_port(coordinator.clone(), scheduler.clone()); + let session_close = scheduled_session_close_port(coordinator.clone(), scheduler.clone()); let session_mode: Arc = coordinator.clone(); let session_model: Arc = coordinator.clone(); let session_restore: Arc = coordinator.clone(); @@ -869,6 +965,7 @@ impl CoreServiceAgentRuntime { cancellation, interaction_response, )? + .with_session_close_port(session_close) .with_dialog_turn_port(dialog_turn) .with_lifecycle_delivery_port(lifecycle_delivery) .build() @@ -1022,6 +1119,30 @@ impl CoreServiceAgentRuntime { ) } + pub(crate) fn sdk_host_product_agent_runtime( + coordinator: Arc, + scheduler: Arc, + event_source: AgentEventSource, + session_fork: Arc, + session_usage: Arc, + turn_settlement: Arc, + services: bitfun_runtime_services::RuntimeServices, + harness_registry: bitfun_harness::HarnessRegistry, + ) -> Result { + let dialog_turn: Arc = scheduler.clone(); + Self::product_agent_runtime_with_dialog_turn( + coordinator, + scheduler, + dialog_turn, + Some(event_source), + Some(session_fork), + Some(session_usage), + Some(turn_settlement), + services, + harness_registry, + ) + } + fn product_agent_runtime_with_dialog_turn( coordinator: Arc, scheduler: Arc, @@ -1036,6 +1157,7 @@ impl CoreServiceAgentRuntime { let submission: Arc = coordinator.clone(); let session_management = scheduled_session_management_port(coordinator.clone(), scheduler.clone()); + let session_close = scheduled_session_close_port(coordinator.clone(), scheduler.clone()); let session_mode: Arc = coordinator.clone(); let session_model: Arc = coordinator.clone(); let session_restore: Arc = coordinator.clone(); @@ -1059,6 +1181,7 @@ impl CoreServiceAgentRuntime { cancellation, interaction_response, )? + .with_session_close_port(session_close) .with_dialog_turn_port(dialog_turn) .with_lifecycle_delivery_port(lifecycle_delivery); let builder = match event_source { @@ -1772,7 +1895,10 @@ mod tests { { } + fn assert_session_lifecycle_port() {} + assert_scheduler_ports::(); + assert_session_lifecycle_port::(); } #[test] diff --git a/src/crates/assembly/product-capabilities/src/lib.rs b/src/crates/assembly/product-capabilities/src/lib.rs index 1c3a692439..2f7bf17076 100644 --- a/src/crates/assembly/product-capabilities/src/lib.rs +++ b/src/crates/assembly/product-capabilities/src/lib.rs @@ -257,7 +257,7 @@ const PRODUCT_DELIVERY_PROFILE_ENTRIES: &[ProductDeliveryProfileEntry] = &[ ), ProductDeliveryProfileEntry::new( DeliveryProfile::Sdk, - ProductCoreDependencyMode::NoDirectCoreDependency, + ProductCoreDependencyMode::ProductFullCompatibility, ), ]; @@ -1081,11 +1081,13 @@ fn product_capability_registry_for_profile(profile: DeliveryProfile) -> ProductC DeliveryProfile::ProductFull | DeliveryProfile::Desktop | DeliveryProfile::Cli - | DeliveryProfile::Acp => default_product_capability_registry(), + | DeliveryProfile::Acp + | DeliveryProfile::Sdk => default_product_capability_registry(), DeliveryProfile::Server | DeliveryProfile::Remote | DeliveryProfile::Web - | DeliveryProfile::MobileWeb - | DeliveryProfile::Sdk => ProductCapabilityRegistry::new(EMPTY_PRODUCT_CAPABILITY_PACKS), + | DeliveryProfile::MobileWeb => { + ProductCapabilityRegistry::new(EMPTY_PRODUCT_CAPABILITY_PACKS) + } } } diff --git a/src/crates/assembly/product-capabilities/tests/plugin_product_shape.rs b/src/crates/assembly/product-capabilities/tests/plugin_product_shape.rs index 495d8aa1e9..dfbc19fad1 100644 --- a/src/crates/assembly/product-capabilities/tests/plugin_product_shape.rs +++ b/src/crates/assembly/product-capabilities/tests/plugin_product_shape.rs @@ -159,7 +159,7 @@ fn non_p0_surfaces_cannot_inherit_executable_plugin_host() { DeliveryProfile::MobileWeb, DeliveryProfile::Sdk, ] { - let services = if matches!(profile, DeliveryProfile::Acp) { + let services = if matches!(profile, DeliveryProfile::Acp | DeliveryProfile::Sdk) { product_full_services() } else { baseline_services() @@ -219,6 +219,7 @@ fn default_assembled_product_shapes_keep_profile_specific_plugin_availability() | DeliveryProfile::Desktop | DeliveryProfile::Cli | DeliveryProfile::Acp + | DeliveryProfile::Sdk ) { product_full_services() } else { diff --git a/src/crates/assembly/product-capabilities/tests/product_capabilities.rs b/src/crates/assembly/product-capabilities/tests/product_capabilities.rs index 2fbb30cf35..755e7cd099 100644 --- a/src/crates/assembly/product-capabilities/tests/product_capabilities.rs +++ b/src/crates/assembly/product-capabilities/tests/product_capabilities.rs @@ -274,7 +274,6 @@ fn no_direct_core_profiles_do_not_select_product_full_runtime_capabilities() { DeliveryProfile::Remote, DeliveryProfile::Web, DeliveryProfile::MobileWeb, - DeliveryProfile::Sdk, ] { let plan = product_assembly_plan_for_profile(profile); @@ -358,7 +357,7 @@ fn product_delivery_profile_matrix_documents_current_core_dependency_shape() { ), ( DeliveryProfile::Sdk, - ProductCoreDependencyMode::NoDirectCoreDependency, + ProductCoreDependencyMode::ProductFullCompatibility, ), ] ); @@ -724,7 +723,6 @@ fn product_assembler_allows_no_direct_core_profiles_without_product_services() { DeliveryProfile::Remote, DeliveryProfile::Web, DeliveryProfile::MobileWeb, - DeliveryProfile::Sdk, ] { let services = FakeRuntimeServicesProvider::with_all_required() .build_services() diff --git a/src/crates/assembly/product-capabilities/tests/product_sdk_assembly.rs b/src/crates/assembly/product-capabilities/tests/product_sdk_assembly.rs index 48a55a9bfb..69434608dd 100644 --- a/src/crates/assembly/product-capabilities/tests/product_sdk_assembly.rs +++ b/src/crates/assembly/product-capabilities/tests/product_sdk_assembly.rs @@ -52,12 +52,6 @@ impl AgentSubmissionPort for ProductSdkAgentProvider { } } -fn baseline_sdk_services() -> RuntimeServices { - FakeRuntimeServicesProvider::with_all_required() - .build_services() - .expect("baseline SDK services should build") -} - fn product_full_compatible_services() -> RuntimeServices { FakeRuntimeServicesProvider::with_all_required() .register(RuntimeServicesBuilder::new()) @@ -69,28 +63,34 @@ fn product_full_compatible_services() -> RuntimeServices { } #[tokio::test] -async fn sdk_delivery_profile_builds_minimal_agent_runtime_without_product_full_capabilities() { +async fn sdk_delivery_profile_builds_shared_runtime_owner_ceiling_without_bitfun_core() { let parts = ProductAssembler::new() .assemble(ProductAssemblyInput::new( DeliveryProfile::Sdk, - baseline_sdk_services(), + product_full_compatible_services(), )) - .expect("SDK delivery profile should assemble without product-full services"); + .expect("SDK delivery profile should assemble with its shared runtime services"); + let cli_plan = + bitfun_product_capabilities::product_assembly_plan_for_profile(DeliveryProfile::Cli); assert_eq!(parts.plan().profile(), DeliveryProfile::Sdk); - assert!(parts.plan().capability_set().ids().is_empty()); - assert!(parts.service_availability().is_empty()); + assert_eq!( + parts.plan().capability_set().ids(), + cli_plan.capability_set().ids(), + "SDK and Headless CLI currently select the same assembly-plan ceiling without sharing product identity" + ); assert!(parts.missing_service_requirements().is_empty()); - assert!(parts.harness_registry().provider_ids().is_empty()); - assert!(!parts - .services() - .has_capability(RuntimeServiceCapability::Terminal)); - assert!(!parts - .services() - .has_capability(RuntimeServiceCapability::Git)); - assert!(!parts - .services() - .has_capability(RuntimeServiceCapability::Network)); + assert_eq!( + parts.harness_registry().provider_ids(), + vec!["core.deep_review", "core.deep_research", "core.miniapp"] + ); + for capability in [ + RuntimeServiceCapability::Terminal, + RuntimeServiceCapability::Git, + RuntimeServiceCapability::Network, + ] { + assert!(parts.services().has_capability(capability)); + } let (services, harness_registry, plugin_runtime) = parts.into_runtime_parts(); let provider = Arc::new(ProductSdkAgentProvider::default()); @@ -100,7 +100,7 @@ async fn sdk_delivery_profile_builds_minimal_agent_runtime_without_product_full_ .with_harness_registry(Arc::new(harness_registry)) .with_plugin_runtime(plugin_runtime) .build() - .expect("SDK profile parts should build a minimal runtime"); + .expect("SDK profile parts should build a runtime from the shared owner contracts"); let handle = runtime .run(AgentRunRequest::new( @@ -108,12 +108,15 @@ async fn sdk_delivery_profile_builds_minimal_agent_runtime_without_product_full_ "hello from sdk profile", )) .await - .expect("SDK delivery profile runtime should accept a minimal run"); + .expect("SDK delivery profile runtime should accept a run"); assert_eq!(handle.session_id, "product-sdk-session"); assert_eq!(handle.turn_id, "product-sdk-turn"); assert!(handle.accepted); - assert!(runtime.harness_provider_ids().is_empty()); + assert_eq!( + runtime.harness_provider_ids(), + vec!["core.deep_review", "core.deep_research", "core.miniapp"] + ); } #[tokio::test] @@ -162,7 +165,7 @@ async fn product_runtime_parts_can_build_agent_runtime_sdk_without_core() { "hello from product assembly", ) .with_turn_id("product-sdk-turn") - .with_source(AgentSubmissionSource::Cli), + .with_source(AgentSubmissionSource::SdkHost), ) .await .expect("product assembly runtime should accept an SDK run"); diff --git a/src/crates/contracts/runtime-ports/src/lib.rs b/src/crates/contracts/runtime-ports/src/lib.rs index f722a3398f..2ec9a61c1b 100644 --- a/src/crates/contracts/runtime-ports/src/lib.rs +++ b/src/crates/contracts/runtime-ports/src/lib.rs @@ -1271,6 +1271,7 @@ pub enum AgentSubmissionSource { RemoteRelay, Bot, Cli, + SdkHost, } pub type DialogTriggerSource = AgentSubmissionSource; @@ -1307,7 +1308,8 @@ impl DialogSubmissionPolicy { DialogTriggerSource::ScheduledJob => DialogQueuePriority::Low, DialogTriggerSource::DesktopUi | DialogTriggerSource::DesktopApi - | DialogTriggerSource::Cli => DialogQueuePriority::Normal, + | DialogTriggerSource::Cli + | DialogTriggerSource::SdkHost => DialogQueuePriority::Normal, DialogTriggerSource::RemoteRelay | DialogTriggerSource::Bot => { DialogQueuePriority::Normal } @@ -1796,6 +1798,24 @@ pub trait AgentSubmissionPort: Send + Sync { )) } + /// Creates one caller-identified connection-scoped Session. + /// + /// The Session uses the normal Runtime owners but must not become durable + /// product state. This narrow operation lets process adapters provide + /// bounded cleanup without pretending that crash-safe durable creation has + /// already been specified. + async fn create_transient_session_with_id( + &self, + session_id: String, + request: AgentSessionCreateRequest, + ) -> PortResult { + let _ = (session_id, request); + Err(PortError::new( + PortErrorKind::NotAvailable, + "transient exact session creation is not supported by this provider", + )) + } + async fn submit_message( &self, request: AgentSubmissionRequest, @@ -1855,6 +1875,36 @@ pub trait AgentSessionManagementPort: Send + Sync { ) -> PortResult>; } +/// Deadline-bearing request for discarding a connection-scoped transient +/// Session. This is separate from [`AgentSessionDeleteRequest`] so adding Host +/// cleanup policy cannot break the established Rust Session-management API. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AgentTransientSessionDiscardRequest { + pub workspace_path: String, + pub session_id: String, + pub remote_connection_id: Option, + pub remote_ssh_host: Option, + pub wait_timeout_ms: u64, +} + +/// Runtime lifecycle owner for connection-scoped Session cleanup. +#[async_trait::async_trait] +pub trait AgentSessionClosePort: Send + Sync { + /// Quiesces and discards only a loaded transient Session owned by the + /// caller. Implementations must reject durable Sessions and must never + /// remove persisted Session storage through this operation. + async fn discard_transient_session( + &self, + request: AgentTransientSessionDiscardRequest, + ) -> PortResult { + let _ = request; + Err(PortError::new( + PortErrorKind::NotAvailable, + "transient session discard is not supported by this provider", + )) + } +} + #[async_trait::async_trait] pub trait AgentLocalCommandTurnPort: Send + Sync { async fn record_completed_local_command_turn( @@ -2378,6 +2428,10 @@ mod tests { .expect("serialize dialog trigger source"); assert_eq!(json, serde_json::json!("cli")); + + let sdk_host = serde_json::to_value(DialogTriggerSource::SdkHost) + .expect("serialize SDK Host trigger source"); + assert_eq!(sdk_host, serde_json::json!("sdk_host")); } #[test] @@ -2393,6 +2447,9 @@ mod tests { let cli = DialogSubmissionPolicy::for_source(DialogTriggerSource::Cli); assert_eq!(cli.queue_priority, DialogQueuePriority::Normal); + + let sdk_host = DialogSubmissionPolicy::for_source(DialogTriggerSource::SdkHost); + assert_eq!(sdk_host.queue_priority, DialogQueuePriority::Normal); } #[test] diff --git a/src/crates/execution/agent-runtime/src/permission.rs b/src/crates/execution/agent-runtime/src/permission.rs index cb54f803b1..9549fa4d2a 100644 --- a/src/crates/execution/agent-runtime/src/permission.rs +++ b/src/crates/execution/agent-runtime/src/permission.rs @@ -76,6 +76,7 @@ pub enum PermissionRequestManagerError { #[derive(Debug)] struct PendingPermission { request: PermissionRequest, + dialog_turn_id: Option, sender: oneshot::Sender, interactive: bool, registration_sequence: u64, @@ -205,19 +206,43 @@ impl PermissionRequestManager { &self, requests: Vec, ) -> Result, PermissionRequestManagerError> { - self.register_batch_with_visibility(requests, true).await + self.register_batch_with_visibility(requests, None, true) + .await + } + + /// Registers a batch with the exact owning Dialog Turn kept as internal + /// coordination state. The public permission DTO remains stable because + /// turn routing is a runtime concern, not an authorization fact. + pub async fn register_batch_for_turn( + &self, + requests: Vec, + dialog_turn_id: impl Into, + ) -> Result, PermissionRequestManagerError> { + self.register_batch_with_visibility(requests, Some(dialog_turn_id.into()), true) + .await } pub async fn register_batch_non_interactive( &self, requests: Vec, ) -> Result, PermissionRequestManagerError> { - self.register_batch_with_visibility(requests, false).await + self.register_batch_with_visibility(requests, None, false) + .await + } + + pub async fn register_batch_non_interactive_for_turn( + &self, + requests: Vec, + dialog_turn_id: impl Into, + ) -> Result, PermissionRequestManagerError> { + self.register_batch_with_visibility(requests, Some(dialog_turn_id.into()), false) + .await } async fn register_batch_with_visibility( &self, requests: Vec, + dialog_turn_id: Option, interactive: bool, ) -> Result, PermissionRequestManagerError> { if requests.is_empty() { @@ -250,6 +275,7 @@ impl PermissionRequestManager { request.request_id.clone(), PendingPermission { request: request.clone(), + dialog_turn_id: dialog_turn_id.clone(), sender, interactive, registration_sequence, @@ -296,6 +322,14 @@ impl PermissionRequestManager { self.ordered_pending_requests(|pending| pending.interactive) } + /// Returns the process-local owning Dialog Turn for exact event routing. + /// This fact is intentionally absent from the persisted/public request DTO. + pub fn pending_request_dialog_turn_id(&self, request_id: &str) -> Option { + self.pending + .get(request_id) + .and_then(|pending| pending.dialog_turn_id.clone()) + } + fn ordered_pending_requests( &self, include: impl Fn(&PendingPermission) -> bool, diff --git a/src/crates/execution/agent-runtime/src/runtime.rs b/src/crates/execution/agent-runtime/src/runtime.rs index c88e52212a..6974779e3e 100644 --- a/src/crates/execution/agent-runtime/src/runtime.rs +++ b/src/crates/execution/agent-runtime/src/runtime.rs @@ -13,21 +13,21 @@ use bitfun_runtime_ports::{ AgentBackgroundResultRequest, AgentDialogTurnPort, AgentDialogTurnRequest, AgentInputAttachment, AgentLifecycleDeliveryPort, AgentLocalCommandTurnPort, AgentLocalCommandTurnRecordRequest, AgentSessionArchiveRequest, - AgentSessionArchiveStateRequest, AgentSessionCreateRequest, AgentSessionCreateResult, - AgentSessionDeleteRequest, AgentSessionForkAtTurnRequest, AgentSessionForkPort, - AgentSessionForkRequest, AgentSessionForkResult, AgentSessionListRequest, + AgentSessionArchiveStateRequest, AgentSessionClosePort, AgentSessionCreateRequest, + AgentSessionCreateResult, AgentSessionDeleteRequest, AgentSessionForkAtTurnRequest, + AgentSessionForkPort, AgentSessionForkRequest, AgentSessionForkResult, AgentSessionListRequest, AgentSessionManagementPort, AgentSessionModePort, AgentSessionModeUpdateRequest, AgentSessionModelPort, AgentSessionModelUpdateRequest, AgentSessionRenameRequest, AgentSessionSummary, AgentSessionUsagePort, AgentSessionUsageRequest, AgentSessionWorkspaceBinding, AgentSessionWorkspaceRequest, AgentSubmissionPort, AgentSubmissionRequest, AgentSubmissionResult, AgentSubmissionSource, AgentThreadGoalCreateRequest, AgentThreadGoalDeliveryRequest, AgentThreadGoalGetRequest, - AgentThreadGoalManagementPort, AgentThreadGoalUpdateStatusRequest, AgentTurnCancellationPort, - AgentTurnCancellationRequest, AgentTurnCancellationResult, AgentTurnSettlementPort, - AgentTurnSettlementRequest, DialogSubmitOutcome, PermissionAuditRecord, PermissionGrant, - PermissionGrantKey, PluginRuntimeBinding, PortError, PortErrorKind, PortResult, - RuntimeEventEnvelope, SessionTranscript, SessionTranscriptReader, SessionTranscriptRequest, - ThreadGoal, + AgentThreadGoalManagementPort, AgentThreadGoalUpdateStatusRequest, + AgentTransientSessionDiscardRequest, AgentTurnCancellationPort, AgentTurnCancellationRequest, + AgentTurnCancellationResult, AgentTurnSettlementPort, AgentTurnSettlementRequest, + DialogSubmitOutcome, PermissionAuditRecord, PermissionGrant, PermissionGrantKey, + PluginRuntimeBinding, PortError, PortErrorKind, PortResult, RuntimeEventEnvelope, + SessionTranscript, SessionTranscriptReader, SessionTranscriptRequest, ThreadGoal, }; use bitfun_runtime_services::RuntimeServices; @@ -184,6 +184,7 @@ pub trait RuntimeAgentRegistry: Send + Sync { pub struct AgentRuntime { submission: Arc, session_management: Option>, + session_close: Option>, session_mode: Option>, session_model: Option>, session_fork: Option>, @@ -219,6 +220,13 @@ impl std::fmt::Debug for AgentRuntime { .as_ref() .map(|_| ""), ) + .field( + "session_close", + &self + .session_close + .as_ref() + .map(|_| ""), + ) .field( "session_mode", &self @@ -360,6 +368,7 @@ where pub struct AgentRuntimeBuilder { submission: Option>, session_management: Option>, + session_close: Option>, session_mode: Option>, session_model: Option>, session_fork: Option>, @@ -402,6 +411,11 @@ impl AgentRuntimeBuilder { self } + pub fn with_session_close_port(mut self, port: Arc) -> Self { + self.session_close = Some(port); + self + } + pub fn with_session_model_port(mut self, port: Arc) -> Self { self.session_model = Some(port); self @@ -534,6 +548,7 @@ impl AgentRuntimeBuilder { let Self { submission, session_management, + session_close, session_mode, session_model, session_fork, @@ -565,6 +580,7 @@ impl AgentRuntimeBuilder { Ok(AgentRuntime { submission: submission.ok_or(RuntimeBuildError::MissingSubmissionPort)?, session_management, + session_close, session_mode, session_model, session_fork, @@ -712,6 +728,16 @@ impl AgentRuntime { .ok_or(RuntimeError::MissingPermissionRequestManager) } + pub fn permission_request_dialog_turn_id( + &self, + request_id: &str, + ) -> Result, RuntimeError> { + self.permission_requests + .as_ref() + .map(|manager| manager.pending_request_dialog_turn_id(request_id)) + .ok_or(RuntimeError::MissingPermissionRequestManager) + } + pub fn subscribe_permission_requests( &self, ) -> Result { @@ -877,6 +903,46 @@ impl AgentRuntime { Ok(result) } + pub async fn create_transient_session_with_id( + &self, + session_id: String, + request: AgentSessionCreateRequest, + ) -> Result { + let result = self + .submission + .create_transient_session_with_id(session_id.clone(), request) + .await + .map_err(RuntimeError::from)?; + if result.session_id != session_id { + return Err(PortError::new( + PortErrorKind::Backend, + format!( + "agent submission provider returned session_id '{}' for requested transient session_id '{}'", + result.session_id, session_id + ), + ) + .into()); + } + Ok(result) + } + + pub async fn discard_transient_session( + &self, + request: AgentTransientSessionDiscardRequest, + ) -> Result { + self.session_close + .as_ref() + .ok_or_else(|| { + RuntimeError::Port(PortError::new( + PortErrorKind::NotAvailable, + "agent session close port is not registered", + )) + })? + .discard_transient_session(request) + .await + .map_err(RuntimeError::from) + } + pub async fn list_sessions( &self, request: AgentSessionListRequest, @@ -1106,14 +1172,30 @@ impl AgentRuntime { &self, request: AgentDialogTurnRequest, ) -> Result { + let requested_session_id = request.session_id.clone(); let dialog_turn = self .dialog_turn .as_ref() .ok_or(RuntimeError::MissingDialogTurnPort)?; - dialog_turn + let outcome = dialog_turn .submit_dialog_turn(request) .await - .map_err(RuntimeError::from) + .map_err(RuntimeError::from)?; + let returned_session_id = match &outcome { + DialogSubmitOutcome::Started { session_id, .. } + | DialogSubmitOutcome::Queued { session_id, .. } => session_id, + }; + if returned_session_id != &requested_session_id { + return Err(PortError::new( + PortErrorKind::Backend, + format!( + "agent dialog provider returned session_id '{}' for requested session_id '{}'", + returned_session_id, requested_session_id + ), + ) + .into()); + } + Ok(outcome) } pub async fn deliver_background_result( @@ -2442,6 +2524,59 @@ mod tests { ); } + #[tokio::test] + async fn submit_dialog_turn_rejects_provider_session_identity_mismatch() { + #[derive(Debug)] + struct MismatchedDialogTurnPort; + + #[async_trait::async_trait] + impl bitfun_runtime_ports::AgentDialogTurnPort for MismatchedDialogTurnPort { + async fn submit_dialog_turn( + &self, + request: AgentDialogTurnRequest, + ) -> PortResult { + Ok(DialogSubmitOutcome::Started { + session_id: "different-session".to_string(), + turn_id: request.turn_id.unwrap_or_else(|| "generated".to_string()), + }) + } + } + + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(Arc::new(FakeAgentRuntimePorts::default())) + .with_dialog_turn_port(Arc::new(MismatchedDialogTurnPort)) + .build() + .expect("runtime"); + + let error = runtime + .submit_dialog_turn(AgentDialogTurnRequest { + session_id: "requested-session".to_string(), + message: "hello".to_string(), + original_message: None, + turn_id: Some("turn-1".to_string()), + agent_type: "agentic".to_string(), + workspace_path: Some("/workspace/project".to_string()), + remote_connection_id: None, + remote_ssh_host: None, + policy: DialogSubmissionPolicy::for_source(AgentSubmissionSource::SdkHost), + reply_route: None, + prepended_reminders: Vec::new(), + attachments: Vec::new(), + metadata: serde_json::Map::new(), + }) + .await + .expect_err("provider session mismatch must fail closed"); + + assert!(matches!( + error, + RuntimeError::Port(PortError { + kind: PortErrorKind::Backend, + .. + }) + )); + assert!(error.into_message().contains("requested-session")); + } + #[tokio::test] async fn deliver_background_result_requires_registered_lifecycle_port() { let ports = Arc::new(FakeAgentRuntimePorts::default()); diff --git a/src/crates/execution/agent-runtime/src/sdk.rs b/src/crates/execution/agent-runtime/src/sdk.rs index 426b93db2c..ffb565f2db 100644 --- a/src/crates/execution/agent-runtime/src/sdk.rs +++ b/src/crates/execution/agent-runtime/src/sdk.rs @@ -61,24 +61,24 @@ pub use bitfun_runtime_ports::{ AgentBackgroundResultRequest, AgentDialogTurnPort, AgentDialogTurnRequest, AgentInputAttachment, AgentLifecycleDeliveryPort, AgentLocalCommandTurnPort, AgentLocalCommandTurnRecordRequest, AgentSessionArchiveRequest, - AgentSessionArchiveStateRequest, AgentSessionCreateRequest, AgentSessionCreateResult, - AgentSessionDeleteRequest, AgentSessionForkAtTurnRequest, AgentSessionForkPort, - AgentSessionForkRequest, AgentSessionForkResult, AgentSessionListRequest, + AgentSessionArchiveStateRequest, AgentSessionClosePort, AgentSessionCreateRequest, + AgentSessionCreateResult, AgentSessionDeleteRequest, AgentSessionForkAtTurnRequest, + AgentSessionForkPort, AgentSessionForkRequest, AgentSessionForkResult, AgentSessionListRequest, AgentSessionManagementPort, AgentSessionModePort, AgentSessionModeUpdateRequest, AgentSessionModelPort, AgentSessionModelUpdateRequest, AgentSessionRenameRequest, AgentSessionSummary, AgentSessionUsagePort, AgentSessionUsageRequest, AgentSessionWorkspaceBinding, AgentSessionWorkspaceRequest, AgentSubmissionPort, AgentSubmissionRequest, AgentSubmissionResult, AgentSubmissionSource, AgentThreadGoalCreateRequest, AgentThreadGoalDeliveryRequest, AgentThreadGoalGetRequest, - AgentThreadGoalManagementPort, AgentThreadGoalUpdateStatusRequest, AgentTurnCancellationPort, - AgentTurnCancellationRequest, AgentTurnCancellationResult, AgentTurnSettlementPort, - AgentTurnSettlementRequest, ClockPort, DialogSubmissionPolicy, DialogSubmitOutcome, - FileSystemPort, GitPort, McpCatalogPort, NetworkPort, PermissionAuditRecord, - PermissionDelegationContext, PermissionGrant, PermissionGrantKey, PermissionReply, - PermissionReplySource, PermissionRequest, PermissionRequestEvent, PermissionRequestSource, - PermissionRequestSourceKind, PortError, PortErrorKind, PortResult, - RemoteAssistantWorkspaceFacts, RemoteCapabilityPort, RemoteConnectionPort, - RemoteProjectionPort, RemoteRecentWorkspaceFacts, RemoteWorkspaceFacts, + AgentThreadGoalManagementPort, AgentThreadGoalUpdateStatusRequest, + AgentTransientSessionDiscardRequest, AgentTurnCancellationPort, AgentTurnCancellationRequest, + AgentTurnCancellationResult, AgentTurnSettlementPort, AgentTurnSettlementRequest, ClockPort, + DialogSubmissionPolicy, DialogSubmitOutcome, FileSystemPort, GitPort, McpCatalogPort, + NetworkPort, PermissionAuditRecord, PermissionDelegationContext, PermissionGrant, + PermissionGrantKey, PermissionReply, PermissionReplySource, PermissionRequest, + PermissionRequestEvent, PermissionRequestSource, PermissionRequestSourceKind, PortError, + PortErrorKind, PortResult, RemoteAssistantWorkspaceFacts, RemoteCapabilityPort, + RemoteConnectionPort, RemoteProjectionPort, RemoteRecentWorkspaceFacts, RemoteWorkspaceFacts, RemoteWorkspaceFileRuntimeHost, RemoteWorkspaceKind, RemoteWorkspacePort, RemoteWorkspaceRuntimeHost, RemoteWorkspaceUpdate, RuntimeEventEnvelope, RuntimeEventSink, RuntimeEventType, RuntimeServiceCapability, RuntimeServicePort, SessionStorageKind, @@ -127,6 +127,11 @@ impl AgentRuntimeBuilder { self } + pub fn with_session_close_port(mut self, port: Arc) -> Self { + self.inner = self.inner.with_session_close_port(port); + self + } + pub fn with_session_model_port(mut self, port: Arc) -> Self { self.inner = self.inner.with_session_model_port(port); self @@ -271,6 +276,13 @@ impl AgentRuntime { self.inner.pending_permission_requests() } + pub fn permission_request_dialog_turn_id( + &self, + request_id: &str, + ) -> Result, RuntimeError> { + self.inner.permission_request_dialog_turn_id(request_id) + } + pub fn subscribe_permission_requests( &self, ) -> Result { @@ -371,6 +383,26 @@ impl AgentRuntime { self.inner.create_session_with_id(session_id, request).await } + /// Creates one connection-scoped Session through the same Runtime owners. + /// It is intentionally separate from durable Session creation so process + /// adapters cannot silently weaken persistence semantics. + pub async fn create_transient_session_with_id( + &self, + session_id: String, + request: AgentSessionCreateRequest, + ) -> Result { + self.inner + .create_transient_session_with_id(session_id, request) + .await + } + + pub async fn discard_transient_session( + &self, + request: AgentTransientSessionDiscardRequest, + ) -> Result { + self.inner.discard_transient_session(request).await + } + pub async fn list_sessions( &self, request: AgentSessionListRequest, diff --git a/src/crates/execution/agent-runtime/tests/permission_contracts.rs b/src/crates/execution/agent-runtime/tests/permission_contracts.rs index a562b4ca23..344d629726 100644 --- a/src/crates/execution/agent-runtime/tests/permission_contracts.rs +++ b/src/crates/execution/agent-runtime/tests/permission_contracts.rs @@ -704,6 +704,37 @@ async fn pending_requests_are_process_local_and_not_restored_by_a_new_manager() assert!(restarted.pending_requests().is_empty()); } +#[tokio::test] +async fn exact_turn_routing_stays_runtime_local_and_is_removed_with_the_request() { + let (manager, _) = manager(); + let request_id = "request-turn-owned"; + let receiver = manager + .register_batch_for_turn( + vec![request(request_id, "session-turn-owned")], + "turn-owned", + ) + .await + .expect("register turn-owned permission") + .pop() + .expect("permission receiver"); + + assert_eq!( + manager + .pending_request_dialog_turn_id(request_id) + .as_deref(), + Some("turn-owned") + ); + manager + .cancel_request(request_id, "test cleanup") + .await + .expect("cancel request"); + assert!(matches!( + receiver.wait().await, + PermissionWaitOutcome::Cancelled { .. } + )); + assert_eq!(manager.pending_request_dialog_turn_id(request_id), None); +} + #[tokio::test] async fn grant_management_is_project_scoped_and_audit_remains_append_only() { let (manager, store) = manager(); diff --git a/src/crates/execution/agent-runtime/tests/sdk_smoke.rs b/src/crates/execution/agent-runtime/tests/sdk_smoke.rs index bc75d0f30d..26b52280a4 100644 --- a/src/crates/execution/agent-runtime/tests/sdk_smoke.rs +++ b/src/crates/execution/agent-runtime/tests/sdk_smoke.rs @@ -4,15 +4,16 @@ use std::sync::{Arc, Mutex}; use async_trait::async_trait; use bitfun_agent_runtime::sdk::{ build_descriptor_harness_registry, AgentEventStream, AgentRunRequest, AgentRuntimeBuilder, - AgentRuntimeSdkCompatibility, AgentRuntimeSdkStability, AgentSessionCreateRequest, - AgentSessionCreateResult, AgentSubmissionPort, AgentSubmissionRequest, AgentSubmissionResult, - AgentSubmissionSource, ClockPort, FileSystemPort, HarnessCapability, HarnessProviderDescriptor, - HarnessWorkflow, PortResult, RuntimeAgentRegistry, RuntimeAgentRegistryQuery, - RuntimeEventEnvelope, RuntimeEventSink, RuntimeEventType, RuntimeHookErrorPolicy, - RuntimeHookKind, RuntimeHookPlan, RuntimeHookRegistry, RuntimeServiceCapability, - RuntimeServicePort, RuntimeServices, RuntimeServicesBuilder, SessionSelector, - SessionStorageKind, SessionStoragePathRequest, SessionStoragePathResolution, SessionStorePort, - ToolRegistry, ToolRegistryItem, WorkspacePort, + AgentRuntimeSdkCompatibility, AgentRuntimeSdkStability, AgentSessionClosePort, + AgentSessionCreateRequest, AgentSessionCreateResult, AgentSubmissionPort, + AgentSubmissionRequest, AgentSubmissionResult, AgentSubmissionSource, + AgentTransientSessionDiscardRequest, ClockPort, FileSystemPort, HarnessCapability, + HarnessProviderDescriptor, HarnessWorkflow, PortErrorKind, PortResult, RuntimeAgentRegistry, + RuntimeAgentRegistryQuery, RuntimeError, RuntimeEventEnvelope, RuntimeEventSink, + RuntimeEventType, RuntimeHookErrorPolicy, RuntimeHookKind, RuntimeHookPlan, + RuntimeHookRegistry, RuntimeServiceCapability, RuntimeServicePort, RuntimeServices, + RuntimeServicesBuilder, SessionSelector, SessionStorageKind, SessionStoragePathRequest, + SessionStoragePathResolution, SessionStorePort, ToolRegistry, ToolRegistryItem, WorkspacePort, }; use serde_json::{json, Value}; @@ -38,6 +39,11 @@ struct FakeSdkRuntimePort { #[derive(Debug, Default)] struct FakeSdkRuntimeEventSink; +#[derive(Debug, Default)] +struct FakeSessionClosePort { + requests: Mutex>, +} + #[test] fn sdk_facade_exposes_versioned_preview_compatibility_contract() { let compatibility = AgentRuntimeSdkCompatibility::current(); @@ -72,6 +78,17 @@ impl RuntimeServicePort for FakeSdkRuntimePort { impl FileSystemPort for FakeSdkRuntimePort {} impl WorkspacePort for FakeSdkRuntimePort {} +#[async_trait] +impl AgentSessionClosePort for FakeSessionClosePort { + async fn discard_transient_session( + &self, + request: AgentTransientSessionDiscardRequest, + ) -> PortResult { + self.requests.lock().unwrap().push(request.clone()); + Ok(true) + } +} + #[async_trait] impl SessionStorePort for FakeSdkRuntimePort { async fn resolve_session_storage_path( @@ -282,3 +299,57 @@ async fn sdk_facade_accepts_fake_services_tools_harnesses_and_hooks_without_core .expect("services should be injected") .has_capability(RuntimeServiceCapability::SessionStore)); } + +#[tokio::test] +async fn sdk_facade_delegates_connection_scoped_session_discard() { + let provider = Arc::new(FakeSdkAgentProvider::default()); + let close_port = Arc::new(FakeSessionClosePort::default()); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(provider) + .with_session_close_port(close_port.clone()) + .build() + .expect("sdk runtime"); + let request = AgentTransientSessionDiscardRequest { + workspace_path: "/workspace/project".to_string(), + session_id: "sdk-session-1".to_string(), + remote_connection_id: None, + remote_ssh_host: None, + wait_timeout_ms: 5_000, + }; + + let result = runtime + .discard_transient_session(request.clone()) + .await + .expect("discard transient session through SDK facade"); + + assert_eq!(close_port.requests.lock().unwrap().as_slice(), &[request]); + assert!(result); +} + +#[tokio::test] +async fn sdk_facade_reports_missing_session_close_capability() { + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(Arc::new(FakeSdkAgentProvider::default())) + .build() + .expect("sdk runtime"); + + let error = runtime + .discard_transient_session(AgentTransientSessionDiscardRequest { + workspace_path: "/workspace/project".to_string(), + session_id: "sdk-session-1".to_string(), + remote_connection_id: None, + remote_ssh_host: None, + wait_timeout_ms: 5_000, + }) + .await + .expect_err("missing transient cleanup port must fail closed"); + + assert!(matches!( + &error, + RuntimeError::Port(port) if port.kind == PortErrorKind::NotAvailable + )); + assert_eq!( + error.into_message(), + "agent session close port is not registered" + ); +} diff --git a/src/crates/interfaces/AGENTS-CN.md b/src/crates/interfaces/AGENTS-CN.md index 5bc291527b..3faaaeb2aa 100644 --- a/src/crates/interfaces/AGENTS-CN.md +++ b/src/crates/interfaces/AGENTS-CN.md @@ -9,6 +9,7 @@ | Crate | 职责 | 本地文档 | |---|---|---| | `acp` | 基于已组装产品 runtime 的 Agent Client Protocol 接口 | [AGENTS.md](acp/AGENTS.md) | +| `sdk-host` | 版本化的本地 Agent SDK Host 协议与连接用例；进程启动和 stdio framing 仍由 `src/apps/sdk-host` 负责 | — | ## 放置规则 @@ -19,4 +20,7 @@ ## 依赖边界 - interface crate 可以依赖 `assembly/core` 暴露选定交付形态。 +- 可移植的 `sdk-host` 协议 crate 边界更窄：只能依赖稳定 Runtime/合同，不得依赖 + `bitfun-core`、`terminal-core`、具体 service、SDK Host app 或 CLI；具体 Host 组装保留在 + `src/apps/sdk-host`。 - interface crate 不拥有产品策略、可复用服务、协议传输内部实现或执行原语。 diff --git a/src/crates/interfaces/AGENTS.md b/src/crates/interfaces/AGENTS.md index cb39fa0d22..663a667f59 100644 --- a/src/crates/interfaces/AGENTS.md +++ b/src/crates/interfaces/AGENTS.md @@ -12,6 +12,7 @@ product behavior. UI apps and delivery hosts remain under `src/apps`, | Crate | Responsibility | Local doc | |---|---|---| | `acp` | Agent Client Protocol interface over the assembled product runtime | [AGENTS.md](acp/AGENTS.md) | +| `sdk-host` | Versioned local Agent SDK Host protocol and connection use cases; process bootstrap and stdio framing remain in `src/apps/sdk-host` | — | ## Placement Rules @@ -25,5 +26,9 @@ product behavior. UI apps and delivery hosts remain under `src/apps`, - Interface crates may depend on `assembly/core` to expose a selected delivery profile. +- The portable `sdk-host` protocol crate is narrower: it depends only on stable + Runtime/contracts and must not depend on `bitfun-core`, `terminal-core`, + concrete services, the SDK Host app, or CLI. Concrete Host assembly stays in + `src/apps/sdk-host`. - Interface crates must not own product policy, reusable services, protocol transport internals, or execution primitives. diff --git a/src/crates/interfaces/sdk-host/Cargo.toml b/src/crates/interfaces/sdk-host/Cargo.toml new file mode 100644 index 0000000000..c3672d6f7a --- /dev/null +++ b/src/crates/interfaces/sdk-host/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "bitfun-sdk-host" +version.workspace = true +authors.workspace = true +edition.workspace = true +description = "Versioned local SDK Host adapter for the BitFun Agent Runtime" + +[lib] +name = "bitfun_sdk_host" + +[dependencies] +async-trait = { workspace = true } +bitfun-agent-runtime = { path = "../../execution/agent-runtime" } +bitfun-core-types = { path = "../../contracts/core-types" } +bitfun-events = { path = "../../contracts/events" } +bitfun-runtime-ports = { path = "../../contracts/runtime-ports" } +futures-util = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +tokio = { workspace = true } +tokio-util = { workspace = true, features = ["codec"] } +tracing = { workspace = true } +uuid = { workspace = true } + +[lints] +workspace = true diff --git a/src/crates/interfaces/sdk-host/src/host.rs b/src/crates/interfaces/sdk-host/src/host.rs new file mode 100644 index 0000000000..f5bdd1a0f6 --- /dev/null +++ b/src/crates/interfaces/sdk-host/src/host.rs @@ -0,0 +1,2209 @@ +//! Connection-scoped SDK Host request and Query lifecycle. + +use std::collections::{HashMap, HashSet}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use bitfun_agent_runtime::sdk::{ + AgentDialogTurnRequest, AgentRuntime, AgentSessionCreateRequest, AgentSessionCreateResult, + AgentSubmissionSource, AgentTransientSessionDiscardRequest, AgentTurnCancellationRequest, + AgentTurnSettlementRequest, DialogSubmissionPolicy, DialogSubmitOutcome, PermissionReply, + PermissionReplySource, PermissionRequest, PermissionRequestEvent, PortErrorKind, RuntimeError, + AUTO_APPROVE_ASK_CONTEXT_KEY, +}; +use bitfun_agent_runtime::user_questions::USER_INPUT_AVAILABLE_CONTEXT_KEY; +use bitfun_core_types::ErrorCategory; +use bitfun_events::AgenticEvent; +use futures_util::{stream::FuturesUnordered, FutureExt, StreamExt}; +use tokio::sync::{mpsc, oneshot, Mutex, OwnedSemaphorePermit, Semaphore}; +use tokio::task::JoinHandle; +use tokio::time::{timeout, Instant}; +use tokio_util::sync::CancellationToken; + +use crate::protocol::{ + ErrorCode, ErrorData, ErrorStage, InitializeParams, InitializeResult, JsonRpcErrorResponse, + JsonRpcNotification, JsonRpcRequest, JsonRpcSuccessResponse, QueryCancelParams, + QueryCancelResult, QueryEvent, QueryEventParams, QueryResultError, QueryResultParams, + QueryStartParams, QueryStartResult, QueryTerminalStatus, RecoveryAction, RequestId, + SessionCloseParams, SessionCloseResult, SessionCreateParams, SessionCreateResult, + SessionLifetime, ShutdownParams, ShutdownResult, JSON_RPC_VERSION, METHOD_INITIALIZE, + METHOD_QUERY_CANCEL, METHOD_QUERY_START, METHOD_SESSION_CLOSE, METHOD_SESSION_CREATE, + METHOD_SHUTDOWN, NOTIFICATION_QUERY_EVENT, NOTIFICATION_QUERY_RESULT, PROTOCOL_VERSION, +}; + +const DEFAULT_SESSION_NAME: &str = "BitFun SDK query"; +const DEFAULT_AGENT: &str = "agentic"; +const DEFAULT_TURN_SETTLEMENT_TIMEOUT_MS: u64 = 5_000; +const PERMISSION_REJECTION_TIMEOUT_MS: u64 = 2_000; +const MAX_SESSION_CLOSE_TIMEOUT_MS: u64 = 30_000; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ConnectionControl { + Continue, + Shutdown, +} + +#[derive(Debug, Clone)] +pub struct SdkHostConfig { + pub max_in_flight_requests: usize, + pub max_in_flight_control_requests: usize, + pub max_active_queries: usize, + pub max_leased_sessions: usize, +} + +impl Default for SdkHostConfig { + fn default() -> Self { + Self { + max_in_flight_requests: 32, + max_in_flight_control_requests: 4, + max_active_queries: 16, + max_leased_sessions: 64, + } + } +} + +#[derive(Clone)] +pub struct SdkHostConnection { + inner: Arc, +} + +#[async_trait::async_trait] +pub trait HostOutput: Send + Sync { + async fn send(&self, value: serde_json::Value) -> Result<(), ()>; +} + +struct ChannelHostOutput(mpsc::Sender); + +#[async_trait::async_trait] +impl HostOutput for ChannelHostOutput { + async fn send(&self, value: serde_json::Value) -> Result<(), ()> { + self.0.send(value).await.map_err(|_| ()) + } +} + +struct ConnectionInner { + runtime: AgentRuntime, + runtime_version: &'static str, + default_cwd: String, + output: Arc, + state: Arc>, + request_budget: Arc, + control_request_budget: Arc, + query_budget: Arc, + session_budget: Arc, + shutdown_started: CancellationToken, +} + +#[derive(Default)] +struct ConnectionState { + initialized: bool, + shutting_down: bool, + cleanup_failed: bool, + sessions: HashMap, + queries: HashMap>, + starting_query_sessions: HashSet, + active_query_sessions: HashSet, + closing_sessions: HashSet, + poisoned_sessions: HashSet, + pending_session_tasks: Vec, + untracked_transient_cleanups: HashMap, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum QueryReservationError { + Unavailable, + Poisoned, +} + +#[derive(Clone)] +struct SessionLease { + workspace_path: String, + remote_connection_id: Option, + remote_ssh_host: Option, + exposed: bool, + _budget: Arc, +} + +struct PendingSessionTask { + transient_cleanup: Option, + task: JoinHandle<()>, +} + +#[derive(Clone)] +struct TransientSessionCleanup { + session_id: String, + workspace_path: String, +} + +struct QueryLease { + query_id: String, + session_id: String, + turn_id: String, + terminal: AtomicBool, + stop_forwarding: CancellationToken, + emit_output: bool, + _budget: OwnedSemaphorePermit, +} + +impl QueryLease { + fn finish_once(&self) -> bool { + !self.terminal.swap(true, Ordering::AcqRel) + } +} + +impl SdkHostConnection { + pub fn new( + runtime: AgentRuntime, + default_cwd: impl Into, + output: mpsc::Sender, + config: SdkHostConfig, + ) -> Self { + Self::with_output( + runtime, + default_cwd, + Arc::new(ChannelHostOutput(output)), + config, + ) + } + + pub fn with_output( + runtime: AgentRuntime, + default_cwd: impl Into, + output: Arc, + config: SdkHostConfig, + ) -> Self { + Self { + inner: Arc::new(ConnectionInner { + runtime, + runtime_version: env!("CARGO_PKG_VERSION"), + default_cwd: default_cwd.into(), + output, + state: Arc::new(Mutex::new(ConnectionState::default())), + request_budget: Arc::new(Semaphore::new(config.max_in_flight_requests.max(1))), + control_request_budget: Arc::new(Semaphore::new( + config.max_in_flight_control_requests.max(1), + )), + query_budget: Arc::new(Semaphore::new(config.max_active_queries.max(1))), + session_budget: Arc::new(Semaphore::new(config.max_leased_sessions.max(1))), + shutdown_started: CancellationToken::new(), + }), + } + } + + pub async fn handle_request(&self, request: JsonRpcRequest) -> ConnectionControl { + let request_id = request.id.clone(); + if request.jsonrpc != JSON_RPC_VERSION { + self.send_rpc_error( + request_id, + -32600, + ErrorCode::InvalidRequest, + ErrorStage::Protocol, + false, + None, + "jsonrpc must be 2.0", + ) + .await; + return ConnectionControl::Continue; + } + + // Resource and request lifecycle methods require an addressable + // response. Executing them as notifications would create Sessions or + // Queries that the caller can neither identify nor close. Shutdown is + // the only fire-and-forget client method in this protocol version. + if request.id.is_none() && request.method != METHOD_SHUTDOWN { + return ConnectionControl::Continue; + } + + if request.method != METHOD_SHUTDOWN { + self.reap_finished_pending_session_tasks().await; + } + + let _request_permit = if request.method == METHOD_SHUTDOWN { + None + } else { + let budget = if matches!( + request.method.as_str(), + METHOD_QUERY_CANCEL | METHOD_SESSION_CLOSE + ) { + self.inner.control_request_budget.clone() + } else { + self.inner.request_budget.clone() + }; + let Ok(permit) = budget.try_acquire_owned() else { + self.send_error( + request_id, + ErrorCode::Overloaded, + ErrorStage::Protocol, + true, + Some(RecoveryAction::Retry), + "SDK Host request capacity is exhausted", + ) + .await; + return ConnectionControl::Continue; + }; + Some(permit) + }; + + if request.method == METHOD_INITIALIZE { + self.handle_initialize(request).await; + return ConnectionControl::Continue; + } + + let (initialized, shutting_down, cleanup_failed) = { + let state = self.inner.state.lock().await; + ( + state.initialized, + state.shutting_down, + state.cleanup_failed || !state.untracked_transient_cleanups.is_empty(), + ) + }; + if !initialized { + self.send_error( + request_id, + ErrorCode::NotInitialized, + ErrorStage::Protocol, + true, + Some(RecoveryAction::Initialize), + "initialize must complete before this method", + ) + .await; + return ConnectionControl::Continue; + } + if cleanup_failed && request.method != METHOD_SHUTDOWN { + self.send_error( + request_id, + ErrorCode::CleanupRequired, + ErrorStage::Protocol, + false, + Some(RecoveryAction::RestartHost), + "SDK Host cleanup is incomplete; shut down and restart the Host before retrying", + ) + .await; + return ConnectionControl::Continue; + } + if shutting_down && request.method != METHOD_SHUTDOWN { + self.send_error( + request_id, + ErrorCode::Cancelled, + ErrorStage::Protocol, + false, + None, + "SDK Host connection is shutting down", + ) + .await; + return ConnectionControl::Continue; + } + + match request.method.as_str() { + METHOD_SESSION_CREATE => self.handle_session_create(request).await, + METHOD_QUERY_START => self.handle_query_start(request).await, + METHOD_QUERY_CANCEL => self.handle_query_cancel(request).await, + METHOD_SESSION_CLOSE => self.handle_session_close(request).await, + METHOD_SHUTDOWN => { + if self + .parse_params::(&request, ErrorStage::Shutdown) + .await + .is_none() + { + return ConnectionControl::Continue; + } + { + let mut state = self.inner.state.lock().await; + state.shutting_down = true; + } + self.send_success(request.id.clone(), ShutdownResult { accepted: true }) + .await; + return ConnectionControl::Shutdown; + } + _ => { + self.send_rpc_error( + request.id.clone(), + -32601, + ErrorCode::CapabilityUnavailable, + ErrorStage::Protocol, + false, + None, + "method is not supported by this SDK Host", + ) + .await; + } + } + ConnectionControl::Continue + } + + /// Emits the protocol-owned overload response when a transport reaches its + /// bounded in-flight request capacity. + pub async fn reject_overloaded(&self, request_id: Option) { + self.send_error( + request_id, + ErrorCode::Overloaded, + ErrorStage::Protocol, + true, + Some(RecoveryAction::Retry), + "SDK Host request capacity is exhausted", + ) + .await; + } + + /// Reports whether this connection has completed the required initialize + /// handshake so a transport can serialize re-initialization safely. + pub async fn is_initialized(&self) -> bool { + self.inner.state.lock().await.initialized + } + + async fn reap_finished_pending_session_tasks(&self) { + let mut state = self.inner.state.lock().await; + let mut active = Vec::with_capacity(state.pending_session_tasks.len()); + for mut pending in std::mem::take(&mut state.pending_session_tasks) { + if !pending.task.is_finished() { + active.push(pending); + continue; + } + match (&mut pending.task).now_or_never() { + Some(Ok(())) => {} + Some(Err(error)) => { + tracing::warn!( + error = %error, + "SDK Host Session ownership task failed" + ); + if let Some(cleanup) = pending.transient_cleanup { + state + .untracked_transient_cleanups + .insert(cleanup.session_id.clone(), cleanup); + } + } + None => active.push(pending), + } + } + state.pending_session_tasks = active; + } + + pub async fn shutdown_connection(&self) { + self.shutdown_connection_inner(None).await; + } + + /// Shuts down one connection without allowing a transient Session create + /// task to keep the Host process alive indefinitely. + pub async fn shutdown_connection_bounded(&self, total_timeout: Duration) -> bool { + self.shutdown_connection_inner(Some(total_timeout)).await + } + + async fn shutdown_connection_inner(&self, total_timeout: Option) -> bool { + self.inner.shutdown_started.cancel(); + let (pending_session_tasks, prior_cleanup_failed) = { + let mut state = self.inner.state.lock().await; + state.shutting_down = true; + ( + std::mem::take(&mut state.pending_session_tasks), + state.cleanup_failed, + ) + }; + let started_at = Instant::now(); + let deadline = total_timeout.map(|timeout| started_at + timeout); + let graceful_deadline = total_timeout.map(|timeout| started_at + timeout / 2); + let mut cleanup_complete = !prior_cleanup_failed; + for pending in pending_session_tasks { + self.settle_pending_session_task(pending, graceful_deadline) + .await; + } + if !self + .compensate_registered_transient_sessions(deadline) + .await + { + cleanup_complete = false; + } + let (queries, sessions) = { + let mut state = self.inner.state.lock().await; + for query in state.queries.values() { + query.stop_forwarding.cancel(); + } + let queries = std::mem::take(&mut state.queries) + .into_values() + .collect::>(); + state.active_query_sessions.clear(); + state.starting_query_sessions.clear(); + (queries, std::mem::take(&mut state.sessions)) + }; + + let mut cancellations = queries + .into_iter() + .map(|query| { + let runtime = self.inner.runtime.clone(); + let cancellation_timeout = deadline + .map(|deadline| { + deadline + .saturating_duration_since(Instant::now()) + .min(Duration::from_millis(2_500)) + }) + .unwrap_or(Duration::from_millis(2_500)); + async move { + timeout( + cancellation_timeout, + runtime.cancel_turn(AgentTurnCancellationRequest { + session_id: query.session_id.clone(), + turn_id: Some(query.turn_id.clone()), + source: Some(AgentSubmissionSource::SdkHost), + requester_session_id: None, + reason: Some("sdk_host_connection_shutdown".to_string()), + wait_timeout_ms: Some(2_000), + }), + ) + .await + } + }) + .collect::>(); + while let Some(result) = cancellations.next().await { + match result { + Ok(Ok(_)) => {} + Ok(Err(error)) => { + cleanup_complete = false; + tracing::warn!( + error_kind = runtime_error_kind(&error), + "Failed to cancel SDK Host Query during connection shutdown" + ); + } + Err(_) => { + cleanup_complete = false; + tracing::warn!( + "SDK Host Query cancellation timed out during connection shutdown" + ); + } + } + } + + let mut cleanup = sessions + .into_iter() + .map(|(session_id, session)| { + let runtime = self.inner.runtime.clone(); + let session_cleanup_timeout = deadline + .map(|deadline| { + deadline + .saturating_duration_since(Instant::now()) + .min(Duration::from_millis(5_500)) + }) + .unwrap_or(Duration::from_millis(5_500)); + async move { + let reported_session_id = session_id.clone(); + let cleanup = async move { + runtime + .discard_transient_session(AgentTransientSessionDiscardRequest { + workspace_path: session.workspace_path, + session_id: session_id.clone(), + remote_connection_id: session.remote_connection_id, + remote_ssh_host: session.remote_ssh_host, + wait_timeout_ms: duration_ms( + session_cleanup_timeout + .saturating_sub(Duration::from_millis(500)), + ), + }) + .await?; + Ok(()) + }; + ( + reported_session_id, + timeout(session_cleanup_timeout, cleanup).await, + ) + } + }) + .collect::>(); + while let Some((session_id, result)) = cleanup.next().await { + match result { + Ok(Ok(())) => {} + Ok(Err(error)) => { + cleanup_complete = false; + tracing::warn!( + session_id = %session_id, + error_kind = runtime_error_kind(&error), + "Failed to clean up SDK Host Session during connection shutdown" + ); + } + Err(_) => { + cleanup_complete = false; + tracing::warn!( + session_id = %session_id, + "SDK Host Session cleanup timed out during connection shutdown" + ); + } + } + } + cleanup_complete + } + + async fn settle_pending_session_task( + &self, + pending: PendingSessionTask, + wait_deadline: Option, + ) { + let PendingSessionTask { + transient_cleanup, + mut task, + } = pending; + let completed = if task.is_finished() { + Some((&mut task).await) + } else { + match wait_deadline { + Some(deadline) => { + let remaining = deadline.saturating_duration_since(Instant::now()); + timeout(remaining, &mut task).await.ok() + } + None => Some((&mut task).await), + } + }; + + match completed { + Some(Ok(())) => {} + Some(Err(error)) => { + tracing::warn!( + error = %error, + "SDK Host Session ownership task failed" + ); + if let Some(cleanup) = transient_cleanup { + self.register_untracked_transient_cleanup(cleanup).await; + } + } + None => { + task.abort(); + let _ = task.await; + if let Some(cleanup) = transient_cleanup { + self.register_untracked_transient_cleanup(cleanup).await; + } + } + } + } + + async fn register_untracked_transient_cleanup(&self, cleanup: TransientSessionCleanup) { + self.inner + .state + .lock() + .await + .untracked_transient_cleanups + .insert(cleanup.session_id.clone(), cleanup); + } + + async fn compensate_registered_transient_sessions( + &self, + cleanup_deadline: Option, + ) -> bool { + let cleanups = self + .inner + .state + .lock() + .await + .untracked_transient_cleanups + .values() + .cloned() + .collect::>(); + let mut compensations = cleanups + .into_iter() + .map(|cleanup| { + let connection = self.clone(); + async move { + let session_id = cleanup.session_id.clone(); + let completed = connection + .compensate_untracked_transient_session(cleanup, cleanup_deadline) + .await; + (session_id, completed) + } + }) + .collect::>(); + let mut cleanup_complete = true; + while let Some((session_id, completed)) = compensations.next().await { + if completed { + self.inner + .state + .lock() + .await + .untracked_transient_cleanups + .remove(&session_id); + } else { + cleanup_complete = false; + } + } + cleanup_complete + } + + async fn compensate_untracked_transient_session( + &self, + cleanup: TransientSessionCleanup, + cleanup_deadline: Option, + ) -> bool { + let cleanup_timeout = cleanup_deadline + .map(|deadline| deadline.saturating_duration_since(Instant::now())) + .unwrap_or(Duration::from_secs(5)); + let result = timeout( + cleanup_timeout, + self.inner + .runtime + .discard_transient_session(AgentTransientSessionDiscardRequest { + workspace_path: cleanup.workspace_path, + session_id: cleanup.session_id.clone(), + remote_connection_id: None, + remote_ssh_host: None, + wait_timeout_ms: duration_ms( + cleanup_timeout.saturating_sub(Duration::from_millis(500)), + ), + }), + ) + .await; + match result { + Ok(Ok(_)) + | Ok(Err(RuntimeError::Port(bitfun_runtime_ports::PortError { + kind: PortErrorKind::NotFound, + .. + }))) => true, + Ok(Err(error)) => { + tracing::warn!( + session_id = %cleanup.session_id, + error_kind = runtime_error_kind(&error), + "Failed to compensate an untracked transient SDK Host Session" + ); + false + } + Err(_) => { + tracing::warn!( + session_id = %cleanup.session_id, + "Transient SDK Host Session compensation timed out" + ); + false + } + } + } + + async fn handle_initialize(&self, request: JsonRpcRequest) { + let Some(params) = self + .parse_params::(&request, ErrorStage::Initialize) + .await + else { + return; + }; + if params.protocol_version != PROTOCOL_VERSION { + self.send_error( + request.id.clone(), + ErrorCode::VersionMismatch, + ErrorStage::Initialize, + false, + Some(RecoveryAction::UpdateSdk), + "SDK Host protocol version is incompatible", + ) + .await; + return; + } + if !params.capabilities.server_notifications { + self.send_error( + request.id.clone(), + ErrorCode::CapabilityUnavailable, + ErrorStage::Initialize, + false, + None, + "query event notifications are required", + ) + .await; + return; + } + { + let mut state = self.inner.state.lock().await; + if state.initialized { + drop(state); + self.send_error( + request.id.clone(), + ErrorCode::AlreadyInitialized, + ErrorStage::Initialize, + false, + None, + "SDK Host connection is already initialized", + ) + .await; + return; + } + state.initialized = true; + } + self.send_success( + request.id.clone(), + InitializeResult::current(self.inner.runtime_version), + ) + .await; + } + + async fn handle_session_create(&self, request: JsonRpcRequest) { + let Some(params) = self + .parse_params::(&request, ErrorStage::Session) + .await + else { + return; + }; + let Ok(session_budget) = self.inner.session_budget.clone().try_acquire_owned() else { + self.send_error( + request.id.clone(), + ErrorCode::Overloaded, + ErrorStage::Session, + true, + Some(RecoveryAction::Retry), + "SDK Host Session capacity is exhausted", + ) + .await; + return; + }; + let workspace_path = params.cwd.unwrap_or_else(|| self.inner.default_cwd.clone()); + let result = self + .create_leased_session( + AgentSessionCreateRequest { + session_name: params + .session_name + .unwrap_or_else(|| DEFAULT_SESSION_NAME.to_string()), + agent_type: params.agent.unwrap_or_else(|| DEFAULT_AGENT.to_string()), + workspace_path: Some(workspace_path.clone()), + workspace_id: None, + remote_connection_id: None, + remote_ssh_host: None, + model_id: params.model, + metadata: serde_json::Map::new(), + }, + workspace_path, + session_budget, + ) + .await; + match result { + Ok(created) => { + let session_id = created.session_id.clone(); + self.deliver_session_create_response(request.id.clone(), created, session_id) + .await; + } + Err(error) => { + self.send_runtime_error(request.id.clone(), ErrorStage::Session, error) + .await + } + } + } + + async fn handle_query_start(&self, request: JsonRpcRequest) { + let emit_output = request.id.is_some(); + let Some(params) = self + .parse_params::(&request, ErrorStage::Query) + .await + else { + return; + }; + if params.prompt.trim().is_empty() { + self.send_invalid_params( + request.id.clone(), + ErrorStage::Query, + "prompt must not be empty", + ) + .await; + return; + } + if params.session_id.is_some() + && (params.session_name.is_some() + || params.agent.is_some() + || params.cwd.is_some() + || params.model.is_some()) + { + self.send_invalid_params( + request.id.clone(), + ErrorStage::Query, + "sessionName, agent, cwd, and model are only valid when creating a Session", + ) + .await; + return; + } + let Ok(query_budget) = self.inner.query_budget.clone().try_acquire_owned() else { + self.send_error( + request.id.clone(), + ErrorCode::Overloaded, + ErrorStage::Query, + true, + Some(RecoveryAction::Retry), + "SDK Host active Query capacity is exhausted", + ) + .await; + return; + }; + + let (session_id, agent_type, created_session) = match params.session_id.clone() { + Some(session_id) => { + let lease = match self.ensure_session_lease(&session_id).await { + Ok(lease) => lease, + Err(error) => { + self.send_runtime_error(request.id.clone(), ErrorStage::Session, error) + .await; + return; + } + }; + let agent_type = match self + .inner + .runtime + .resolve_session_agent_type(&session_id) + .await + { + Ok(Some(agent_type)) => agent_type, + Ok(None) => DEFAULT_AGENT.to_string(), + Err(error) => { + drop(lease); + self.send_runtime_error(request.id.clone(), ErrorStage::Session, error) + .await; + return; + } + }; + (session_id, agent_type, false) + } + None => { + let Ok(session_budget) = self.inner.session_budget.clone().try_acquire_owned() + else { + self.send_error( + request.id.clone(), + ErrorCode::Overloaded, + ErrorStage::Session, + true, + Some(RecoveryAction::Retry), + "SDK Host Session capacity is exhausted", + ) + .await; + return; + }; + let workspace_path = params + .cwd + .clone() + .unwrap_or_else(|| self.inner.default_cwd.clone()); + match self + .create_leased_session( + AgentSessionCreateRequest { + session_name: params + .session_name + .clone() + .unwrap_or_else(|| DEFAULT_SESSION_NAME.to_string()), + agent_type: params + .agent + .clone() + .unwrap_or_else(|| DEFAULT_AGENT.to_string()), + workspace_path: Some(workspace_path.clone()), + workspace_id: None, + remote_connection_id: None, + remote_ssh_host: None, + model_id: params.model.clone(), + metadata: serde_json::Map::new(), + }, + workspace_path, + session_budget, + ) + .await + { + Ok(created) => { + let agent_type = created.agent_type.clone(); + (created.session_id, agent_type, true) + } + Err(error) => { + self.send_runtime_error(request.id.clone(), ErrorStage::Session, error) + .await; + return; + } + } + } + }; + + let session = match self.reserve_query_session(&session_id).await { + Ok(session) => session, + Err(reservation_error) => { + if created_session && self.delete_unexposed_session(&session_id).await.is_err() { + self.send_cleanup_required(request.id.clone(), ErrorStage::Query, &session_id) + .await; + return; + } + match reservation_error { + QueryReservationError::Unavailable => { + self.send_error( + request.id.clone(), + ErrorCode::Overloaded, + ErrorStage::Query, + true, + Some(RecoveryAction::Retry), + "Session cannot accept a new Query while another Query, close, or shutdown is active", + ) + .await; + } + QueryReservationError::Poisoned => { + self.send_cleanup_required( + request.id.clone(), + ErrorStage::Query, + &session_id, + ) + .await; + } + } + return; + } + }; + let session_lifetime = SessionLifetime::Connection; + + let mut events = match self.inner.runtime.subscribe_session_events(&session_id) { + Ok(events) => events, + Err(error) => { + self.release_query_session(&session_id).await; + if created_session && self.delete_unexposed_session(&session_id).await.is_err() { + self.send_cleanup_required(request.id.clone(), ErrorStage::Query, &session_id) + .await; + return; + } + self.send_runtime_error(request.id.clone(), ErrorStage::Query, error) + .await; + return; + } + }; + let mut permission_events = match self.inner.runtime.subscribe_permission_requests() { + Ok(events) => events, + Err(error) => { + self.release_query_session(&session_id).await; + if created_session && self.delete_unexposed_session(&session_id).await.is_err() { + self.send_cleanup_required(request.id.clone(), ErrorStage::Query, &session_id) + .await; + return; + } + self.send_runtime_error(request.id.clone(), ErrorStage::Query, error) + .await; + return; + } + }; + let mut submission_metadata = serde_json::Map::new(); + submission_metadata.insert( + USER_INPUT_AVAILABLE_CONTEXT_KEY.to_string(), + serde_json::Value::Bool(false), + ); + submission_metadata.insert( + AUTO_APPROVE_ASK_CONTEXT_KEY.to_string(), + serde_json::Value::Bool(false), + ); + let submitted = match self + .inner + .runtime + .submit_dialog_turn(AgentDialogTurnRequest { + session_id: session_id.clone(), + message: params.prompt, + original_message: None, + turn_id: None, + agent_type, + workspace_path: Some(session.workspace_path.clone()), + remote_connection_id: session.remote_connection_id.clone(), + remote_ssh_host: session.remote_ssh_host.clone(), + policy: DialogSubmissionPolicy::for_source(AgentSubmissionSource::SdkHost), + reply_route: None, + prepended_reminders: Vec::new(), + attachments: Vec::new(), + metadata: submission_metadata, + }) + .await + { + Ok(outcome) => outcome, + Err(error) => { + self.release_query_session(&session_id).await; + if created_session && self.delete_unexposed_session(&session_id).await.is_err() { + self.send_cleanup_required(request.id.clone(), ErrorStage::Query, &session_id) + .await; + return; + } + self.send_runtime_error(request.id.clone(), ErrorStage::Query, error) + .await; + return; + } + }; + let (submitted_session_id, turn_id) = match submitted { + DialogSubmitOutcome::Started { + session_id, + turn_id, + } => (session_id, turn_id), + DialogSubmitOutcome::Queued { + session_id, + turn_id, + } => (session_id, turn_id), + }; + let query_id = format!("query_{}", uuid::Uuid::new_v4()); + let lease = Arc::new(QueryLease { + query_id: query_id.clone(), + session_id: submitted_session_id.clone(), + turn_id: turn_id.clone(), + terminal: AtomicBool::new(false), + stop_forwarding: CancellationToken::new(), + emit_output, + _budget: query_budget, + }); + { + let mut state = self.inner.state.lock().await; + state.queries.insert(query_id.clone(), lease.clone()); + state.starting_query_sessions.remove(&session_id); + state.active_query_sessions.insert(session_id.clone()); + } + + let start_delivered = self + .deliver_query_start_response( + request.id.clone(), + QueryStartResult { + query_id: query_id.clone(), + session_id: submitted_session_id.clone(), + turn_id: turn_id.clone(), + accepted: true, + created_session, + session_lifetime, + }, + lease.clone(), + created_session, + ) + .await; + if !start_delivered { + return; + } + + let connection = self.clone(); + tokio::spawn(async move { + let mut sequence = 0u64; + loop { + let envelope = tokio::select! { + _ = lease.stop_forwarding.cancelled() => return, + permission = permission_events.recv() => { + match permission { + Ok(PermissionRequestEvent::Asked { request }) + if permission_request_targets_query( + &request, + connection + .inner + .runtime + .permission_request_dialog_turn_id(&request.request_id) + .ok() + .flatten() + .as_deref(), + &lease, + ) => + { + connection.reject_permission_and_finish(&lease, &request).await; + return; + } + Ok(_) => continue, + Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { + match connection.inner.runtime.pending_permission_requests() { + Ok(pending) => { + if let Some(request) = pending + .into_iter() + .find(|request| { + permission_request_targets_query( + request, + connection + .inner + .runtime + .permission_request_dialog_turn_id( + &request.request_id, + ) + .ok() + .flatten() + .as_deref(), + &lease, + ) + }) + { + connection.reject_permission_and_finish(&lease, &request).await; + return; + } + continue; + } + Err(error) => { + connection.cancel_and_finish( + &lease, + query_error_from_runtime( + &lease.query_id, + error, + "SDK Host could not recover permission requests after event lag", + ), + true, + ).await; + return; + } + } + } + Err(_) => { + connection.cancel_and_finish( + &lease, + QueryResultError::new( + ErrorCode::Internal, + true, + Some(RecoveryAction::RestartHost), + &lease.query_id, + "SDK Host permission event stream is unavailable", + ), + true, + ).await; + return; + } + } + } + event = events.recv() => match event { + Ok(event) => event, + Err(tokio::sync::broadcast::error::RecvError::Lagged(_)) => { + connection.cancel_and_finish( + &lease, + QueryResultError::new( + ErrorCode::Internal, + true, + Some(RecoveryAction::RestartHost), + &lease.query_id, + "SDK Host event stream lagged", + ), + true, + ).await; + return; + } + Err(tokio::sync::broadcast::error::RecvError::Closed) => { + connection.cancel_and_finish( + &lease, + QueryResultError::new( + ErrorCode::Internal, + true, + Some(RecoveryAction::RestartHost), + &lease.query_id, + "SDK Host event stream closed", + ), + true, + ).await; + return; + } + } + }; + if event_turn_id(&envelope.event) != Some(lease.turn_id.as_str()) { + continue; + } + let terminal = terminal_fact(&envelope.event, &lease.turn_id, &lease.query_id); + if lease.emit_output { + if let Some(projected) = project_query_event(&envelope.event) { + sequence += 1; + if !connection + .send_notification( + NOTIFICATION_QUERY_EVENT, + QueryEventParams { + query_id: lease.query_id.clone(), + session_id: lease.session_id.clone(), + turn_id: lease.turn_id.clone(), + sequence, + event: projected, + }, + ) + .await + { + connection + .cancel_and_finish( + &lease, + QueryResultError::new( + ErrorCode::ProcessLost, + false, + None, + &lease.query_id, + "SDK Host output is unavailable", + ), + false, + ) + .await; + return; + } + } + } + if let Some((status, error)) = terminal { + connection.finish_query(&lease, status, error, true).await; + return; + } + } + }); + } + + async fn handle_query_cancel(&self, request: JsonRpcRequest) { + let Some(params) = self + .parse_params::(&request, ErrorStage::Query) + .await + else { + return; + }; + let lease = self + .inner + .state + .lock() + .await + .queries + .get(¶ms.query_id) + .cloned(); + let Some(lease) = lease else { + self.send_error( + request.id.clone(), + ErrorCode::NotFound, + ErrorStage::Query, + false, + None, + "Query was not found", + ) + .await; + return; + }; + match timeout( + Duration::from_millis(2_500), + self.inner + .runtime + .cancel_turn(AgentTurnCancellationRequest { + session_id: lease.session_id.clone(), + turn_id: Some(lease.turn_id.clone()), + source: Some(AgentSubmissionSource::SdkHost), + requester_session_id: None, + reason: Some("sdk_query_cancel".to_string()), + wait_timeout_ms: Some(2_000), + }), + ) + .await + { + Ok(Ok(result)) => { + self.send_success( + request.id.clone(), + QueryCancelResult { + query_id: lease.query_id.clone(), + session_id: lease.session_id.clone(), + turn_id: lease.turn_id.clone(), + requested: result.requested, + }, + ) + .await; + } + Ok(Err(error)) => { + self.send_runtime_error(request.id.clone(), ErrorStage::Query, error) + .await + } + Err(_) => { + self.send_error( + request.id.clone(), + ErrorCode::Timeout, + ErrorStage::Query, + true, + Some(RecoveryAction::Retry), + "SDK Host Query cancellation timed out", + ) + .await; + } + } + } + + async fn handle_session_close(&self, request: JsonRpcRequest) { + let Some(params) = self + .parse_params::(&request, ErrorStage::Session) + .await + else { + return; + }; + if params + .wait_timeout_ms + .is_some_and(|timeout| timeout == 0 || timeout > MAX_SESSION_CLOSE_TIMEOUT_MS) + { + self.send_invalid_params( + request.id.clone(), + ErrorStage::Session, + "waitTimeoutMs must be between 1 and 30000", + ) + .await; + return; + } + let session = { + let mut state = self.inner.state.lock().await; + if state.closing_sessions.contains(¶ms.session_id) { + drop(state); + self.send_error( + request.id.clone(), + ErrorCode::Overloaded, + ErrorStage::Session, + true, + Some(RecoveryAction::Retry), + "Session close is already in progress", + ) + .await; + return; + } + if state.starting_query_sessions.contains(¶ms.session_id) { + drop(state); + self.send_error( + request.id.clone(), + ErrorCode::Overloaded, + ErrorStage::Session, + true, + Some(RecoveryAction::Retry), + "Session has a Query start in progress", + ) + .await; + return; + } + let session = state.sessions.get(¶ms.session_id).cloned(); + if session.is_some() { + state.closing_sessions.insert(params.session_id.clone()); + } + session + }; + let Some(session) = session else { + self.send_error( + request.id.clone(), + ErrorCode::NotFound, + ErrorStage::Session, + false, + None, + "Session is not owned by this SDK Host connection", + ) + .await; + return; + }; + let close_timeout_ms = params.wait_timeout_ms.unwrap_or(5_000); + let session_id = params.session_id.clone(); + let operation = async { + self.inner + .runtime + .discard_transient_session(AgentTransientSessionDiscardRequest { + workspace_path: session.workspace_path, + session_id: session_id.clone(), + remote_connection_id: session.remote_connection_id, + remote_ssh_host: session.remote_ssh_host, + wait_timeout_ms: close_timeout_ms, + }) + .await + }; + match timeout(Duration::from_millis(close_timeout_ms + 500), operation).await { + Ok(Ok(unloaded)) => { + let queries = { + let mut state = self.inner.state.lock().await; + state.sessions.remove(¶ms.session_id); + state.closing_sessions.remove(¶ms.session_id); + state.poisoned_sessions.remove(¶ms.session_id); + state.starting_query_sessions.remove(¶ms.session_id); + state.active_query_sessions.remove(¶ms.session_id); + let query_ids = state + .queries + .iter() + .filter(|(_, lease)| lease.session_id == params.session_id) + .map(|(query_id, _)| query_id.clone()) + .collect::>(); + query_ids + .into_iter() + .filter_map(|query_id| state.queries.remove(&query_id)) + .collect::>() + }; + for query in queries { + query.stop_forwarding.cancel(); + self.finish_query(&query, QueryTerminalStatus::Cancelled, None, true) + .await; + } + self.send_success( + request.id.clone(), + SessionCloseResult { + session_id: params.session_id, + unloaded, + }, + ) + .await; + } + Ok(Err(error)) => { + tracing::warn!( + session_id = %params.session_id, + error_kind = runtime_error_kind(&error), + "SDK Host Session close ended with uncertain cleanup" + ); + self.mark_session_cleanup_failed(¶ms.session_id).await; + self.send_error( + request.id.clone(), + ErrorCode::CleanupRequired, + ErrorStage::Session, + false, + Some(RecoveryAction::RestartHost), + "SDK Host Session cleanup is incomplete; restart the Host before retrying", + ) + .await; + } + Err(_) => { + tracing::warn!( + session_id = %params.session_id, + "SDK Host Session close timed out with uncertain cleanup" + ); + self.mark_session_cleanup_failed(¶ms.session_id).await; + self.send_error( + request.id.clone(), + ErrorCode::CleanupRequired, + ErrorStage::Session, + false, + Some(RecoveryAction::RestartHost), + "SDK Host Session cleanup timed out; restart the Host before retrying", + ) + .await; + } + } + } + + async fn mark_session_cleanup_failed(&self, session_id: &str) { + let mut state = self.inner.state.lock().await; + state.closing_sessions.remove(session_id); + state.poisoned_sessions.insert(session_id.to_string()); + state.cleanup_failed = true; + } + + async fn ensure_session_lease(&self, session_id: &str) -> Result { + self.inner + .state + .lock() + .await + .sessions + .get(session_id) + .cloned() + .ok_or_else(|| { + bitfun_runtime_ports::PortError::new( + PortErrorKind::NotAvailable, + "sessionId must belong to the same SDK Host connection; durable Session resume is not available in this protocol version", + ) + .into() + }) + } + + async fn create_leased_session( + &self, + request: AgentSessionCreateRequest, + workspace_path: String, + session_budget: OwnedSemaphorePermit, + ) -> Result { + let session_id = uuid::Uuid::new_v4().to_string(); + let runtime = self.inner.runtime.clone(); + let state = self.inner.state.clone(); + let task_state = state.clone(); + let (result_tx, result_rx) = oneshot::channel(); + let mut connection_state = state.lock().await; + if connection_state.shutting_down { + return Err(bitfun_runtime_ports::PortError::new( + PortErrorKind::Cancelled, + "SDK Host connection is shutting down", + ) + .into()); + } + let task_session_id = session_id.clone(); + let cleanup_workspace_path = workspace_path.clone(); + let creation = tokio::spawn(async move { + let result = runtime + .create_transient_session_with_id(task_session_id, request) + .await; + if let Ok(created) = &result { + task_state.lock().await.sessions.insert( + created.session_id.clone(), + SessionLease { + workspace_path, + remote_connection_id: None, + remote_ssh_host: None, + exposed: false, + _budget: Arc::new(session_budget), + }, + ); + } + let _ = result_tx.send(result); + }); + connection_state + .pending_session_tasks + .push(PendingSessionTask { + transient_cleanup: Some(TransientSessionCleanup { + session_id, + workspace_path: cleanup_workspace_path, + }), + task: creation, + }); + drop(connection_state); + result_rx.await.map_err(|_| { + RuntimeError::from(bitfun_runtime_ports::PortError::new( + PortErrorKind::Backend, + "SDK Host Session creation task ended without a result", + )) + })? + } + + async fn deliver_session_create_response( + &self, + request_id: Option, + created: AgentSessionCreateResult, + session_id: String, + ) -> bool { + let connection = self.clone(); + let (delivered_tx, delivered_rx) = oneshot::channel(); + let mut state = self.inner.state.lock().await; + if state.shutting_down { + return false; + } + let delivery = tokio::spawn(async move { + let delivered = connection + .send_success( + request_id, + SessionCreateResult { + session_id: created.session_id, + session_name: created.session_name, + agent: created.agent_type, + lifetime: SessionLifetime::Connection, + }, + ) + .await; + if delivered { + connection.mark_session_exposed(&session_id).await; + } else { + tokio::select! { + _ = connection.inner.shutdown_started.cancelled() => {} + _ = connection.delete_unexposed_session(&session_id) => {} + } + } + let _ = delivered_tx.send(delivered); + }); + state.pending_session_tasks.push(PendingSessionTask { + transient_cleanup: None, + task: delivery, + }); + drop(state); + delivered_rx.await.unwrap_or(false) + } + + async fn deliver_query_start_response( + &self, + request_id: Option, + result: QueryStartResult, + lease: Arc, + created_session: bool, + ) -> bool { + let connection = self.clone(); + let (delivered_tx, delivered_rx) = oneshot::channel(); + let mut state = self.inner.state.lock().await; + if state.shutting_down { + return false; + } + let delivery = tokio::spawn(async move { + let delivered = connection.send_success(request_id, result).await; + if delivered { + if created_session { + connection.mark_session_exposed(&lease.session_id).await; + } + } else { + let cleanup = async { + connection + .cancel_and_finish( + &lease, + QueryResultError::new( + ErrorCode::ProcessLost, + false, + None, + &lease.query_id, + "SDK Host output is unavailable", + ), + false, + ) + .await; + if created_session { + let _ = connection.delete_unexposed_session(&lease.session_id).await; + } + }; + tokio::select! { + _ = connection.inner.shutdown_started.cancelled() => {} + _ = cleanup => {} + } + } + let _ = delivered_tx.send(delivered); + }); + state.pending_session_tasks.push(PendingSessionTask { + transient_cleanup: None, + task: delivery, + }); + drop(state); + delivered_rx.await.unwrap_or(false) + } + + async fn reserve_query_session( + &self, + session_id: &str, + ) -> Result { + let mut state = self.inner.state.lock().await; + if state.poisoned_sessions.contains(session_id) { + return Err(QueryReservationError::Poisoned); + } + if state.shutting_down + || state.closing_sessions.contains(session_id) + || state.starting_query_sessions.contains(session_id) + || state.active_query_sessions.contains(session_id) + { + return Err(QueryReservationError::Unavailable); + } + let session = state + .sessions + .get(session_id) + .cloned() + .ok_or(QueryReservationError::Unavailable)?; + state.starting_query_sessions.insert(session_id.to_string()); + Ok(session) + } + + async fn release_query_session(&self, session_id: &str) { + let mut state = self.inner.state.lock().await; + state.starting_query_sessions.remove(session_id); + state.active_query_sessions.remove(session_id); + } + + async fn mark_session_exposed(&self, session_id: &str) { + if let Some(session) = self.inner.state.lock().await.sessions.get_mut(session_id) { + session.exposed = true; + } + } + + async fn delete_unexposed_session(&self, session_id: &str) -> Result<(), ()> { + let lease = self + .inner + .state + .lock() + .await + .sessions + .get(session_id) + .cloned(); + let Some(lease) = lease else { + return Ok(()); + }; + match timeout( + Duration::from_millis(5_000), + self.inner + .runtime + .discard_transient_session(AgentTransientSessionDiscardRequest { + workspace_path: lease.workspace_path, + session_id: session_id.to_string(), + remote_connection_id: lease.remote_connection_id, + remote_ssh_host: lease.remote_ssh_host, + wait_timeout_ms: 4_500, + }), + ) + .await + { + Ok(Ok(_)) => { + self.inner.state.lock().await.sessions.remove(session_id); + Ok(()) + } + Ok(Err(error)) => { + tracing::warn!( + session_id = %session_id, + error_kind = runtime_error_kind(&error), + "Failed to delete an unexposed SDK Host Session" + ); + self.inner.state.lock().await.cleanup_failed = true; + Err(()) + } + Err(_) => { + tracing::warn!( + session_id = %session_id, + "Timed out while deleting an unexposed SDK Host Session" + ); + self.inner.state.lock().await.cleanup_failed = true; + Err(()) + } + } + } + + async fn send_cleanup_required( + &self, + request_id: Option, + stage: ErrorStage, + session_id: &str, + ) { + let message = format!( + "SDK Host could not remove unexposed Session {session_id}; restart the Host before retrying" + ); + self.send_error( + request_id, + ErrorCode::CleanupRequired, + stage, + false, + Some(RecoveryAction::RestartHost), + &message, + ) + .await; + } + + async fn finish_query( + &self, + lease: &Arc, + mut status: QueryTerminalStatus, + mut error: Option, + emit_result: bool, + ) { + if !lease.finish_once() { + return; + } + let settlement = timeout( + Duration::from_millis(DEFAULT_TURN_SETTLEMENT_TIMEOUT_MS + 500), + self.inner + .runtime + .wait_for_turn_settlement(AgentTurnSettlementRequest { + session_id: lease.session_id.clone(), + turn_id: lease.turn_id.clone(), + wait_timeout_ms: DEFAULT_TURN_SETTLEMENT_TIMEOUT_MS, + }), + ) + .await; + let mut poison_session = false; + match settlement { + Ok(Ok(())) => {} + Ok(Err(_settlement_error)) => { + poison_session = true; + status = QueryTerminalStatus::Failed; + error = Some(QueryResultError::new( + ErrorCode::CleanupRequired, + false, + Some(RecoveryAction::RestartHost), + &lease.query_id, + "SDK Host could not confirm Turn settlement; restart the Host before retrying", + )); + } + Err(_) => { + poison_session = true; + status = QueryTerminalStatus::Failed; + error = Some(QueryResultError::new( + ErrorCode::CleanupRequired, + false, + Some(RecoveryAction::RestartHost), + &lease.query_id, + "SDK Host Turn settlement timed out; restart the Host before retrying", + )); + } + } + { + let mut state = self.inner.state.lock().await; + state.queries.remove(&lease.query_id); + state.starting_query_sessions.remove(&lease.session_id); + state.active_query_sessions.remove(&lease.session_id); + if poison_session { + state.poisoned_sessions.insert(lease.session_id.clone()); + } + } + if emit_result { + self.send_query_result(lease, status, error).await; + } + } + + async fn send_query_result( + &self, + lease: &QueryLease, + status: QueryTerminalStatus, + error: Option, + ) -> bool { + if !lease.emit_output { + return true; + } + self.send_notification( + NOTIFICATION_QUERY_RESULT, + QueryResultParams { + query_id: lease.query_id.clone(), + session_id: lease.session_id.clone(), + turn_id: lease.turn_id.clone(), + status, + error, + }, + ) + .await + } + + async fn cancel_and_finish( + &self, + lease: &Arc, + mut error: QueryResultError, + emit_result: bool, + ) { + let cancellation = timeout( + Duration::from_millis(2_500), + self.inner + .runtime + .cancel_turn(AgentTurnCancellationRequest { + session_id: lease.session_id.clone(), + turn_id: Some(lease.turn_id.clone()), + source: Some(AgentSubmissionSource::SdkHost), + requester_session_id: None, + reason: Some("sdk_host_fail_closed".to_string()), + wait_timeout_ms: Some(2_000), + }), + ) + .await; + match cancellation { + Ok(Ok(_)) => {} + Ok(Err(cancel_error)) => { + error = query_error_from_runtime( + &lease.query_id, + cancel_error, + "SDK Host could not cancel the Turn after a Host failure", + ); + } + Err(_) => { + error = QueryResultError::new( + ErrorCode::Timeout, + true, + Some(RecoveryAction::RestartHost), + &lease.query_id, + "SDK Host cancellation timed out after a Host failure", + ); + } + } + self.finish_query(lease, QueryTerminalStatus::Failed, Some(error), emit_result) + .await; + } + + async fn reject_permission_and_finish( + &self, + lease: &Arc, + request: &PermissionRequest, + ) { + let reply = PermissionReply::Reject { + feedback: Some( + "Non-interactive SDK execution requires an explicit permission callback" + .to_string(), + ), + }; + match timeout( + Duration::from_millis(PERMISSION_REJECTION_TIMEOUT_MS), + self.inner.runtime.respond_permission_with_source( + &request.request_id, + reply, + PermissionReplySource::System, + ), + ) + .await + { + Ok(Ok(_)) => {} + Ok(Err(error)) => { + self.cancel_and_finish( + lease, + query_error_from_runtime( + &lease.query_id, + error, + "SDK Host could not reject a pending permission request", + ), + true, + ) + .await; + return; + } + Err(_) => { + self.cancel_and_finish( + lease, + QueryResultError::new( + ErrorCode::Timeout, + true, + Some(RecoveryAction::RestartHost), + &lease.query_id, + "SDK Host permission rejection timed out", + ), + true, + ) + .await; + return; + } + } + self.cancel_and_finish( + lease, + QueryResultError::new( + ErrorCode::ActionRequired, + false, + None, + &lease.query_id, + "Permission approval is required but permission callbacks are unavailable", + ), + true, + ) + .await; + } + + async fn parse_params(&self, request: &JsonRpcRequest, stage: ErrorStage) -> Option + where + T: serde::de::DeserializeOwned, + { + match request.params_as() { + Ok(params) => Some(params), + Err(_) => { + self.send_invalid_params(request.id.clone(), stage, "invalid method parameters") + .await; + None + } + } + } + + async fn send_invalid_params( + &self, + id: Option, + stage: ErrorStage, + message: &'static str, + ) { + self.send_rpc_error( + id, + -32602, + ErrorCode::InvalidRequest, + stage, + false, + None, + message, + ) + .await; + } + + async fn send_runtime_error( + &self, + id: Option, + stage: ErrorStage, + error: RuntimeError, + ) { + let (code, retryable, recovery) = runtime_error_facts(&error); + self.send_error(id, code, stage, retryable, recovery, &error.into_message()) + .await; + } + + async fn send_error( + &self, + id: Option, + code: ErrorCode, + stage: ErrorStage, + retryable: bool, + recovery: Option, + message: &str, + ) { + self.send_rpc_error(id, -32000, code, stage, retryable, recovery, message) + .await; + } + + #[allow(clippy::too_many_arguments)] + async fn send_rpc_error( + &self, + id: Option, + rpc_code: i32, + code: ErrorCode, + stage: ErrorStage, + retryable: bool, + recovery: Option, + message: &str, + ) { + let Some(id) = id else { + return; + }; + let correlation_id = id.correlation_id(); + self.send_value(JsonRpcErrorResponse::new( + id, + rpc_code, + message, + ErrorData { + code, + stage, + retryable, + correlation_id, + recovery, + }, + )) + .await; + } + + async fn send_success(&self, id: Option, result: T) -> bool { + match id { + Some(id) => { + self.send_value(JsonRpcSuccessResponse::new(id, result)) + .await + } + None => true, + } + } + + async fn send_notification( + &self, + method: &'static str, + params: T, + ) -> bool { + self.send_value(JsonRpcNotification::new(method, params)) + .await + } + + async fn send_value(&self, value: T) -> bool { + let Ok(value) = serde_json::to_value(value) else { + return false; + }; + self.inner.output.send(value).await.is_ok() + } +} + +fn event_turn_id(event: &AgenticEvent) -> Option<&str> { + match event { + AgenticEvent::DialogTurnCompleted { turn_id, .. } + | AgenticEvent::DialogTurnCancelled { turn_id, .. } + | AgenticEvent::DialogTurnFailed { turn_id, .. } + | AgenticEvent::TextChunk { turn_id, .. } => Some(turn_id), + _ => None, + } +} + +fn duration_ms(duration: Duration) -> u64 { + duration.as_millis().min(u64::MAX as u128) as u64 +} + +fn project_query_event(event: &AgenticEvent) -> Option { + match event { + AgenticEvent::TextChunk { text, .. } => { + Some(QueryEvent::AssistantTextDelta { text: text.clone() }) + } + _ => None, + } +} + +fn terminal_fact( + event: &AgenticEvent, + expected_turn_id: &str, + query_id: &str, +) -> Option<(QueryTerminalStatus, Option)> { + match event { + AgenticEvent::DialogTurnCompleted { + turn_id, + success, + finish_reason, + has_final_response, + .. + } if turn_id == expected_turn_id => { + if success == &Some(false) || has_final_response == &Some(false) { + Some(( + QueryTerminalStatus::Failed, + Some(QueryResultError::new( + ErrorCode::Internal, + false, + None, + query_id, + format!( + "Query completed unsuccessfully: {}", + finish_reason + .as_deref() + .unwrap_or("unsuccessful_completion") + ), + )), + )) + } else { + Some((QueryTerminalStatus::Completed, None)) + } + } + AgenticEvent::DialogTurnCancelled { turn_id, .. } if turn_id == expected_turn_id => { + Some((QueryTerminalStatus::Cancelled, None)) + } + AgenticEvent::DialogTurnFailed { + turn_id, + error, + error_category, + error_detail, + .. + } if turn_id == expected_turn_id => Some(( + QueryTerminalStatus::Failed, + Some(query_error_from_failure( + query_id, + error, + error_category.as_ref(), + error_detail.as_ref().and_then(|detail| detail.retryable), + )), + )), + _ => None, + } +} + +fn query_error_from_failure( + correlation_id: &str, + message: &str, + category: Option<&ErrorCategory>, + explicit_retryable: Option, +) -> QueryResultError { + let (code, default_retryable, recovery) = match category { + Some(ErrorCategory::Network | ErrorCategory::ProviderUnavailable) => ( + ErrorCode::ProviderUnavailable, + true, + Some(RecoveryAction::Retry), + ), + Some(ErrorCategory::Auth) => (ErrorCode::Authentication, false, None), + Some(ErrorCategory::RateLimit) => { + (ErrorCode::RateLimited, true, Some(RecoveryAction::Retry)) + } + Some(ErrorCategory::ContextOverflow) => (ErrorCode::ContextOverflow, false, None), + Some(ErrorCategory::Timeout) => (ErrorCode::Timeout, true, Some(RecoveryAction::Retry)), + Some(ErrorCategory::ProviderQuota) => (ErrorCode::ProviderQuota, false, None), + Some(ErrorCategory::ProviderBilling) => (ErrorCode::ProviderBilling, false, None), + Some(ErrorCategory::Permission) => (ErrorCode::PermissionDenied, false, None), + Some(ErrorCategory::InvalidRequest) => (ErrorCode::InvalidRequest, false, None), + Some(ErrorCategory::ContentPolicy) => (ErrorCode::ContentPolicy, false, None), + Some(ErrorCategory::ModelError | ErrorCategory::Unknown) | None => { + (ErrorCode::Internal, false, None) + } + }; + QueryResultError::new( + code, + explicit_retryable.unwrap_or(default_retryable), + recovery, + correlation_id, + message, + ) +} + +fn query_error_from_runtime( + query_id: &str, + error: RuntimeError, + context: &str, +) -> QueryResultError { + let (code, retryable, recovery) = runtime_error_facts(&error); + QueryResultError::new( + code, + retryable, + recovery, + query_id, + format!("{context}: {}", error.into_message()), + ) +} + +fn runtime_error_facts(error: &RuntimeError) -> (ErrorCode, bool, Option) { + match error { + RuntimeError::Port(port) => match port.kind { + PortErrorKind::NotAvailable => (ErrorCode::CapabilityUnavailable, false, None), + PortErrorKind::NotFound => (ErrorCode::NotFound, false, None), + PortErrorKind::InvalidRequest => (ErrorCode::InvalidRequest, false, None), + PortErrorKind::PermissionDenied => (ErrorCode::PermissionDenied, false, None), + PortErrorKind::Cancelled => (ErrorCode::Cancelled, false, None), + PortErrorKind::Timeout => (ErrorCode::Timeout, true, Some(RecoveryAction::Retry)), + PortErrorKind::CleanupRequired => ( + ErrorCode::CleanupRequired, + false, + Some(RecoveryAction::RestartHost), + ), + PortErrorKind::Backend => { + (ErrorCode::Internal, true, Some(RecoveryAction::RestartHost)) + } + }, + RuntimeError::PermissionRequest(_) => (ErrorCode::PermissionDenied, false, None), + _ => (ErrorCode::CapabilityUnavailable, false, None), + } +} + +fn permission_request_targets_query( + request: &PermissionRequest, + dialog_turn_id: Option<&str>, + lease: &QueryLease, +) -> bool { + (request.session_id == lease.session_id && dialog_turn_id == Some(lease.turn_id.as_str())) + || request.delegation.as_ref().is_some_and(|delegation| { + delegation.parent_session_id == lease.session_id + && delegation.parent_dialog_turn_id.as_deref() == Some(lease.turn_id.as_str()) + }) +} + +fn runtime_error_kind(error: &RuntimeError) -> &'static str { + match error { + RuntimeError::Port(port) => match port.kind { + PortErrorKind::NotAvailable => "not_available", + PortErrorKind::NotFound => "not_found", + PortErrorKind::InvalidRequest => "invalid_request", + PortErrorKind::PermissionDenied => "permission_denied", + PortErrorKind::Cancelled => "cancelled", + PortErrorKind::Timeout => "timeout", + PortErrorKind::CleanupRequired => "cleanup_required", + PortErrorKind::Backend => "backend", + }, + RuntimeError::MissingDialogTurnPort + | RuntimeError::MissingLifecycleDeliveryPort + | RuntimeError::MissingCancellationPort + | RuntimeError::MissingSessionManagementPort + | RuntimeError::MissingSessionRestorePort + | RuntimeError::MissingLocalCommandTurnPort + | RuntimeError::MissingSessionTranscriptReader + | RuntimeError::MissingThreadGoalManagementPort + | RuntimeError::MissingInteractionResponsePort + | RuntimeError::MissingEventSink + | RuntimeError::MissingEventSource + | RuntimeError::MissingPermissionRequestManager => "capability_unavailable", + RuntimeError::PermissionRequest(_) => "permission_request", + } +} diff --git a/src/crates/interfaces/sdk-host/src/lib.rs b/src/crates/interfaces/sdk-host/src/lib.rs new file mode 100644 index 0000000000..53d4130a79 --- /dev/null +++ b/src/crates/interfaces/sdk-host/src/lib.rs @@ -0,0 +1,8 @@ +//! Versioned local SDK Host adapter for the shared BitFun Agent Runtime. +//! +//! This crate owns only protocol and connection lifecycle. Agent execution, +//! Session persistence, Tool/MCP, Permission, and Hook behavior remain in the +//! existing runtime owners supplied through [`bitfun_agent_runtime`]. + +pub mod host; +pub mod protocol; diff --git a/src/crates/interfaces/sdk-host/src/protocol.rs b/src/crates/interfaces/sdk-host/src/protocol.rs new file mode 100644 index 0000000000..cde91ea0b0 --- /dev/null +++ b/src/crates/interfaces/sdk-host/src/protocol.rs @@ -0,0 +1,469 @@ +//! Versioned JSON-RPC contracts for the local SDK Host. + +use serde::de::DeserializeOwned; +use serde::{Deserialize, Serialize}; + +pub const JSON_RPC_VERSION: &str = "2.0"; +pub const PROTOCOL_VERSION: u32 = 1; + +pub const METHOD_INITIALIZE: &str = "initialize"; +pub const METHOD_SESSION_CREATE: &str = "session/create"; +pub const METHOD_QUERY_START: &str = "query/start"; +pub const METHOD_QUERY_CANCEL: &str = "query/cancel"; +pub const METHOD_SESSION_CLOSE: &str = "session/close"; +pub const METHOD_SHUTDOWN: &str = "shutdown"; +pub const NOTIFICATION_QUERY_EVENT: &str = "query/event"; +pub const NOTIFICATION_QUERY_RESULT: &str = "query/result"; + +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +#[serde(untagged)] +pub enum RequestId { + Number(i64), + String(String), +} + +impl RequestId { + pub fn correlation_id(&self) -> String { + match self { + Self::Number(value) => format!("request:number:{value}"), + Self::String(value) => format!("request:string:{value}"), + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct JsonRpcRequest { + pub jsonrpc: String, + /// Absent for a JSON-RPC notification. `null` remains invalid at the + /// transport envelope boundary and is not treated as a notification. + #[serde( + default, + deserialize_with = "deserialize_optional_request_id", + skip_serializing_if = "Option::is_none" + )] + pub id: Option, + pub method: String, + #[serde(default = "empty_object")] + pub params: serde_json::Value, +} + +impl JsonRpcRequest { + pub fn params_as(&self) -> Result { + serde_json::from_value(self.params.clone()) + } +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct JsonRpcSuccessResponse { + pub jsonrpc: &'static str, + pub id: RequestId, + pub result: T, +} + +impl JsonRpcSuccessResponse { + pub fn new(id: RequestId, result: T) -> Self { + Self { + jsonrpc: JSON_RPC_VERSION, + id, + result, + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct JsonRpcErrorResponse { + pub jsonrpc: &'static str, + pub id: Option, + pub error: JsonRpcErrorObject, +} + +impl JsonRpcErrorResponse { + pub fn new(id: RequestId, rpc_code: i32, message: impl Into, data: ErrorData) -> Self { + Self { + jsonrpc: JSON_RPC_VERSION, + id: Some(id), + error: JsonRpcErrorObject { + code: rpc_code, + message: message.into(), + data, + }, + } + } + + pub fn parse_error(message: impl Into, correlation_id: impl Into) -> Self { + Self { + jsonrpc: JSON_RPC_VERSION, + id: None, + error: JsonRpcErrorObject { + code: -32700, + message: message.into(), + data: ErrorData { + code: ErrorCode::InvalidRequest, + stage: ErrorStage::Protocol, + retryable: false, + correlation_id: correlation_id.into(), + recovery: None, + }, + }, + } + } + + pub fn invalid_request( + id: Option, + message: impl Into, + correlation_id: impl Into, + ) -> Self { + Self { + jsonrpc: JSON_RPC_VERSION, + id, + error: JsonRpcErrorObject { + code: -32600, + message: message.into(), + data: ErrorData { + code: ErrorCode::InvalidRequest, + stage: ErrorStage::Protocol, + retryable: false, + correlation_id: correlation_id.into(), + recovery: None, + }, + }, + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct JsonRpcErrorObject { + pub code: i32, + pub message: String, + pub data: ErrorData, +} + +#[derive(Debug, Clone, PartialEq, Serialize)] +pub struct JsonRpcNotification { + pub jsonrpc: &'static str, + pub method: &'static str, + pub params: T, +} + +impl JsonRpcNotification { + pub fn new(method: &'static str, params: T) -> Self { + Self { + jsonrpc: JSON_RPC_VERSION, + method, + params, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields, rename_all = "camelCase")] +pub struct InitializeParams { + pub protocol_version: u32, + pub client_info: ClientInfo, + pub capabilities: ClientCapabilities, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ClientInfo { + pub name: String, + pub version: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields, rename_all = "camelCase")] +pub struct ClientCapabilities { + pub server_notifications: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct InitializeResult { + pub protocol_version: u32, + pub runtime_version: String, + pub stability: Stability, + pub capabilities: HostCapabilities, +} + +impl InitializeResult { + pub fn current(runtime_version: impl Into) -> Self { + Self { + protocol_version: PROTOCOL_VERSION, + runtime_version: runtime_version.into(), + stability: Stability::NotDelivered, + capabilities: HostCapabilities::current(), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum Stability { + /// Internal implementation candidate. It is not a supported SDK surface. + NotDelivered, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct HostCapabilities { + pub session_create: bool, + pub session_create_lifetime: SessionLifetime, + pub query: bool, + pub query_cancel: bool, + pub session_close: bool, + pub event_stream: bool, + pub structured_output: bool, + pub usage: bool, + pub custom_tools: bool, + pub permission_callbacks: bool, + pub hooks: bool, + pub mcp_configuration: bool, + pub prestarted_transport: bool, +} + +impl HostCapabilities { + pub const fn current() -> Self { + Self { + session_create: true, + session_create_lifetime: SessionLifetime::Connection, + query: true, + query_cancel: true, + session_close: true, + event_stream: true, + structured_output: false, + usage: false, + custom_tools: false, + permission_callbacks: false, + hooks: false, + mcp_configuration: false, + prestarted_transport: false, + } + } +} + +/// Persistence boundary of a Session visible through the internal Host candidate. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum SessionLifetime { + /// Created and deleted by this Host connection. + Connection, +} + +fn deserialize_optional_request_id<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + RequestId::deserialize(deserializer).map(Some) +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ErrorCode { + InvalidRequest, + NotInitialized, + AlreadyInitialized, + VersionMismatch, + CapabilityUnavailable, + NotFound, + PermissionDenied, + ActionRequired, + Authentication, + RateLimited, + ProviderQuota, + ProviderBilling, + ProviderUnavailable, + ContextOverflow, + ContentPolicy, + Overloaded, + Timeout, + Cancelled, + ProcessLost, + CleanupRequired, + Internal, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ErrorStage { + Protocol, + Initialize, + Session, + Query, + Shutdown, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RecoveryAction { + Initialize, + Retry, + UpdateSdk, + RestartHost, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ErrorData { + pub code: ErrorCode, + pub stage: ErrorStage, + pub retryable: bool, + pub correlation_id: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub recovery: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields, rename_all = "camelCase")] +pub struct SessionCreateParams { + #[serde(default)] + pub session_name: Option, + #[serde(default)] + pub agent: Option, + #[serde(default)] + pub cwd: Option, + #[serde(default)] + pub model: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionCreateResult { + pub session_id: String, + pub session_name: String, + pub agent: String, + pub lifetime: SessionLifetime, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields, rename_all = "camelCase")] +pub struct QueryStartParams { + pub prompt: String, + #[serde(default)] + pub session_id: Option, + #[serde(default)] + pub session_name: Option, + #[serde(default)] + pub agent: Option, + #[serde(default)] + pub cwd: Option, + #[serde(default)] + pub model: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct QueryStartResult { + pub query_id: String, + pub session_id: String, + pub turn_id: String, + pub accepted: bool, + pub created_session: bool, + pub session_lifetime: SessionLifetime, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields, rename_all = "camelCase")] +pub struct QueryCancelParams { + pub query_id: String, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct QueryCancelResult { + pub query_id: String, + pub session_id: String, + pub turn_id: String, + pub requested: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields, rename_all = "camelCase")] +pub struct SessionCloseParams { + pub session_id: String, + #[serde(default)] + pub wait_timeout_ms: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct SessionCloseResult { + pub session_id: String, + pub unloaded: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct ShutdownParams {} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ShutdownResult { + pub accepted: bool, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct QueryEventParams { + pub query_id: String, + pub session_id: String, + pub turn_id: String, + pub sequence: u64, + pub event: QueryEvent, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum QueryEvent { + AssistantTextDelta { text: String }, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum QueryTerminalStatus { + Completed, + Failed, + Cancelled, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct QueryResultParams { + pub query_id: String, + pub session_id: String, + pub turn_id: String, + pub status: QueryTerminalStatus, + #[serde(skip_serializing_if = "Option::is_none")] + pub error: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct QueryResultError { + pub message: String, + pub data: ErrorData, +} + +impl QueryResultError { + pub fn new( + code: ErrorCode, + retryable: bool, + recovery: Option, + query_id: &str, + message: impl Into, + ) -> Self { + Self { + message: message.into(), + data: ErrorData { + code, + stage: ErrorStage::Query, + retryable, + correlation_id: format!("query:{query_id}"), + recovery, + }, + } + } +} + +fn empty_object() -> serde_json::Value { + serde_json::Value::Object(serde_json::Map::new()) +} diff --git a/src/crates/interfaces/sdk-host/tests/host_lifecycle.rs b/src/crates/interfaces/sdk-host/tests/host_lifecycle.rs new file mode 100644 index 0000000000..45a704efe3 --- /dev/null +++ b/src/crates/interfaces/sdk-host/tests/host_lifecycle.rs @@ -0,0 +1,2240 @@ +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use async_trait::async_trait; +use bitfun_agent_runtime::event_queue::{EventQueue, EventQueueConfig}; +use bitfun_agent_runtime::sdk::{ + AgentDialogTurnPort, AgentDialogTurnRequest, AgentEventSource, AgentRuntimeBuilder, + AgentSessionClosePort, AgentSessionCreateRequest, AgentSessionCreateResult, + AgentSessionDeleteRequest, AgentSessionListRequest, AgentSessionManagementPort, + AgentSessionSummary, AgentSessionWorkspaceBinding, AgentSessionWorkspaceRequest, + AgentSubmissionPort, AgentSubmissionRequest, AgentSubmissionResult, + AgentTransientSessionDiscardRequest, AgentTurnCancellationPort, AgentTurnCancellationRequest, + AgentTurnCancellationResult, AgentTurnSettlementPort, AgentTurnSettlementRequest, + DialogSubmitOutcome, PermissionRequest, PermissionRequestManager, PermissionRequestSource, + PermissionRequestSourceKind, PortError, PortErrorKind, PortResult, +}; +use bitfun_core_types::ErrorCategory; +use bitfun_events::AgenticEvent; +use bitfun_runtime_ports::{ + ClockPort, PermissionAuditRecord, PermissionAuditStorePort, PermissionGrant, + PermissionReplyStorePort, RuntimeServiceCapability, RuntimeServicePort, +}; +use bitfun_sdk_host::host::{ConnectionControl, HostOutput, SdkHostConfig, SdkHostConnection}; +use bitfun_sdk_host::protocol::{JsonRpcRequest, PROTOCOL_VERSION}; +use tokio::sync::{mpsc, Notify}; + +#[derive(Default)] +struct FakeOwner { + queue: Mutex>>, + created_session_ids: Mutex>, + cancel_requests: Mutex>, + discard_requests: Mutex>, + settlement_requests: Mutex>, + dialog_metadata: Mutex>>, + emit_terminal: bool, + fail_dialog_submit: bool, + fail_delete: bool, + fail_settlement: bool, + queue_dialog: bool, + dialog_session_override: Option, + block_dialog_submit: bool, + block_agent_resolution: bool, + block_first_cancel: bool, + block_delete: bool, + block_session_create: bool, + panic_after_session_create: bool, + dialog_submit_started: Notify, + release_dialog_submit: Notify, + agent_resolution_started: Notify, + release_agent_resolution: Notify, + first_cancel_started: Notify, + release_first_cancel: Notify, + delete_started: Notify, + release_delete: Notify, + session_create_started: Notify, + release_session_create: Notify, +} + +impl FakeOwner { + fn owns_session(&self, session_id: &str) -> bool { + session_id == "session-fixture" + || session_id == "transient-fixture" + || self + .created_session_ids + .lock() + .unwrap() + .iter() + .any(|created| created == session_id) + } + + fn last_created_session_id(&self) -> String { + self.created_session_ids + .lock() + .unwrap() + .last() + .expect("fixture must have created a Session") + .clone() + } + + fn with_queue(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + emit_terminal: true, + ..Self::default() + } + } + + fn without_terminal(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + emit_terminal: false, + ..Self::default() + } + } + + fn failing_dialog(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + fail_dialog_submit: true, + ..Self::default() + } + } + + fn failing_dialog_and_delete(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + fail_dialog_submit: true, + fail_delete: true, + ..Self::default() + } + } + + fn failing_settlement(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + emit_terminal: true, + fail_settlement: true, + ..Self::default() + } + } + + fn queued_dialog(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + queue_dialog: true, + ..Self::default() + } + } + + fn mismatched_dialog(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + dialog_session_override: Some("different-session".to_string()), + ..Self::default() + } + } + + fn blocking_dialog(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + block_dialog_submit: true, + ..Self::default() + } + } + + fn blocking_agent_resolution(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + emit_terminal: true, + block_agent_resolution: true, + ..Self::default() + } + } + + fn blocking_first_cancel(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + emit_terminal: false, + block_first_cancel: true, + ..Self::default() + } + } + + fn blocking_session_create(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + emit_terminal: false, + block_session_create: true, + ..Self::default() + } + } + + fn panicking_session_create(queue: Arc, fail_delete: bool) -> Self { + Self { + queue: Mutex::new(Some(queue)), + panic_after_session_create: true, + fail_delete, + ..Self::default() + } + } + + fn blocking_delete(queue: Arc) -> Self { + Self { + queue: Mutex::new(Some(queue)), + block_delete: true, + ..Self::default() + } + } +} + +#[async_trait] +impl AgentSubmissionPort for FakeOwner { + async fn create_session( + &self, + request: AgentSessionCreateRequest, + ) -> PortResult { + if self.block_session_create { + self.session_create_started.notify_one(); + self.release_session_create.notified().await; + } + let session_id = "session-fixture".to_string(); + self.created_session_ids + .lock() + .unwrap() + .push(session_id.clone()); + Ok(AgentSessionCreateResult { + session_id, + session_name: request.session_name, + agent_type: request.agent_type, + }) + } + + async fn create_session_with_id( + &self, + session_id: String, + request: AgentSessionCreateRequest, + ) -> PortResult { + if self.block_session_create { + self.session_create_started.notify_one(); + self.release_session_create.notified().await; + } + self.created_session_ids + .lock() + .unwrap() + .push(session_id.clone()); + if self.panic_after_session_create { + panic!("fixture panics after creating the transient Session"); + } + Ok(AgentSessionCreateResult { + session_id, + session_name: request.session_name, + agent_type: request.agent_type, + }) + } + + async fn create_transient_session_with_id( + &self, + session_id: String, + request: AgentSessionCreateRequest, + ) -> PortResult { + self.create_session_with_id(session_id, request).await + } + + async fn submit_message( + &self, + request: AgentSubmissionRequest, + ) -> PortResult { + Ok(AgentSubmissionResult { + turn_id: request + .turn_id + .unwrap_or_else(|| "submission-turn-fixture".to_string()), + accepted: true, + }) + } + + async fn resolve_session_agent_type(&self, session_id: &str) -> PortResult> { + if self.block_agent_resolution { + self.agent_resolution_started.notify_one(); + self.release_agent_resolution.notified().await; + } + if self.owns_session(session_id) { + Ok(Some("agentic".to_string())) + } else { + Err(PortError::new(PortErrorKind::NotFound, "session not found")) + } + } +} + +#[async_trait] +impl AgentDialogTurnPort for FakeOwner { + async fn submit_dialog_turn( + &self, + request: AgentDialogTurnRequest, + ) -> PortResult { + self.dialog_metadata + .lock() + .unwrap() + .push(request.metadata.clone()); + if self.fail_dialog_submit { + return Err(PortError::new( + PortErrorKind::Backend, + "dialog submission failed", + )); + } + if self.block_dialog_submit { + self.dialog_submit_started.notify_one(); + self.release_dialog_submit.notified().await; + } + let session_id = self + .dialog_session_override + .clone() + .unwrap_or_else(|| request.session_id.clone()); + let turn_id = request + .turn_id + .clone() + .unwrap_or_else(|| "turn-fixture".to_string()); + if self.queue_dialog { + return Ok(DialogSubmitOutcome::Queued { + session_id: request.session_id, + turn_id, + }); + } + let queue = self.queue.lock().unwrap().clone().unwrap(); + queue + .enqueue( + AgenticEvent::TextChunk { + session_id: request.session_id.clone(), + turn_id: turn_id.clone(), + round_id: "round-fixture".to_string(), + attempt_id: Some("attempt-fixture".to_string()), + attempt_index: Some(0), + text: "fixture result".to_string(), + }, + None, + ) + .await + .unwrap(); + if self.emit_terminal { + queue + .enqueue( + AgenticEvent::DialogTurnCompleted { + session_id: session_id.clone(), + turn_id: turn_id.clone(), + total_rounds: 1, + total_tools: 0, + duration_ms: 1, + partial_recovery_reason: None, + success: Some(true), + finish_reason: Some("stop".to_string()), + has_final_response: Some(true), + }, + None, + ) + .await + .unwrap(); + } + Ok(DialogSubmitOutcome::Started { + session_id, + turn_id, + }) + } +} + +#[async_trait] +impl AgentTurnSettlementPort for FakeOwner { + async fn wait_for_turn_settlement( + &self, + request: AgentTurnSettlementRequest, + ) -> PortResult<()> { + self.settlement_requests.lock().unwrap().push(request); + if self.fail_settlement { + return Err(PortError::new( + PortErrorKind::Backend, + "turn settlement is unknown", + )); + } + Ok(()) + } +} + +#[async_trait] +impl AgentSessionManagementPort for FakeOwner { + async fn list_sessions( + &self, + _request: AgentSessionListRequest, + ) -> PortResult> { + Ok(Vec::new()) + } + + async fn delete_session(&self, _request: AgentSessionDeleteRequest) -> PortResult<()> { + if self.fail_delete { + return Err(PortError::new( + PortErrorKind::CleanupRequired, + "session deletion failed", + )); + } + Ok(()) + } + + async fn resolve_session_workspace_binding( + &self, + request: AgentSessionWorkspaceRequest, + ) -> PortResult> { + if !self.owns_session(&request.session_id) { + return Ok(None); + } + Ok(Some(AgentSessionWorkspaceBinding { + workspace_id: None, + workspace_path: "D:/workspace/project".to_string(), + remote_connection_id: None, + remote_ssh_host: None, + })) + } +} + +#[derive(Default)] +struct PermissionStore { + audit: Mutex>, +} + +#[derive(Default)] +struct BlockingPermissionReplyStore { + audit: Mutex>, +} + +impl RuntimeServicePort for BlockingPermissionReplyStore { + fn capability(&self) -> RuntimeServiceCapability { + RuntimeServiceCapability::Permission + } +} + +#[async_trait] +impl PermissionAuditStorePort for BlockingPermissionReplyStore { + async fn append_permission_audit(&self, record: PermissionAuditRecord) -> PortResult<()> { + self.audit.lock().unwrap().push(record); + Ok(()) + } + + async fn list_project_permission_audit( + &self, + project_id: &str, + ) -> PortResult> { + Ok(self + .audit + .lock() + .unwrap() + .iter() + .filter(|record| record.request.project_id == project_id) + .cloned() + .collect()) + } +} + +#[async_trait] +impl PermissionReplyStorePort for BlockingPermissionReplyStore { + async fn commit_permission_reply( + &self, + _grants: Vec, + _audit: Vec, + ) -> PortResult<()> { + std::future::pending::<()>().await; + unreachable!("blocking permission reply store must be cancelled by the Host deadline") + } +} + +impl RuntimeServicePort for PermissionStore { + fn capability(&self) -> RuntimeServiceCapability { + RuntimeServiceCapability::Permission + } +} + +#[async_trait] +impl PermissionAuditStorePort for PermissionStore { + async fn append_permission_audit(&self, record: PermissionAuditRecord) -> PortResult<()> { + self.audit.lock().unwrap().push(record); + Ok(()) + } + + async fn list_project_permission_audit( + &self, + project_id: &str, + ) -> PortResult> { + Ok(self + .audit + .lock() + .unwrap() + .iter() + .filter(|record| record.request.project_id == project_id) + .cloned() + .collect()) + } +} + +#[async_trait] +impl PermissionReplyStorePort for PermissionStore { + async fn commit_permission_reply( + &self, + _grants: Vec, + audit: Vec, + ) -> PortResult<()> { + self.audit.lock().unwrap().extend(audit); + Ok(()) + } +} + +struct FixedClock; + +impl RuntimeServicePort for FixedClock { + fn capability(&self) -> RuntimeServiceCapability { + RuntimeServiceCapability::Clock + } +} + +impl ClockPort for FixedClock { + fn now_unix_millis(&self) -> i64 { + 1_720_000_000_000 + } +} + +fn permission_manager() -> Arc { + let store = Arc::new(PermissionStore::default()); + Arc::new(PermissionRequestManager::new( + store.clone(), + store, + Arc::new(FixedClock), + )) +} + +fn blocking_permission_manager() -> Arc { + let store = Arc::new(BlockingPermissionReplyStore::default()); + Arc::new(PermissionRequestManager::new( + store.clone(), + store, + Arc::new(FixedClock), + )) +} + +async fn host_with_query_limit( + max_active_queries: usize, +) -> ( + SdkHostConnection, + Arc, + mpsc::Receiver, +) { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::without_terminal(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (output, receiver) = mpsc::channel(32); + ( + SdkHostConnection::new( + runtime, + "D:/workspace/project", + output, + SdkHostConfig { + max_active_queries, + ..SdkHostConfig::default() + }, + ), + owner, + receiver, + ) +} + +#[async_trait] +impl AgentTurnCancellationPort for FakeOwner { + async fn cancel_turn( + &self, + request: AgentTurnCancellationRequest, + ) -> PortResult { + let cancel_index = { + let mut requests = self.cancel_requests.lock().unwrap(); + requests.push(request.clone()); + requests.len() + }; + if self.block_first_cancel && cancel_index == 1 { + self.first_cancel_started.notify_one(); + self.release_first_cancel.notified().await; + } + Ok(AgentTurnCancellationResult { + session_id: request.session_id, + turn_id: request.turn_id, + requested: true, + }) + } +} + +struct FailQueryStartOutput { + output: mpsc::Sender, +} + +#[async_trait] +impl HostOutput for FailQueryStartOutput { + async fn send(&self, value: serde_json::Value) -> Result<(), ()> { + if value + .get("result") + .and_then(|result| result.get("queryId")) + .is_some() + { + return Err(()); + } + self.output.send(value).await.map_err(|_| ()) + } +} + +struct BlockingSessionCreateOutput { + output: mpsc::Sender, + response_visible: Arc, + release_response: Arc, +} + +#[async_trait] +impl HostOutput for BlockingSessionCreateOutput { + async fn send(&self, value: serde_json::Value) -> Result<(), ()> { + let is_session_create = value + .get("result") + .and_then(|result| result.get("sessionId")) + .is_some() + && value + .get("result") + .and_then(|result| result.get("queryId")) + .is_none(); + self.output.send(value).await.map_err(|_| ())?; + if is_session_create { + self.response_visible.notify_one(); + self.release_response.notified().await; + } + Ok(()) + } +} + +#[async_trait] +impl AgentSessionClosePort for FakeOwner { + async fn discard_transient_session( + &self, + request: AgentTransientSessionDiscardRequest, + ) -> PortResult { + self.discard_requests.lock().unwrap().push(request); + if self.block_delete { + self.delete_started.notify_one(); + self.release_delete.notified().await; + } + if self.fail_delete { + return Err(PortError::new( + PortErrorKind::CleanupRequired, + "session discard failed", + )); + } + Ok(true) + } +} + +fn request(value: serde_json::Value) -> JsonRpcRequest { + serde_json::from_value(value).unwrap() +} + +async fn host() -> ( + SdkHostConnection, + Arc, + mpsc::Receiver, +) { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::with_queue(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (output, receiver) = mpsc::channel(32); + ( + SdkHostConnection::new( + runtime, + "D:/workspace/project", + output, + SdkHostConfig::default(), + ), + owner, + receiver, + ) +} + +async fn initialize(host: &SdkHostConnection, output: &mut mpsc::Receiver) { + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": PROTOCOL_VERSION, + "clientInfo": { "name": "fixture", "version": "0.1.0" }, + "capabilities": { "serverNotifications": true } + } + }))) + .await; + assert_eq!(output.recv().await.unwrap()["id"], 1); +} + +#[tokio::test] +async fn resource_lifecycle_notifications_do_not_create_unaddressable_sessions() { + let (host, owner, mut output) = host().await; + initialize(&host, &mut output).await; + + for _ in 0..64 { + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "method": "session/create", + "params": {} + }))) + .await; + } + + assert!(owner.created_session_ids.lock().unwrap().is_empty()); + assert!(output.try_recv().is_err()); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "create-after-notifications", + "method": "session/create", + "params": {} + }))) + .await; + let created = output.recv().await.unwrap(); + assert_eq!(created["id"], "create-after-notifications"); + assert_eq!(created["result"]["lifetime"], "connection"); + + host.shutdown_connection().await; +} + +#[tokio::test] +async fn initialize_is_required_and_version_mismatch_fails_closed() { + let (host, _, mut output) = host().await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-before-init", + "method": "query/start", + "params": { "prompt": "hello" } + }))) + .await; + let error = output.recv().await.unwrap(); + assert_eq!(error["error"]["data"]["code"], "not_initialized"); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "bad-version", + "method": "initialize", + "params": { + "protocolVersion": 99, + "clientInfo": { "name": "fixture", "version": "0.1.0" }, + "capabilities": { "serverNotifications": true } + } + }))) + .await; + let error = output.recv().await.unwrap(); + assert_eq!(error["error"]["data"]["code"], "version_mismatch"); + assert_eq!(error["error"]["data"]["recovery"], "update_sdk"); +} + +#[tokio::test] +async fn query_streams_existing_events_and_one_terminal_result() { + let (host, _, mut output) = host().await; + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-1", + "method": "query/start", + "params": { "prompt": "hello" } + }))) + .await; + + let accepted = output.recv().await.unwrap(); + assert_eq!(accepted["result"]["accepted"], true); + assert_eq!(accepted["result"]["createdSession"], true); + assert_eq!(accepted["result"]["sessionLifetime"], "connection"); + let query_id = accepted["result"]["queryId"].as_str().unwrap().to_string(); + + let event = output.recv().await.unwrap(); + assert_eq!(event["method"], "query/event"); + assert_eq!(event["params"]["queryId"], query_id); + assert_eq!(event["params"]["event"]["type"], "assistant_text_delta"); + assert_eq!(event["params"]["event"]["text"], "fixture result"); + + let result = output.recv().await.unwrap(); + assert_eq!(result["method"], "query/result"); + assert_eq!(result["params"]["queryId"], query_id); + assert_eq!(result["params"]["status"], "completed"); + assert!(output.try_recv().is_err(), "terminal result must be unique"); +} + +#[tokio::test] +async fn query_on_created_transient_session_preserves_connection_lifetime() { + let (host, _, mut output) = host().await; + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "create-1", + "method": "session/create", + "params": {} + }))) + .await; + let created = output.recv().await.unwrap(); + assert_eq!(created["result"]["lifetime"], "connection"); + let session_id = created["result"]["sessionId"].as_str().unwrap(); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-1", + "method": "query/start", + "params": { + "prompt": "hello", + "sessionId": session_id + } + }))) + .await; + + let accepted = output.recv().await.unwrap(); + assert_eq!(accepted["id"], "query-1"); + assert_eq!(accepted["result"]["createdSession"], false); + assert_eq!(accepted["result"]["sessionLifetime"], "connection"); + + host.shutdown_connection().await; +} + +#[tokio::test] +async fn dialog_session_identity_mismatch_releases_the_requested_session_reservation() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::mismatched_dialog(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "create-mismatch", + "method": "session/create", + "params": {} + }))) + .await; + let created = output.recv().await.unwrap(); + let session_id = created["result"]["sessionId"].as_str().unwrap(); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-mismatch", + "method": "query/start", + "params": { "prompt": "hello", "sessionId": session_id } + }))) + .await; + let rejected = output.recv().await.unwrap(); + assert_eq!(rejected["id"], "query-mismatch"); + assert_eq!(rejected["error"]["data"]["code"], "internal"); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "close-after-mismatch", + "method": "session/close", + "params": { "sessionId": session_id } + }))) + .await; + let closed = output.recv().await.unwrap(); + assert_eq!(closed["id"], "close-after-mismatch"); + assert_eq!(closed["result"]["unloaded"], true); +} + +#[tokio::test] +async fn cancel_close_and_shutdown_use_existing_runtime_owners() { + let (host, owner, mut output) = host().await; + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-1", + "method": "query/start", + "params": { "prompt": "hello" } + }))) + .await; + let accepted = output.recv().await.unwrap(); + let query_id = accepted["result"]["queryId"].as_str().unwrap(); + let session_id = accepted["result"]["sessionId"] + .as_str() + .unwrap() + .to_string(); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "cancel-1", + "method": "query/cancel", + "params": { "queryId": query_id } + }))) + .await; + while output.recv().await.unwrap()["id"] != "cancel-1" {} + assert_eq!(owner.cancel_requests.lock().unwrap().len(), 1); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "close-1", + "method": "session/close", + "params": { "sessionId": session_id, "waitTimeoutMs": 1234 } + }))) + .await; + while output.recv().await.unwrap()["id"] != "close-1" {} + let discard_requests = owner.discard_requests.lock().unwrap(); + assert_eq!(discard_requests.len(), 1); + assert_eq!(discard_requests[0].wait_timeout_ms, 1234); + drop(discard_requests); + + let control = host + .handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "shutdown-1", + "method": "shutdown", + "params": {} + }))) + .await; + assert_eq!(control, ConnectionControl::Shutdown); + assert_eq!(output.recv().await.unwrap()["id"], "shutdown-1"); +} + +#[tokio::test] +async fn uncertain_session_close_cleanup_requires_host_restart() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::blocking_delete(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "create-before-close-timeout", + "method": "session/create", + "params": {} + }))) + .await; + let created = output.recv().await.unwrap(); + let session_id = created["result"]["sessionId"].as_str().unwrap(); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "close-timeout", + "method": "session/close", + "params": { "sessionId": session_id, "waitTimeoutMs": 1 } + }))) + .await; + let close_error = output.recv().await.unwrap(); + assert_eq!(close_error["error"]["data"]["code"], "cleanup_required"); + assert_eq!(close_error["error"]["data"]["retryable"], false); + assert_eq!(close_error["error"]["data"]["recovery"], "restart_host"); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "rejected-after-close-timeout", + "method": "session/create", + "params": {} + }))) + .await; + assert_eq!( + output.recv().await.unwrap()["error"]["data"]["code"], + "cleanup_required" + ); +} + +#[tokio::test] +async fn active_query_capacity_fails_closed_with_typed_overload() { + let (host, _, mut output) = host_with_query_limit(1).await; + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-1", + "method": "query/start", + "params": { "prompt": "first" } + }))) + .await; + assert_eq!(output.recv().await.unwrap()["id"], "query-1"); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-2", + "method": "query/start", + "params": { "prompt": "second" } + }))) + .await; + let mut response = output.recv().await.unwrap(); + while response.get("id").is_none() { + response = output.recv().await.unwrap(); + } + assert_eq!(response["id"], "query-2"); + assert_eq!(response["error"]["data"]["code"], "overloaded"); + assert_eq!(response["error"]["data"]["stage"], "query"); + assert_eq!(response["error"]["data"]["recovery"], "retry"); +} + +#[tokio::test] +async fn cancellation_remains_available_when_data_request_capacity_is_exhausted() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::blocking_session_create(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig { + max_in_flight_requests: 1, + max_in_flight_control_requests: 1, + ..SdkHostConfig::default() + }, + ); + initialize(&host, &mut output).await; + + let create_host = host.clone(); + let create = tokio::spawn(async move { + create_host + .handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "blocked-create", + "method": "session/create", + "params": {} + }))) + .await + }); + owner.session_create_started.notified().await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "cancel-while-data-busy", + "method": "query/cancel", + "params": { "queryId": "missing-query" } + }))) + .await; + let cancellation = output.recv().await.unwrap(); + assert_eq!(cancellation["id"], "cancel-while-data-busy"); + assert_eq!(cancellation["error"]["data"]["code"], "not_found"); + + owner.release_session_create.notify_one(); + assert_eq!(create.await.unwrap(), ConnectionControl::Continue); + assert_eq!(output.recv().await.unwrap()["id"], "blocked-create"); + host.shutdown_connection().await; +} + +#[tokio::test] +async fn connection_loss_discards_owned_transient_sessions_through_core_port() { + let (host, owner, mut output) = host().await; + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "create-1", + "method": "session/create", + "params": { "sessionName": "owned by connection" } + }))) + .await; + let created = output.recv().await.unwrap(); + assert_eq!(created["id"], "create-1"); + assert_eq!(created["result"]["lifetime"], "connection"); + let session_id = created["result"]["sessionId"].as_str().unwrap(); + + host.shutdown_connection().await; + + let requests = owner.discard_requests.lock().unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].session_id, session_id); +} + +#[tokio::test] +async fn existing_durable_session_is_not_adopted_without_cross_process_fencing() { + let (host, owner, mut output) = host().await; + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-durable", + "method": "query/start", + "params": { + "prompt": "use an existing durable Session", + "sessionId": "session-fixture" + } + }))) + .await; + + let rejected = output.recv().await.unwrap(); + assert_eq!(rejected["id"], "query-durable"); + assert_eq!(rejected["error"]["data"]["code"], "capability_unavailable"); + + host.shutdown_connection().await; + + assert!(owner.discard_requests.lock().unwrap().is_empty()); +} + +#[tokio::test] +async fn visible_session_create_response_is_exposed_before_shutdown_cleanup() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::with_queue(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let response_visible = Arc::new(Notify::new()); + let release_response = Arc::new(Notify::new()); + let host = SdkHostConnection::with_output( + runtime, + "D:/workspace/project", + Arc::new(BlockingSessionCreateOutput { + output: sender, + response_visible: response_visible.clone(), + release_response: release_response.clone(), + }), + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + + let create_host = host.clone(); + let create = tokio::spawn(async move { + create_host + .handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "visible-create", + "method": "session/create", + "params": {} + }))) + .await + }); + response_visible.notified().await; + assert_eq!(output.recv().await.unwrap()["id"], "visible-create"); + + let shutdown_host = host.clone(); + let mut shutdown = tokio::spawn(async move { shutdown_host.shutdown_connection().await }); + assert!( + tokio::time::timeout(Duration::from_millis(50), &mut shutdown) + .await + .is_err(), + "shutdown must wait until response visibility is committed" + ); + release_response.notify_one(); + + assert_eq!(create.await.unwrap(), ConnectionControl::Continue); + shutdown.await.unwrap(); + assert_eq!(owner.discard_requests.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn shutdown_waits_for_in_flight_session_creation_then_cleans_it() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::blocking_session_create(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + + let create_host = host.clone(); + let create = tokio::spawn(async move { + create_host + .handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "late-create", + "method": "session/create", + "params": {} + }))) + .await + }); + owner.session_create_started.notified().await; + + let shutdown_host = host.clone(); + let mut shutdown = tokio::spawn(async move { shutdown_host.shutdown_connection().await }); + assert!( + tokio::time::timeout(Duration::from_millis(50), &mut shutdown) + .await + .is_err(), + "shutdown must not abandon the Core Session creation transaction" + ); + owner.release_session_create.notify_one(); + + assert_eq!(create.await.unwrap(), ConnectionControl::Continue); + shutdown.await.unwrap(); + let deleted = owner.discard_requests.lock().unwrap(); + assert_eq!(deleted.len(), 1); + assert_eq!(deleted[0].session_id, owner.last_created_session_id()); +} + +#[tokio::test] +async fn shutdown_compensates_a_session_creation_task_that_panics_after_creation() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::panicking_session_create(queue.clone(), false)); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "panic-create", + "method": "session/create", + "params": {} + }))) + .await; + let _creation_failure = output.recv().await.unwrap(); + + assert!( + host.shutdown_connection_bounded(Duration::from_secs(1)) + .await + ); + let discarded = owner.discard_requests.lock().unwrap(); + assert_eq!(discarded.len(), 1); + assert_eq!(discarded[0].session_id, owner.last_created_session_id()); +} + +#[tokio::test] +async fn shutdown_reports_failure_when_post_panic_session_compensation_fails() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::panicking_session_create(queue.clone(), true)); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "panic-create-failed-cleanup", + "method": "session/create", + "params": {} + }))) + .await; + let _creation_failure = output.recv().await.unwrap(); + + assert!( + !host + .shutdown_connection_bounded(Duration::from_secs(1)) + .await + ); + assert_eq!(owner.discard_requests.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn a_later_request_registers_panicked_session_cleanup_for_shutdown() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::panicking_session_create(queue.clone(), false)); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "panic-create-before-reap", + "method": "session/create", + "params": {} + }))) + .await; + let _creation_failure = output.recv().await.unwrap(); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "after-panic", + "method": "session/close", + "params": { "sessionId": "missing" } + }))) + .await; + let cleanup_required = output.recv().await.unwrap(); + assert_eq!(cleanup_required["id"], "after-panic"); + assert_eq!( + cleanup_required["error"]["data"]["code"], + "cleanup_required" + ); + assert!(owner.discard_requests.lock().unwrap().is_empty()); + + assert!( + host.shutdown_connection_bounded(Duration::from_secs(1)) + .await + ); + let discarded = owner.discard_requests.lock().unwrap(); + assert_eq!(discarded.len(), 1); + assert_eq!(discarded[0].session_id, owner.last_created_session_id()); +} + +#[tokio::test] +async fn shutdown_does_not_forget_cleanup_registered_by_a_later_request() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::panicking_session_create(queue.clone(), true)); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "panic-create-before-failed-reap", + "method": "session/create", + "params": {} + }))) + .await; + let _creation_failure = output.recv().await.unwrap(); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "after-panic-failed-reap", + "method": "session/close", + "params": { "sessionId": "missing" } + }))) + .await; + let cleanup_required = output.recv().await.unwrap(); + assert_eq!( + cleanup_required["error"]["data"]["code"], + "cleanup_required" + ); + + assert!( + !host + .shutdown_connection_bounded(Duration::from_secs(1)) + .await + ); + let discarded = owner.discard_requests.lock().unwrap(); + assert_eq!(discarded.len(), 1); + assert_eq!(discarded[0].session_id, owner.last_created_session_id()); +} + +#[tokio::test] +async fn existing_session_rejects_create_only_query_options() { + let (host, _, mut output) = host().await; + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "invalid-query-options", + "method": "query/start", + "params": { + "prompt": "hello", + "sessionId": "session-fixture", + "model": "model-for-a-new-session" + } + }))) + .await; + + let error = output.recv().await.unwrap(); + assert_eq!(error["error"]["code"], -32602); + assert_eq!(error["error"]["data"]["code"], "invalid_request"); +} + +#[tokio::test] +async fn existing_transient_session_cannot_be_adopted_as_durable() { + let (host, _, mut output) = host().await; + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "transient-adopt", + "method": "query/start", + "params": { + "prompt": "do not adopt another connection's transient Session", + "sessionId": "transient-fixture" + } + }))) + .await; + + let error = output.recv().await.unwrap(); + assert_eq!(error["id"], "transient-adopt"); + assert_eq!(error["error"]["data"]["code"], "capability_unavailable"); + assert!(error["error"]["message"] + .as_str() + .is_some_and(|message| message.contains("same SDK Host connection"))); +} + +#[tokio::test] +async fn failed_implicit_query_submission_deletes_the_unexposed_session() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::failing_dialog(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "failed-query", + "method": "query/start", + "params": { "prompt": "hello" } + }))) + .await; + + let error = output.recv().await.unwrap(); + assert_eq!(error["error"]["data"]["code"], "internal"); + let deleted = owner.discard_requests.lock().unwrap(); + assert_eq!(deleted.len(), 1); + assert_eq!(deleted[0].session_id, owner.last_created_session_id()); +} + +#[tokio::test] +async fn shutdown_takes_over_failed_query_start_cleanup_within_its_total_budget() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::blocking_first_cancel(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::with_output( + runtime, + "D:/workspace/project", + Arc::new(FailQueryStartOutput { output: sender }), + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + + let query_host = host.clone(); + let query = tokio::spawn(async move { + query_host + .handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "lost-query-start", + "method": "query/start", + "params": { "prompt": "hello" } + }))) + .await + }); + owner.first_cancel_started.notified().await; + + let shutdown = + tokio::time::timeout(Duration::from_millis(500), host.shutdown_connection()).await; + if shutdown.is_err() { + owner.release_first_cancel.notify_waiters(); + host.shutdown_connection().await; + } + assert!( + shutdown.is_ok(), + "connection shutdown must take over a failed response's slower cleanup path" + ); + assert_eq!(query.await.unwrap(), ConnectionControl::Continue); + assert!(owner.cancel_requests.lock().unwrap().len() >= 2); + assert_eq!(owner.discard_requests.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn failed_unexposed_session_cleanup_poison_connection_and_allows_shutdown() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::failing_dialog_and_delete(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "failed-query-cleanup", + "method": "query/start", + "params": { "prompt": "hello" } + }))) + .await; + let cleanup_error = output.recv().await.unwrap(); + assert_eq!(cleanup_error["error"]["data"]["code"], "cleanup_required"); + assert_eq!(cleanup_error["error"]["data"]["retryable"], false); + assert_eq!(cleanup_error["error"]["data"]["recovery"], "restart_host"); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "rejected-after-cleanup", + "method": "session/create", + "params": {} + }))) + .await; + assert_eq!( + output.recv().await.unwrap()["error"]["data"]["code"], + "cleanup_required" + ); + + assert_eq!( + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "shutdown-after-cleanup", + "method": "shutdown", + "params": {} + }))) + .await, + ConnectionControl::Shutdown + ); + assert_eq!(output.recv().await.unwrap()["result"]["accepted"], true); +} + +#[tokio::test] +async fn terminal_failure_is_typed_and_emitted_after_settlement() { + let (host, owner, mut output) = host_with_query_limit(1).await; + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-failure", + "method": "query/start", + "params": { "prompt": "hello" } + }))) + .await; + let accepted = output.recv().await.unwrap(); + let query_id = accepted["result"]["queryId"].as_str().unwrap(); + let turn_id = accepted["result"]["turnId"].as_str().unwrap().to_string(); + let session_id = accepted["result"]["sessionId"] + .as_str() + .unwrap() + .to_string(); + let queue = owner.queue.lock().unwrap().clone().unwrap(); + queue + .enqueue( + AgenticEvent::DialogTurnFailed { + session_id, + turn_id, + error: "provider unavailable".to_string(), + error_category: Some(ErrorCategory::ProviderUnavailable), + error_detail: None, + }, + None, + ) + .await + .unwrap(); + + let result = loop { + let value = output.recv().await.unwrap(); + if value["method"] == "query/result" { + break value; + } + }; + assert_eq!(result["params"]["queryId"], query_id); + assert_eq!(result["params"]["status"], "failed"); + assert_eq!( + result["params"]["error"]["data"]["code"], + "provider_unavailable" + ); + assert_eq!(result["params"]["error"]["data"]["retryable"], true); + assert_eq!(owner.settlement_requests.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn uncertain_turn_settlement_poisons_the_session_against_retry() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::failing_settlement(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "uncertain-query", + "method": "query/start", + "params": { "prompt": "hello" } + }))) + .await; + let accepted = output.recv().await.unwrap(); + let session_id = accepted["result"]["sessionId"] + .as_str() + .unwrap() + .to_string(); + let result = loop { + let value = output.recv().await.unwrap(); + if value["method"] == "query/result" { + break value; + } + }; + assert_eq!( + result["params"]["error"]["data"]["code"], + "cleanup_required" + ); + assert_eq!(result["params"]["error"]["data"]["retryable"], false); + assert_eq!( + result["params"]["error"]["data"]["recovery"], + "restart_host" + ); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "retry-uncertain-session", + "method": "query/start", + "params": { "prompt": "do not duplicate", "sessionId": session_id } + }))) + .await; + let retry = output.recv().await.unwrap(); + assert_eq!(retry["error"]["data"]["code"], "cleanup_required"); + assert_eq!(owner.dialog_metadata.lock().unwrap().len(), 1); +} + +#[tokio::test] +async fn query_submission_disables_unavailable_interactive_callbacks() { + let (host, owner, mut output) = host_with_query_limit(1).await; + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "noninteractive-query", + "method": "query/start", + "params": { "prompt": "hello" } + }))) + .await; + assert_eq!(output.recv().await.unwrap()["result"]["accepted"], true); + + let metadata = owner.dialog_metadata.lock().unwrap(); + assert_eq!(metadata.len(), 1); + assert_eq!(metadata[0]["user_input_available"], false); + assert_eq!(metadata[0]["auto_approve_ask"], false); + drop(metadata); + host.shutdown_connection().await; +} + +#[tokio::test] +async fn provider_quota_and_billing_keep_distinct_wire_codes() { + for (category, expected_code) in [ + (ErrorCategory::ProviderQuota, "provider_quota"), + (ErrorCategory::ProviderBilling, "provider_billing"), + ] { + let (host, owner, mut output) = host_with_query_limit(1).await; + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": expected_code, + "method": "query/start", + "params": { "prompt": "hello" } + }))) + .await; + let accepted = output.recv().await.unwrap(); + let turn_id = accepted["result"]["turnId"].as_str().unwrap().to_string(); + let session_id = accepted["result"]["sessionId"] + .as_str() + .unwrap() + .to_string(); + owner + .queue + .lock() + .unwrap() + .clone() + .unwrap() + .enqueue( + AgenticEvent::DialogTurnFailed { + session_id, + turn_id, + error: format!("{expected_code} fixture"), + error_category: Some(category), + error_detail: None, + }, + None, + ) + .await + .unwrap(); + + let result = loop { + let value = output.recv().await.unwrap(); + if value["method"] == "query/result" { + break value; + } + }; + assert_eq!(result["params"]["error"]["data"]["code"], expected_code); + assert_eq!(result["params"]["error"]["data"]["retryable"], false); + } +} + +#[tokio::test] +async fn queued_query_is_accepted_and_tracked_by_its_exact_turn() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::queued_dialog(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "create-for-queued-query", + "method": "session/create", + "params": {} + }))) + .await; + let created = output.recv().await.unwrap(); + assert_eq!(created["result"]["lifetime"], "connection"); + let session_id = created["result"]["sessionId"] + .as_str() + .expect("created Session id") + .to_string(); + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "queued-query", + "method": "query/start", + "params": { + "prompt": "hello", + "sessionId": session_id + } + }))) + .await; + + let accepted = output.recv().await.unwrap(); + assert_eq!(accepted["result"]["accepted"], true); + assert_eq!(accepted["result"]["turnId"], "turn-fixture"); + assert!(owner.cancel_requests.lock().unwrap().is_empty()); + + let queue = owner.queue.lock().unwrap().clone().unwrap(); + queue + .enqueue( + AgenticEvent::DialogTurnCompleted { + session_id: session_id.clone(), + turn_id: "another-surface-turn".to_string(), + total_rounds: 1, + total_tools: 0, + duration_ms: 1, + partial_recovery_reason: None, + success: Some(true), + finish_reason: Some("stop".to_string()), + has_final_response: Some(true), + }, + None, + ) + .await + .unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(50), output.recv()) + .await + .is_err() + ); + queue + .enqueue( + AgenticEvent::DialogTurnCompleted { + session_id, + turn_id: "turn-fixture".to_string(), + total_rounds: 1, + total_tools: 0, + duration_ms: 1, + partial_recovery_reason: None, + success: Some(true), + finish_reason: Some("stop".to_string()), + has_final_response: Some(true), + }, + None, + ) + .await + .unwrap(); + let result = output.recv().await.unwrap(); + assert_eq!(result["method"], "query/result"); + assert_eq!(result["params"]["turnId"], "turn-fixture"); +} + +#[tokio::test] +async fn session_close_rejects_while_query_start_is_in_flight() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::blocking_dialog(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "create-session", + "method": "session/create", + "params": {} + }))) + .await; + let created = output.recv().await.unwrap(); + assert_eq!(created["id"], "create-session"); + let session_id = created["result"]["sessionId"].as_str().unwrap().to_string(); + + let query_host = host.clone(); + let query_session_id = session_id.clone(); + let query = tokio::spawn(async move { + query_host + .handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "slow-query-start", + "method": "query/start", + "params": { + "prompt": "hello", + "sessionId": query_session_id + } + }))) + .await + }); + owner.dialog_submit_started.notified().await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "racing-close", + "method": "session/close", + "params": { "sessionId": session_id } + }))) + .await; + let close_error = output.recv().await.unwrap(); + assert_eq!(close_error["id"], "racing-close"); + assert_eq!(close_error["error"]["data"]["code"], "overloaded"); + owner.release_dialog_submit.notify_one(); + assert_eq!(query.await.unwrap(), ConnectionControl::Continue); + let started = output.recv().await.unwrap(); + assert_eq!(started["id"], "slow-query-start"); + assert_eq!(started["result"]["accepted"], true); + host.shutdown_connection().await; +} + +#[tokio::test] +async fn query_start_rejects_if_session_close_finishes_before_reservation() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::blocking_agent_resolution(queue.clone())); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permission_manager()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "create-session", + "method": "session/create", + "params": {} + }))) + .await; + let created = output.recv().await.unwrap(); + assert_eq!(created["id"], "create-session"); + let session_id = created["result"]["sessionId"].as_str().unwrap().to_string(); + + let query_host = host.clone(); + let query_session_id = session_id.clone(); + let query = tokio::spawn(async move { + query_host + .handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-after-close", + "method": "query/start", + "params": { + "prompt": "hello", + "sessionId": query_session_id + } + }))) + .await + }); + owner.agent_resolution_started.notified().await; + + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "close-before-reserve", + "method": "session/close", + "params": { "sessionId": session_id } + }))) + .await; + let closed = output.recv().await.unwrap(); + assert_eq!(closed["id"], "close-before-reserve"); + assert_eq!(closed["result"]["unloaded"], true); + + owner.release_agent_resolution.notify_one(); + assert_eq!(query.await.unwrap(), ConnectionControl::Continue); + let rejected = output.recv().await.unwrap(); + assert_eq!(rejected["id"], "query-after-close"); + assert_eq!(rejected["error"]["data"]["code"], "overloaded"); +} + +#[tokio::test] +async fn permission_without_callback_is_rejected_and_finishes_action_required() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::without_terminal(queue.clone())); + let permissions = permission_manager(); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permissions.clone()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-permission", + "method": "query/start", + "params": { "prompt": "edit a file" } + }))) + .await; + let accepted = output.recv().await.unwrap(); + assert_eq!(accepted["result"]["accepted"], true); + let session_id = accepted["result"]["sessionId"] + .as_str() + .unwrap() + .to_string(); + assert_eq!(output.recv().await.unwrap()["method"], "query/event"); + + let permission_request = PermissionRequest { + request_id: "permission-fixture".to_string(), + round_id: "round-fixture".to_string(), + order: 0, + tool_call_id: Some("tool-fixture".to_string()), + project_path: Some("D:/workspace/project".to_string()), + project_id: "project-fixture".to_string(), + session_id, + agent_id: "agentic".to_string(), + action: "edit".to_string(), + resources: vec!["src/lib.rs".to_string()], + save_resources: Vec::new(), + source: PermissionRequestSource { + kind: PermissionRequestSourceKind::ToolCall, + identity: "edit".to_string(), + }, + delegation: None, + display_metadata: serde_json::Map::new(), + }; + let unrelated = permissions + .register_batch_for_turn( + vec![PermissionRequest { + request_id: "permission-other-turn".to_string(), + ..permission_request.clone() + }], + "another-turn", + ) + .await + .unwrap() + .pop() + .unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(50), output.recv()) + .await + .is_err() + ); + permissions + .cancel_request("permission-other-turn", "test cleanup") + .await + .unwrap(); + assert!(matches!( + unrelated.wait().await, + bitfun_agent_runtime::permission::PermissionWaitOutcome::Cancelled { .. } + )); + + let pending = permissions + .register_batch_for_turn(vec![permission_request], "turn-fixture") + .await + .unwrap() + .pop() + .unwrap(); + + let result = loop { + let value = output.recv().await.unwrap(); + if value["method"] == "query/result" { + break value; + } + }; + assert_eq!(result["params"]["status"], "failed"); + assert_eq!(result["params"]["error"]["data"]["code"], "action_required"); + let resolution = pending.wait().await; + assert!(matches!( + resolution, + bitfun_agent_runtime::permission::PermissionWaitOutcome::Replied( + bitfun_agent_runtime::sdk::PermissionReply::Reject { .. } + ) + )); +} + +#[tokio::test] +async fn stalled_permission_rejection_is_bounded_and_cancels_the_exact_turn() { + let queue = Arc::new(EventQueue::new(EventQueueConfig::default())); + let owner = Arc::new(FakeOwner::without_terminal(queue.clone())); + let permissions = blocking_permission_manager(); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(owner.clone()) + .with_dialog_turn_port(owner.clone()) + .with_cancellation_port(owner.clone()) + .with_turn_settlement_port(owner.clone()) + .with_session_management_port(owner.clone()) + .with_session_close_port(owner.clone()) + .with_permission_request_manager(permissions.clone()) + .with_event_source(AgentEventSource::new(queue)) + .build() + .unwrap(); + let (sender, mut output) = mpsc::channel(16); + let host = SdkHostConnection::new( + runtime, + "D:/workspace/project", + sender, + SdkHostConfig::default(), + ); + initialize(&host, &mut output).await; + host.handle_request(request(serde_json::json!({ + "jsonrpc": "2.0", + "id": "query-stalled-permission", + "method": "query/start", + "params": { "prompt": "edit a file" } + }))) + .await; + let accepted = output.recv().await.unwrap(); + assert_eq!(accepted["result"]["accepted"], true); + let session_id = accepted["result"]["sessionId"] + .as_str() + .unwrap() + .to_string(); + + let _pending = permissions + .register_batch_for_turn( + vec![PermissionRequest { + request_id: "permission-stalled".to_string(), + round_id: "round-fixture".to_string(), + order: 0, + tool_call_id: Some("tool-fixture".to_string()), + project_path: Some("D:/workspace/project".to_string()), + project_id: "project-fixture".to_string(), + session_id, + agent_id: "agentic".to_string(), + action: "edit".to_string(), + resources: vec!["src/lib.rs".to_string()], + save_resources: Vec::new(), + source: PermissionRequestSource { + kind: PermissionRequestSourceKind::ToolCall, + identity: "edit".to_string(), + }, + delegation: None, + display_metadata: serde_json::Map::new(), + }], + "turn-fixture", + ) + .await + .unwrap() + .pop() + .unwrap(); + + let result = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let value = output.recv().await.unwrap(); + if value["method"] == "query/result" { + break value; + } + } + }) + .await + .expect("permission rejection must remain bounded"); + assert_eq!(result["params"]["error"]["data"]["code"], "timeout"); + assert_eq!( + owner.cancel_requests.lock().unwrap()[0].turn_id.as_deref(), + Some("turn-fixture") + ); +} diff --git a/src/crates/interfaces/sdk-host/tests/protocol_contracts.rs b/src/crates/interfaces/sdk-host/tests/protocol_contracts.rs new file mode 100644 index 0000000000..d06eaa5db6 --- /dev/null +++ b/src/crates/interfaces/sdk-host/tests/protocol_contracts.rs @@ -0,0 +1,180 @@ +use bitfun_sdk_host::protocol::{ + ErrorCode, ErrorData, ErrorStage, HostCapabilities, InitializeParams, InitializeResult, + JsonRpcErrorResponse, JsonRpcRequest, JsonRpcSuccessResponse, QueryEvent, QueryResultError, + QueryResultParams, QueryTerminalStatus, RecoveryAction, RequestId, SessionLifetime, Stability, + PROTOCOL_VERSION, +}; + +#[test] +fn initialize_contract_is_versioned_and_uses_familiar_capability_names() { + let request: JsonRpcRequest = serde_json::from_value(serde_json::json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "initialize", + "params": { + "protocolVersion": 1, + "clientInfo": { "name": "fixture", "version": "0.1.0" }, + "capabilities": { "serverNotifications": true } + } + })) + .unwrap(); + let params: InitializeParams = request.params_as().unwrap(); + + assert_eq!(request.id, Some(RequestId::Number(1))); + assert_eq!(params.protocol_version, PROTOCOL_VERSION); + assert!(params.capabilities.server_notifications); + + let result = InitializeResult::current("0.2.13"); + assert_eq!(result.protocol_version, PROTOCOL_VERSION); + assert_eq!(result.stability, Stability::NotDelivered); + assert_eq!( + result.capabilities, + HostCapabilities { + session_create: true, + session_create_lifetime: SessionLifetime::Connection, + query: true, + query_cancel: true, + session_close: true, + event_stream: true, + structured_output: false, + usage: false, + custom_tools: false, + permission_callbacks: false, + hooks: false, + mcp_configuration: false, + prestarted_transport: false, + } + ); +} + +#[test] +fn current_host_capabilities_are_a_deliberate_subset_of_the_headless_cli_target() { + let capabilities = HostCapabilities::current(); + + assert!(capabilities.session_create); + assert!(capabilities.query); + assert!(capabilities.query_cancel); + assert!(capabilities.session_close); + assert!(capabilities.event_stream); + + assert_eq!( + capabilities.session_create_lifetime, + SessionLifetime::Connection + ); + assert!(!capabilities.structured_output); + assert!(!capabilities.usage); + assert!(!capabilities.custom_tools); + assert!(!capabilities.permission_callbacks); + assert!(!capabilities.hooks); + assert!(!capabilities.mcp_configuration); + assert!(!capabilities.prestarted_transport); +} + +#[test] +fn query_events_and_terminal_errors_are_closed_protocol_values() { + let event = serde_json::to_value(QueryEvent::AssistantTextDelta { + text: "hello".to_string(), + }) + .unwrap(); + assert_eq!( + event, + serde_json::json!({ "type": "assistant_text_delta", "text": "hello" }) + ); + + let result = serde_json::to_value(QueryResultParams { + query_id: "query-1".to_string(), + session_id: "session-1".to_string(), + turn_id: "turn-1".to_string(), + status: QueryTerminalStatus::Failed, + error: Some(QueryResultError { + message: "Permission approval is required".to_string(), + data: ErrorData { + code: ErrorCode::ActionRequired, + stage: ErrorStage::Query, + retryable: false, + correlation_id: "query:query-1".to_string(), + recovery: None, + }, + }), + }) + .unwrap(); + assert_eq!(result["error"]["data"]["code"], "action_required"); + assert_eq!(result["error"]["data"]["stage"], "query"); + assert_eq!( + result["error"]["message"], + "Permission approval is required" + ); + assert_eq!( + serde_json::to_value(ErrorCode::ProviderQuota).unwrap(), + "provider_quota" + ); + assert_eq!( + serde_json::to_value(ErrorCode::ProviderBilling).unwrap(), + "provider_billing" + ); + assert_eq!( + serde_json::to_value(ErrorCode::CleanupRequired).unwrap(), + "cleanup_required" + ); +} + +#[test] +fn success_and_error_envelopes_are_strict_json_rpc() { + let success = JsonRpcSuccessResponse::new( + RequestId::String("request-1".to_string()), + serde_json::json!({ "accepted": true }), + ); + assert_eq!( + serde_json::to_value(success).unwrap(), + serde_json::json!({ + "jsonrpc": "2.0", + "id": "request-1", + "result": { "accepted": true } + }) + ); + + let error = JsonRpcErrorResponse::new( + RequestId::Number(2), + -32003, + "SDK Host is overloaded", + ErrorData { + code: ErrorCode::Overloaded, + stage: ErrorStage::Query, + retryable: true, + correlation_id: "request:2".to_string(), + recovery: Some(RecoveryAction::Retry), + }, + ); + let value = serde_json::to_value(error).unwrap(); + assert_eq!(value["error"]["data"]["code"], "overloaded"); + assert_eq!(value["error"]["data"]["stage"], "query"); + assert_eq!(value["error"]["data"]["recovery"], "retry"); + assert_eq!(value["error"]["data"]["retryable"], true); +} + +#[test] +fn request_ids_reject_null_fractional_and_structured_values() { + for id in [ + serde_json::Value::Null, + serde_json::json!(1.5), + serde_json::json!({ "nested": true }), + serde_json::json!([1]), + ] { + let request = serde_json::json!({ + "jsonrpc": "2.0", + "id": id, + "method": "initialize", + "params": {} + }); + assert!(serde_json::from_value::(request).is_err()); + } +} + +#[test] +fn request_correlation_ids_preserve_json_rpc_id_type() { + assert_eq!(RequestId::Number(1).correlation_id(), "request:number:1"); + assert_eq!( + RequestId::String("1".to_string()).correlation_id(), + "request:string:1" + ); +} diff --git a/src/crates/services/services-core/Cargo.toml b/src/crates/services/services-core/Cargo.toml index 62b30af328..0369d14475 100644 --- a/src/crates/services/services-core/Cargo.toml +++ b/src/crates/services/services-core/Cargo.toml @@ -21,6 +21,7 @@ serde_json = { workspace = true } serde_yaml = { workspace = true, optional = true } base64 = { workspace = true } chrono = { workspace = true } +dunce = { workspace = true, optional = true } zip = { workspace = true, optional = true } thiserror = { workspace = true } log = { workspace = true } @@ -49,7 +50,7 @@ libc = { workspace = true } default = ["lsp"] lsp = ["dep:anyhow", "dep:notify", "dep:zip"] markdown = ["dep:serde_yaml"] -workspace-runtime = ["dep:anyhow", "dep:async-trait", "dep:bitfun-runtime-ports"] +workspace-runtime = ["dep:anyhow", "dep:async-trait", "dep:bitfun-runtime-ports", "dep:dunce"] permission = ["dep:async-trait", "dep:bitfun-runtime-ports", "dep:rusqlite", "bitfun-runtime-ports/permission"] [dev-dependencies] diff --git a/src/crates/services/services-core/src/lib.rs b/src/crates/services/services-core/src/lib.rs index f652159759..757cd92dd8 100644 --- a/src/crates/services/services-core/src/lib.rs +++ b/src/crates/services/services-core/src/lib.rs @@ -7,6 +7,8 @@ pub mod diagnostics; pub mod diff; pub mod filesystem; pub mod json_store; +#[cfg(feature = "workspace-runtime")] +pub mod local_runtime_ports; #[cfg(feature = "lsp")] pub mod lsp; pub mod managed_runtime; diff --git a/src/crates/services/services-core/src/local_runtime_ports.rs b/src/crates/services/services-core/src/local_runtime_ports.rs new file mode 100644 index 0000000000..db1ac55101 --- /dev/null +++ b/src/crates/services/services-core/src/local_runtime_ports.rs @@ -0,0 +1,147 @@ +//! Reusable local implementations of the required runtime service ports. +//! +//! Product composition roots select these ports. This module only owns local +//! workspace identity, the system clock, and the in-process event sink. + +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::time::{SystemTime, UNIX_EPOCH}; + +use bitfun_runtime_ports::{ + ClockPort, FileSystemPort, PortResult, RuntimeEventEnvelope, RuntimeEventSink, + RuntimeServiceCapability, RuntimeServicePort, WorkspacePort, +}; +use tokio::sync::broadcast; + +#[derive(Debug)] +struct LocalFileSystemPort { + _workspace_root: PathBuf, +} + +impl RuntimeServicePort for LocalFileSystemPort { + fn capability(&self) -> RuntimeServiceCapability { + RuntimeServiceCapability::FileSystem + } +} + +impl FileSystemPort for LocalFileSystemPort {} + +#[derive(Debug)] +struct LocalWorkspacePort { + _workspace_root: PathBuf, +} + +impl RuntimeServicePort for LocalWorkspacePort { + fn capability(&self) -> RuntimeServiceCapability { + RuntimeServiceCapability::Workspace + } +} + +impl WorkspacePort for LocalWorkspacePort {} + +#[derive(Debug, Clone, Copy, Default)] +struct SystemClock; + +impl RuntimeServicePort for SystemClock { + fn capability(&self) -> RuntimeServiceCapability { + RuntimeServiceCapability::Clock + } +} + +impl ClockPort for SystemClock { + fn now_unix_millis(&self) -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_millis().min(i64::MAX as u128) as i64) + .unwrap_or_default() + } +} + +#[derive(Debug, Clone)] +struct LocalRuntimeEventSink { + tx: broadcast::Sender, +} + +impl LocalRuntimeEventSink { + fn new(capacity: usize) -> Self { + let (tx, _) = broadcast::channel(capacity.max(1)); + Self { tx } + } +} + +#[async_trait::async_trait] +impl RuntimeEventSink for LocalRuntimeEventSink { + async fn publish_runtime_event(&self, event: RuntimeEventEnvelope) -> PortResult<()> { + let _ = self.tx.send(event); + Ok(()) + } +} + +/// Local runtime ports bound to one canonical workspace. +#[derive(Clone)] +pub struct LocalRuntimePorts { + workspace_root: PathBuf, + filesystem: Arc, + workspace: Arc, + events: Arc, + clock: Arc, +} + +impl std::fmt::Debug for LocalRuntimePorts { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("LocalRuntimePorts") + .field("workspace_root", &self.workspace_root) + .finish_non_exhaustive() + } +} + +impl LocalRuntimePorts { + pub fn new(workspace_root: impl AsRef, event_capacity: usize) -> anyhow::Result { + let requested_root = workspace_root.as_ref(); + let canonical_root = dunce::canonicalize(requested_root).map_err(|error| { + anyhow::anyhow!( + "workspace root is not available ({}): {error}", + requested_root.display() + ) + })?; + if !canonical_root.is_dir() { + anyhow::bail!( + "workspace root is not a directory: {}", + canonical_root.display() + ); + } + + Ok(Self { + workspace_root: canonical_root.clone(), + filesystem: Arc::new(LocalFileSystemPort { + _workspace_root: canonical_root.clone(), + }), + workspace: Arc::new(LocalWorkspacePort { + _workspace_root: canonical_root, + }), + events: Arc::new(LocalRuntimeEventSink::new(event_capacity)), + clock: Arc::new(SystemClock), + }) + } + + pub fn workspace_root(&self) -> &Path { + &self.workspace_root + } + + pub fn filesystem(&self) -> Arc { + self.filesystem.clone() + } + + pub fn workspace(&self) -> Arc { + self.workspace.clone() + } + + pub fn events(&self) -> Arc { + self.events.clone() + } + + pub fn clock(&self) -> Arc { + self.clock.clone() + } +} diff --git a/src/crates/services/services-core/tests/local_runtime_ports.rs b/src/crates/services/services-core/tests/local_runtime_ports.rs new file mode 100644 index 0000000000..bd28c71430 --- /dev/null +++ b/src/crates/services/services-core/tests/local_runtime_ports.rs @@ -0,0 +1,44 @@ +#![cfg(feature = "workspace-runtime")] + +use bitfun_runtime_ports::{RuntimeEventEnvelope, RuntimeEventType, RuntimeServiceCapability}; +use bitfun_services_core::local_runtime_ports::LocalRuntimePorts; + +#[tokio::test] +async fn local_runtime_ports_bind_one_canonical_workspace_and_runtime_facts() { + let workspace = tempfile::tempdir().expect("workspace"); + let ports = LocalRuntimePorts::new(workspace.path(), 8).expect("local runtime ports"); + + assert_eq!( + ports.workspace_root(), + dunce::canonicalize(workspace.path()).unwrap() + ); + assert_eq!( + ports.filesystem().capability(), + RuntimeServiceCapability::FileSystem + ); + assert_eq!( + ports.workspace().capability(), + RuntimeServiceCapability::Workspace + ); + assert!(ports.clock().now_unix_millis() > 0); + ports + .events() + .publish_runtime_event(RuntimeEventEnvelope { + session_id: "local-runtime-test".to_string(), + turn_id: None, + source: None, + event_type: RuntimeEventType::SessionStateChanged, + payload: serde_json::json!({ "status": "ready" }), + }) + .await + .expect("publish local runtime event"); +} + +#[test] +fn local_runtime_ports_reject_a_missing_workspace() { + let temp = tempfile::tempdir().expect("tempdir"); + let error = LocalRuntimePorts::new(temp.path().join("missing"), 8) + .expect_err("missing workspace must fail"); + + assert!(error.to_string().contains("workspace"), "{error}"); +} diff --git a/src/crates/services/terminal/src/pty/process.rs b/src/crates/services/terminal/src/pty/process.rs index 876df82d6a..c2a7b3a8e5 100644 --- a/src/crates/services/terminal/src/pty/process.rs +++ b/src/crates/services/terminal/src/pty/process.rs @@ -23,8 +23,8 @@ use std::thread; #[cfg(windows)] use log::debug; use log::{error, warn}; -use portable_pty::{native_pty_system, CommandBuilder, PtySize}; -use tokio::sync::mpsc; +use portable_pty::{native_pty_system, Child, CommandBuilder, PtySize}; +use tokio::sync::{mpsc, oneshot}; use crate::config::ShellConfig; use crate::shell::ShellType; @@ -36,6 +36,10 @@ use super::flow_control::{HIGH_WATER_MARK, LOW_WATER_MARK}; mod shutdown { /// Time to wait for data flush after exit is queued pub(super) const DATA_FLUSH_TIMEOUT_MS: u64 = 250; + /// Maximum time to wait for the managed child to report termination. + pub(super) const EXIT_CONFIRM_TIMEOUT_MS: u64 = 5_000; + /// Poll interval while confirming termination after a kill request. + pub(super) const EXIT_CONFIRM_POLL_MS: u64 = 10; } /// Resize constants for Windows ConPTY @@ -74,7 +78,10 @@ enum InternalCommand { /// Send a signal to the process Signal(String), /// Shutdown the process - Shutdown { immediate: bool }, + Shutdown { + immediate: bool, + completion: oneshot::Sender>, + }, } /// Events emitted by the PTY process @@ -193,10 +200,17 @@ impl PtyController { /// Shutdown the process pub async fn shutdown(&self, immediate: bool) -> TerminalResult<()> { + let (completion, completed) = oneshot::channel(); self.command_tx - .send(InternalCommand::Shutdown { immediate }) + .send(InternalCommand::Shutdown { + immediate, + completion, + }) .await - .map_err(|_| TerminalError::ProcessNotRunning) + .map_err(|_| TerminalError::ProcessNotRunning)?; + completed + .await + .map_err(|_| TerminalError::ProcessNotRunning)? } /// Check if process is still running @@ -570,29 +584,20 @@ pub fn spawn_pty( } } } - InternalCommand::Shutdown { immediate } => { - has_exited_cmd.store(true, Ordering::Relaxed); - - if !immediate { - // Wait for data flush - tokio::time::sleep(tokio::time::Duration::from_millis( - shutdown::DATA_FLUSH_TIMEOUT_MS, - )) - .await; + InternalCommand::Shutdown { + immediate, + completion, + } => match confirm_child_shutdown(child.as_mut(), immediate).await { + Ok(exit_code) => { + has_exited_cmd.store(true, Ordering::Relaxed); + let _ = event_tx.try_send(PtyEvent::Exit { exit_code }); + let _ = completion.send(Ok(())); + break; } - - // Kill the process - let code = match child.try_wait() { - Ok(Some(status)) => Some(status.exit_code()), - _ => { - let _ = child.kill(); - child.try_wait().ok().flatten().map(|s| s.exit_code()) - } - }; - - let _ = event_tx.send(PtyEvent::Exit { exit_code: code }).await; - break; - } + Err(error) => { + let _ = completion.send(Err(error)); + } + }, } } }); @@ -655,6 +660,49 @@ fn is_tauri_host_env(key: &OsStr) -> bool { || key.starts_with("TAURI_ANDROID_PACKAGE_NAME_") } +async fn confirm_child_shutdown( + child: &mut (dyn Child + Send), + immediate: bool, +) -> TerminalResult> { + if !immediate { + tokio::time::sleep(tokio::time::Duration::from_millis( + shutdown::DATA_FLUSH_TIMEOUT_MS, + )) + .await; + } + + match child.try_wait() { + Ok(Some(status)) => return Ok(Some(status.exit_code())), + Ok(None) => {} + Err(error) => return Err(TerminalError::Io(error)), + } + + let kill_error = child.kill().err().map(|error| error.to_string()); + let deadline = tokio::time::Instant::now() + + tokio::time::Duration::from_millis(shutdown::EXIT_CONFIRM_TIMEOUT_MS); + loop { + match child.try_wait() { + Ok(Some(status)) => return Ok(Some(status.exit_code())), + Ok(None) if tokio::time::Instant::now() < deadline => { + tokio::time::sleep(tokio::time::Duration::from_millis( + shutdown::EXIT_CONFIRM_POLL_MS, + )) + .await; + } + Ok(None) => { + let detail = kill_error + .as_deref() + .map(|error| format!("; kill request failed: {error}")) + .unwrap_or_default(); + return Err(TerminalError::Timeout(format!( + "PTY process did not exit after shutdown request{detail}" + ))); + } + Err(error) => return Err(TerminalError::Io(error)), + } + } +} + // ============================================================================ // Legacy compatibility - PtyCommand enum (for external use if needed) // ============================================================================ @@ -674,7 +722,9 @@ pub enum PtyCommand { #[cfg(test)] mod tests { - use super::is_tauri_host_env; + use super::{is_tauri_host_env, spawn_pty, PtyEvent}; + use crate::config::ShellConfig; + use crate::shell::ShellType; #[test] fn strips_tauri_host_configuration_from_parent_env() { @@ -691,4 +741,34 @@ mod tests { assert!(!is_tauri_host_env("PATH".as_ref())); assert!(!is_tauri_host_env("TERMINAL_NONCE".as_ref())); } + + #[tokio::test(flavor = "current_thread")] + async fn shutdown_returns_only_after_process_exit_is_confirmed() { + let shell = if cfg!(windows) { + ("cmd.exe", ShellType::Cmd) + } else { + ("/bin/sh", ShellType::Sh) + }; + let mut spawned = spawn_pty( + 1, + &ShellConfig { + executable: shell.0.to_string(), + ..ShellConfig::default() + }, + shell.1, + 80, + 24, + ) + .expect("spawn test PTY"); + + spawned + .controller + .shutdown(true) + .await + .expect("shutdown test PTY"); + + assert!(!spawned.controller.is_running()); + assert!(std::iter::from_fn(|| spawned.events.try_recv()) + .any(|event| matches!(event, PtyEvent::Exit { .. }))); + } } diff --git a/src/crates/services/terminal/src/pty/service.rs b/src/crates/services/terminal/src/pty/service.rs index 9648e13259..37a7b1efeb 100644 --- a/src/crates/services/terminal/src/pty/service.rs +++ b/src/crates/services/terminal/src/pty/service.rs @@ -292,13 +292,25 @@ impl PtyService { /// Shutdown a PTY process pub async fn shutdown(&self, id: u32, immediate: bool) -> TerminalResult<()> { - let process = { - let mut processes = self.processes.write().await; - processes.remove(&id) + let controller = { + let processes = self.processes.read().await; + processes.get(&id).map(|process| process.controller.clone()) }; - if let Some(process) = process { - process.controller.shutdown(immediate).await?; + if let Some(controller) = controller { + if !controller.is_running() { + self.processes.write().await.remove(&id); + return Ok(()); + } + match controller.shutdown(immediate).await { + Ok(()) => { + self.processes.write().await.remove(&id); + } + Err(TerminalError::ProcessNotRunning) if !controller.is_running() => { + self.processes.write().await.remove(&id); + } + Err(error) => return Err(error), + } } Ok(()) @@ -376,3 +388,52 @@ impl Drop for PtyService { // Note: Processes should be shut down explicitly before dropping } } + +#[cfg(test)] +mod tests { + use super::PtyService; + use crate::config::{ShellConfig, TerminalConfig}; + use crate::shell::ShellType; + + #[tokio::test(flavor = "current_thread")] + async fn shutdown_evicts_a_process_whose_controller_already_confirmed_exit() { + let service = PtyService::new(TerminalConfig::default()); + let (executable, shell_type) = if cfg!(windows) { + ("cmd.exe", ShellType::Cmd) + } else { + ("/bin/sh", ShellType::Sh) + }; + let process_id = service + .create_process( + ShellConfig { + executable: executable.to_string(), + ..ShellConfig::default() + }, + shell_type, + 80, + 24, + ) + .await + .expect("create test PTY"); + let controller = service + .processes + .read() + .await + .get(&process_id) + .expect("managed process") + .controller + .clone(); + + controller + .shutdown(true) + .await + .expect("confirm child exit before service eviction"); + assert!(service.has_process(process_id).await); + + service + .shutdown(process_id, true) + .await + .expect("stale process eviction must be idempotent"); + assert!(!service.has_process(process_id).await); + } +} diff --git a/src/crates/services/terminal/src/session/binding.rs b/src/crates/services/terminal/src/session/binding.rs index 26222b544d..a4fad97365 100644 --- a/src/crates/services/terminal/src/session/binding.rs +++ b/src/crates/services/terminal/src/session/binding.rs @@ -13,6 +13,7 @@ use std::collections::HashMap; use std::sync::Arc; +use dashmap::mapref::entry::Entry; use dashmap::DashMap; use log::warn; @@ -222,38 +223,55 @@ impl TerminalSessionBinding { pub async fn remove(&self, owner_id: &str) -> TerminalResult<()> { let session_manager = get_session_manager() .ok_or_else(|| TerminalError::Session("SessionManager not initialized".to_string()))?; + let mut first_error = None; // Close primary session - if let Some(terminal_session_id) = self.unbind(owner_id) { - if let Err(e) = session_manager + if let Some(terminal_session_id) = self.get(owner_id) { + if let Err(error) = session_manager .close_session(&terminal_session_id, false) .await { warn!( "Failed to close terminal session {}: {}", - terminal_session_id, e + terminal_session_id, error ); + first_error = Some(error); + } else if let Entry::Occupied(entry) = self.bindings.entry(owner_id.to_string()) { + if entry.get() == &terminal_session_id { + entry.remove(); + } } } // Close all background sessions - if let Some((_, bg_sessions)) = self.background_bindings.remove(owner_id) { - for bg_session_id in bg_sessions { - if let Err(e) = session_manager.close_session(&bg_session_id, false).await { - warn!( - "Failed to close background terminal session {}: {}", - bg_session_id, e - ); + for bg_session_id in self.list_background_sessions(owner_id) { + if let Err(error) = session_manager.close_session(&bg_session_id, false).await { + warn!( + "Failed to close background terminal session {}: {}", + bg_session_id, error + ); + if first_error.is_none() { + first_error = Some(error); + } + } else if let Entry::Occupied(mut entry) = + self.background_bindings.entry(owner_id.to_string()) + { + entry.get_mut().retain(|current| current != &bg_session_id); + if entry.get().is_empty() { + entry.remove(); } } } - Ok(()) + match first_error { + Some(error) => Err(error), + None => Ok(()), + } } /// Check if a binding exists for the given owner pub fn has(&self, owner_id: &str) -> bool { - self.bindings.contains_key(owner_id) + self.bindings.contains_key(owner_id) || self.background_bindings.contains_key(owner_id) } /// List all current bindings @@ -338,4 +356,14 @@ mod tests { binding.clear(); assert_eq!(binding.count(), 0); } + + #[test] + fn background_only_binding_is_owned_by_the_session() { + let binding = TerminalSessionBinding::new(); + binding + .background_bindings + .insert("owner1".to_string(), vec!["background-session".to_string()]); + + assert!(binding.has("owner1")); + } } diff --git a/src/crates/services/terminal/src/session/manager.rs b/src/crates/services/terminal/src/session/manager.rs index 89d1933b0a..4ad7eb785a 100644 --- a/src/crates/services/terminal/src/session/manager.rs +++ b/src/crates/services/terminal/src/session/manager.rs @@ -1488,13 +1488,11 @@ impl SessionManager { // Shutdown PTY if exists if let Some(pty_id) = pty_id { - // Remove mapping - { - let mut mapping = self.pty_to_session.write().await; - mapping.remove(&pty_id); - } - self.pty_service.shutdown(pty_id, immediate).await?; + + // Keep the mapping available for retry until shutdown is confirmed. + let mut mapping = self.pty_to_session.write().await; + mapping.remove(&pty_id); } if is_manual_session {