diff --git a/docs/architecture/agent-runtime-services-design.md b/docs/architecture/agent-runtime-services-design.md index 97ebf6e119..3901efcb96 100644 --- a/docs/architecture/agent-runtime-services-design.md +++ b/docs/architecture/agent-runtime-services-design.md @@ -693,9 +693,9 @@ pub struct HarnessExecutionContext { 通过 `assemble(ProductAssemblyInput)` 产生 `ProductRuntimeParts`，本文件不再为它定义第二套目标接口。 当前 CLI 入口已使用类型化 `RuntimeServices` 构造 `ProductRuntimeParts`，并通过一个调用级上下文把 -Agent Runtime SDK、Harness、能力注册、调用级权限和 Agentic 事件广播交给 TUI、Exec、Session 与 Usage。 -SDK 已承接会话创建/列举/删除、轮次提交和取消；SDK v1 尚未覆盖的恢复视图、消息、分支、用量和 -工具确认由 `assembly/core` 的单一兼容门面转发。`doctor` 与 `health` 校验真实组装结果及必需注册完整性； +Agent Runtime SDK、Harness、能力注册、调用级权限和 Agentic 事件广播交给 TUI、Exec、Session、Usage 与 +交互模式下的 Peer Host。SDK 已承接会话创建/列举/删除、轮次提交和取消；SDK v1 尚未覆盖的恢复视图、 +消息、分支、用量、快照和工具确认由 `assembly/core` 的单一兼容门面转发。`doctor` 与 `health` 校验真实组装结果及必需注册完整性； Core 的 Network、Git 和 MCP Catalog 当前仍含兼容 marker，因此该诊断不等于对这些外部服务做实时探活。 该切换仍是 `product-full` 兼容组装，不是 owner 迁移。协调器、调度器、持久化、工具管线和 Agentic Event Queue @@ -713,7 +713,7 @@ ping 路由。未接入入口的 profile、枚举分支和单元测试仍不能 | 阶段 | 约束 | |---|---| -| 当前 | CLI 消费真实 Runtime Parts 与 SDK，Core 兼容门面只承接 SDK v1 缺口；不扩张字段或再造描述符 | +| 当前 | CLI 与 Peer Host 消费真实 Runtime Parts 与 SDK，Core 兼容门面只承接 SDK v1 缺口；不扩张字段或再造描述符 | | 迁移 | 迁移执行 owner 或继续接入 ACP、Desktop 前，必须分别证明行为等价；`assembly/core -> apps/relay-server` 反向依赖仍需消除 | | 完成 | 每个声称支持的 profile 都由生产入口消费组装结果，并有最小入口验证；无消费方的 profile 不对外宣称可用 | @@ -998,9 +998,10 @@ Product 测试： - `bitfun-runtime-services` 提供类型化服务注入；工具 contracts、provider groups 与 execution 已分层。 - `bitfun-harness` 已提供类型化工作流描述与注册能力。 - `bitfun-core` 可继续作为 `product-full` 兼容门面，避免迁移期间一次性重写入口。 -- CLI 已以 `DeliveryProfile::Cli` 构造真实 Runtime Parts 和 SDK runtime；本地 Agent 入口、会话和用量 - 共用一个调用级上下文与广播事件源，审批策略不再写回全局配置。Peer Host 仍保留既有 Core 兼容路径， - 其协议与生命周期切换由独立变更处理。 +- CLI 已以 `DeliveryProfile::Cli` 构造真实 Runtime Parts 和 SDK runtime；本地 Agent 入口、会话、用量和 + Peer Host 共用一个调用级上下文与广播事件源，审批策略不再写回全局配置。Peer Host 通过 SDK 提交/取消 + turn，通过单一 Core 兼容门面处理会话与快照缺口，不再构造独立调度器、持久化 manager 或事件队列； + wire schema、Relay ACK/重放和重连协议未在该切换中扩张。 - CLI 通过 SDK 处理已覆盖的 session/turn/cancel 操作，并通过一个 Core 兼容门面处理 SDK v1 缺口； 该门面复用现有 owner，不建立第二套状态或事件 schema。 diff --git a/docs/architecture/cli-product-line-design.md b/docs/architecture/cli-product-line-design.md index 510e052a06..866f321a73 100644 --- a/docs/architecture/cli-product-line-design.md +++ b/docs/architecture/cli-product-line-design.md @@ -88,9 +88,10 @@ BitFun CLI 应成为可独立安装和发布的 Agent 产品，而不是 Desktop - BitFun 原生插件目录的发现、内容校验、来源确认，以及 OpenCode custom tool 静态名称预览。 - CLI 本地 Agent 入口以类型化 `RuntimeServices` 调用 `ProductAssembler`，选择 `DeliveryProfile::Cli`， 并把 `ProductRuntimeParts`、Agent Runtime SDK、事件源和调用级审批策略保存在一个 `CliRuntimeContext` 中。 -- TUI、`exec`、会话和用量复用同一上下文。SDK 已承接会话创建/列举/删除、轮次提交和取消； - SDK v1 尚未覆盖的固定 ID、恢复视图、消息、分支、用量和工具确认由一个 Core 兼容门面转发给原 owner。 -- Agentic Event Queue 仍是唯一事件 owner；TUI 与 `exec` 使用独立广播订阅，不互相消费事件。 +- TUI、`exec`、会话、用量和交互模式下的 Peer Host 复用同一上下文。SDK 已承接会话创建/列举/删除、 + 轮次提交和取消；SDK v1 尚未覆盖的固定 ID、恢复视图、消息、分支、用量、快照和工具确认由一个 Core + 兼容门面转发给原 owner。 +- Agentic Event Queue 仍是唯一事件 owner；TUI、`exec` 与 Peer Host 使用独立广播订阅，不互相消费事件。 - 有界旧队列只承担兼容存储；达到容量时不得抑制广播。CLI 保持一个后台 drain，订阅方一旦报告 lag/closed， 必须取消活动 turn 并显式失败，不能在状态不完整时继续报告成功。 - 会话 ID 在进入存储路径前统一校验；运行时索引同时绑定 ID 与规范化存储路径，并以待提交 claim 计数保护 @@ -101,14 +102,24 @@ BitFun CLI 应成为可独立安装和发布的 Agent 产品，而不是 Desktop 关闭输入捕获、关闭 raw mode 并显示光标。真实 PTY/ConPTY 故障注入仍需独立验收。 - 初始化按入口分级：交互模式启动 Peer Host 与 MCP，`exec` 只启动 MCP；本地 session 管理和 usage 查询不启动 Peer Host/MCP。该分级不改变 Agentic/Terminal owner，也不等同于管理命令已有独立轻量 Runtime。 -- Peer Host 的 HostInvoke、Relay、控制器身份、确认和重连协议仍走既有兼容路径，不属于本次本地 Runtime 切换。 +- Peer Host 保持既有 HostInvoke / DeviceEvent wire schema 与 Relay 路由，但执行已接入上述调用级上下文： + 对话提交和精确取消走 SDK，会话与快照缺口走单一 Core 兼容门面。Peer Host 只跟踪由 Peer 提交的根 turn、 + 其子 turn 与待确认工具；可确认工具始终由控制器确认，即使宿主全局策略跳过确认，Agent 也会暂停等待控制器。 + 该 Peer 专属确认要求会沿精确后台结果 follow-up 保留。后台结果按 Core 内部元数据中的精确父 turn 与来源子 turn + 继承 ownership；仅在父 turn 仍运行时注入，否则排在 + 无关 turn 之后并保留 Peer ownership。来源 turn 完成而结果仍等待会话串行化时，仅保留有界、一次性内部 tombstone； + 会话清理或事件流中断会移除它。最后一个控制器离线 + 或分离、事件订阅 lag/closed 时取消这些 turn；事件失步同时投递既有失败 terminal event，终态在实际发送尝试前 + 不提前清理，队列关闭时改走同一直接投递路径。事件以入队时的控制器快照为上界，每个目标发送前再确认仍连接； + 单目标投递租约将分离或离线移除与本地 Relay 入队尝试串行化。显式断开无法确认取消时，本地界面仍安全退出并 + 显示警告。不承诺本次变更范围外的 ACK、重放或重连恢复。 - `doctor` 与 `health` 构造并校验真实 Runtime Parts，区分 assembly-ready、Core compatibility owner 和不可用扩展。 它们证明必需能力已注册，不把 Core 的 Network/Git/MCP compatibility marker 描述为外部服务实时可用。 - 独立 CLI 测试与打包工作流；主 CI 的三平台 workspace check 同时覆盖 `bitfun-cli` 编译。 上述切换不等于运行时 owner 已迁移，也不表示 CLI-P0 全部完成。CLI crate 仍以 `bitfun-core/product-full` -承载协调器、调度器、持久化、工具管线和部分 SDK v1 缺口；ACP stdio 仍走原入口，插件命令仍以来源管理和 -静态预览为主。兼容门面只转发，不重新计算或写入同一事实。 +承载协调器、调度器、持久化、工具管线和部分 SDK v1 缺口，但 Peer Host 不再自行构造这些 owner；ACP stdio +仍走原入口，插件命令仍以来源管理和静态预览为主。兼容门面只转发，不重新计算或写入同一事实。 目标态仍存在以下结构缺口： diff --git a/docs/architecture/peer-device-mode.md b/docs/architecture/peer-device-mode.md index 79e8f8a0e8..e5632c4e18 100644 --- a/docs/architecture/peer-device-mode.md +++ b/docs/architecture/peer-device-mode.md @@ -55,12 +55,37 @@ FS) and must not be mixed with Peer Device Mode. - Peer: decrypt → allow/deny → execute on the peer host: - Desktop: webview bridge `peer-host-invoke://request` → same Tauri handlers as local UI → `peer_host_invoke_complete` - - CLI: Core HostInvoke registry (`WorkspaceService`, FS, config, session, - git, `DialogScheduler`) — no webview. Desktop-only surfaces (MiniApp / - cron / ACP list) return empty or no-op so hydrate does not fail. + - CLI: the invocation-scoped CLI product runtime handles dialog submit/cancel + through the Agent Runtime SDK and session/snapshot gaps through one Core + compatibility facade — no webview and no second scheduler, persistence + manager, or event queue. Desktop-only surfaces (MiniApp / cron / ACP list) + return empty or no-op so hydrate does not fail. - Events: peer agentic projection (and other product events such as terminal / FS / MCP interaction) fan-out as `RemoteCommand::DeviceEvent` to attached controllers; controller re-emits the same event names locally. +- CLI Peer Host forwards only turns submitted through Peer Host and linked + child turns. A background-result follow-up inherits ownership only when its + Core-internal metadata identifies the exact tracked parent and source child + turns; if an unrelated turn is running in the same session, the result queues + behind it without losing Peer ownership. Completed source lineage uses a + bounded, one-shot tombstone while delivery waits on session serialization; + session drain or event-stream interruption clears it. Peer Host + requires an attached controller before submit and binds tool confirmation to + the exact observed tool and turn. Confirmable Peer tools always wait for the + controller even when the host's global policy skips confirmation, so an Agent + pauses until the controller responds; exact background-result follow-ups + retain this Peer-only confirmation requirement. The host cancels tracked turns when the + last controller detaches/goes offline or the agent-event subscription + lags/closes; continuity loss also projects the existing dialog-turn-failed + terminal event. Terminal ownership remains tracked until the event reaches the + delivery attempt, and a closed local delivery queue uses the same direct + DeviceEvent path. Delivery targets are captured when an event is queued and + rechecked against the currently attached set before each send. A per-target + delivery lease serializes detach or offline removal with the local Relay + enqueue attempt. An explicit disconnect still restores the local controller + UI, but reports a warning when host cancellation was not confirmed. This + boundary does not change the Relay envelope or add ACK, replay, or reconnect + recovery. - Relay `POST /api/devices/:id/rpc` waits up to **120s** for the peer response; reverse proxies in front of the relay must use a matching (or higher) read timeout or they will return 504 first. diff --git a/docs/architecture/product-architecture.md b/docs/architecture/product-architecture.md index 56a729b68f..fcc86ea745 100644 --- a/docs/architecture/product-architecture.md +++ b/docs/architecture/product-architecture.md @@ -236,7 +236,7 @@ flowchart LR | 产品形态 | 当前 P0 插件能力 | 入口行为 | |---|---|---| | Desktop / product-full | 生产入口仍直接依赖 `bitfun-core/product-full`；当前没有 managed-plugin 管理或 OpenCode 静态预览的生产 UI/调用方 | 共享代码可编译不等于 Desktop 已消费插件能力 | -| CLI | 入口仍以 `bitfun-core/product-full` 作为执行兼容 owner；只为 BitFun 原生包提供来源审核、启用预览、精确内容确认和停用 | 本地 Agent 路径选择 `DeliveryProfile::Cli`，校验必需 Runtime Service 注册并消费 Runtime Parts/SDK；SDK 缺口由单一 Core 兼容门面转发。部分注册仍是 compatibility marker，不代表实时探活；插件 binding 明确禁用，不执行 OpenCode 插件代码 | +| CLI | 入口仍以 `bitfun-core/product-full` 作为执行兼容 owner；只为 BitFun 原生包提供来源审核、启用预览、精确内容确认和停用 | 本地 Agent 与 Peer Host 路径选择 `DeliveryProfile::Cli`，校验必需 Runtime Service 注册并消费同一 Runtime Parts/SDK；SDK 缺口由单一 Core 兼容门面转发。Peer Host 不再构造第二套调度、持久化或事件 owner。部分注册仍是 compatibility marker，不代表实时探活；插件 binding 明确禁用，不执行 OpenCode 插件代码 | | ACP | 生产入口仍直接依赖 `bitfun-core/product-full` | `DeliveryProfile::Acp` 尚未进入入口组装；不得把测试中的 profile 解释为生产隔离 | | Server / Remote | 当前生产路由没有插件状态消费闭环；Remote 插件执行未实现 | 不在本地替远端项目发现、准备或执行插件；未接入时返回明确不支持 | | Web / Mobile Web | 依赖现有后端入口，不持有插件执行单元 | 对应 profile 当前为空计划或未接入生产，不能据枚举值宣称独立产品能力 | diff --git a/src/apps/cli/README.md b/src/apps/cli/README.md index 97080e7c54..d491d0b294 100644 --- a/src/apps/cli/README.md +++ b/src/apps/cli/README.md @@ -6,7 +6,9 @@ The local Agent paths build the CLI product profile once per invocation. Interac session commands, and usage reports use that invocation-scoped runtime context and event source. Local management queries do not start Peer Host or MCP; `exec` starts MCP but not Peer Host. Core remains the compatibility owner for execution and persistence operations not yet covered by the -Agent Runtime SDK; plugin execution is not enabled by this assembly path. +Agent Runtime SDK. When interactive mode enables Peer Host, Peer dialog submission, cancellation, +and agent-event fan-out reuse the same runtime context; Peer Host does not construct another +scheduler, persistence manager, or event queue. Plugin execution is not enabled by this assembly path. ## Common commands diff --git a/src/apps/cli/src/account.rs b/src/apps/cli/src/account.rs index a2620066a7..96eba67539 100644 --- a/src/apps/cli/src/account.rs +++ b/src/apps/cli/src/account.rs @@ -241,6 +241,9 @@ async fn spawn_device_routing(relay_url: &str, device_name: &str) -> Result<()> while let Some(event) = event_rx.recv().await { handle_relay_event(event, &session_arc, &relay_client_arc).await; } + if is_current_routing_client(&relay_client_arc).await { + crate::peer_host::update_controller_presence(Vec::new()).await; + } tracing::info!("Device routing event loop exited"); }); @@ -249,9 +252,19 @@ async fn spawn_device_routing(relay_url: &str, device_name: &str) -> Result<()> /// Disconnect the device-routing connection (if any). pub(crate) async fn stop_device_routing() { - if let Some(client) = device_relay_client().write().await.take() { + let client = { device_relay_client().write().await.take() }; + if let Some(client) = client { client.disconnect().await; } + crate::peer_host::update_controller_presence(Vec::new()).await; +} + +async fn is_current_routing_client(client: &Arc) -> bool { + device_relay_client() + .read() + .await + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, client)) } /// Log out: tear down routing, revoke the token (best-effort), clear state. @@ -277,6 +290,10 @@ async fn handle_relay_event( session_arc: &Arc>>, relay_client: &Arc, ) { + if !is_current_routing_client(relay_client).await { + tracing::debug!("Ignoring event from a stale device routing client"); + return; + } match event { RelayEvent::AuthOk { user_id, device_id } => { tracing::info!("Device routing auth ok: user={user_id} device={device_id}"); @@ -284,9 +301,14 @@ async fn handle_relay_event( RelayEvent::AuthError { message } => { tracing::warn!("Device routing auth error: {message}"); TOKEN_EXPIRED.store(true, Ordering::Relaxed); + crate::peer_host::update_controller_presence(Vec::new()).await; } RelayEvent::DevicePresence { devices } => { tracing::info!("Device presence updated: {} online", devices.len()); + crate::peer_host::update_controller_presence( + devices.into_iter().map(|device| device.device_id).collect(), + ) + .await; } RelayEvent::DeviceMessageReceived { source_device_id, @@ -365,6 +387,7 @@ async fn handle_relay_event( } RelayEvent::Disconnected => { tracing::info!("Device routing disconnected"); + crate::peer_host::update_controller_presence(Vec::new()).await; } RelayEvent::Reconnected => { tracing::info!("Device routing reconnected"); diff --git a/src/apps/cli/src/main.rs b/src/apps/cli/src/main.rs index 30a25adf8e..a0d42c2c26 100644 --- a/src/apps/cli/src/main.rs +++ b/src/apps/cli/src/main.rs @@ -498,7 +498,7 @@ async fn initialize_core_services( ); if bootstrap_profile.starts_peer_host() { - if let Err(e) = peer_host::ensure_peer_host_ready(runtime.agentic_system()).await { + if let Err(e) = peer_host::ensure_peer_host_ready(runtime.as_ref()).await { tracing::warn!("Failed to initialize CLI peer host services: {e}"); } else { tracing::info!("CLI peer host services initialized"); diff --git a/src/apps/cli/src/peer_host/bootstrap.rs b/src/apps/cli/src/peer_host/bootstrap.rs index b7b39052a2..9081034384 100644 --- a/src/apps/cli/src/peer_host/bootstrap.rs +++ b/src/apps/cli/src/peer_host/bootstrap.rs @@ -3,42 +3,20 @@ use std::sync::Arc; use anyhow::{Context, Result}; -use bitfun_core::agentic::coordination::{self, DialogScheduler}; -use bitfun_core::agentic::system::AgenticSystem; -use bitfun_core::infrastructure::try_get_path_manager_arc; use bitfun_core::service::filesystem::FileSystemServiceFactory; use bitfun_core::service::workspace::{self, WorkspaceService}; +use crate::runtime::CliRuntimeContext; + use super::fanout::start_peer_event_fanout; -use super::state::{set_peer_host_state, try_peer_host_state, PeerHostState}; +use super::state::{set_peer_host_state, try_peer_host_state, PeerHostState, PeerTurnTracker}; /// Ensure Peer Host services are ready. Idempotent. -pub(crate) async fn ensure_peer_host_ready(agentic: &AgenticSystem) -> Result<()> { +pub(crate) async fn ensure_peer_host_ready(runtime: &CliRuntimeContext) -> Result<()> { if try_peer_host_state().is_some() { return Ok(()); } - let path_manager = try_get_path_manager_arc().context("path manager")?; - let persistence = Arc::new( - bitfun_core::agentic::persistence::PersistenceManager::new(path_manager) - .context("persistence manager")?, - ); - - let scheduler = if let Some(existing) = coordination::get_global_scheduler() { - existing - } else { - let session_manager = agentic.coordinator.get_session_manager().clone(); - let scheduler = DialogScheduler::new(agentic.coordinator.clone(), session_manager); - agentic - .coordinator - .set_scheduler_notifier(scheduler.outcome_sender()); - agentic - .coordinator - .set_round_injection_source(scheduler.round_injection_monitor()); - coordination::set_global_scheduler(scheduler.clone()); - scheduler - }; - let workspace_service = if let Some(existing) = workspace::get_global_workspace_service() { existing } else { @@ -54,12 +32,12 @@ pub(crate) async fn ensure_peer_host_ready(agentic: &AgenticSystem) -> Result<() let filesystem_service = Arc::new(FileSystemServiceFactory::create_default()); let state = PeerHostState { - coordinator: agentic.coordinator.clone(), - scheduler, - event_queue: agentic.event_queue.clone(), + agent_runtime: runtime.agent_runtime().clone(), + compatibility: runtime.compatibility().clone(), + agent_events: runtime.agent_events().clone(), + turns: PeerTurnTracker::new(), workspace_service, filesystem_service, - persistence, }; if set_peer_host_state(state.clone()).is_err() { @@ -67,7 +45,7 @@ pub(crate) async fn ensure_peer_host_ready(agentic: &AgenticSystem) -> Result<() return Ok(()); } - start_peer_event_fanout(state.event_queue.clone()); + start_peer_event_fanout(state); tracing::info!("CLI peer host services ready"); Ok(()) } diff --git a/src/apps/cli/src/peer_host/commands/dialog.rs b/src/apps/cli/src/peer_host/commands/dialog.rs index df667d9947..b898a71e55 100644 --- a/src/apps/cli/src/peer_host/commands/dialog.rs +++ b/src/apps/cli/src/peer_host/commands/dialog.rs @@ -2,10 +2,34 @@ use serde_json::{json, Value}; -use bitfun_core::agentic::coordination::{DialogSubmissionPolicy, DialogTriggerSource}; +use bitfun_runtime_ports::{ + AgentDialogTurnRequest, AgentSubmissionSource, AgentTurnCancellationRequest, + DialogSubmissionPolicy, DialogTriggerSource, +}; use crate::peer_host::args::{get_string, optional_string, request_value}; -use crate::peer_host::state::PeerHostState; +use crate::peer_host::control::{attached_controller_lease, is_controller_lease_current}; +use crate::peer_host::state::{PeerHostState, PeerTurnKey}; + +fn peer_dialog_metadata(request: &Value) -> Result, String> { + let mut metadata = match request.get("userMessageMetadata") { + Some(Value::Object(metadata)) => metadata.clone(), + Some(Value::Null) | None => serde_json::Map::new(), + Some(_) => return Err("userMessageMetadata must be an object".to_string()), + }; + for reserved_key in [ + "acp_transport", + "backgroundTaskId", + "parentSessionId", + "parentDialogTurnId", + "subagentSessionId", + "subagentDialogTurnId", + ] { + metadata.remove(reserved_key); + } + metadata.insert("require_tool_confirmation".to_string(), Value::Bool(true)); + Ok(metadata) +} pub(crate) async fn start_dialog_turn( state: &PeerHostState, @@ -19,28 +43,73 @@ pub(crate) async fn start_dialog_turn( let workspace_path = optional_string(request, "workspacePath"); let remote_connection_id = optional_string(request, "remoteConnectionId"); let remote_ssh_host = optional_string(request, "remoteSshHost"); - let turn_id = optional_string(request, "turnId"); - let user_message_metadata = request.get("userMessageMetadata").cloned(); + let controller_lease = attached_controller_lease()?; + let turn_id = + optional_string(request, "turnId").unwrap_or_else(|| uuid::Uuid::new_v4().to_string()); + let metadata = peer_dialog_metadata(request)?; + let turn = PeerTurnKey::new(session_id.clone(), turn_id.clone()); + let stream_generation = state.turns.register_root(turn.clone())?; + if !state + .turns + .is_event_stream_generation_current(stream_generation) + || !is_controller_lease_current(controller_lease) + { + state.turns.finish_turn(&turn); + return Err( + "Peer controller or event stream continuity was lost before dialog submission" + .to_string(), + ); + } let policy = DialogSubmissionPolicy::for_source(DialogTriggerSource::DesktopUi); - state - .scheduler - .submit( - session_id, - user_input, - original_user_input, - turn_id, + let submit_result = state + .agent_runtime + .submit_dialog_turn(AgentDialogTurnRequest { + session_id: session_id.clone(), + message: user_input, + original_message: original_user_input, + turn_id: Some(turn_id.clone()), agent_type, workspace_path, remote_connection_id, remote_ssh_host, policy, - None, - user_message_metadata, - None, - ) - .await - .map_err(|e| format!("Failed to start dialog turn: {e}"))?; + reply_route: None, + prepended_reminders: Vec::new(), + attachments: Vec::new(), + metadata, + }) + .await; + if let Err(error) = submit_result { + state.turns.finish_turn(&turn); + return Err(format!("Failed to start dialog turn: {error}")); + } + if !state + .turns + .is_event_stream_generation_current(stream_generation) + || !is_controller_lease_current(controller_lease) + { + let cancellation = state + .agent_runtime + .cancel_turn(AgentTurnCancellationRequest { + session_id: session_id.clone(), + turn_id: Some(turn_id.clone()), + source: Some(AgentSubmissionSource::Cli), + requester_session_id: None, + reason: Some("Peer controller or event stream lost continuity".to_string()), + wait_timeout_ms: Some(1_500), + }) + .await; + if let Err(error) = cancellation { + return Err(format!( + "Peer continuity was lost after dialog submission and cancellation could not be confirmed: session_id={session_id}, turn_id={turn_id}, error={error}" + )); + } + return Err( + "Peer controller or event stream lost continuity while starting the dialog turn" + .to_string(), + ); + } Ok(json!({ "success": true, "message": "Dialog turn started" })) } @@ -52,9 +121,19 @@ pub(crate) async fn cancel_dialog_turn( let request = request_value(args); let session_id = get_string(request, "sessionId")?; let dialog_turn_id = get_string(request, "dialogTurnId")?; + if !state.turns.owns(&session_id, Some(&dialog_turn_id)) { + return Err("The dialog turn is not owned by the Peer controller".to_string()); + } state - .coordinator - .cancel_dialog_turn(&session_id, &dialog_turn_id) + .agent_runtime + .cancel_turn(AgentTurnCancellationRequest { + session_id, + turn_id: Some(dialog_turn_id), + source: Some(AgentSubmissionSource::Cli), + requester_session_id: None, + reason: Some("Peer controller requested cancellation".to_string()), + wait_timeout_ms: Some(1_500), + }) .await .map_err(|e| format!("Failed to cancel dialog turn: {e}"))?; Ok(json!({ "success": true })) @@ -66,12 +145,27 @@ pub(crate) async fn confirm_tool_execution( ) -> Result { let request = request_value(args); let tool_id = get_string(request, "toolId")?; + let ownership = state + .turns + .claim_confirmation(&tool_id) + .ok_or_else(|| "Tool confirmation is not owned by an active Peer turn".to_string())?; + if optional_string(request, "sessionId") + .is_some_and(|value| value.as_str() != ownership.session_id.as_str()) + || optional_string(request, "dialogTurnId") + .is_some_and(|value| value.as_str() != ownership.turn_id.as_str()) + { + state.turns.restore_confirmation(tool_id, ownership); + return Err("Tool confirmation session or turn does not match its Peer owner".to_string()); + } let updated_input = request.get("updatedInput").cloned(); - state - .coordinator + if let Err(error) = state + .compatibility .confirm_tool(&tool_id, updated_input) .await - .map_err(|e| format!("Confirm tool failed: {e}"))?; + { + state.turns.restore_confirmation(tool_id, ownership); + return Err(format!("Confirm tool failed: {error}")); + } Ok(Value::Null) } @@ -81,11 +175,80 @@ pub(crate) async fn reject_tool_execution( ) -> Result { let request = request_value(args); let tool_id = get_string(request, "toolId")?; + let ownership = state + .turns + .claim_confirmation(&tool_id) + .ok_or_else(|| "Tool confirmation is not owned by an active Peer turn".to_string())?; + if optional_string(request, "sessionId") + .is_some_and(|value| value.as_str() != ownership.session_id.as_str()) + || optional_string(request, "dialogTurnId") + .is_some_and(|value| value.as_str() != ownership.turn_id.as_str()) + { + state.turns.restore_confirmation(tool_id, ownership); + return Err("Tool confirmation session or turn does not match its Peer owner".to_string()); + } let reason = optional_string(request, "reason").unwrap_or_else(|| "User rejected".to_string()); - state - .coordinator - .reject_tool(&tool_id, reason) - .await - .map_err(|e| format!("Reject tool failed: {e}"))?; + if let Err(error) = state.compatibility.reject_tool(&tool_id, reason).await { + state.turns.restore_confirmation(tool_id, ownership); + return Err(format!("Reject tool failed: {error}")); + } Ok(Value::Null) } + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::peer_dialog_metadata; + + #[test] + fn peer_metadata_forces_confirmation_and_cannot_claim_acp_transport() { + let metadata = peer_dialog_metadata(&json!({ + "userMessageMetadata": { + "acp_transport": true, + "kind": "background_result", + "sourceKind": "subagent", + "backgroundTaskId": "background-task", + "parentSessionId": "parent-session", + "parentDialogTurnId": "parent-turn", + "subagentSessionId": "subagent-session", + "subagentDialogTurnId": "subagent-turn", + "require_tool_confirmation": false, + "caller": "desktop", + } + })) + .expect("metadata"); + + assert_eq!( + metadata.get("require_tool_confirmation"), + Some(&json!(true)) + ); + assert!(!metadata.contains_key("acp_transport")); + for reserved_key in [ + "backgroundTaskId", + "parentSessionId", + "parentDialogTurnId", + "subagentSessionId", + "subagentDialogTurnId", + ] { + assert!(!metadata.contains_key(reserved_key)); + } + assert_eq!(metadata.get("kind"), Some(&json!("background_result"))); + assert_eq!(metadata.get("sourceKind"), Some(&json!("subagent"))); + assert_eq!(metadata.get("caller"), Some(&json!("desktop"))); + } + + #[test] + fn peer_metadata_preserves_non_lineage_classification() { + let metadata = peer_dialog_metadata(&json!({ + "userMessageMetadata": { + "kind": "manual_compaction", + "sourceKind": "user", + } + })) + .expect("metadata"); + + assert_eq!(metadata.get("kind"), Some(&json!("manual_compaction"))); + assert_eq!(metadata.get("sourceKind"), Some(&json!("user"))); + } +} diff --git a/src/apps/cli/src/peer_host/commands/mod.rs b/src/apps/cli/src/peer_host/commands/mod.rs index 61ba928b32..3a7d3a0968 100644 --- a/src/apps/cli/src/peer_host/commands/mod.rs +++ b/src/apps/cli/src/peer_host/commands/mod.rs @@ -67,12 +67,12 @@ pub(crate) async fn dispatch( "update_session_model" => session::update_session_model(state, args).await, "ensure_coordinator_session" => session::ensure_coordinator_session(state, args).await, "get_available_modes" => session::get_available_modes().await, - "get_session_stats" => session::get_session_stats(args).await, + "get_session_stats" => session::get_session_stats(state, args).await, "save_session_turn" => session::save_session_turn(state, args).await, // Snapshot / rollback "rollback_to_turn" => snapshot::rollback_to_turn(state, args).await, - "get_session_files" => snapshot::get_session_files(args).await, + "get_session_files" => snapshot::get_session_files(state, args).await, // Dialog / tools "start_dialog_turn" => dialog::start_dialog_turn(state, args).await, diff --git a/src/apps/cli/src/peer_host/commands/session.rs b/src/apps/cli/src/peer_host/commands/session.rs index 4ee3ad4a1d..dcf3f16925 100644 --- a/src/apps/cli/src/peer_host/commands/session.rs +++ b/src/apps/cli/src/peer_host/commands/session.rs @@ -8,13 +8,39 @@ use serde_json::{json, Value}; use bitfun_core::agentic::core::{Session, SessionConfig}; use bitfun_core::agentic::get_agent_registry; use bitfun_core::service::session::SessionStatus; -use bitfun_core::service::snapshot::get_snapshot_manager_for_workspace; +use bitfun_runtime_ports::{AgentSessionDeleteRequest, SessionStoragePathRequest}; use crate::peer_host::args::{get_string, optional_bool, optional_string, request_value}; use crate::peer_host::state::PeerHostState; -fn storage_path(workspace_path: &str) -> PathBuf { - PathBuf::from(workspace_path) +fn session_storage_request(request: &Value) -> Result { + let workspace_path = get_string(request, "workspacePath")?; + let workspace_path = workspace_path.trim(); + if workspace_path.is_empty() { + return Err("workspace_path is required".to_string()); + } + Ok(SessionStoragePathRequest { + workspace_path: PathBuf::from(workspace_path), + remote_connection_id: optional_string(request, "remoteConnectionId"), + remote_ssh_host: optional_string(request, "remoteSshHost"), + }) +} + +pub(super) async fn resolved_session_storage_path( + state: &PeerHostState, + request: &Value, +) -> Result { + state + .compatibility + .resolve_persisted_session_storage_path(session_storage_request(request)?) + .await + .map_err(|error| format!("Failed to resolve session storage path: {error}")) +} + +fn validated_session_id(request: &Value) -> Result { + let session_id = get_string(request, "sessionId")?; + bitfun_agent_runtime::session_control::validate_session_id(&session_id)?; + Ok(session_id) } fn system_time_to_unix_secs(time: SystemTime) -> u64 { @@ -42,10 +68,10 @@ pub(crate) async fn list_persisted_sessions( args: &Value, ) -> Result { let request = request_value(args); - let workspace_path = storage_path(&get_string(request, "workspacePath")?); + let workspace_path = resolved_session_storage_path(state, request).await?; let list = state - .persistence - .list_session_metadata(&workspace_path) + .compatibility + .list_persisted_sessions(&workspace_path) .await .map_err(|e| format!("Failed to list persisted sessions: {e}"))?; serde_json::to_value(list).map_err(|e| format!("serialize sessions: {e}")) @@ -56,12 +82,12 @@ pub(crate) async fn list_persisted_sessions_page( args: &Value, ) -> Result { let request = request_value(args); - let workspace_path = storage_path(&get_string(request, "workspacePath")?); + let workspace_path = resolved_session_storage_path(state, request).await?; let limit = request.get("limit").and_then(|v| v.as_u64()).unwrap_or(50) as usize; let cursor = optional_string(request, "cursor"); let page = state - .persistence - .list_session_metadata_page(&workspace_path, cursor.as_deref(), limit) + .compatibility + .list_persisted_sessions_page(&workspace_path, cursor.as_deref(), limit) .await .map_err(|e| format!("Failed to list persisted session page: {e}"))?; serde_json::to_value(page).map_err(|e| format!("serialize session page: {e}")) @@ -72,10 +98,10 @@ pub(crate) async fn list_persisted_sessions_count( args: &Value, ) -> Result { let request = request_value(args); - let workspace_path = storage_path(&get_string(request, "workspacePath")?); + let workspace_path = resolved_session_storage_path(state, request).await?; let list = state - .persistence - .list_session_metadata(&workspace_path) + .compatibility + .list_persisted_sessions(&workspace_path) .await .map_err(|e| format!("Failed to count persisted sessions: {e}"))?; Ok(json!(list.len())) @@ -86,20 +112,17 @@ pub(crate) async fn load_session_turns( args: &Value, ) -> Result { let request = request_value(args); - let session_id = get_string(request, "sessionId")?; - let workspace_path = storage_path(&get_string(request, "workspacePath")?); - let turns = if let Some(limit) = request.get("limit").and_then(|v| v.as_u64()) { - state - .persistence - .load_recent_turns(&workspace_path, &session_id, limit as usize) - .await - } else { - state - .persistence - .load_session_turns(&workspace_path, &session_id) - .await - } - .map_err(|e| format!("Failed to load session turns: {e}"))?; + let session_id = validated_session_id(request)?; + let workspace_path = resolved_session_storage_path(state, request).await?; + let limit = request + .get("limit") + .and_then(|value| value.as_u64()) + .map(|value| value as usize); + let turns = state + .compatibility + .load_persisted_session_turns(&workspace_path, &session_id, limit) + .await + .map_err(|e| format!("Failed to load session turns: {e}"))?; serde_json::to_value(turns).map_err(|e| format!("serialize turns: {e}")) } @@ -108,8 +131,8 @@ pub(crate) async fn restore_session_view( args: &Value, ) -> Result { let request = request_value(args); - let session_id = get_string(request, "sessionId")?; - let workspace_path = storage_path(&get_string(request, "workspacePath")?); + let session_id = validated_session_id(request)?; + let storage_request = session_storage_request(request)?; let include_internal = request .get("includeInternal") .and_then(|v| v.as_bool()) @@ -121,46 +144,16 @@ pub(crate) async fn restore_session_view( .filter(|n| *n > 0) .map(|n| n.min(16)); - let (session, turns, total_turn_count, timings) = if let Some(tail) = tail_turn_count { - if include_internal { - state - .coordinator - .restore_internal_session_view_from_storage_path_tail_timed( - &workspace_path, - &session_id, - tail, - ) - .await - } else { - state - .coordinator - .restore_session_view_from_storage_path_tail_timed( - &workspace_path, - &session_id, - tail, - ) - .await - } - } else if include_internal { - state - .coordinator - .restore_internal_session_view_from_storage_path_timed(&workspace_path, &session_id) - .await - .map(|(session, turns, timings)| { - let total = turns.len(); - (session, turns, total, timings) - }) - } else { - state - .coordinator - .restore_session_view_from_storage_path_timed(&workspace_path, &session_id) - .await - .map(|(session, turns, timings)| { - let total = turns.len(); - (session, turns, total, timings) - }) - } - .map_err(|e| format!("Failed to restore session view: {e}"))?; + let (session, turns, total_turn_count, timings) = state + .compatibility + .restore_session_view_for_workspace( + storage_request, + &session_id, + include_internal, + tail_turn_count, + ) + .await + .map_err(|e| format!("Failed to restore session view: {e}"))?; let loaded_turn_count = turns.len(); let is_partial = loaded_turn_count < total_turn_count; @@ -180,25 +173,18 @@ pub(crate) async fn restore_session_with_turns( args: &Value, ) -> Result { let request = request_value(args); - let session_id = get_string(request, "sessionId")?; - let workspace_path = storage_path(&get_string(request, "workspacePath")?); + let session_id = validated_session_id(request)?; + let storage_request = session_storage_request(request)?; let include_internal = request .get("includeInternal") .and_then(|v| v.as_bool()) .unwrap_or(false); - let (session, turns) = if include_internal { - state - .coordinator - .restore_internal_session_with_turns_from_storage_path(&workspace_path, &session_id) - .await - } else { - state - .coordinator - .restore_session_with_turns_from_storage_path(&workspace_path, &session_id) - .await - } - .map_err(|e| format!("Failed to restore session with turns: {e}"))?; + let (session, turns) = state + .compatibility + .restore_session_with_turns_for_workspace(storage_request, &session_id, include_internal) + .await + .map_err(|e| format!("Failed to restore session with turns: {e}"))?; let turn_count = turns.len(); Ok(json!({ @@ -209,25 +195,18 @@ pub(crate) async fn restore_session_with_turns( pub(crate) async fn restore_session(state: &PeerHostState, args: &Value) -> Result { let request = request_value(args); - let session_id = get_string(request, "sessionId")?; - let workspace_path = storage_path(&get_string(request, "workspacePath")?); + let session_id = validated_session_id(request)?; + let storage_request = session_storage_request(request)?; let include_internal = request .get("includeInternal") .and_then(|v| v.as_bool()) .unwrap_or(false); - let session = if include_internal { - state - .coordinator - .restore_internal_session_from_storage_path(&workspace_path, &session_id) - .await - } else { - state - .coordinator - .restore_session_from_storage_path(&workspace_path, &session_id) - .await - } - .map_err(|e| format!("Failed to restore session: {e}"))?; + let session = state + .compatibility + .restore_session_for_workspace(storage_request, &session_id, include_internal) + .await + .map_err(|e| format!("Failed to restore session: {e}"))?; let turn_count = session.dialog_turn_ids.len(); Ok(session_to_json(session, turn_count)) @@ -266,7 +245,7 @@ pub(crate) async fn create_session(state: &PeerHostState, args: &Value) -> Resul }; let session = state - .coordinator + .compatibility .create_session_with_workspace(session_id, session_name, agent_type, config, workspace_path) .await .map_err(|e| format!("Failed to create session: {e}"))?; @@ -280,34 +259,31 @@ pub(crate) async fn create_session(state: &PeerHostState, args: &Value) -> Resul pub(crate) async fn delete_session(state: &PeerHostState, args: &Value) -> Result { let request = request_value(args); - let session_id = get_string(request, "sessionId")?; - let workspace_path = storage_path(&get_string(request, "workspacePath")?); - - // Prefer coordinator delete when session may be live; also remove persisted files. - if let Err(e) = state - .coordinator - .delete_session(&workspace_path, &session_id) + let session_id = validated_session_id(request)?; + let workspace_path = get_string(request, "workspacePath")?; + state + .agent_runtime + .delete_session(AgentSessionDeleteRequest { + workspace_path, + session_id, + remote_connection_id: optional_string(request, "remoteConnectionId"), + remote_ssh_host: optional_string(request, "remoteSshHost"), + }) .await - { - tracing::debug!("coordinator delete_session: {e}; falling back to persistence delete"); - state - .persistence - .delete_session(&workspace_path, &session_id) - .await - .map_err(|err| format!("Failed to delete session: {err}"))?; - } + .map_err(|error| format!("Failed to delete session: {error}"))?; Ok(Value::Null) } pub(crate) async fn rename_session(state: &PeerHostState, args: &Value) -> Result { let request = request_value(args); - let session_id = get_string(request, "sessionId")?; + let session_id = validated_session_id(request)?; + let workspace_path = resolved_session_storage_path(state, request).await?; let title = get_string(request, "sessionName") .or_else(|_| get_string(request, "title")) .or_else(|_| get_string(request, "name"))?; state - .coordinator - .update_session_title(&session_id, &title) + .compatibility + .update_session_title_for_storage_path(&workspace_path, &session_id, &title) .await .map_err(|e| format!("Failed to rename session: {e}"))?; Ok(Value::Null) @@ -315,18 +291,23 @@ pub(crate) async fn rename_session(state: &PeerHostState, args: &Value) -> Resul pub(crate) async fn archive_session(state: &PeerHostState, args: &Value) -> Result { let request = request_value(args); - let session_id = get_string(request, "sessionId")?; - let workspace_path = storage_path(&get_string(request, "workspacePath")?); + let session_id = validated_session_id(request)?; + let workspace_path = resolved_session_storage_path(state, request).await?; + let _mutation = state + .compatibility + .begin_persisted_session_mutation(&workspace_path, &session_id) + .await + .map_err(|error| format!("Failed to lock session archive: {error}"))?; let mut metadata = state - .persistence - .load_session_metadata(&workspace_path, &session_id) + .compatibility + .load_persisted_session_metadata(&workspace_path, &session_id) .await .map_err(|e| format!("Failed to load session metadata: {e}"))? .ok_or_else(|| "Session not found".to_string())?; metadata.status = SessionStatus::Archived; state - .persistence - .save_session_metadata(&workspace_path, &metadata) + .compatibility + .save_persisted_session_metadata(&workspace_path, &metadata) .await .map_err(|e| format!("Failed to archive session: {e}"))?; Ok(Value::Null) @@ -337,11 +318,16 @@ pub(crate) async fn touch_session_activity( args: &Value, ) -> Result { let request = request_value(args); - let session_id = get_string(request, "sessionId")?; - let workspace_path = storage_path(&get_string(request, "workspacePath")?); + let session_id = validated_session_id(request)?; + let workspace_path = resolved_session_storage_path(state, request).await?; + let _mutation = state + .compatibility + .begin_persisted_session_mutation(&workspace_path, &session_id) + .await + .map_err(|error| format!("Failed to lock session activity update: {error}"))?; state - .persistence - .touch_session(&workspace_path, &session_id) + .compatibility + .touch_persisted_session(&workspace_path, &session_id) .await .map_err(|e| format!("Failed to update session activity: {e}"))?; Ok(Value::Null) @@ -352,12 +338,14 @@ pub(crate) async fn get_session_thread_goal( args: &Value, ) -> Result { let request = request_value(args); - let session_id = get_string(request, "sessionId")?; - let workspace_path = optional_string(request, "workspacePath") - .map(|p| storage_path(&p)) - .unwrap_or_else(|| PathBuf::from(".")); + let session_id = validated_session_id(request)?; + let workspace_path = if optional_string(request, "workspacePath").is_some() { + resolved_session_storage_path(state, request).await? + } else { + PathBuf::from(".") + }; let goal = state - .coordinator + .compatibility .get_thread_goal(&session_id, workspace_path.as_path()) .await .map_err(|e| e.to_string())?; @@ -369,10 +357,10 @@ pub(crate) async fn update_session_model( args: &Value, ) -> Result { let request = request_value(args); - let session_id = get_string(request, "sessionId")?; + let session_id = validated_session_id(request)?; let model_name = get_string(request, "modelName")?; state - .coordinator + .compatibility .update_session_model(&session_id, &model_name) .await .map_err(|e| format!("Failed to update session model: {e}"))?; @@ -384,41 +372,21 @@ pub(crate) async fn ensure_coordinator_session( args: &Value, ) -> Result { let request = request_value(args); - let session_id = get_string(request, "sessionId")?; - let session_id = session_id.trim(); - if session_id.is_empty() { - return Err("session_id is required".to_string()); - } - + let session_id = validated_session_id(request)?; if state - .coordinator - .get_session_manager() - .get_session(session_id) - .is_some() + .compatibility + .is_session_loaded_in_memory(&session_id) + .map_err(|error| error.to_string())? { return Ok(Value::Null); } - - let workspace_path = get_string(request, "workspacePath")?; - let workspace_path = workspace_path.trim(); - if workspace_path.is_empty() { - return Err("workspace_path is required when the session is not loaded".to_string()); - } - let storage = storage_path(workspace_path); + let storage = resolved_session_storage_path(state, request).await?; let include_internal = optional_bool(request, "includeInternal").unwrap_or(false); - let restore_result = if include_internal { - state - .coordinator - .restore_internal_session_from_storage_path(&storage, session_id) - .await - } else { - state - .coordinator - .restore_session_from_storage_path(&storage, session_id) - .await - }; - restore_result + state + .compatibility + .ensure_session_loaded_from_storage_path(&storage, &session_id, include_internal) + .await .map(|_| Value::Null) .map_err(|e| e.to_string()) } @@ -452,17 +420,22 @@ pub(crate) async fn get_available_modes() -> Result { Ok(Value::Array(dtos)) } -pub(crate) async fn get_session_stats(args: &Value) -> Result { +pub(crate) async fn get_session_stats( + state: &PeerHostState, + args: &Value, +) -> Result { let request = request_value(args); let session_id = get_string(request, "sessionId")?; let workspace_path = get_string(request, "workspacePath")?; let workspace = PathBuf::from(&workspace_path); - if let Some(manager) = get_snapshot_manager_for_workspace(&workspace) { - return manager - .get_session_stats(&session_id) - .await - .map_err(|e| format!("Failed to get session stats: {e}")); + if let Some(stats) = state + .compatibility + .get_session_snapshot_stats(&workspace, &session_id) + .await + .map_err(|e| format!("Failed to get session stats: {e}"))? + { + return Ok(stats); } Ok(json!({ @@ -478,19 +451,31 @@ pub(crate) async fn save_session_turn( args: &Value, ) -> Result { let request = request_value(args); - let workspace_path = storage_path(&get_string(request, "workspacePath")?); + let workspace_path = resolved_session_storage_path(state, request).await?; let turn_data = request .get("turnData") .or_else(|| request.get("turn_data")) .cloned() .ok_or_else(|| "Missing 'turn_data' field".to_string())?; - let turn: bitfun_core::service::session::DialogTurnData = serde_json::from_value(turn_data) - .map_err(|e| format!("Invalid turn_data: {e}"))?; + let turn: bitfun_core::service::session::DialogTurnData = + serde_json::from_value(turn_data).map_err(|e| format!("Invalid turn_data: {e}"))?; + bitfun_agent_runtime::session_control::validate_session_id(&turn.session_id)?; + if let Some(request_session_id) = optional_string(request, "sessionId") { + bitfun_agent_runtime::session_control::validate_session_id(&request_session_id)?; + if request_session_id != turn.session_id { + return Err("turn_data session_id does not match request session_id".to_string()); + } + } + let _mutation = state + .compatibility + .begin_persisted_session_mutation(&workspace_path, &turn.session_id) + .await + .map_err(|error| format!("Failed to lock session turn save: {error}"))?; state - .persistence - .save_dialog_turn(&workspace_path, &turn) + .compatibility + .save_persisted_dialog_turn(&workspace_path, &turn) .await .map_err(|e| format!("Failed to save session turn: {e}"))?; Ok(Value::Null) diff --git a/src/apps/cli/src/peer_host/commands/snapshot.rs b/src/apps/cli/src/peer_host/commands/snapshot.rs index ae6074f8f9..f5c2e710a5 100644 --- a/src/apps/cli/src/peer_host/commands/snapshot.rs +++ b/src/apps/cli/src/peer_host/commands/snapshot.rs @@ -2,84 +2,108 @@ use std::collections::HashSet; use std::path::PathBuf; -use std::sync::Arc; -use std::time::Duration; use serde_json::{json, Value}; -use bitfun_core::service::snapshot::{ - get_snapshot_manager_for_workspace, initialize_snapshot_manager_for_workspace, SnapshotManager, -}; - use crate::peer_host::args::{get_string, get_usize, optional_bool, request_value}; use crate::peer_host::fanout::fanout_peer_device_event; use crate::peer_host::state::PeerHostState; -async fn ensure_snapshot_manager(workspace_path: &str) -> Result, String> { - let workspace_dir = PathBuf::from(workspace_path); - if let Some(manager) = get_snapshot_manager_for_workspace(&workspace_dir) { - return Ok(manager); - } - - initialize_snapshot_manager_for_workspace(workspace_dir.clone(), None) - .await - .map_err(|e| { - format!( - "Failed to initialize snapshot system for workspace {}: {e}", - workspace_dir.display() - ) - })?; - - get_snapshot_manager_for_workspace(&workspace_dir).ok_or_else(|| { - format!( - "Failed to get snapshot manager for workspace {}", - workspace_dir.display() - ) - }) -} +use super::session::resolved_session_storage_path; -pub(crate) async fn get_session_files(args: &Value) -> Result { +pub(crate) async fn get_session_files( + state: &PeerHostState, + args: &Value, +) -> Result { let request = request_value(args); let session_id = get_string(request, "sessionId")?; let workspace_path = get_string(request, "workspacePath")?; - let manager = ensure_snapshot_manager(&workspace_path).await?; - let files = manager - .get_session_files(&session_id) + bitfun_agent_runtime::session_control::validate_session_id(&session_id)?; + let files = state + .compatibility + .get_session_snapshot_files(&PathBuf::from(&workspace_path), &session_id) .await .map_err(|e| format!("Failed to get session files: {e}"))?; - Ok(json!( - files - .iter() - .map(|p| p.to_string_lossy().to_string()) - .collect::>() - )) + Ok(json!(files + .iter() + .map(|p| p.to_string_lossy().to_string()) + .collect::>())) } -pub(crate) async fn rollback_to_turn( - state: &PeerHostState, - args: &Value, -) -> Result { +pub(crate) async fn rollback_to_turn(state: &PeerHostState, args: &Value) -> Result { let request = request_value(args); let session_id = get_string(request, "sessionId")?; let workspace_path = get_string(request, "workspacePath")?; let turn_index = get_usize(request, "turnIndex")?; let delete_turns = optional_bool(request, "deleteTurns").unwrap_or(false); - if let Err(e) = state - .coordinator - .cancel_active_turn_for_session(&session_id, Duration::from_secs(2)) - .await - { - tracing::warn!( - "Failed to cancel active turn before rollback: session_id={session_id}, turn_index={turn_index}, error={e}" - ); + bitfun_agent_runtime::session_control::validate_session_id(&session_id)?; + let workspace = PathBuf::from(&workspace_path); + let session_storage_path = resolved_session_storage_path(state, request).await?; + if delete_turns { + state + .compatibility + .ensure_session_loaded_from_storage_path(&session_storage_path, &session_id, false) + .await + .map_err(|error| format!("Failed to load session before rollback: {error}"))?; } - - let manager = ensure_snapshot_manager(&workspace_path).await?; - let restored_files = manager - .rollback_to_turn(&session_id, turn_index) + let maintenance = state + .compatibility + .begin_session_maintenance(&session_storage_path, &session_id, 2_000) + .await + .map_err(|error| format!("Failed to quiesce session before rollback: {error}"))?; + let mut descendant_cancellation = state.turns.session_turns_for_cancellation(&session_id); + descendant_cancellation + .turns + .retain(|turn| turn.session_id != session_id); + state + .cancel_peer_turns(descendant_cancellation, "Peer session rollback") + .await + .map_err(|error| format!("Failed to cancel Peer descendants before rollback: {error}"))?; + state.turns.drain_session_turns(&session_id); + + let mutation = if delete_turns { + Some( + state + .compatibility + .begin_persisted_session_mutation(&session_storage_path, &session_id) + .await + .map_err(|error| format!("Failed to lock session rollback: {error}"))?, + ) + } else { + None + }; + + let rolled_back_parent_turn_ids = if delete_turns { + let turns = state + .compatibility + .load_persisted_session_turns(&session_storage_path, &session_id, None) + .await + .map_err(|error| format!("Failed to load turns before rollback: {error}"))?; + state + .compatibility + .validate_persisted_session_context_rollback( + mutation + .as_ref() + .expect("mutation exists when deleting turns"), + turn_index, + ) + .await + .map_err(|error| format!("Failed to validate session rollback: {error}"))?; + turns + .into_iter() + .filter(|turn| turn.turn_index >= turn_index) + .map(|turn| turn.turn_id) + .collect::>() + } else { + HashSet::new() + }; + + let restored_files = state + .compatibility + .rollback_workspace_files_to_turn(&workspace, &session_id, turn_index) .await .map_err(|e| format!("Failed to rollback turn: {e}"))?; @@ -88,70 +112,44 @@ pub(crate) async fn rollback_to_turn( .map(|p| p.to_string_lossy().to_string()) .collect(); - let mut deleted_turns_count = 0usize; + let deleted_turns_count = rolled_back_parent_turn_ids.len(); if delete_turns { - let workspace = PathBuf::from(&workspace_path); - let mut rolled_back_parent_turn_ids = HashSet::new(); - - match state - .persistence - .load_session_turns(&workspace, &session_id) + if let Err(error) = state + .compatibility + .rollback_persisted_session_context_to_turn_start( + mutation + .as_ref() + .expect("mutation exists when deleting turns"), + turn_index, + ) .await { - Ok(turns) => { - rolled_back_parent_turn_ids = turns - .into_iter() - .filter(|turn| turn.turn_index >= turn_index) - .map(|turn| turn.turn_id) - .collect(); - } - Err(e) => { - tracing::warn!( - "Failed to load parent turns before rollback cleanup: session_id={session_id}, turn_index={turn_index}, error={e}" - ); - } + return Err(format!( + "Workspace files were rolled back, but session history rollback failed. Reload the session before retrying: {error}" + )); } if !rolled_back_parent_turn_ids.is_empty() { - if let Err(e) = state - .coordinator + if let Err(error) = state + .compatibility .delete_hidden_subagent_sessions_for_parent_turns( - &workspace, + &session_storage_path, &session_id, &rolled_back_parent_turn_ids, ) .await { tracing::warn!( - "Failed to delete hidden subagent sessions during rollback: session_id={session_id}, turn_index={turn_index}, error={e}" + "Failed to delete hidden subagent sessions during rollback: session_id={session_id}, turn_index={turn_index}, error={error}" ); } } + } - if let Err(e) = state - .coordinator - .get_session_manager() - .rollback_context_to_turn_start(&workspace, &session_id, turn_index) - .await - { - tracing::warn!( - "Rollback agentic context failed: session_id={session_id}, turn_index={turn_index}, error={e}" - ); - } - - match state - .persistence - .delete_turns_from(&workspace, &session_id, turn_index) - .await - { - Ok(count) => deleted_turns_count = count, - Err(e) => { - tracing::warn!( - "Failed to delete conversation turns: session_id={session_id}, turn_index={turn_index}, error={e}" - ); - } - } + drop(mutation); + drop(maintenance); + if delete_turns { fanout_peer_device_event( "conversation_turns_deleted".to_string(), json!({ @@ -159,7 +157,8 @@ pub(crate) async fn rollback_to_turn( "remaining_turns": turn_index, "deleted_count": deleted_turns_count, }), - ); + ) + .await; } fanout_peer_device_event( @@ -171,7 +170,8 @@ pub(crate) async fn rollback_to_turn( "deleted_turns": delete_turns, "deleted_turns_count": deleted_turns_count, }), - ); + ) + .await; Ok(json!(restored_files_str)) } diff --git a/src/apps/cli/src/peer_host/control.rs b/src/apps/cli/src/peer_host/control.rs index 87e813ca27..7caa9e7de2 100644 --- a/src/apps/cli/src/peer_host/control.rs +++ b/src/apps/cli/src/peer_host/control.rs @@ -4,35 +4,125 @@ use std::collections::HashSet; use std::sync::{Mutex, OnceLock}; use serde_json::{json, Value}; +use tokio::sync::{RwLock, RwLockReadGuard}; use bitfun_core::service::remote_connect::DeviceIdentity; -static CONTROL_SUBSCRIBERS: OnceLock>> = OnceLock::new(); +#[derive(Default)] +struct ControllerRegistry { + ids: HashSet, + generation: u64, +} + +#[derive(Clone, Copy)] +pub(crate) struct ControllerLease { + generation: u64, +} + +static CONTROL_SUBSCRIBERS: OnceLock> = OnceLock::new(); +static CONTROLLER_DELIVERY: OnceLock> = OnceLock::new(); +const MAX_ATTACHED_CONTROLLERS: usize = 64; -fn control_subscribers() -> &'static Mutex> { - CONTROL_SUBSCRIBERS.get_or_init(|| Mutex::new(HashSet::new())) +fn control_subscribers() -> &'static Mutex { + CONTROL_SUBSCRIBERS.get_or_init(|| Mutex::new(ControllerRegistry::default())) } -pub(crate) fn attach_controller(device_id: String) { +fn controller_delivery() -> &'static RwLock<()> { + CONTROLLER_DELIVERY.get_or_init(|| RwLock::new(())) +} + +pub(crate) async fn attach_controller(device_id: String) -> Result<(), String> { if device_id.trim().is_empty() { - return; + return Err("controller_device_id is required".to_string()); } - if let Ok(mut set) = control_subscribers().lock() { - set.insert(device_id); + let _delivery = controller_delivery().write().await; + let mut registry = control_subscribers() + .lock() + .map_err(|_| "Peer controller registry is unavailable".to_string())?; + if !registry.ids.contains(&device_id) && registry.ids.len() >= MAX_ATTACHED_CONTROLLERS { + return Err("Peer controller capacity is exhausted".to_string()); } + registry.ids.insert(device_id); + Ok(()) } -pub(crate) fn detach_controller(device_id: &str) { - if let Ok(mut set) = control_subscribers().lock() { - set.remove(device_id); +pub(crate) async fn detach_controller(device_id: &str) -> bool { + let _delivery = controller_delivery().write().await; + control_subscribers() + .lock() + .map(|mut registry| detach_from_registry(&mut registry, device_id)) + .unwrap_or(false) +} + +pub(crate) async fn retain_online_controllers<'a>( + online: impl IntoIterator, +) -> bool { + let online = online.into_iter().collect::>(); + let _delivery = controller_delivery().write().await; + control_subscribers() + .lock() + .map(|mut registry| retain_online_in_registry(&mut registry, &online)) + .unwrap_or(false) +} + +pub(crate) async fn controller_delivery_lease( + device_id: &str, +) -> Option> { + let lease = controller_delivery().read().await; + let attached = control_subscribers() + .lock() + .map(|registry| registry.ids.contains(device_id)) + .unwrap_or(false); + attached.then_some(lease) +} + +fn detach_from_registry(registry: &mut ControllerRegistry, device_id: &str) -> bool { + let was_attached = registry.ids.remove(device_id); + let lost_all = was_attached && registry.ids.is_empty(); + if lost_all { + registry.generation = registry.generation.wrapping_add(1); } + lost_all } -pub(crate) fn attached_controllers() -> Vec { +fn retain_online_in_registry(registry: &mut ControllerRegistry, online: &HashSet<&str>) -> bool { + let had_controllers = !registry.ids.is_empty(); + registry + .ids + .retain(|device_id| online.contains(device_id.as_str())); + let lost_all = had_controllers && registry.ids.is_empty(); + if lost_all { + registry.generation = registry.generation.wrapping_add(1); + } + lost_all +} + +pub(crate) fn attached_controller_lease() -> Result { + let registry = control_subscribers() + .lock() + .map_err(|_| "Peer controller registry is unavailable".to_string())?; + if registry.ids.is_empty() { + return Err("A Peer controller must attach before starting a dialog turn".to_string()); + } + Ok(ControllerLease { + generation: registry.generation, + }) +} + +pub(crate) fn is_controller_lease_current(lease: ControllerLease) -> bool { control_subscribers() .lock() - .map(|set| set.iter().cloned().collect()) - .unwrap_or_default() + .map(|registry| !registry.ids.is_empty() && registry.generation == lease.generation) + .unwrap_or(false) +} + +pub(crate) fn attached_controllers() -> Vec { + let mut controllers: Vec = control_subscribers() + .lock() + .map(|registry| registry.ids.iter().cloned().collect()) + .unwrap_or_default(); + controllers.sort(); + controllers } pub(crate) fn peer_mode_ping_value() -> Value { @@ -59,3 +149,91 @@ pub(crate) fn parse_controller_device_id(args: &Value) -> String { .unwrap_or("") .to_string() } + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + use std::sync::Arc; + + use super::{ + attach_controller, controller_delivery_lease, detach_controller, detach_from_registry, + retain_online_in_registry, ControllerLease, ControllerRegistry, + }; + + fn lease_is_current(registry: &ControllerRegistry, lease: ControllerLease) -> bool { + !registry.ids.is_empty() && registry.generation == lease.generation + } + + #[test] + fn only_the_last_detach_reports_loss_of_all_controllers() { + let mut registry = ControllerRegistry { + ids: HashSet::from(["controller-1".to_string(), "controller-2".to_string()]), + generation: 0, + }; + let lease = ControllerLease { generation: 0 }; + + assert!(!detach_from_registry(&mut registry, "controller-1")); + assert!(lease_is_current(®istry, lease)); + assert!(detach_from_registry(&mut registry, "controller-2")); + assert!(!lease_is_current(®istry, lease)); + assert!(!detach_from_registry(&mut registry, "controller-2")); + } + + #[test] + fn presence_removal_reports_when_the_last_controller_goes_offline() { + let mut registry = ControllerRegistry { + ids: HashSet::from(["controller-1".to_string(), "controller-2".to_string()]), + generation: 0, + }; + let first_online = HashSet::from(["controller-1"]); + assert!(!retain_online_in_registry(&mut registry, &first_online)); + + assert!(retain_online_in_registry(&mut registry, &HashSet::new())); + } + + #[test] + fn reattach_does_not_revalidate_a_lease_from_before_the_last_detach() { + let mut registry = ControllerRegistry { + ids: HashSet::from(["controller-1".to_string()]), + generation: 7, + }; + let old_lease = ControllerLease { generation: 7 }; + + assert!(detach_from_registry(&mut registry, "controller-1")); + registry.ids.insert("controller-2".to_string()); + + assert!(!lease_is_current(®istry, old_lease)); + assert!(lease_is_current( + ®istry, + ControllerLease { + generation: registry.generation, + } + )); + } + + #[tokio::test] + async fn detach_waits_for_an_in_flight_delivery_lease() { + let controller_id = "delivery-lease-controller".to_string(); + attach_controller(controller_id.clone()) + .await + .expect("attach controller"); + let delivery_lease = controller_delivery_lease(&controller_id) + .await + .expect("delivery lease"); + let started = Arc::new(tokio::sync::Barrier::new(2)); + let detach_started = Arc::clone(&started); + let detach_id = controller_id.clone(); + let detach_task = tokio::spawn(async move { + detach_started.wait().await; + detach_controller(&detach_id).await + }); + + started.wait().await; + tokio::task::yield_now().await; + assert!(!detach_task.is_finished()); + + drop(delivery_lease); + detach_task.await.expect("detach task"); + assert!(controller_delivery_lease(&controller_id).await.is_none()); + } +} diff --git a/src/apps/cli/src/peer_host/dispatch.rs b/src/apps/cli/src/peer_host/dispatch.rs index 1f22eb422b..614b643b57 100644 --- a/src/apps/cli/src/peer_host/dispatch.rs +++ b/src/apps/cli/src/peer_host/dispatch.rs @@ -9,7 +9,7 @@ use super::control::{ attach_controller, detach_controller, parse_controller_device_id, peer_mode_ping_value, }; use super::deny::{is_cli_unsupported_command, is_local_only_command}; -use super::state::peer_host_state; +use super::state::{peer_host_state, try_peer_host_state}; #[derive(Debug, Clone)] struct HostInvokeBridgeResult { @@ -62,12 +62,25 @@ async fn handle_host_invoke_inner(command: &str, args: Value) -> HostInvokeBridg if controller_id.trim().is_empty() { return HostInvokeBridgeResult::err("controller_device_id is required"); } - attach_controller(controller_id); + if let Err(error) = attach_controller(controller_id).await { + return HostInvokeBridgeResult::err(error); + } return HostInvokeBridgeResult::ok_value(json!({ "attached": true })); } if command == "peer_control_detach" { let controller_id = parse_controller_device_id(&args); - detach_controller(&controller_id); + if detach_controller(&controller_id).await { + if let Some(state) = try_peer_host_state() { + if let Err(error) = state + .cancel_and_drain_peer_turns("last Peer controller detached") + .await + { + return HostInvokeBridgeResult::err(format!( + "Peer controller detached, but active work was not fully cancelled: {error}" + )); + } + } + } return HostInvokeBridgeResult::ok_value(json!({ "detached": true })); } if command == "peer_mode_ping" { diff --git a/src/apps/cli/src/peer_host/fanout.rs b/src/apps/cli/src/peer_host/fanout.rs index 785e18e8bb..a6e9a161a1 100644 --- a/src/apps/cli/src/peer_host/fanout.rs +++ b/src/apps/cli/src/peer_host/fanout.rs @@ -1,99 +1,768 @@ //! DeviceEvent fan-out to attached Peer Mode controllers. -use std::sync::{Arc, OnceLock}; +use std::collections::HashSet; +use std::sync::OnceLock; -use bitfun_core::agentic::events::EventQueue; use bitfun_core::service::remote_connect::encryption::encrypt_to_base64; use bitfun_core::service::remote_connect::remote_server::RemoteCommand; -use bitfun_events::project_agentic_frontend_event; -use tokio::sync::mpsc; +use bitfun_events::{project_agentic_frontend_event, AgenticEvent, ToolEventData}; +use tokio::sync::{broadcast, mpsc}; -use super::control::attached_controllers; +use super::control::{attached_controllers, controller_delivery_lease}; +use super::state::{PeerHostState, PeerTurnKey}; -static PEER_EVENT_FANOUT_TX: OnceLock> = - OnceLock::new(); +const PEER_EVENT_DELIVERY_CAPACITY: usize = 512; -/// Start a second EventQueue subscriber that fans agentic events to controllers. -pub(crate) fn start_peer_event_fanout(event_queue: Arc) { +struct QueuedPeerDeviceEvent { + targets: Vec, + event: String, + payload: serde_json::Value, + continuity: Option<(super::state::PeerTurnTracker, u64)>, + terminal: Option<(super::state::PeerTurnTracker, u64, PeerTurnKey)>, +} + +impl QueuedPeerDeviceEvent { + fn new(targets: Vec, event: String, payload: serde_json::Value) -> Self { + Self { + targets, + event, + payload, + continuity: None, + terminal: None, + } + } + + fn for_agent_event( + targets: Vec, + event: String, + payload: serde_json::Value, + turns: super::state::PeerTurnTracker, + generation: u64, + terminal_turn: Option, + ) -> Self { + let terminal = terminal_turn.map(|turn| (turns.clone(), generation, turn)); + Self { + targets, + event, + payload, + continuity: Some((turns, generation)), + terminal, + } + } +} + +fn continuity_is_current(continuity: &Option<(super::state::PeerTurnTracker, u64)>) -> bool { + continuity + .as_ref() + .is_none_or(|(turns, generation)| turns.is_event_stream_generation_current(*generation)) +} + +static PEER_EVENT_FANOUT_TX: OnceLock> = OnceLock::new(); + +fn peer_event_sender() -> &'static mpsc::Sender { + PEER_EVENT_FANOUT_TX.get_or_init(|| { + let (tx, mut rx) = mpsc::channel::(PEER_EVENT_DELIVERY_CAPACITY); + tokio::spawn(async move { + while let Some(queued) = rx.recv().await { + fanout_peer_device_event_once(queued).await; + } + }); + tx + }) +} + +/// Subscribe to the invocation-scoped event source and forward only Peer-owned turns. +pub(crate) fn start_peer_event_fanout(state: PeerHostState) { + let mut rx = state.agent_events.subscribe(); + state.turns.mark_event_stream_ready(); tokio::spawn(async move { - let mut rx = event_queue.subscribe(); loop { match rx.recv().await { Ok(envelope) => { - if attached_controllers().is_empty() { - continue; + if let Err(error) = handle_agentic_event(&state, envelope.event).await { + tracing::warn!("CLI Peer event fanout lost continuity: {error}"); + interrupt_and_fail_peer_turns( + &state, + false, + "Peer event fanout lost continuity", + ) + .await; + if drain_broadcast_receiver(&mut rx) { + state.turns.interrupt_event_stream(true); + break; + } + state.turns.mark_event_stream_ready(); } - if let Some(projected) = project_agentic_frontend_event(envelope.event) { - fanout_peer_device_event(projected.event_name, projected.payload); + } + Err(broadcast::error::RecvError::Lagged(skipped)) => { + tracing::warn!("CLI Peer event fanout lagged by {skipped} events"); + interrupt_and_fail_peer_turns(&state, false, "Peer event stream lagged").await; + if drain_broadcast_receiver(&mut rx) { + state.turns.interrupt_event_stream(true); + break; } + state.turns.mark_event_stream_ready(); } - Err(tokio::sync::broadcast::error::RecvError::Lagged(skipped)) => { - tracing::warn!("CLI peer event fanout lagged by {skipped} events"); + Err(broadcast::error::RecvError::Closed) => { + interrupt_and_fail_peer_turns(&state, true, "Peer event stream closed").await; + break; } - Err(tokio::sync::broadcast::error::RecvError::Closed) => break, } } }); } -/// Queue a DeviceEvent for sequential delivery to attached controllers. -pub(crate) fn fanout_peer_device_event(event: String, payload: serde_json::Value) { - if attached_controllers().is_empty() { - return; +async fn interrupt_and_fail_peer_turns(state: &PeerHostState, closed: bool, reason: &'static str) { + let drain = state.turns.interrupt_event_stream(closed); + let interrupted_turns = drain.turns.clone(); + if let Err(error) = state.cancel_peer_turns(drain, reason).await { + tracing::warn!("Peer turn cancellation after event interruption was incomplete: {error}"); } - let tx = PEER_EVENT_FANOUT_TX.get_or_init(|| { - let (tx, mut rx) = mpsc::unbounded_channel::<(String, serde_json::Value)>(); - tokio::spawn(async move { - while let Some((event, payload)) = rx.recv().await { - fanout_peer_device_event_once(event, payload).await; + for turn in interrupted_turns { + let (event, payload) = interrupted_turn_failure_projection(&turn, reason); + fanout_peer_device_event(event, payload).await; + } +} + +/// Returns true when the sender side closed while the stale backlog was drained. +fn drain_broadcast_receiver( + rx: &mut broadcast::Receiver, +) -> bool { + loop { + match rx.try_recv() { + Ok(_) | Err(broadcast::error::TryRecvError::Lagged(_)) => continue, + Err(broadcast::error::TryRecvError::Empty) => return false, + Err(broadcast::error::TryRecvError::Closed) => return true, + } + } +} + +async fn handle_agentic_event(state: &PeerHostState, event: AgenticEvent) -> Result<(), String> { + let event_turn = event_turn_key(&event); + let terminal_turn = terminal_turn_key(&event); + if terminal_turn + .as_ref() + .is_some_and(|turn| state.turns.is_interrupted_terminal(turn)) + { + return Ok(()); + } + if let AgenticEvent::DialogTurnStarted { + user_message_metadata, + .. + } = &event + { + let background_lineage = user_message_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .filter(|metadata| { + metadata.get("kind").and_then(serde_json::Value::as_str) + == Some("background_result") + && metadata + .get("sourceKind") + .and_then(serde_json::Value::as_str) + == Some("subagent") + }) + .and_then(|metadata| { + Some(( + PeerTurnKey::new( + metadata.get("parentSessionId")?.as_str()?, + metadata.get("parentDialogTurnId")?.as_str()?, + ), + PeerTurnKey::new( + metadata.get("subagentSessionId")?.as_str()?, + metadata.get("subagentDialogTurnId")?.as_str()?, + ), + )) + }); + if let Some((parent, source_child)) = background_lineage { + if let Some(turn) = event_turn.as_ref() { + state + .turns + .register_background_follow_up(&parent, &source_child, turn.clone())?; } - }); - tx - }); - if let Err(e) = tx.send((event, payload)) { - tracing::debug!("peer event fanout queue closed: {e}"); + } + } + + if let AgenticEvent::SubagentSessionLinked { + session_id, + subagent_dialog_turn_id, + parent_session_id, + parent_dialog_turn_id, + parent_tool_call_id, + .. + } = &event + { + state.turns.register_linked_child( + &PeerTurnKey::new(parent_session_id, parent_dialog_turn_id), + PeerTurnKey::new(session_id, subagent_dialog_turn_id), + parent_tool_call_id, + )?; + } + + if matches!(&event, AgenticEvent::DialogTurnStarted { .. }) { + let Some(turn) = event_turn.as_ref() else { + return Ok(()); + }; + if !state.turns.mark_started(turn) { + return Ok(()); + } + } + + let Some(session_id) = event.session_id() else { + return Ok(()); + }; + if !state.turns.owns( + session_id, + event_turn.as_ref().map(|turn| turn.turn_id.as_str()), + ) { + return Ok(()); + } + if let Some(turn) = terminal_turn.as_ref() { + if !state.turns.claim_terminal_delivery(turn)? { + return Ok(()); + } + } + + if let AgenticEvent::UserSteeringInjected { + session_id, + turn_id, + steering_id, + .. + } = &event + { + state + .turns + .finish_background_injection(&PeerTurnKey::new(session_id, turn_id), steering_id); + } + + if let AgenticEvent::ToolEvent { + session_id, + turn_id, + tool_event: + ToolEventData::Started { + tool_id, + tool_name, + params, + .. + }, + .. + } = &event + { + if tool_name == "Task" + && params + .get("run_in_background") + .and_then(serde_json::Value::as_bool) + == Some(true) + { + state.turns.record_background_task_call( + &PeerTurnKey::new(session_id, turn_id), + tool_id.clone(), + )?; + } else if tool_name == "Task" + && params.get("action").and_then(serde_json::Value::as_str) == Some("cancel") + { + if let Some(target_session_id) = + params.get("session_id").and_then(serde_json::Value::as_str) + { + state.turns.record_background_task_cancellation( + &PeerTurnKey::new(session_id, turn_id), + tool_id.clone(), + target_session_id.to_string(), + )?; + } + } + } + + if let AgenticEvent::ToolEvent { + session_id, + turn_id, + tool_event, + .. + } = &event + { + let terminal_task_call = match tool_event { + ToolEventData::Completed { + tool_id, + tool_name, + result, + .. + } if tool_name == "Task" => Some(( + tool_id.as_str(), + result + .get("background_task_id") + .and_then(serde_json::Value::as_str), + result + .get("cancelled_background_tasks") + .and_then(serde_json::Value::as_u64), + )), + ToolEventData::Failed { + tool_id, tool_name, .. + } + | ToolEventData::Cancelled { + tool_id, tool_name, .. + } if tool_name == "Task" => Some((tool_id.as_str(), None, None)), + _ => None, + }; + if let Some((tool_id, background_task_id, cancelled_background_tasks)) = terminal_task_call + { + state.turns.finish_task_call( + &PeerTurnKey::new(session_id, turn_id), + tool_id, + background_task_id, + cancelled_background_tasks, + ); + } + } + + if let AgenticEvent::ToolEvent { + session_id, + turn_id, + tool_event: ToolEventData::ConfirmationNeeded { tool_id, .. }, + .. + } = &event + { + state + .turns + .record_confirmation(&PeerTurnKey::new(session_id, turn_id), tool_id.clone())?; + } + + let Some(projected) = project_agentic_frontend_event(event) else { + if let Some(turn) = terminal_turn { + state.turns.finish_turn(&turn); + } + return Ok(()); + }; + let targets = attached_controllers(); + if targets.is_empty() { + return Err("no attached Peer controller can receive Agent events".to_string()); + } + let generation = state.turns.current_event_stream_generation()?; + enqueue_peer_device_event( + peer_event_sender(), + QueuedPeerDeviceEvent::for_agent_event( + targets, + projected.event_name, + projected.payload, + state.turns.clone(), + generation, + terminal_turn, + ), + ) + .await + .map_err(|_| "Peer event delivery queue is closed".to_string())?; + Ok(()) +} + +fn event_turn_key(event: &AgenticEvent) -> Option { + let (session_id, turn_id) = match event { + AgenticEvent::DialogTurnStarted { + session_id, + turn_id, + .. + } + | AgenticEvent::DialogTurnCompleted { + session_id, + turn_id, + .. + } + | AgenticEvent::DialogTurnCancelled { + session_id, + turn_id, + .. + } + | AgenticEvent::DialogTurnFailed { + session_id, + turn_id, + .. + } + | AgenticEvent::TokenUsageUpdated { + session_id, + turn_id, + .. + } + | AgenticEvent::ContextCompressionStarted { + session_id, + turn_id, + .. + } + | AgenticEvent::ContextCompressionCompleted { + session_id, + turn_id, + .. + } + | AgenticEvent::ContextCompressionFailed { + session_id, + turn_id, + .. + } + | AgenticEvent::ModelRoundStarted { + session_id, + turn_id, + .. + } + | AgenticEvent::ModelRoundCompleted { + session_id, + turn_id, + .. + } + | AgenticEvent::TextChunk { + session_id, + turn_id, + .. + } + | AgenticEvent::ThinkingChunk { + session_id, + turn_id, + .. + } + | AgenticEvent::ToolEvent { + session_id, + turn_id, + .. + } + | AgenticEvent::DeepReviewQueueStateChanged { + session_id, + turn_id, + .. + } + | AgenticEvent::UserSteeringInjected { + session_id, + turn_id, + .. + } => (session_id, turn_id), + AgenticEvent::SubagentSessionLinked { + session_id, + subagent_dialog_turn_id, + .. + } => (session_id, subagent_dialog_turn_id), + _ => return None, + }; + Some(PeerTurnKey::new(session_id, turn_id)) +} + +fn terminal_turn_key(event: &AgenticEvent) -> Option { + match event { + AgenticEvent::DialogTurnCompleted { + session_id, + turn_id, + .. + } + | AgenticEvent::DialogTurnCancelled { + session_id, + turn_id, + .. + } + | AgenticEvent::DialogTurnFailed { + session_id, + turn_id, + .. + } => Some(PeerTurnKey::new(session_id, turn_id)), + _ => None, } } -async fn fanout_peer_device_event_once(event: String, payload: serde_json::Value) { +/// Queue an explicit Peer command event with its current delivery targets. +pub(crate) async fn fanout_peer_device_event(event: String, payload: serde_json::Value) { let targets = attached_controllers(); if targets.is_empty() { return; } + let queued = QueuedPeerDeviceEvent::new(targets, event, payload); + if let Err(queued) = enqueue_peer_device_event(peer_event_sender(), queued).await { + tracing::warn!( + "Peer event delivery queue closed before accepting command event; using direct delivery" + ); + fanout_peer_device_event_once(queued).await; + } +} + +async fn enqueue_peer_device_event( + sender: &mpsc::Sender, + queued: QueuedPeerDeviceEvent, +) -> Result<(), QueuedPeerDeviceEvent> { + sender.send(queued).await.map_err(|error| error.0) +} + +async fn fanout_peer_device_event_once(queued: QueuedPeerDeviceEvent) { + let QueuedPeerDeviceEvent { + targets, + event, + payload, + continuity, + terminal, + } = queued; + if !continuity_is_current(&continuity) { + return; + } + let _terminal_delivery = TerminalDeliveryGuard::new(terminal); + if targets.is_empty() { + return; + } let (session, relay_client) = match crate::account::peer_fanout_context().await { Ok(ctx) => ctx, - Err(e) => { - tracing::debug!("peer event fanout skipped: {e}"); + Err(error) => { + tracing::debug!("Peer event fanout skipped: {error}"); return; } }; - let envelope = match serde_json::to_string(&RemoteCommand::DeviceEvent { - event: event.clone(), - payload, - }) { - Ok(s) => s, - Err(e) => { - tracing::warn!("peer event fanout serialize failed: {e}"); + let envelope = match serde_json::to_string(&RemoteCommand::DeviceEvent { event, payload }) { + Ok(envelope) => envelope, + Err(error) => { + tracing::warn!("Peer event fanout serialization failed: {error}"); return; } }; let (encrypted_data, nonce) = match encrypt_to_base64(&session.master_key, &envelope) { - Ok(v) => v, - Err(e) => { - tracing::warn!("peer event fanout encrypt failed: {e}"); + Ok(encrypted) => encrypted, + Err(error) => { + tracing::warn!("Peer event fanout encryption failed: {error}"); return; } }; + let targets = retained_delivery_targets(&targets, &attached_controllers()); + if targets.is_empty() { + return; + } - for target in targets { + for target in &targets { + if !continuity_is_current(&continuity) { + break; + } + let Some(_delivery_lease) = controller_delivery_lease(target).await else { + continue; + }; let correlation_id = uuid::Uuid::new_v4().to_string(); - if let Err(e) = relay_client + if let Err(error) = relay_client .send_device_message(&target, &correlation_id, &encrypted_data, &nonce) .await { - tracing::debug!("peer event fanout to {target} failed: {e}"); + tracing::debug!("Peer event fanout to {target} failed: {error}"); + } + } +} + +fn retained_delivery_targets(snapshot: &[String], currently_attached: &[String]) -> Vec { + let currently_attached = currently_attached.iter().collect::>(); + snapshot + .iter() + .filter(|target| currently_attached.contains(target)) + .cloned() + .collect() +} + +struct TerminalDeliveryGuard { + terminal: Option<(super::state::PeerTurnTracker, u64, PeerTurnKey)>, +} + +impl TerminalDeliveryGuard { + fn new(terminal: Option<(super::state::PeerTurnTracker, u64, PeerTurnKey)>) -> Self { + Self { terminal } + } +} + +impl Drop for TerminalDeliveryGuard { + fn drop(&mut self) { + complete_terminal_delivery(self.terminal.take()); + } +} + +fn complete_terminal_delivery(terminal: Option<(super::state::PeerTurnTracker, u64, PeerTurnKey)>) { + if let Some((turns, generation, turn)) = terminal { + turns.complete_terminal_delivery(generation, &turn); + } +} + +fn interrupted_turn_failure_projection( + turn: &PeerTurnKey, + reason: &str, +) -> (String, serde_json::Value) { + let projected = project_agentic_frontend_event(AgenticEvent::DialogTurnFailed { + session_id: turn.session_id.clone(), + turn_id: turn.turn_id.clone(), + error: reason.to_string(), + error_category: None, + error_detail: None, + }) + .expect("DialogTurnFailed must have a frontend projection"); + (projected.event_name, projected.payload) +} + +#[cfg(test)] +mod tests { + use bitfun_events::{AgenticEvent, AgenticEventEnvelope, AgenticEventPriority}; + + use super::{ + continuity_is_current, drain_broadcast_receiver, enqueue_peer_device_event, event_turn_key, + interrupted_turn_failure_projection, retained_delivery_targets, QueuedPeerDeviceEvent, + TerminalDeliveryGuard, + }; + use crate::peer_host::state::{PeerTurnKey, PeerTurnTracker}; + + #[test] + fn queued_events_keep_the_target_snapshot_from_enqueue_time() { + let mut current_targets = vec!["controller-1".to_string()]; + let queued = QueuedPeerDeviceEvent::new( + current_targets.clone(), + "dialog_turn_started".to_string(), + serde_json::json!({}), + ); + current_targets.push("controller-2".to_string()); + + assert_eq!(queued.targets, vec!["controller-1"]); + assert_eq!(current_targets, vec!["controller-1", "controller-2"]); + } + + #[test] + fn queued_events_exclude_controllers_that_detached_after_enqueue() { + let queued_targets = vec!["controller-1".to_string(), "controller-2".to_string()]; + let currently_attached = vec!["controller-2".to_string(), "controller-3".to_string()]; + + assert_eq!( + retained_delivery_targets(&queued_targets, ¤tly_attached), + vec!["controller-2"] + ); + assert!(retained_delivery_targets(&queued_targets, &[]).is_empty()); + } + + #[test] + fn queued_peer_turn_does_not_authorize_session_scoped_events() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let turn = PeerTurnKey::new("session-1", "peer-turn"); + tracker.register_root(turn.clone()).expect("register root"); + let event = AgenticEvent::ImageAnalysisStarted { + session_id: turn.session_id.clone(), + image_count: 1, + user_input: "local input".to_string(), + image_metadata: None, + }; + + assert!(event_turn_key(&event).is_none()); + assert!(!tracker.owns(&turn.session_id, None)); + + assert!(tracker.mark_started(&turn)); + assert!(tracker.owns(&turn.session_id, None)); + tracker.finish_turn(&turn); + assert!(!tracker.owns(&turn.session_id, None)); + } + + #[tokio::test] + async fn closed_delivery_queue_returns_the_event_for_direct_fallback() { + let (tx, rx) = tokio::sync::mpsc::channel(1); + drop(rx); + let queued = QueuedPeerDeviceEvent::new( + vec!["controller-1".to_string()], + "agentic://dialog-turn-failed".to_string(), + serde_json::json!({ "turnId": "turn-1" }), + ); + + let recovered = enqueue_peer_device_event(&tx, queued) + .await + .expect_err("closed queue must return the undelivered event"); + assert_eq!(recovered.event, "agentic://dialog-turn-failed"); + } + + #[test] + fn terminal_turn_stays_owned_until_delivery_completion() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let turn = PeerTurnKey::new("session-1", "turn-1"); + let generation = tracker.register_root(turn.clone()).expect("register root"); + + assert!(tracker.owns("session-1", Some("turn-1"))); + { + let _delivery = TerminalDeliveryGuard::new(Some((tracker.clone(), generation, turn))); + assert!(tracker.owns("session-1", Some("turn-1"))); } + assert!(!tracker.owns("session-1", Some("turn-1"))); + } + + #[test] + fn stale_terminal_delivery_cannot_release_an_interrupted_turn_key_after_reset() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let turn = PeerTurnKey::new("session-1", "turn-1"); + let stale_generation = tracker + .register_root(turn.clone()) + .expect("register old root"); + let stale_delivery = + TerminalDeliveryGuard::new(Some((tracker.clone(), stale_generation, turn.clone()))); + + tracker.interrupt_event_stream(false); + tracker.mark_event_stream_ready(); + assert!(tracker.register_root(turn.clone()).is_err()); + drop(stale_delivery); + assert!(tracker.register_root(turn.clone()).is_err()); + assert!(!tracker.owns("session-1", Some("turn-1"))); + } + + #[test] + fn interrupted_turn_projects_an_existing_failed_terminal_event() { + let (event, payload) = interrupted_turn_failure_projection( + &PeerTurnKey::new("session-1", "turn-1"), + "Peer event stream lagged", + ); + + assert_eq!(event, "agentic://dialog-turn-failed"); + assert_eq!(payload["sessionId"], "session-1"); + assert_eq!(payload["turnId"], "turn-1"); + assert_eq!(payload["error"], "Peer event stream lagged"); + } + + #[test] + fn queued_agent_events_are_invalidated_by_stream_interruption() { + let turns = PeerTurnTracker::new(); + turns.mark_event_stream_ready(); + let generation = turns + .current_event_stream_generation() + .expect("ready generation"); + let queued = QueuedPeerDeviceEvent::for_agent_event( + vec!["controller-1".to_string()], + "dialog_turn_started".to_string(), + serde_json::json!({}), + turns.clone(), + generation, + None, + ); + + assert!(continuity_is_current(&queued.continuity)); + turns.interrupt_event_stream(false); + turns.mark_event_stream_ready(); + assert!(!continuity_is_current(&queued.continuity)); + } + + #[test] + fn draining_distinguishes_an_empty_live_stream_from_a_closed_stream() { + let (tx, mut live_rx) = tokio::sync::broadcast::channel(1); + assert!(!drain_broadcast_receiver(&mut live_rx)); + + let mut closed_rx = tx.subscribe(); + drop(tx); + assert!(drain_broadcast_receiver(&mut closed_rx)); + } + + #[test] + fn draining_stale_backlog_does_not_release_interrupted_turn_quarantine() { + let turns = PeerTurnTracker::new(); + turns.mark_event_stream_ready(); + let turn = PeerTurnKey::new("session-1", "turn-1"); + turns.register_root(turn.clone()).expect("register root"); + turns.interrupt_event_stream(false); + turns.mark_event_stream_ready(); + + let (tx, mut rx) = tokio::sync::broadcast::channel(2); + tx.send(AgenticEventEnvelope::new( + AgenticEvent::DialogTurnCancelled { + session_id: turn.session_id.clone(), + turn_id: turn.turn_id.clone(), + }, + AgenticEventPriority::Normal, + )) + .expect("queue stale terminal"); + + assert!(!drain_broadcast_receiver(&mut rx)); + assert!(turns.register_root(turn).is_err()); } } diff --git a/src/apps/cli/src/peer_host/mod.rs b/src/apps/cli/src/peer_host/mod.rs index 7734aa5536..c116b80b4d 100644 --- a/src/apps/cli/src/peer_host/mod.rs +++ b/src/apps/cli/src/peer_host/mod.rs @@ -16,3 +16,19 @@ mod workspace_dto; pub(crate) use bootstrap::ensure_peer_host_ready; pub(crate) use dispatch::{handle_device_event_command, handle_host_invoke}; + +pub(crate) async fn update_controller_presence(online_device_ids: Vec) { + let lost_last_controller = + control::retain_online_controllers(online_device_ids.iter().map(String::as_str)).await; + if !lost_last_controller { + return; + } + if let Some(state) = state::try_peer_host_state() { + if let Err(error) = state + .cancel_and_drain_peer_turns("last Peer controller went offline") + .await + { + tracing::warn!("Peer work was not fully cancelled after controller loss: {error}"); + } + } +} diff --git a/src/apps/cli/src/peer_host/state.rs b/src/apps/cli/src/peer_host/state.rs index e9efe84996..2206583f21 100644 --- a/src/apps/cli/src/peer_host/state.rs +++ b/src/apps/cli/src/peer_host/state.rs @@ -1,21 +1,1879 @@ -//! Shared Peer Host service handles. +//! Shared Peer Host service handles and bounded Peer-owned turn state. -use std::sync::{Arc, OnceLock}; +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex, OnceLock}; -use bitfun_core::agentic::coordination::{ConversationCoordinator, DialogScheduler}; -use bitfun_core::agentic::events::EventQueue; -use bitfun_core::agentic::persistence::PersistenceManager; +use bitfun_agent_runtime::sdk::AgentRuntime; +use bitfun_core::product_runtime::CoreAgentRuntimeCompatibility; use bitfun_core::service::filesystem::FileSystemService; use bitfun_core::service::workspace::WorkspaceService; +use bitfun_runtime_ports::{AgentSubmissionSource, AgentTurnCancellationRequest}; + +use crate::runtime::events::CliAgentEventSource; + +const MAX_TRACKED_PEER_TURNS: usize = 256; +const MAX_BACKGROUND_PEER_AUTHORIZATIONS: usize = 256; +const MAX_PENDING_PEER_TASK_CANCELLATIONS: usize = 256; +const MAX_PENDING_PEER_CONFIRMATIONS: usize = 512; + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub(crate) struct PeerTurnKey { + pub(crate) session_id: String, + pub(crate) turn_id: String, +} + +impl PeerTurnKey { + pub(crate) fn new(session_id: impl Into, turn_id: impl Into) -> Self { + Self { + session_id: session_id.into(), + turn_id: turn_id.into(), + } + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub(crate) struct PeerBackgroundSubagentLink { + pub(crate) parent_session_id: String, + pub(crate) subagent_session_id: String, +} + +#[derive(Default)] +pub(crate) struct PeerTurnDrain { + pub(crate) turns: Vec, + pub(crate) background_subagents: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum PeerEventStreamState { + Recovering, + Ready, + Closed, +} + +struct PeerTurnTrackerInner { + stream: PeerEventStreamState, + stream_generation: u64, + parents: HashMap>, + active: HashSet, + started: HashSet, + terminal_deliveries: HashSet, + interrupted_turns: HashSet, + background_task_calls: HashMap<(PeerTurnKey, String), PeerBackgroundTaskCall>, + background_task_cancellations: HashMap<(PeerTurnKey, String), String>, + background_source_children: HashSet, + background_source_tasks: HashMap, + background_follow_ups: HashSet, + early_background_follow_ups: HashSet, + completed_background_sources: HashMap, + confirmations: HashMap, +} + +#[derive(Default)] +struct PeerBackgroundTaskCall { + background_task_id: Option, + source_child: Option, +} + +#[derive(Clone)] +pub(crate) struct PeerTurnTracker { + inner: Arc>, +} + +impl PeerTurnTracker { + pub(crate) fn new() -> Self { + Self { + inner: Arc::new(Mutex::new(PeerTurnTrackerInner { + stream: PeerEventStreamState::Recovering, + stream_generation: 0, + parents: HashMap::new(), + active: HashSet::new(), + started: HashSet::new(), + terminal_deliveries: HashSet::new(), + interrupted_turns: HashSet::new(), + background_task_calls: HashMap::new(), + background_task_cancellations: HashMap::new(), + background_source_children: HashSet::new(), + background_source_tasks: HashMap::new(), + background_follow_ups: HashSet::new(), + early_background_follow_ups: HashSet::new(), + completed_background_sources: HashMap::new(), + confirmations: HashMap::new(), + })), + } + } + + pub(crate) fn mark_event_stream_ready(&self) { + if let Ok(mut inner) = self.inner.lock() { + if inner.stream != PeerEventStreamState::Closed { + inner.stream = PeerEventStreamState::Ready; + } + } + } + + pub(crate) fn register_root(&self, key: PeerTurnKey) -> Result { + let mut inner = self + .inner + .lock() + .map_err(|_| "Peer turn tracker is unavailable".to_string())?; + match inner.stream { + PeerEventStreamState::Ready => {} + PeerEventStreamState::Recovering => { + return Err("Peer event stream is recovering; retry the request".to_string()) + } + PeerEventStreamState::Closed => return Err("Peer event stream is closed".to_string()), + } + if inner.parents.contains_key(&key) { + return Err("Peer turn is already tracked".to_string()); + } + reject_interrupted_turn_id(&inner, &key)?; + if inner.active.len() >= MAX_TRACKED_PEER_TURNS + || inner.parents.len() >= MAX_TRACKED_PEER_TURNS + { + return Err("Peer turn tracking capacity is exhausted".to_string()); + } + inner.parents.insert(key.clone(), None); + inner.active.insert(key); + Ok(inner.stream_generation) + } + + pub(crate) fn claim_terminal_delivery(&self, key: &PeerTurnKey) -> Result { + let mut inner = self + .inner + .lock() + .map_err(|_| "Peer turn tracker is unavailable".to_string())?; + if !inner.active.contains(key) { + return Ok(false); + } + Ok(inner.terminal_deliveries.insert(key.clone())) + } + + pub(crate) fn is_interrupted_terminal(&self, key: &PeerTurnKey) -> bool { + self.inner + .lock() + .map(|inner| inner.interrupted_turns.contains(key)) + .unwrap_or(false) + } + + pub(crate) fn complete_terminal_delivery(&self, generation: u64, key: &PeerTurnKey) { + if let Ok(mut inner) = self.inner.lock() { + if inner.stream == PeerEventStreamState::Ready && inner.stream_generation == generation + { + finish_turn_locked(&mut inner, key); + } + } + } + + pub(crate) fn is_event_stream_generation_current(&self, generation: u64) -> bool { + self.inner + .lock() + .map(|inner| { + inner.stream == PeerEventStreamState::Ready && inner.stream_generation == generation + }) + .unwrap_or(false) + } + + pub(crate) fn current_event_stream_generation(&self) -> Result { + let inner = self + .inner + .lock() + .map_err(|_| "Peer turn tracker is unavailable".to_string())?; + if inner.stream != PeerEventStreamState::Ready { + return Err("Peer event stream is not ready".to_string()); + } + Ok(inner.stream_generation) + } + + #[cfg(test)] + pub(crate) fn register_child( + &self, + parent: &PeerTurnKey, + child: PeerTurnKey, + ) -> Result { + let mut inner = self + .inner + .lock() + .map_err(|_| "Peer turn tracker is unavailable".to_string())?; + register_child_locked(&mut inner, parent, child, false) + } + + pub(crate) fn record_background_task_call( + &self, + parent: &PeerTurnKey, + tool_call_id: String, + ) -> Result { + let mut inner = self + .inner + .lock() + .map_err(|_| "Peer turn tracker is unavailable".to_string())?; + if !inner.active.contains(parent) { + return Ok(false); + } + let call_key = (parent.clone(), tool_call_id); + if inner.background_task_calls.contains_key(&call_key) { + return Ok(true); + } + if background_authorization_len(&inner) >= MAX_BACKGROUND_PEER_AUTHORIZATIONS { + return Err("Peer background authorization capacity is exhausted".to_string()); + } + inner + .background_task_calls + .insert(call_key, PeerBackgroundTaskCall::default()); + Ok(true) + } + + pub(crate) fn record_background_task_cancellation( + &self, + parent: &PeerTurnKey, + tool_call_id: String, + target_session_id: String, + ) -> Result { + let mut inner = self + .inner + .lock() + .map_err(|_| "Peer turn tracker is unavailable".to_string())?; + if !inner.active.contains(parent) { + return Ok(false); + } + let cancellation_key = (parent.clone(), tool_call_id); + if let Some(existing_target) = inner.background_task_cancellations.get(&cancellation_key) { + if existing_target == &target_session_id { + return Ok(true); + } + return Err("Peer Task cancellation is already bound to another session".to_string()); + } + if inner.background_task_cancellations.len() >= MAX_PENDING_PEER_TASK_CANCELLATIONS { + return Err("Peer Task cancellation tracking capacity is exhausted".to_string()); + } + inner + .background_task_cancellations + .insert(cancellation_key, target_session_id); + Ok(true) + } + + pub(crate) fn finish_task_call( + &self, + parent: &PeerTurnKey, + tool_call_id: &str, + background_task_id: Option<&str>, + cancelled_background_tasks: Option, + ) { + if let Ok(mut inner) = self.inner.lock() { + let call_key = (parent.clone(), tool_call_id.to_string()); + if let Some(background_task_id) = background_task_id { + if let Some(call) = inner.background_task_calls.get_mut(&call_key) { + call.background_task_id = Some(background_task_id.to_string()); + } + bind_background_source_task(&mut inner, &call_key); + } else { + inner.background_task_calls.remove(&call_key); + } + let cancellation = inner + .background_task_cancellations + .remove(&(parent.clone(), tool_call_id.to_string())); + if cancelled_background_tasks.is_some_and(|count| count > 0) { + if let Some(target_session_id) = cancellation { + release_background_sources_for_subagent( + &mut inner, + &parent.session_id, + &target_session_id, + ); + } + } + } + } + + pub(crate) fn register_linked_child( + &self, + parent: &PeerTurnKey, + child: PeerTurnKey, + parent_tool_call_id: &str, + ) -> Result { + let mut inner = self + .inner + .lock() + .map_err(|_| "Peer turn tracker is unavailable".to_string())?; + let background_call_key = (parent.clone(), parent_tool_call_id.to_string()); + let is_background_source = inner + .background_task_calls + .contains_key(&background_call_key); + let registered = register_child_locked(&mut inner, parent, child.clone(), false)?; + if registered && is_background_source { + if let Some(call) = inner.background_task_calls.get_mut(&background_call_key) { + call.source_child = Some(child.clone()); + } + inner.background_source_children.insert(child); + bind_background_source_task(&mut inner, &background_call_key); + } + Ok(registered) + } + + pub(crate) fn finish_background_injection( + &self, + parent: &PeerTurnKey, + background_task_id: &str, + ) -> bool { + let Ok(mut inner) = self.inner.lock() else { + return false; + }; + let Some((mapped_parent, source)) = inner + .background_source_tasks + .get(background_task_id) + .cloned() + else { + return false; + }; + if &mapped_parent != parent { + return false; + } + inner.background_source_tasks.remove(background_task_id); + inner.background_source_children.remove(&source); + inner.early_background_follow_ups.remove(&source); + take_completed_background_source(&mut inner, &source); + prune_completed_branch(&mut inner, &source); + true + } + + pub(crate) fn register_background_follow_up( + &self, + parent: &PeerTurnKey, + source_child: &PeerTurnKey, + follow_up: PeerTurnKey, + ) -> Result { + let mut inner = self + .inner + .lock() + .map_err(|_| "Peer turn tracker is unavailable".to_string())?; + reject_interrupted_turn_id(&inner, &follow_up)?; + if inner.active.contains(&follow_up) { + return Ok(true); + } + let parent_is_tracked = inner.parents.contains_key(parent); + let completed_source_matches = + inner.completed_background_sources.get(source_child) == Some(parent); + if !parent_is_tracked && !completed_source_matches { + return Ok(false); + } + let required_parent_slots = usize::from(!parent_is_tracked) + 1; + if inner.active.len().saturating_add(1) > MAX_TRACKED_PEER_TURNS + || inner.parents.len().saturating_add(required_parent_slots) > MAX_TRACKED_PEER_TURNS + { + return Err("Peer turn tracking capacity is exhausted".to_string()); + } + let tracked_source_matches = inner.parents.get(source_child).and_then(Option::as_ref) + == Some(parent) + && inner.background_source_children.contains(source_child); + if (!tracked_source_matches && !completed_source_matches) + || inner.background_follow_ups.contains(source_child) + { + return Err( + "Peer background follow-up source child is not owned by its parent".to_string(), + ); + } + if completed_source_matches { + take_completed_background_source(&mut inner, source_child); + } + if !parent_is_tracked { + inner.parents.insert(parent.clone(), None); + } + + if !completed_source_matches { + if !inner.active.contains(source_child) { + return Err( + "Peer background follow-up source child is not actively owned by its parent" + .to_string(), + ); + } + inner + .early_background_follow_ups + .insert(source_child.clone()); + inner.background_source_children.remove(source_child); + } + inner + .parents + .insert(follow_up.clone(), Some(parent.clone())); + inner.active.insert(follow_up.clone()); + inner.background_follow_ups.insert(follow_up); + Ok(true) + } + + pub(crate) fn owns(&self, session_id: &str, turn_id: Option<&str>) -> bool { + self.inner + .lock() + .map(|inner| match turn_id { + Some(turn_id) => inner + .active + .contains(&PeerTurnKey::new(session_id, turn_id)), + None => inner.started.iter().any(|key| key.session_id == session_id), + }) + .unwrap_or(false) + } + + pub(crate) fn mark_started(&self, key: &PeerTurnKey) -> bool { + self.inner + .lock() + .map(|mut inner| { + if !inner.active.contains(key) { + return false; + } + inner.started.insert(key.clone()); + true + }) + .unwrap_or(false) + } + + pub(crate) fn record_confirmation( + &self, + key: &PeerTurnKey, + tool_id: String, + ) -> Result<(), String> { + let mut inner = self + .inner + .lock() + .map_err(|_| "Peer turn tracker is unavailable".to_string())?; + if !inner.active.contains(key) { + return Err("Tool confirmation does not belong to a Peer-owned turn".to_string()); + } + if let Some(existing_key) = inner.confirmations.get(&tool_id) { + if existing_key == key { + return Ok(()); + } + return Err("Tool confirmation is already owned by another Peer turn".to_string()); + } + if inner.confirmations.len() >= MAX_PENDING_PEER_CONFIRMATIONS { + return Err("Peer tool confirmation capacity is exhausted".to_string()); + } + inner.confirmations.insert(tool_id, key.clone()); + Ok(()) + } + + pub(crate) fn claim_confirmation(&self, tool_id: &str) -> Option { + self.inner + .lock() + .ok() + .and_then(|mut inner| inner.confirmations.remove(tool_id)) + } + + pub(crate) fn restore_confirmation(&self, tool_id: String, key: PeerTurnKey) { + if let Ok(mut inner) = self.inner.lock() { + if inner.active.contains(&key) + && inner.confirmations.len() < MAX_PENDING_PEER_CONFIRMATIONS + { + inner.confirmations.insert(tool_id, key); + } + } + } + + pub(crate) fn finish_turn(&self, key: &PeerTurnKey) { + if let Ok(mut inner) = self.inner.lock() { + finish_turn_locked(&mut inner, key); + } + } + + pub(crate) fn drain_session_turns(&self, session_id: &str) -> PeerTurnDrain { + self.inner + .lock() + .map(|mut inner| { + let removed = session_tree_keys(&inner, session_id); + if !try_quarantine_active_turns(&mut inner, &removed) { + inner.interrupted_turns.clear(); + inner.stream = PeerEventStreamState::Closed; + } + let mut drain = peer_turn_drain_for_keys(&inner, &removed); + merge_completed_background_subagents(&inner, &mut drain, Some(session_id)); + remove_completed_background_sources_for_session(&mut inner, session_id); + remove_tracked_turns(&mut inner, &removed); + drain + }) + .unwrap_or_default() + } + + pub(crate) fn session_turns_for_cancellation(&self, session_id: &str) -> PeerTurnDrain { + self.inner + .lock() + .map(|inner| { + let keys = session_tree_keys(&inner, session_id); + let mut drain = peer_turn_drain_for_keys(&inner, &keys); + merge_completed_background_subagents(&inner, &mut drain, Some(session_id)); + drain + }) + .unwrap_or_default() + } + + pub(crate) fn interrupt_event_stream(&self, closed: bool) -> PeerTurnDrain { + let Ok(mut inner) = self.inner.lock() else { + return PeerTurnDrain::default(); + }; + let interrupted = inner.active.clone(); + let quarantine_fits = try_quarantine_active_turns(&mut inner, &interrupted); + let stream_closed = closed || !quarantine_fits; + if stream_closed { + inner.interrupted_turns.clear(); + } + let turns = drain_peer_turns(&mut inner); + inner.stream_generation = inner.stream_generation.wrapping_add(1); + inner.stream = if stream_closed { + PeerEventStreamState::Closed + } else { + PeerEventStreamState::Recovering + }; + turns + } + + pub(crate) fn drain_peer_turns(&self) -> PeerTurnDrain { + self.inner + .lock() + .map(|mut inner| { + let active = inner.active.clone(); + if !try_quarantine_active_turns(&mut inner, &active) { + inner.interrupted_turns.clear(); + inner.stream = PeerEventStreamState::Closed; + } + drain_peer_turns(&mut inner) + }) + .unwrap_or_default() + } + + pub(crate) fn peer_turns_for_cancellation(&self) -> PeerTurnDrain { + self.inner + .lock() + .map(|inner| { + let keys = inner.parents.keys().cloned().collect::>(); + let mut drain = peer_turn_drain_for_keys(&inner, &keys); + merge_completed_background_subagents(&inner, &mut drain, None); + drain + }) + .unwrap_or_default() + } +} + +fn register_child_locked( + inner: &mut PeerTurnTrackerInner, + parent: &PeerTurnKey, + child: PeerTurnKey, + is_background_source: bool, +) -> Result { + reject_interrupted_turn_id(inner, &child)?; + if !inner.parents.contains_key(parent) { + return Ok(false); + } + if let Some(existing_parent) = inner.parents.get(&child) { + if existing_parent.as_ref() == Some(parent) { + return Ok(true); + } + return Err("Peer child turn is already owned by another parent".to_string()); + } + if inner.active.len() >= MAX_TRACKED_PEER_TURNS || inner.parents.len() >= MAX_TRACKED_PEER_TURNS + { + return Err("Peer turn tracking capacity is exhausted".to_string()); + } + inner.parents.insert(child.clone(), Some(parent.clone())); + inner.active.insert(child.clone()); + if is_background_source { + inner.background_source_children.insert(child); + } + Ok(true) +} + +fn reject_interrupted_turn_id( + inner: &PeerTurnTrackerInner, + key: &PeerTurnKey, +) -> Result<(), String> { + if inner.interrupted_turns.contains(key) { + return Err( + "Peer turn ID was interrupted and cannot be reused; retry with a new turn ID" + .to_string(), + ); + } + Ok(()) +} + +fn drain_peer_turns(inner: &mut PeerTurnTrackerInner) -> PeerTurnDrain { + let all_keys = inner.parents.keys().cloned().collect::>(); + let mut drain = peer_turn_drain_for_keys(inner, &all_keys); + merge_completed_background_subagents(inner, &mut drain, None); + inner.parents.clear(); + inner.active.clear(); + inner.started.clear(); + inner.terminal_deliveries.clear(); + inner.background_task_calls.clear(); + inner.background_task_cancellations.clear(); + inner.background_source_children.clear(); + inner.background_source_tasks.clear(); + inner.background_follow_ups.clear(); + inner.early_background_follow_ups.clear(); + inner.completed_background_sources.clear(); + inner.confirmations.clear(); + drain +} + +fn try_quarantine_active_turns( + inner: &mut PeerTurnTrackerInner, + keys: &HashSet, +) -> bool { + let new_keys = keys + .iter() + .filter(|key| inner.active.contains(*key) && !inner.interrupted_turns.contains(*key)) + .cloned() + .collect::>(); + if inner.interrupted_turns.len().saturating_add(new_keys.len()) > MAX_TRACKED_PEER_TURNS { + return false; + } + inner.interrupted_turns.extend(new_keys); + true +} + +fn finish_turn_locked(inner: &mut PeerTurnTrackerInner, key: &PeerTurnKey) { + if !inner.active.remove(key) { + return; + } + inner.started.remove(key); + inner.terminal_deliveries.remove(key); + inner + .background_task_calls + .retain(|(parent, _), _| parent != key); + inner + .background_task_cancellations + .retain(|(parent, _), _| parent != key); + let parent = inner.parents.get(key).cloned().flatten(); + let is_background_follow_up = inner.background_follow_ups.remove(key); + let is_background_source = inner.background_source_children.remove(key); + let follow_up_already_registered = inner.early_background_follow_ups.remove(key); + + if !is_background_follow_up && is_background_source && !follow_up_already_registered { + if let Some(parent) = parent.as_ref() { + remember_completed_background_source(inner, key, parent); + } + } + + prune_completed_branch(inner, key); + if let Some(root) = root_for(inner, key).or_else(|| parent.clone()) { + prune_idle_tree(inner, &root); + } + let owned = inner.active.clone(); + inner + .confirmations + .retain(|_, confirmation_key| owned.contains(confirmation_key)); +} + +fn bind_background_source_task(inner: &mut PeerTurnTrackerInner, call_key: &(PeerTurnKey, String)) { + let Some(call) = inner.background_task_calls.get(call_key) else { + return; + }; + let (Some(background_task_id), Some(source_child)) = + (call.background_task_id.as_ref(), call.source_child.as_ref()) + else { + return; + }; + inner.background_source_tasks.insert( + background_task_id.clone(), + (call_key.0.clone(), source_child.clone()), + ); + inner.background_task_calls.remove(call_key); +} + +fn merge_completed_background_subagents( + inner: &PeerTurnTrackerInner, + drain: &mut PeerTurnDrain, + session_id: Option<&str>, +) { + let mut links = drain.background_subagents.drain(..).collect::>(); + links.extend( + inner + .completed_background_sources + .iter() + .filter(|(source, parent)| { + source.session_id != parent.session_id + && session_id.is_none_or(|session_id| { + source.session_id == session_id || parent.session_id == session_id + }) + }) + .map(|(source, parent)| PeerBackgroundSubagentLink { + parent_session_id: parent.session_id.clone(), + subagent_session_id: source.session_id.clone(), + }), + ); + drain.background_subagents = links.into_iter().collect(); +} + +fn peer_turn_drain_for_keys( + inner: &PeerTurnTrackerInner, + keys: &HashSet, +) -> PeerTurnDrain { + let turns = keys + .iter() + .filter(|key| inner.active.contains(*key)) + .cloned() + .collect(); + let background_subagents = keys + .iter() + .filter_map(|key| { + let parent = inner.parents.get(key)?.as_ref()?; + (parent.session_id != key.session_id).then(|| PeerBackgroundSubagentLink { + parent_session_id: parent.session_id.clone(), + subagent_session_id: key.session_id.clone(), + }) + }) + .collect::>() + .into_iter() + .collect(); + PeerTurnDrain { + turns, + background_subagents, + } +} + +fn session_tree_keys(inner: &PeerTurnTrackerInner, session_id: &str) -> HashSet { + let mut keys = inner + .parents + .iter() + .filter(|(key, parent)| { + key.session_id == session_id + || parent + .as_ref() + .is_some_and(|parent| parent.session_id == session_id) + }) + .map(|(key, _)| key.clone()) + .collect::>(); + loop { + let descendants = inner + .parents + .iter() + .filter_map(|(key, parent)| { + parent + .as_ref() + .filter(|parent| keys.contains(*parent)) + .map(|_| key.clone()) + }) + .filter(|key| !keys.contains(key)) + .collect::>(); + if descendants.is_empty() { + break; + } + keys.extend(descendants); + } + keys +} + +fn root_for(inner: &PeerTurnTrackerInner, key: &PeerTurnKey) -> Option { + let mut current = key.clone(); + let mut remaining = inner.parents.len().saturating_add(1); + while remaining > 0 { + remaining -= 1; + match inner.parents.get(¤t) { + Some(Some(parent)) => current = parent.clone(), + Some(None) => return Some(current), + None => return None, + } + } + None +} + +fn tree_keys(inner: &PeerTurnTrackerInner, root: &PeerTurnKey) -> HashSet { + inner + .parents + .keys() + .filter(|key| root_for(inner, key).as_ref() == Some(root)) + .cloned() + .collect() +} + +fn remember_completed_background_source( + inner: &mut PeerTurnTrackerInner, + source: &PeerTurnKey, + parent: &PeerTurnKey, +) { + if inner.completed_background_sources.contains_key(source) { + return; + } + inner + .completed_background_sources + .insert(source.clone(), parent.clone()); +} + +fn take_completed_background_source(inner: &mut PeerTurnTrackerInner, source: &PeerTurnKey) { + inner.completed_background_sources.remove(source); + inner + .background_source_tasks + .retain(|_, (_, mapped_source)| mapped_source != source); +} + +fn remove_completed_background_sources_for_session( + inner: &mut PeerTurnTrackerInner, + session_id: &str, +) { + inner.completed_background_sources.retain(|source, parent| { + source.session_id != session_id && parent.session_id != session_id + }); + inner.background_source_tasks.retain(|_, (parent, source)| { + source.session_id != session_id && parent.session_id != session_id + }); +} + +fn release_background_sources_for_subagent( + inner: &mut PeerTurnTrackerInner, + parent_session_id: &str, + subagent_session_id: &str, +) { + let tracked_sources = inner + .parents + .iter() + .filter_map(|(source, parent)| { + parent.as_ref().filter(|parent| { + parent.session_id == parent_session_id && source.session_id == subagent_session_id + })?; + (inner.background_source_children.contains(source) + || inner.early_background_follow_ups.contains(source)) + .then(|| source.clone()) + }) + .collect::>(); + let completed_sources = inner + .completed_background_sources + .iter() + .filter(|(source, parent)| { + parent.session_id == parent_session_id && source.session_id == subagent_session_id + }) + .map(|(source, _)| source.clone()) + .collect::>(); + + for source in &tracked_sources { + inner.background_source_children.remove(source); + inner.early_background_follow_ups.remove(source); + } + for source in &completed_sources { + take_completed_background_source(inner, source); + } + for source in tracked_sources { + prune_completed_branch(inner, &source); + } + inner.background_source_tasks.retain(|_, (parent, source)| { + parent.session_id != parent_session_id || source.session_id != subagent_session_id + }); +} + +fn background_authorization_len(inner: &PeerTurnTrackerInner) -> usize { + inner + .background_task_calls + .values() + .filter(|call| call.source_child.is_none()) + .count() + .saturating_add(inner.background_source_children.len()) + .saturating_add(inner.completed_background_sources.len()) +} + +fn remove_tracked_turns(inner: &mut PeerTurnTrackerInner, removed: &HashSet) { + inner.parents.retain(|key, _| !removed.contains(key)); + inner.active.retain(|key| !removed.contains(key)); + inner.started.retain(|key| !removed.contains(key)); + inner + .terminal_deliveries + .retain(|key| !removed.contains(key)); + inner.background_task_calls.retain(|(parent, _), call| { + !removed.contains(parent) + && call + .source_child + .as_ref() + .is_none_or(|source| !removed.contains(source)) + }); + inner + .background_task_cancellations + .retain(|(parent, _), _| !removed.contains(parent)); + inner + .background_source_children + .retain(|key| !removed.contains(key)); + let completed_background_sources = &inner.completed_background_sources; + inner.background_source_tasks.retain(|_, (parent, source)| { + completed_background_sources.get(source) == Some(parent) + || (!removed.contains(parent) && !removed.contains(source)) + }); + inner + .background_follow_ups + .retain(|key| !removed.contains(key)); + inner + .early_background_follow_ups + .retain(|key| !removed.contains(key)); + inner.confirmations.retain(|_, key| !removed.contains(key)); +} + +fn prune_completed_branch(inner: &mut PeerTurnTrackerInner, key: &PeerTurnKey) { + let branch = inner + .parents + .keys() + .filter(|candidate| { + let mut current = (*candidate).clone(); + let mut remaining = inner.parents.len().saturating_add(1); + while remaining > 0 { + remaining -= 1; + if ¤t == key { + return true; + } + match inner.parents.get(¤t).and_then(Clone::clone) { + Some(parent) => current = parent, + None => return false, + } + } + false + }) + .cloned() + .collect::>(); + if branch.iter().any(|key| inner.active.contains(key)) + || branch + .iter() + .any(|key| inner.early_background_follow_ups.contains(key)) + { + return; + } + remove_tracked_turns(inner, &branch); +} + +fn prune_idle_tree(inner: &mut PeerTurnTrackerInner, key: &PeerTurnKey) { + let Some(root) = root_for(inner, key) else { + return; + }; + let tree = tree_keys(inner, &root); + if tree.iter().any(|key| inner.active.contains(key)) + || tree + .iter() + .any(|key| inner.early_background_follow_ups.contains(key)) + { + return; + } + remove_tracked_turns(inner, &tree); +} #[derive(Clone)] pub(crate) struct PeerHostState { - pub(crate) coordinator: Arc, - pub(crate) scheduler: Arc, - pub(crate) event_queue: Arc, + pub(crate) agent_runtime: AgentRuntime, + pub(crate) compatibility: CoreAgentRuntimeCompatibility, + pub(crate) agent_events: CliAgentEventSource, + pub(crate) turns: PeerTurnTracker, pub(crate) workspace_service: Arc, pub(crate) filesystem_service: Arc, - pub(crate) persistence: Arc, +} + +impl PeerHostState { + pub(crate) async fn cancel_and_drain_peer_turns( + &self, + reason: &'static str, + ) -> Result<(), String> { + let initial = self.turns.peer_turns_for_cancellation(); + let initial_result = self.cancel_peer_turns(initial, reason).await; + tokio::task::yield_now().await; + let raced = self.turns.drain_peer_turns(); + let raced_result = self.cancel_peer_turns(raced, reason).await; + aggregate_cancellation_results(initial_result, raced_result) + } + + pub(crate) async fn cancel_peer_turns( + &self, + drain: PeerTurnDrain, + reason: &'static str, + ) -> Result<(), String> { + const MAX_CONCURRENT_CANCELLATIONS: usize = 32; + + let mut failure_count = 0usize; + let mut pending_background = drain.background_subagents.into_iter(); + let mut background_tasks = tokio::task::JoinSet::new(); + for _ in 0..MAX_CONCURRENT_CANCELLATIONS { + let Some(link) = pending_background.next() else { + break; + }; + spawn_background_subagent_cancellation( + &mut background_tasks, + self.compatibility.clone(), + link, + ); + } + while let Some(joined) = background_tasks.join_next().await { + match joined { + Ok((link, Err(error))) => { + failure_count += 1; + tracing::warn!( + "Failed to cancel Peer-owned background subagent: parent_session_id={}, subagent_session_id={}, reason={}, error={}", + link.parent_session_id, + link.subagent_session_id, + reason, + error + ); + } + Err(error) => { + failure_count += 1; + tracing::warn!( + "Peer-owned background subagent cancellation task failed: reason={}, error={}", + reason, + error + ); + } + Ok((_, Ok(_))) => {} + } + if let Some(link) = pending_background.next() { + spawn_background_subagent_cancellation( + &mut background_tasks, + self.compatibility.clone(), + link, + ); + } + } + + let mut pending = drain.turns.into_iter(); + let mut tasks = tokio::task::JoinSet::new(); + for _ in 0..MAX_CONCURRENT_CANCELLATIONS { + let Some(turn) = pending.next() else { + break; + }; + spawn_turn_cancellation(&mut tasks, self.agent_runtime.clone(), turn, reason); + } + + while let Some(joined) = tasks.join_next().await { + match joined { + Ok((turn, Err(error))) => { + failure_count += 1; + tracing::warn!( + "Failed to cancel Peer-owned turn: session_id={}, turn_id={}, reason={}, error={}", + turn.session_id, + turn.turn_id, + reason, + error + ); + } + Err(error) => { + failure_count += 1; + tracing::warn!( + "Peer-owned turn cancellation task failed: reason={}, error={}", + reason, + error + ); + } + Ok((_, Ok(_))) => {} + } + if let Some(turn) = pending.next() { + spawn_turn_cancellation(&mut tasks, self.agent_runtime.clone(), turn, reason); + } + } + if failure_count == 0 { + Ok(()) + } else { + Err(format!( + "{failure_count} Peer-owned cancellation operation(s) failed" + )) + } + } +} + +fn aggregate_cancellation_results( + initial: Result<(), String>, + raced: Result<(), String>, +) -> Result<(), String> { + let failures = [initial, raced] + .into_iter() + .filter_map(Result::err) + .collect::>(); + if failures.is_empty() { + Ok(()) + } else { + Err(failures.join("; ")) + } +} + +fn spawn_background_subagent_cancellation( + tasks: &mut tokio::task::JoinSet<(PeerBackgroundSubagentLink, Result<(), String>)>, + compatibility: CoreAgentRuntimeCompatibility, + link: PeerBackgroundSubagentLink, +) { + tasks.spawn(async move { + let result = compatibility + .cancel_background_subagents_for_parent( + &link.parent_session_id, + &link.subagent_session_id, + ) + .await + .map(|_| ()) + .map_err(|error| error.to_string()); + (link, result) + }); +} + +fn spawn_turn_cancellation( + tasks: &mut tokio::task::JoinSet<(PeerTurnKey, Result<(), String>)>, + runtime: AgentRuntime, + turn: PeerTurnKey, + reason: &'static str, +) { + tasks.spawn(async move { + let result = runtime + .cancel_turn(AgentTurnCancellationRequest { + session_id: turn.session_id.clone(), + turn_id: Some(turn.turn_id.clone()), + source: Some(AgentSubmissionSource::Cli), + requester_session_id: None, + reason: Some(reason.to_string()), + wait_timeout_ms: Some(1_500), + }) + .await; + (turn, result.map(|_| ()).map_err(|error| error.to_string())) + }); +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use super::{aggregate_cancellation_results, PeerTurnKey, PeerTurnTracker}; + + fn register_background_child( + tracker: &PeerTurnTracker, + parent: &PeerTurnKey, + child: PeerTurnKey, + ) { + let tool_call_id = format!("task-{}", child.turn_id); + assert!(tracker + .record_background_task_call(parent, tool_call_id.clone()) + .expect("record background Task call")); + assert!(tracker + .register_linked_child(parent, child, &tool_call_id) + .expect("register background child")); + } + + #[test] + fn detach_reports_any_unconfirmed_cancellation_round() { + assert!(aggregate_cancellation_results(Ok(()), Ok(())).is_ok()); + + let error = + aggregate_cancellation_results(Err("initial cancellation failed".to_string()), Ok(())) + .expect_err("detach must not hide an unconfirmed cancellation"); + assert!(error.contains("initial cancellation failed"), "{error}"); + + let error = + aggregate_cancellation_results(Ok(()), Err("raced cancellation failed".to_string())) + .expect_err("detach must report a raced cancellation failure"); + assert!(error.contains("raced cancellation failed"), "{error}"); + } + + #[test] + fn tracker_rejects_turns_until_event_stream_is_ready_and_after_close() { + let tracker = PeerTurnTracker::new(); + let turn = PeerTurnKey::new("session-1", "turn-1"); + assert!(tracker.register_root(turn.clone()).is_err()); + + tracker.mark_event_stream_ready(); + tracker.register_root(turn.clone()).expect("register turn"); + assert_eq!( + tracker.interrupt_event_stream(true).turns, + vec![turn.clone()] + ); + assert!(tracker.register_root(turn).is_err()); + } + + #[test] + fn finishing_a_root_preserves_an_active_child_and_its_confirmation() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session-1", "turn-1"); + let child = PeerTurnKey::new("session-2", "turn-2"); + tracker.register_root(root.clone()).expect("register root"); + assert!(tracker + .register_child(&root, child.clone()) + .expect("register child")); + tracker + .record_confirmation(&child, "tool-1".to_string()) + .expect("record confirmation"); + + tracker.finish_turn(&root); + + assert!(!tracker.owns("session-1", Some("turn-1"))); + assert!(tracker.owns("session-2", Some("turn-2"))); + assert_eq!(tracker.claim_confirmation("tool-1"), Some(child)); + } + + #[test] + fn confirmation_claim_is_bound_to_the_exact_turn_and_can_be_restored() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let turn = PeerTurnKey::new("session-1", "turn-1"); + tracker.register_root(turn.clone()).expect("register turn"); + tracker + .record_confirmation(&turn, "tool-1".to_string()) + .expect("record confirmation"); + + let claimed = tracker + .claim_confirmation("tool-1") + .expect("claim confirmation"); + assert_eq!(claimed, turn); + assert!(tracker.claim_confirmation("tool-1").is_none()); + + tracker.restore_confirmation("tool-1".to_string(), claimed); + assert_eq!(tracker.claim_confirmation("tool-1"), Some(turn)); + } + + #[test] + fn draining_a_parent_session_after_root_completion_returns_the_active_child_only() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session-1", "turn-1"); + let child = PeerTurnKey::new("session-2", "turn-2"); + let other = PeerTurnKey::new("session-3", "turn-3"); + tracker.register_root(root.clone()).expect("register root"); + tracker + .register_child(&root, child.clone()) + .expect("register child"); + tracker + .register_root(other.clone()) + .expect("register other root"); + + tracker.finish_turn(&root); + let drained = tracker + .drain_session_turns("session-1") + .turns + .into_iter() + .collect::>(); + + assert_eq!(drained, HashSet::from([child])); + assert!(!tracker.owns("session-1", Some("turn-1"))); + assert!(!tracker.owns("session-2", Some("turn-2"))); + assert!(tracker.owns("session-3", Some("turn-3"))); + } + + #[test] + fn background_result_follow_up_inherits_peer_ownership() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session-1", "turn-1"); + let child = PeerTurnKey::new("session-2", "turn-2"); + let follow_up = PeerTurnKey::new("session-1", "turn-3"); + tracker.register_root(root.clone()).expect("register root"); + register_background_child(&tracker, &root, child.clone()); + + tracker.finish_turn(&root); + tracker.finish_turn(&child); + let pending = tracker.peer_turns_for_cancellation(); + assert!(pending.turns.is_empty()); + assert_eq!( + pending + .background_subagents + .into_iter() + .collect::>(), + HashSet::from([super::PeerBackgroundSubagentLink { + parent_session_id: "session-1".to_string(), + subagent_session_id: "session-2".to_string(), + }]) + ); + assert!(!tracker + .register_background_follow_up( + &PeerTurnKey::new("session-1", "local-turn"), + &child, + PeerTurnKey::new("session-1", "local-follow-up") + ) + .expect("reject another owner's follow-up")); + assert!(tracker + .register_background_follow_up(&root, &child, follow_up.clone()) + .expect("register follow-up")); + assert!(tracker.owns("session-1", Some("turn-3"))); + + tracker.finish_turn(&follow_up); + assert!(!tracker.owns("session-1", None)); + assert!(!tracker + .register_background_follow_up(&root, &child, PeerTurnKey::new("session-1", "turn-4")) + .expect("reject unrelated follow-up")); + } + + #[test] + fn background_follow_up_can_start_before_terminal_fanout_finishes() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session-1", "turn-1"); + let child = PeerTurnKey::new("session-2", "turn-2"); + let follow_up = PeerTurnKey::new("session-1", "turn-3"); + tracker.register_root(root.clone()).expect("register root"); + register_background_child(&tracker, &root, child.clone()); + + assert!(tracker + .register_background_follow_up(&root, &child, follow_up.clone()) + .expect("register exact follow-up while terminal events are queued")); + assert!(tracker.owns("session-1", Some("turn-3"))); + + tracker.finish_turn(&root); + tracker.finish_turn(&PeerTurnKey::new("session-2", "turn-2")); + tracker.finish_turn(&follow_up); + assert!(!tracker.owns("session-1", None)); + assert!(!tracker + .register_background_follow_up( + &root, + &child, + PeerTurnKey::new("session-1", "unrelated-follow-up") + ) + .expect("reject a follow-up after the exact tree is complete")); + } + + #[test] + fn completed_source_child_can_register_its_delayed_exact_follow_up() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session-1", "turn-1"); + let child = PeerTurnKey::new("session-2", "turn-2"); + let follow_up = PeerTurnKey::new("session-1", "turn-3"); + tracker.register_root(root.clone()).expect("register root"); + register_background_child(&tracker, &root, child.clone()); + + tracker.finish_turn(&child); + tracker.finish_turn(&root); + assert!(tracker + .register_background_follow_up(&root, &child, follow_up.clone()) + .expect("register delayed exact follow-up")); + assert!(tracker.owns("session-1", Some("turn-3"))); + + tracker.finish_turn(&follow_up); + assert!(!tracker.owns("session-1", None)); + assert!(!tracker + .register_background_follow_up( + &root, + &child, + PeerTurnKey::new("session-1", "unrelated-follow-up") + ) + .expect("consumed tombstone must not authorize another follow-up")); + } + + #[test] + fn tombstone_only_lineage_is_included_in_cancellation_drain() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session-1", "turn-1"); + let child = PeerTurnKey::new("session-2", "turn-2"); + tracker.register_root(root.clone()).expect("register root"); + register_background_child(&tracker, &root, child.clone()); + + tracker.finish_turn(&child); + tracker.finish_turn(&root); + + let pending = tracker.peer_turns_for_cancellation(); + assert!(pending.turns.is_empty()); + assert_eq!( + pending + .background_subagents + .into_iter() + .collect::>(), + HashSet::from([super::PeerBackgroundSubagentLink { + parent_session_id: "session-1".to_string(), + subagent_session_id: "session-2".to_string(), + }]) + ); + assert_eq!( + tracker + .drain_peer_turns() + .background_subagents + .into_iter() + .collect::>(), + HashSet::from([super::PeerBackgroundSubagentLink { + parent_session_id: "session-1".to_string(), + subagent_session_id: "session-2".to_string(), + }]) + ); + assert!(!tracker + .register_background_follow_up( + &root, + &child, + PeerTurnKey::new("session-1", "late-follow-up") + ) + .expect("drained tombstone must not authorize a late follow-up")); + } + + #[test] + fn completed_background_authorizations_do_not_reduce_live_turn_capacity() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + for index in 0..super::MAX_BACKGROUND_PEER_AUTHORIZATIONS { + let root = PeerTurnKey::new(format!("parent-{index}"), format!("root-{index}")); + let child = PeerTurnKey::new(format!("child-{index}"), format!("child-turn-{index}")); + tracker.register_root(root.clone()).expect("register root"); + register_background_child(&tracker, &root, child.clone()); + tracker.finish_turn(&child); + tracker.finish_turn(&root); + } + + let mut live_roots = Vec::new(); + for index in 0..super::MAX_TRACKED_PEER_TURNS { + let root = PeerTurnKey::new( + format!("live-session-{index}"), + format!("live-turn-{index}"), + ); + tracker + .register_root(root.clone()) + .expect("background history must not reduce live capacity"); + live_roots.push(root); + } + assert!(tracker + .record_background_task_call(&live_roots[0], "overflow-task".to_string()) + .expect_err("the next background authorization must fail closed") + .contains("background authorization capacity")); + } + + #[test] + fn active_capacity_rejects_a_new_root_without_requesting_a_reset() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + for index in 0..super::MAX_TRACKED_PEER_TURNS { + tracker + .register_root(PeerTurnKey::new( + format!("session-{index}"), + format!("turn-{index}"), + )) + .expect("register active root"); + } + + assert!(tracker + .register_root(PeerTurnKey::new("overflow-session", "overflow-turn")) + .expect_err("live capacity must reject one more root") + .contains("capacity is exhausted")); + } + + #[test] + fn unlinked_background_and_cancellation_task_markers_are_bounded() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session", "turn"); + tracker.register_root(root.clone()).expect("register root"); + for index in 0..super::MAX_BACKGROUND_PEER_AUTHORIZATIONS { + tracker + .record_background_task_call(&root, format!("background-{index}")) + .expect("reserve background authorization"); + } + assert!(tracker + .record_background_task_call(&root, "background-overflow".to_string()) + .expect_err("unlinked background calls must be bounded") + .contains("background authorization capacity")); + + for index in 0..super::MAX_PENDING_PEER_TASK_CANCELLATIONS { + tracker + .record_background_task_cancellation( + &root, + format!("cancel-{index}"), + format!("subagent-{index}"), + ) + .expect("track Task cancellation"); + } + assert!(tracker + .record_background_task_cancellation( + &root, + "cancel-overflow".to_string(), + "subagent-overflow".to_string(), + ) + .expect_err("Task cancellation markers must be bounded") + .contains("cancellation tracking capacity")); + } + + #[test] + fn foreground_children_do_not_consume_background_lineage_capacity() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + for index in 0..(super::MAX_TRACKED_PEER_TURNS * 2) { + let root = PeerTurnKey::new(format!("parent-{index}"), format!("root-{index}")); + let child = PeerTurnKey::new(format!("child-{index}"), format!("child-turn-{index}")); + tracker.register_root(root.clone()).expect("register root"); + tracker + .register_child(&root, child.clone()) + .expect("register foreground child"); + tracker.finish_turn(&child); + tracker.finish_turn(&root); + } + + tracker + .register_root(PeerTurnKey::new("final-parent", "final-root")) + .expect("foreground Task history must not exhaust lineage capacity"); + } + + #[test] + fn successful_task_cancel_releases_the_exact_background_source() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("parent-session", "root-turn"); + let child = PeerTurnKey::new("subagent-session", "child-turn"); + tracker.register_root(root.clone()).expect("register root"); + register_background_child(&tracker, &root, child.clone()); + tracker.finish_turn(&child); + tracker.finish_turn(&root); + + let cancel_turn = PeerTurnKey::new("parent-session", "cancel-turn"); + tracker + .register_root(cancel_turn.clone()) + .expect("register cancel turn"); + assert!(tracker + .record_background_task_cancellation( + &cancel_turn, + "cancel-tool".to_string(), + "subagent-session".to_string(), + ) + .expect("record Task cancellation")); + tracker.finish_task_call(&cancel_turn, "cancel-tool", None, Some(1)); + + assert!(!tracker + .register_background_follow_up( + &root, + &child, + PeerTurnKey::new("parent-session", "late-follow-up"), + ) + .expect("cancelled background source must not authorize a follow-up")); + } + + #[test] + fn zero_count_task_cancel_keeps_a_completed_background_source() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("parent-session", "root-turn"); + let child = PeerTurnKey::new("subagent-session", "child-turn"); + tracker.register_root(root.clone()).expect("register root"); + register_background_child(&tracker, &root, child.clone()); + tracker.finish_turn(&child); + tracker.finish_turn(&root); + + let cancel_turn = PeerTurnKey::new("parent-session", "cancel-turn"); + tracker + .register_root(cancel_turn.clone()) + .expect("register cancel turn"); + tracker + .record_background_task_cancellation( + &cancel_turn, + "cancel-tool".to_string(), + "subagent-session".to_string(), + ) + .expect("record Task cancellation"); + tracker.finish_task_call(&cancel_turn, "cancel-tool", None, Some(0)); + + assert!(tracker + .register_background_follow_up( + &root, + &child, + PeerTurnKey::new("parent-session", "follow-up"), + ) + .expect("zero-count cancellation must not consume the source")); + } + + #[test] + fn one_tombstone_follow_up_does_not_consume_a_sibling_authorization() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session-1", "turn-1"); + let child_a = PeerTurnKey::new("session-2", "turn-2"); + let child_b = PeerTurnKey::new("session-3", "turn-3"); + let follow_up_a = PeerTurnKey::new("session-1", "turn-4"); + let follow_up_b = PeerTurnKey::new("session-1", "turn-5"); + tracker.register_root(root.clone()).expect("register root"); + register_background_child(&tracker, &root, child_a.clone()); + register_background_child(&tracker, &root, child_b.clone()); + + tracker.finish_turn(&child_a); + tracker.finish_turn(&root); + tracker.finish_turn(&child_b); + assert!(tracker + .register_background_follow_up(&root, &child_a, follow_up_a.clone()) + .expect("register child A tombstone follow-up")); + assert!(tracker + .register_background_follow_up(&root, &child_b, follow_up_b.clone()) + .expect("register child B follow-up")); + + tracker.finish_turn(&follow_up_a); + tracker.finish_turn(&follow_up_b); + assert!(!tracker.owns("session-1", None)); + } + + #[test] + fn draining_a_child_session_releases_its_early_follow_up_reservation() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session-1", "turn-1"); + let child = PeerTurnKey::new("session-2", "turn-2"); + let follow_up = PeerTurnKey::new("session-1", "turn-3"); + tracker.register_root(root.clone()).expect("register root"); + register_background_child(&tracker, &root, child.clone()); + assert!(tracker + .register_background_follow_up(&root, &child, follow_up.clone()) + .expect("register exact early follow-up")); + + assert_eq!( + tracker.drain_session_turns("session-2").turns, + vec![child.clone()] + ); + tracker.finish_turn(&root); + tracker.finish_turn(&follow_up); + + assert!(!tracker + .register_background_follow_up( + &root, + &child, + PeerTurnKey::new("session-1", "unrelated-follow-up") + ) + .expect("completed lineage must be pruned")); + } + + #[test] + fn draining_a_sibling_child_does_not_release_another_childs_reservation() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session-1", "turn-1"); + let child_a = PeerTurnKey::new("session-2", "turn-2"); + let child_b = PeerTurnKey::new("session-3", "turn-3"); + let follow_up = PeerTurnKey::new("session-1", "turn-4"); + tracker.register_root(root.clone()).expect("register root"); + register_background_child(&tracker, &root, child_a.clone()); + tracker + .register_child(&root, child_b.clone()) + .expect("register child B"); + assert!(tracker + .register_background_follow_up(&root, &child_a, follow_up.clone()) + .expect("register child A follow-up")); + + assert_eq!( + tracker.drain_session_turns("session-3").turns, + vec![child_b] + ); + tracker.finish_turn(&root); + tracker.finish_turn(&child_a); + tracker.finish_turn(&follow_up); + + assert!(!tracker + .register_background_follow_up( + &root, + &child_a, + PeerTurnKey::new("session-1", "unrelated-follow-up") + ) + .expect("completed lineage must be pruned")); + } + + #[test] + fn unrelated_running_turn_does_not_consume_background_follow_up_authorization() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session-1", "turn-1"); + let child = PeerTurnKey::new("session-2", "turn-2"); + tracker.register_root(root.clone()).expect("register root"); + register_background_child(&tracker, &root, child.clone()); + + tracker.finish_turn(&root); + tracker.finish_turn(&child); + + let follow_up = PeerTurnKey::new("session-1", "turn-3"); + assert!(tracker + .register_background_follow_up(&root, &child, follow_up.clone()) + .expect("register queued follow-up")); + assert!(tracker.owns("session-1", Some("turn-3"))); + } + + #[test] + fn event_stream_interruption_returns_every_active_turn_once() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("session-1", "turn-1"); + let child = PeerTurnKey::new("session-2", "turn-2"); + tracker.register_root(root.clone()).expect("register root"); + tracker + .register_child(&root, child.clone()) + .expect("register child"); + + assert_eq!( + tracker + .interrupt_event_stream(false) + .turns + .into_iter() + .collect::>(), + HashSet::from([root, child]) + ); + assert!(tracker.interrupt_event_stream(false).turns.is_empty()); + } + + #[test] + fn completion_keeps_the_stream_generation_but_interruption_invalidates_it() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let turn = PeerTurnKey::new("session-1", "turn-1"); + let generation = tracker.register_root(turn.clone()).expect("register root"); + + tracker.finish_turn(&turn); + assert!(tracker.is_event_stream_generation_current(generation)); + + tracker.interrupt_event_stream(false); + tracker.mark_event_stream_ready(); + assert!(!tracker.is_event_stream_generation_current(generation)); + } + + #[test] + fn duplicate_root_ids_are_rejected() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let turn = PeerTurnKey::new("session-1", "turn-1"); + tracker.register_root(turn.clone()).expect("register root"); + + assert!(tracker.register_root(turn).is_err()); + } + + #[test] + fn terminal_delivery_is_claimed_once_and_cleared_with_the_turn() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let turn = PeerTurnKey::new("session-1", "turn-1"); + tracker.register_root(turn.clone()).expect("register root"); + + assert!(tracker + .claim_terminal_delivery(&turn) + .expect("claim terminal")); + assert!(!tracker + .claim_terminal_delivery(&turn) + .expect("reject duplicate terminal")); + tracker.finish_turn(&turn); + assert!(!tracker + .claim_terminal_delivery(&turn) + .expect("finished turn is not claimable")); + } + + #[test] + fn interrupted_turn_id_remains_quarantined_after_late_terminals() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let interrupted = PeerTurnKey::new("session-1", "turn-1"); + tracker + .register_root(interrupted.clone()) + .expect("register root"); + + tracker.interrupt_event_stream(false); + tracker.mark_event_stream_ready(); + assert!(tracker.register_root(interrupted.clone()).is_err()); + assert!(tracker.is_interrupted_terminal(&interrupted)); + assert!(tracker.register_root(interrupted).is_err()); + } + + #[test] + fn completed_background_source_is_released_by_its_exact_injection() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let parent = PeerTurnKey::new("parent-session", "parent-turn"); + let source = PeerTurnKey::new("subagent-session", "subagent-turn"); + tracker + .register_root(parent.clone()) + .expect("register parent"); + tracker + .record_background_task_call(&parent, "task-tool".to_string()) + .expect("record background Task"); + tracker + .register_linked_child(&parent, source.clone(), "task-tool") + .expect("link source"); + tracker.finish_task_call(&parent, "task-tool", Some("background-task"), None); + tracker.finish_turn(&source); + + assert!(tracker.finish_background_injection(&parent, "background-task")); + assert!(!tracker.finish_background_injection(&parent, "background-task")); + assert!(!tracker + .register_background_follow_up( + &parent, + &source, + PeerTurnKey::new("parent-session", "late-follow-up"), + ) + .unwrap_or(false)); + } + + #[test] + fn injected_background_results_do_not_accumulate_authorization() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let parent = PeerTurnKey::new("parent-session", "parent-turn"); + tracker + .register_root(parent.clone()) + .expect("register parent"); + + for index in 0..(super::MAX_BACKGROUND_PEER_AUTHORIZATIONS * 2) { + let tool_call_id = format!("task-tool-{index}"); + let background_task_id = format!("background-task-{index}"); + let source = PeerTurnKey::new("subagent-session", format!("subagent-turn-{index}")); + tracker + .record_background_task_call(&parent, tool_call_id.clone()) + .expect("record background Task"); + tracker + .register_linked_child(&parent, source.clone(), &tool_call_id) + .expect("link source"); + tracker.finish_task_call(&parent, &tool_call_id, Some(&background_task_id), None); + tracker.finish_turn(&source); + assert!(tracker.finish_background_injection(&parent, &background_task_id)); + } + } + + #[test] + fn interrupted_turn_quarantine_does_not_reduce_live_turn_capacity() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + tracker + .register_root(PeerTurnKey::new("old-session", "old-turn")) + .expect("register interrupted root"); + tracker.interrupt_event_stream(false); + tracker.mark_event_stream_ready(); + + for index in 0..super::MAX_TRACKED_PEER_TURNS { + tracker + .register_root(PeerTurnKey::new( + format!("session-{index}"), + format!("turn-{index}"), + )) + .expect("quarantine must not reduce live capacity"); + } + } + + #[test] + fn interrupted_child_turn_id_cannot_be_reused() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let old_root = PeerTurnKey::new("old-parent", "old-root"); + let interrupted_child = PeerTurnKey::new("subagent", "shared-child-turn"); + tracker + .register_root(old_root.clone()) + .expect("register old root"); + tracker + .register_child(&old_root, interrupted_child.clone()) + .expect("register old child"); + tracker.interrupt_event_stream(false); + tracker.mark_event_stream_ready(); + + let new_root = PeerTurnKey::new("new-parent", "new-root"); + tracker + .register_root(new_root.clone()) + .expect("register new root"); + assert!(tracker + .register_child(&new_root, interrupted_child) + .is_err()); + } + + #[test] + fn interrupted_follow_up_turn_id_cannot_be_reused() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let interrupted_follow_up = PeerTurnKey::new("parent-session", "shared-follow-up"); + tracker + .register_root(interrupted_follow_up.clone()) + .expect("register interrupted turn"); + tracker.interrupt_event_stream(false); + tracker.mark_event_stream_ready(); + + let parent = PeerTurnKey::new("parent-session", "new-parent"); + let source = PeerTurnKey::new("subagent-session", "source-turn"); + tracker + .register_root(parent.clone()) + .expect("register new parent"); + register_background_child(&tracker, &parent, source.clone()); + assert!(tracker + .register_background_follow_up(&parent, &source, interrupted_follow_up) + .is_err()); + } + + #[test] + fn explicit_drains_quarantine_removed_turn_ids() { + let tracker = PeerTurnTracker::new(); + tracker.mark_event_stream_ready(); + let root = PeerTurnKey::new("parent-session", "root-turn"); + let child = PeerTurnKey::new("child-session", "child-turn"); + tracker.register_root(root.clone()).expect("register root"); + tracker + .register_child(&root, child.clone()) + .expect("register child"); + + tracker.drain_session_turns(&child.session_id); + assert!(tracker.register_child(&root, child).is_err()); + + tracker.drain_peer_turns(); + assert!(tracker.register_root(root).is_err()); + } } static PEER_HOST_STATE: OnceLock = OnceLock::new(); diff --git a/src/apps/cli/src/runtime/mod.rs b/src/apps/cli/src/runtime/mod.rs index d7ade23d93..0268bff8aa 100644 --- a/src/apps/cli/src/runtime/mod.rs +++ b/src/apps/cli/src/runtime/mod.rs @@ -51,7 +51,6 @@ impl CliProductRuntimeState { #[derive(Clone)] pub(crate) struct CliRuntimeContext { workspace_root: PathBuf, - agentic_system: AgenticSystem, agent_runtime: AgentRuntime, compatibility: CoreAgentRuntimeCompatibility, agent_events: CliAgentEventSource, @@ -101,6 +100,7 @@ impl CliRuntimeContext { .context("Failed to build CLI Agent Runtime SDK")?; let compatibility = CoreAgentRuntimeCompatibility::build( agentic_system.coordinator.clone(), + scheduler, agentic_system.token_usage_service.clone(), ); @@ -116,7 +116,6 @@ impl CliRuntimeContext { Ok(Self { workspace_root, agent_events: CliAgentEventSource::new(agentic_system.event_queue.clone()), - agentic_system, agent_runtime, compatibility, services, @@ -130,10 +129,6 @@ impl CliRuntimeContext { &self.workspace_root } - pub(crate) fn agentic_system(&self) -> &AgenticSystem { - &self.agentic_system - } - pub(crate) fn agent_runtime(&self) -> &AgentRuntime { &self.agent_runtime } diff --git a/src/apps/cli/tests/product_assembly_cli.rs b/src/apps/cli/tests/product_assembly_cli.rs index b153886a16..a99b8900b3 100644 --- a/src/apps/cli/tests/product_assembly_cli.rs +++ b/src/apps/cli/tests/product_assembly_cli.rs @@ -143,12 +143,20 @@ fn doctor_rejects_incomplete_e2e_storage_roots() { fn cli_local_persistence_stays_behind_core_compatibility_facade() { const ACCOUNT_SYNC: &str = include_str!("../src/account_sync.rs"); const STARTUP_PAGE: &str = include_str!("../src/ui/startup.rs"); + const PEER_BOOTSTRAP: &str = include_str!("../src/peer_host/bootstrap.rs"); + const PEER_STATE: &str = include_str!("../src/peer_host/state.rs"); + const PEER_SESSION_COMMANDS: &str = include_str!("../src/peer_host/commands/session.rs"); + const PEER_SNAPSHOT_COMMANDS: &str = include_str!("../src/peer_host/commands/snapshot.rs"); const CORE_RUNTIME_SERVICES: &str = include_str!("../../../crates/assembly/core/src/product_runtime/runtime_services.rs"); for (path, source) in [ ("account_sync.rs", ACCOUNT_SYNC), ("ui/startup.rs", STARTUP_PAGE), + ("peer_host/bootstrap.rs", PEER_BOOTSTRAP), + ("peer_host/state.rs", PEER_STATE), + ("peer_host/commands/session.rs", PEER_SESSION_COMMANDS), + ("peer_host/commands/snapshot.rs", PEER_SNAPSHOT_COMMANDS), ] { assert!( !source.contains("PersistenceManager"), @@ -168,4 +176,17 @@ fn cli_local_persistence_stays_behind_core_compatibility_facade() { !CORE_RUNTIME_SERVICES.contains("pub fn persistence_manager"), "runtime services provider must not expose a concrete persistence factory" ); + assert!( + !PEER_BOOTSTRAP.contains("DialogScheduler::new") + && !PEER_BOOTSTRAP.contains("get_global_scheduler"), + "Peer Host must consume the invocation-scoped scheduler instead of assembling one" + ); + assert!( + !PEER_STATE.contains("pub(crate) persistence") + && !PEER_SESSION_COMMANDS.contains("state.persistence") + && !PEER_SNAPSHOT_COMMANDS.contains("state.persistence") + && !PEER_SESSION_COMMANDS.contains("get_snapshot_manager_for_workspace") + && !PEER_SNAPSHOT_COMMANDS.contains("get_snapshot_manager_for_workspace"), + "Peer Host persistence operations must stay behind the Core compatibility facade" + ); } diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index 961ca7bd25..ae49241e77 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -139,7 +139,8 @@ fn should_require_tool_confirmation( policy: DialogSubmissionPolicy, user_message_metadata: Option<&serde_json::Value>, ) -> bool { - policy.requires_tool_confirmation() + (policy.requires_tool_confirmation() + || metadata_bool(user_message_metadata, "require_tool_confirmation") == Some(true)) && metadata_bool(user_message_metadata, "acp_transport") != Some(true) } @@ -297,6 +298,53 @@ fn format_background_subagent_display_text( } } +fn background_subagent_delivery_metadata( + background_task_id: &str, + parent: &SubagentParentInfo, + subagent_session_id: &str, + subagent_dialog_turn_id: &str, + require_tool_confirmation: bool, + agent_type: &str, + task_description: &str, +) -> serde_json::Map { + let mut metadata = serde_json::Map::from_iter([ + ("kind".to_string(), serde_json::json!("background_result")), + ("sourceKind".to_string(), serde_json::json!("subagent")), + ( + "backgroundTaskId".to_string(), + serde_json::json!(background_task_id), + ), + ( + "parentSessionId".to_string(), + serde_json::json!(parent.session_id.as_str()), + ), + ( + "parentDialogTurnId".to_string(), + serde_json::json!(parent.dialog_turn_id.as_str()), + ), + ( + "subagentSessionId".to_string(), + serde_json::json!(subagent_session_id), + ), + ( + "subagentDialogTurnId".to_string(), + serde_json::json!(subagent_dialog_turn_id), + ), + ("subagentType".to_string(), serde_json::json!(agent_type)), + ( + "taskDescription".to_string(), + serde_json::json!(task_description), + ), + ]); + if require_tool_confirmation { + metadata.insert( + "require_tool_confirmation".to_string(), + serde_json::Value::Bool(true), + ); + } + metadata +} + fn build_subagent_session_relationship( parent_info: Option<&SubagentParentInfo>, agent_type: &str, @@ -659,6 +707,9 @@ pub struct ConversationCoordinator { active_subagent_executions: Arc>, /// Background Task runs keyed by background_task_id. background_subagent_tasks: Arc>, + /// Cancelled deliveries consumed by the scheduler after it acquires the + /// parent session's operation lock. + background_subagent_delivery_suppressions: Arc>, /// Notifies DialogScheduler of turn outcomes; injected after construction scheduler_notify_tx: OnceLock>, /// Round-boundary user steering source (mid-turn user message injection); injected after construction @@ -1210,6 +1261,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet subagent_timeout_registry: Arc::new(RwLock::new(HashMap::new())), active_subagent_executions: Arc::new(DashMap::new()), background_subagent_tasks: Arc::new(DashMap::new()), + background_subagent_delivery_suppressions: Arc::new(DashMap::new()), scheduler_notify_tx: OnceLock::new(), round_injection_source: OnceLock::new(), active_turns_per_session: Arc::new(DashMap::new()), @@ -6542,6 +6594,41 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet suppress_delivery } + #[cfg(test)] + pub(crate) fn register_background_subagent_task_for_test( + &self, + background_task_id: &str, + parent_session_id: &str, + subagent_session_id: &str, + ) { + self.register_background_subagent_task( + background_task_id.to_string(), + parent_session_id.to_string(), + subagent_session_id.to_string(), + BackgroundSubagentCancelTarget::Direct(CancellationToken::new()), + ); + } + + #[cfg(test)] + pub(crate) fn has_background_subagent_task_for_test(&self, background_task_id: &str) -> bool { + self.background_subagent_tasks + .contains_key(background_task_id) + } + + #[cfg(test)] + pub(crate) fn suppress_background_subagent_task_for_test( + &self, + background_task_id: &str, + ) -> bool { + self.background_subagent_tasks + .remove_if(background_task_id, |task_id, control| { + control.suppress_delivery.store(true, Ordering::SeqCst); + self.mark_background_subagent_delivery_suppression(task_id.clone()); + true + }) + .is_some() + } + pub(crate) async fn cancel_background_subagents_for_parent( &self, parent_session_id: &str, @@ -6550,22 +6637,109 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet self.ensure_subagent_session_loaded_for_reuse(subagent_session_id, parent_session_id, None) .await?; - let controls: Vec<(String, BackgroundSubagentTaskControl)> = self + let controls = self.claim_background_subagent_controls(|control| { + control.parent_session_id == parent_session_id + && control.subagent_session_id == subagent_session_id + }); + let background_task_ids = controls + .iter() + .map(|(background_task_id, _)| background_task_id.clone()) + .collect::>(); + let cancelled = self.cancel_background_subagent_controls(controls).await?; + if let Some(scheduler) = get_global_scheduler() { + scheduler + .cancel_background_result_deliveries(parent_session_id, &background_task_ids) + .await + .map_err(BitFunError::tool)?; + } + Ok(cancelled) + } + + pub(crate) async fn cancel_background_subagents_for_parent_session( + &self, + parent_session_id: &str, + ) -> BitFunResult { + let controls = self.claim_background_subagent_controls(|control| { + control.parent_session_id == parent_session_id + }); + + self.cancel_background_subagent_controls(controls).await + } + + pub(crate) fn take_background_subagent_delivery_suppression( + &self, + background_task_id: &str, + ) -> bool { + self.background_subagent_delivery_suppressions + .remove(background_task_id) + .is_some() + } + + pub(crate) fn mark_background_subagent_delivery_suppression(&self, background_task_id: String) { + self.background_subagent_delivery_suppressions + .insert(background_task_id, ()); + } + + pub(crate) fn finish_background_subagent_delivery(&self, background_task_id: &str) { + self.background_subagent_tasks.remove(background_task_id); + self.background_subagent_delivery_suppressions + .remove(background_task_id); + } + + pub(crate) fn background_subagent_control_available_for_injection( + &self, + background_task_id: &str, + ) -> bool { + self.background_subagent_tasks + .get(background_task_id) + .is_some_and(|control| !control.suppress_delivery.load(Ordering::SeqCst)) + } + + pub(crate) fn claim_background_subagent_control_for_injection( + &self, + background_task_id: &str, + ) -> bool { + self.background_subagent_tasks + .remove_if(background_task_id, |_, control| { + !control.suppress_delivery.load(Ordering::SeqCst) + }) + .is_some() + } + + fn claim_background_subagent_controls( + &self, + matches: impl Fn(&BackgroundSubagentTaskControl) -> bool, + ) -> Vec<(String, BackgroundSubagentTaskControl)> { + let candidate_ids = self .background_subagent_tasks .iter() - .filter(|entry| { - entry.parent_session_id == parent_session_id - && entry.subagent_session_id == subagent_session_id + .filter(|entry| matches(entry.value())) + .map(|entry| entry.key().clone()) + .collect::>(); + candidate_ids + .into_iter() + .filter_map(|background_task_id| { + self.background_subagent_tasks + .remove_if(&background_task_id, |task_id, control| { + if !matches(control) { + return false; + } + control.suppress_delivery.store(true, Ordering::SeqCst); + self.mark_background_subagent_delivery_suppression(task_id.clone()); + true + }) }) - .map(|entry| (entry.key().clone(), entry.value().clone())) - .collect(); + .collect() + } + async fn cancel_background_subagent_controls( + &self, + controls: Vec<(String, BackgroundSubagentTaskControl)>, + ) -> BitFunResult { for (background_task_id, control) in &controls { - control.suppress_delivery.store(true, Ordering::SeqCst); - self.background_subagent_tasks.remove(background_task_id); debug!( "Cancelling background subagent task: background_task_id={}, parent_session_id={}, subagent_session_id={}", - background_task_id, parent_session_id, subagent_session_id + background_task_id, control.parent_session_id, control.subagent_session_id ); match &control.cancel_target { BackgroundSubagentCancelTarget::Scheduler(handle) => { @@ -6574,7 +6748,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet } else { warn!( "Cannot cancel scheduler-backed background subagent because scheduler is unavailable: background_task_id={}, subagent_session_id={}", - background_task_id, subagent_session_id + background_task_id, control.subagent_session_id ); } } @@ -6697,6 +6871,15 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet ) })? .to_string(); + let subagent_dialog_turn_id = request + .dialog_turn_id + .as_deref() + .ok_or_else(|| { + BitFunError::Validation( + "prepared hidden subagent request is missing dialog_turn_id".to_string(), + ) + })? + .to_string(); let agent_type = request.agent_type.clone(); let subagent_parent_info = match request.subagent_parent_info.clone() { Some(info) => info, @@ -6727,6 +6910,12 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet let parent_workspace_path = parent_session.config.workspace_path.clone(); let parent_remote_connection_id = parent_session.config.remote_connection_id.clone(); let parent_remote_ssh_host = parent_session.config.remote_ssh_host.clone(); + let parent_requires_tool_confirmation = get_global_scheduler().is_some_and(|scheduler| { + scheduler.active_turn_requires_tool_confirmation( + &subagent_parent_info.session_id, + &subagent_parent_info.dialog_turn_id, + ) + }); let background_task_id = format!("bg-subagent-{}", uuid::Uuid::new_v4()); let background_task_id_for_delivery = background_task_id.clone(); let task_description = request.user_input_text.clone(); @@ -6767,7 +6956,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet BackgroundSubagentCancelTarget::Scheduler(cancel_handle.clone()), ); let background_subagent_tasks = self.background_subagent_tasks.clone(); + let background_delivery_suppressions = + self.background_subagent_delivery_suppressions.clone(); let subagent_session_id_for_delivery = subagent_session_id.clone(); + let subagent_dialog_turn_id_for_delivery = subagent_dialog_turn_id.clone(); tokio::spawn(async move { let result = match parent_cancel_token.as_ref() { @@ -6789,8 +6981,9 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet } None => Self::await_hidden_subagent_receiver(receiver).await, }; - background_subagent_tasks.remove(&background_task_id_for_delivery); if suppress_delivery.load(Ordering::SeqCst) { + background_subagent_tasks.remove(&background_task_id_for_delivery); + background_delivery_suppressions.remove(&background_task_id_for_delivery); debug!( "Suppressing cancelled background subagent result delivery: background_task_id={}, parent_session_id={}", background_task_id_for_delivery, subagent_parent_info.session_id @@ -6815,32 +7008,23 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet ), }; - let mut metadata = serde_json::Map::new(); - metadata.insert( - "kind".to_string(), - serde_json::Value::String("background_result".to_string()), - ); - metadata.insert( - "sourceKind".to_string(), - serde_json::Value::String("subagent".to_string()), - ); - metadata.insert( - "backgroundTaskId".to_string(), - serde_json::Value::String(background_task_id_for_delivery.clone()), - ); - metadata.insert( - "subagentType".to_string(), - serde_json::Value::String(agent_type), - ); - metadata.insert( - "taskDescription".to_string(), - serde_json::Value::String(task_description), + let metadata = background_subagent_delivery_metadata( + &background_task_id_for_delivery, + &subagent_parent_info, + &subagent_session_id_for_delivery, + &subagent_dialog_turn_id_for_delivery, + parent_requires_tool_confirmation, + &agent_type, + &task_description, ); let runtime = match CoreServiceAgentRuntime::global_agent_runtime_with_lifecycle_delivery() { Ok(runtime) => runtime, Err(error) => { + background_subagent_tasks.remove(&background_task_id_for_delivery); + background_delivery_suppressions + .remove(&background_task_id_for_delivery); warn!( "Agent runtime lifecycle delivery is not available; background subagent result dropped: background_task_id={}, parent_session_id={}, error={}", background_task_id_for_delivery, @@ -6851,6 +7035,16 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet } }; + if suppress_delivery.load(Ordering::SeqCst) { + background_subagent_tasks.remove(&background_task_id_for_delivery); + background_delivery_suppressions.remove(&background_task_id_for_delivery); + debug!( + "Suppressing cancelled background subagent result delivery: background_task_id={}, parent_session_id={}", + background_task_id_for_delivery, subagent_parent_info.session_id + ); + return; + } + if let Err(error) = runtime .deliver_background_result(AgentBackgroundResultRequest { session_id: subagent_parent_info.session_id.clone(), @@ -6864,6 +7058,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet }) .await { + background_subagent_tasks.remove(&background_task_id_for_delivery); + background_delivery_suppressions.remove(&background_task_id_for_delivery); warn!( "Failed to deliver background subagent result through scheduler path: background_task_id={}, parent_session_id={}, error={}", background_task_id_for_delivery, @@ -6906,7 +7102,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet BackgroundSubagentCancelTarget::Direct(background_cancel_token), ); let background_subagent_tasks = self.background_subagent_tasks.clone(); + let background_delivery_suppressions = + self.background_subagent_delivery_suppressions.clone(); let subagent_session_id_for_delivery = subagent_session_id.clone(); + let subagent_dialog_turn_id_for_delivery = subagent_dialog_turn_id.clone(); tokio::spawn(async move { let (delivery_text, display_text) = match coordinator @@ -6933,8 +7132,9 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet ), }; cancel_bridge_handle.abort(); - background_subagent_tasks.remove(&background_task_id_for_delivery); if suppress_delivery.load(Ordering::SeqCst) { + background_subagent_tasks.remove(&background_task_id_for_delivery); + background_delivery_suppressions.remove(&background_task_id_for_delivery); debug!( "Suppressing cancelled background subagent result delivery: background_task_id={}, parent_session_id={}", background_task_id_for_delivery, subagent_parent_info.session_id @@ -6942,32 +7142,22 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet return; } - let mut metadata = serde_json::Map::new(); - metadata.insert( - "kind".to_string(), - serde_json::Value::String("background_result".to_string()), - ); - metadata.insert( - "sourceKind".to_string(), - serde_json::Value::String("subagent".to_string()), - ); - metadata.insert( - "backgroundTaskId".to_string(), - serde_json::Value::String(background_task_id_for_delivery.clone()), - ); - metadata.insert( - "subagentType".to_string(), - serde_json::Value::String(agent_type), - ); - metadata.insert( - "taskDescription".to_string(), - serde_json::Value::String(task_description), + let metadata = background_subagent_delivery_metadata( + &background_task_id_for_delivery, + &subagent_parent_info, + &subagent_session_id_for_delivery, + &subagent_dialog_turn_id_for_delivery, + parent_requires_tool_confirmation, + &agent_type, + &task_description, ); let runtime = match CoreServiceAgentRuntime::global_agent_runtime_with_lifecycle_delivery() { Ok(runtime) => runtime, Err(error) => { + background_subagent_tasks.remove(&background_task_id_for_delivery); + background_delivery_suppressions.remove(&background_task_id_for_delivery); warn!( "Agent runtime lifecycle delivery is not available; background subagent result dropped: background_task_id={}, parent_session_id={}, error={}", background_task_id_for_delivery, subagent_parent_info.session_id, error @@ -6976,6 +7166,16 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet } }; + if suppress_delivery.load(Ordering::SeqCst) { + background_subagent_tasks.remove(&background_task_id_for_delivery); + background_delivery_suppressions.remove(&background_task_id_for_delivery); + debug!( + "Suppressing cancelled background subagent result delivery: background_task_id={}, parent_session_id={}", + background_task_id_for_delivery, subagent_parent_info.session_id + ); + return; + } + if let Err(error) = runtime .deliver_background_result(AgentBackgroundResultRequest { session_id: subagent_parent_info.session_id.clone(), @@ -6989,6 +7189,8 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet }) .await { + background_subagent_tasks.remove(&background_task_id_for_delivery); + background_delivery_suppressions.remove(&background_task_id_for_delivery); warn!( "Failed to deliver background subagent result: background_task_id={}, parent_session_id={}, error={}", background_task_id_for_delivery, @@ -7786,10 +7988,11 @@ fn merge_prepended_messages_for_turn( #[cfg(test)] mod tests { use super::{ - merge_prepended_messages_for_turn, normalize_subagent_max_concurrency, - resolve_agent_session_create_created_by, resolve_agent_submission_turn_id, - should_require_tool_confirmation, turn_review_manifest_for_agent, - validate_background_subagent_delivery, ConversationCoordinator, SubagentExecutionRequest, + background_subagent_delivery_metadata, merge_prepended_messages_for_turn, + normalize_subagent_max_concurrency, resolve_agent_session_create_created_by, + resolve_agent_submission_turn_id, should_require_tool_confirmation, + turn_review_manifest_for_agent, validate_background_subagent_delivery, + ConversationCoordinator, SubagentExecutionRequest, }; use crate::agentic::agents::{CustomSubagent, CustomSubagentKind, UserContextPolicy}; use crate::agentic::core::{ @@ -7897,6 +8100,147 @@ mod tests { )); } + #[test] + fn peer_metadata_can_only_strengthen_tool_confirmation() { + let desktop_policy = DialogSubmissionPolicy::new( + AgentSubmissionSource::DesktopUi, + DialogQueuePriority::Normal, + false, + ); + let peer_metadata = serde_json::json!({ "require_tool_confirmation": true }); + + assert!(should_require_tool_confirmation( + desktop_policy, + Some(&peer_metadata) + )); + assert!(!should_require_tool_confirmation( + desktop_policy, + Some(&serde_json::json!({ + "require_tool_confirmation": true, + "acp_transport": true, + })) + )); + } + + #[test] + fn background_subagent_delivery_identifies_its_exact_parent_turn() { + let parent = SubagentParentInfo { + session_id: "parent-session".to_string(), + dialog_turn_id: "parent-turn".to_string(), + tool_call_id: "tool-call".to_string(), + }; + + let metadata = background_subagent_delivery_metadata( + "background-task", + &parent, + "subagent-session", + "subagent-turn", + true, + "agentic", + "Investigate", + ); + + assert_eq!(metadata["kind"], serde_json::json!("background_result")); + assert_eq!(metadata["sourceKind"], serde_json::json!("subagent")); + assert_eq!( + metadata["parentSessionId"], + serde_json::json!("parent-session") + ); + assert_eq!( + metadata["parentDialogTurnId"], + serde_json::json!("parent-turn") + ); + assert_eq!( + metadata["subagentSessionId"], + serde_json::json!("subagent-session") + ); + assert_eq!( + metadata["subagentDialogTurnId"], + serde_json::json!("subagent-turn") + ); + assert_eq!( + metadata["require_tool_confirmation"], + serde_json::json!(true) + ); + assert!(should_require_tool_confirmation( + DialogSubmissionPolicy::for_source(AgentSubmissionSource::AgentSession), + Some(&serde_json::Value::Object(metadata)) + )); + + let non_peer_metadata = background_subagent_delivery_metadata( + "background-task", + &parent, + "subagent-session", + "subagent-turn", + false, + "agentic", + "Investigate", + ); + assert!(!non_peer_metadata.contains_key("require_tool_confirmation")); + } + + #[tokio::test] + async fn completed_background_control_cannot_be_claimed_for_late_cancellation() { + let (coordinator, _) = test_coordinator(); + coordinator.register_background_subagent_task( + "background-task".to_string(), + "parent-session".to_string(), + "subagent-session".to_string(), + super::BackgroundSubagentCancelTarget::Direct( + tokio_util::sync::CancellationToken::new(), + ), + ); + coordinator.finish_background_subagent_delivery("background-task"); + + let claimed = coordinator.claim_background_subagent_controls(|control| { + control.parent_session_id == "parent-session" + && control.subagent_session_id == "subagent-session" + }); + + assert!(claimed.is_empty()); + assert!(!coordinator.take_background_subagent_delivery_suppression("background-task")); + } + + #[tokio::test] + async fn injection_and_cancellation_claim_the_same_background_control_once() { + let (coordinator, _) = test_coordinator(); + coordinator.register_background_subagent_task( + "delivery-wins".to_string(), + "parent-session".to_string(), + "subagent-session".to_string(), + super::BackgroundSubagentCancelTarget::Direct( + tokio_util::sync::CancellationToken::new(), + ), + ); + assert!(coordinator.claim_background_subagent_control_for_injection("delivery-wins")); + assert!(coordinator + .claim_background_subagent_controls(|control| { + control.parent_session_id == "parent-session" + && control.subagent_session_id == "subagent-session" + }) + .is_empty()); + + coordinator.register_background_subagent_task( + "cancellation-wins".to_string(), + "parent-session".to_string(), + "subagent-session".to_string(), + super::BackgroundSubagentCancelTarget::Direct( + tokio_util::sync::CancellationToken::new(), + ), + ); + assert_eq!( + coordinator + .claim_background_subagent_controls(|control| { + control.parent_session_id == "parent-session" + && control.subagent_session_id == "subagent-session" + }) + .len(), + 1 + ); + assert!(!coordinator.claim_background_subagent_control_for_injection("cancellation-wins")); + assert!(coordinator.take_background_subagent_delivery_suppression("cancellation-wins")); + } + #[tokio::test] async fn coordinator_test_fixture_injects_terminal_port() { let (coordinator, _) = test_coordinator(); diff --git a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs index b57d825692..6c933feea1 100644 --- a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs +++ b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs @@ -41,8 +41,9 @@ use uuid::Uuid; use bitfun_agent_runtime::scheduler::{ build_thread_goal_objective_updated_delivery_plan, build_thread_goal_resumed_delivery_plan, - resolve_agent_session_reply_action, resolve_background_delivery_action, - resolve_background_delivery_injection, resolve_dialog_start_route, + is_background_result_injection, resolve_agent_session_reply_action, + resolve_background_delivery_action, resolve_background_delivery_injection, + resolve_background_delivery_injection_for_turn, resolve_dialog_start_route, resolve_dialog_steering_action, resolve_turn_outcome_lifecycle_plan, ActiveDialogTurn, ActiveDialogTurnStore, ActiveDialogTurnTakeResult, AgentSessionReplyAction, AgentSessionReplyPlan, BackgroundDeliveryAction, BackgroundDeliveryFacts, @@ -57,7 +58,7 @@ use bitfun_runtime_ports::{ AgentThreadGoalDeliveryKind, AgentThreadGoalDeliveryRequest, AgentTurnCancellationPort, AgentTurnCancellationRequest, AgentTurnCancellationResult, DialogSessionStateFact, DialogSubmitQueueAction, DialogSubmitQueueFacts, PortError, PortErrorKind, PortResult, - SessionStoragePathRequest, SessionStorePort, + RoundInjection, RoundInjectionKind, SessionStoragePathRequest, SessionStorePort, }; pub use bitfun_runtime_ports::{ AgentSessionReplyRoute, DialogQueuePriority, DialogSteerOutcome, DialogSubmissionPolicy, @@ -99,6 +100,41 @@ fn remove_queued_turn_by_id( queues.remove_first_matching(session_id, |turn| turn.turn_id.as_deref() == Some(turn_id)) } +fn background_task_id_from_metadata(metadata: Option<&serde_json::Value>) -> Option<&str> { + let metadata = metadata.and_then(serde_json::Value::as_object)?; + (metadata.get("kind").and_then(serde_json::Value::as_str) == Some("background_result") + && metadata + .get("sourceKind") + .and_then(serde_json::Value::as_str) + == Some("subagent")) + .then(|| { + metadata + .get("backgroundTaskId") + .and_then(serde_json::Value::as_str) + }) + .flatten() +} + +fn background_result_injection_id(background_task_id: Option) -> String { + background_task_id.unwrap_or_else(|| Uuid::new_v4().to_string()) +} + +fn remove_queued_background_result_by_task_id( + queues: &DialogTurnQueue, + session_id: &str, + background_task_id: &str, +) -> Option { + queues.remove_first_matching(session_id, |turn| { + queued_background_task_id(turn) == Some(background_task_id) + }) +} + +fn queued_background_task_id(turn: &QueuedTurn) -> Option<&str> { + (turn.policy.trigger_source == DialogTriggerSource::AgentSession) + .then(|| background_task_id_from_metadata(turn.user_message_metadata.as_ref())) + .flatten() +} + #[derive(Debug, Clone)] pub(crate) struct HiddenSubagentQueuedExecution { request: HiddenSubagentExecutionRequest, @@ -176,6 +212,81 @@ enum ActiveInternalTurn { HiddenSubagent, } +#[derive(Clone)] +struct PendingBackgroundResultDelivery { + session_id: String, + agent_type: String, + workspace_path: Option, + remote_connection_id: Option, + remote_ssh_host: Option, + content: String, + display_content: Option, + user_message_metadata: Option, +} + +struct SchedulerRoundInjectionSource { + buffer: Arc, + coordinator: Arc, + pending_background_results: Arc>, +} + +impl DialogRoundInjectionSource for SchedulerRoundInjectionSource { + fn has_pending(&self, session_id: &str, turn_id: &str) -> bool { + self.buffer.has_pending_for_turn(session_id, turn_id) + } + + fn pending_tool_preemption( + &self, + session_id: &str, + turn_id: &str, + ) -> bitfun_runtime_ports::RoundInjectionToolPreemption { + self.buffer + .pending_tool_preemption_for_turn(session_id, turn_id) + } + + fn take_pending(&self, session_id: &str, turn_id: &str) -> Vec { + self.buffer + .drain_for_turn(session_id, turn_id) + .into_iter() + .filter(|injection| { + if !is_background_result_injection(injection.kind) + || !self.pending_background_results.contains_key(&injection.id) + { + return true; + } + if self + .coordinator + .claim_background_subagent_control_for_injection(&injection.id) + { + return true; + } + self.pending_background_results.remove(&injection.id); + self.coordinator + .finish_background_subagent_delivery(&injection.id); + false + }) + .collect() + } + + fn acknowledge_consumed( + &self, + _session_id: &str, + _turn_id: &str, + injection_id: &str, + kind: RoundInjectionKind, + ) { + if is_background_result_injection(kind) + && self + .pending_background_results + .remove(injection_id) + .is_some() + { + self.coordinator + .finish_background_subagent_delivery(injection_id); + } + } +} + /// Message queue manager for dialog turns. /// /// All user-facing callers (frontend Tauri commands, remote server, bot router) @@ -196,11 +307,11 @@ pub struct DialogScheduler { /// Turns whose cancelled auto-reply should be suppressed because the source /// agent explicitly cancelled its own outstanding SessionMessage request. suppressed_cancelled_replies: Arc, - /// Exact outcomes retired by destructive session deletion. The outcome - /// channel may receive them only after the deletion permit releases its + /// Exact outcomes retired by destructive session maintenance. The outcome + /// channel may receive them only after the maintenance permit releases its /// per-session operation lock; tombstoning prevents them from mutating a /// newly created session that reuses the same explicit ID. - retired_deletion_outcomes: Arc, + retired_maintenance_outcomes: Arc, /// Set when the user cancels an in-flight turn; aborts goal-continuation submit retries. goal_continuation_abort: Arc, /// Cloneable sender given to ConversationCoordinator for turn outcome notifications @@ -208,6 +319,10 @@ pub struct DialogScheduler { /// Per-session FIFO buffer of round injections drained at round boundaries /// by the engine and injected into the running dialog turn. round_injection_buffer: Arc, + round_injection_source: Arc, + pending_background_results: Arc>, + #[cfg(test)] + background_delivery_before_lock: std::sync::Mutex>>, } /// Holds the scheduler's exclusive session-operation boundary while a caller @@ -218,11 +333,11 @@ pub(crate) struct SessionMaintenancePermit { fn take_active_turn_for_outcome( active_turns: &ActiveDialogTurnStore, - retired_deletion_outcomes: &DialogReplySuppressionSet, + retired_maintenance_outcomes: &DialogReplySuppressionSet, session_id: &str, turn_id: &str, ) -> Option { - if retired_deletion_outcomes.take(session_id, turn_id) { + if retired_maintenance_outcomes.take(session_id, turn_id) { None } else { Some(active_turns.take_for_outcome(session_id, turn_id)) @@ -257,6 +372,13 @@ impl DialogScheduler { session_manager: Arc, ) -> Arc { let (outcome_tx, outcome_rx) = mpsc::channel(128); + let round_injection_buffer = Arc::new(SessionRoundInjectionBuffer::default()); + let pending_background_results = Arc::new(dashmap::DashMap::new()); + let round_injection_source = Arc::new(SchedulerRoundInjectionSource { + buffer: round_injection_buffer.clone(), + coordinator: coordinator.clone(), + pending_background_results: pending_background_results.clone(), + }); let scheduler = Arc::new(Self { coordinator, @@ -266,10 +388,14 @@ impl DialogScheduler { active_turns: Arc::new(ActiveDialogTurnStore::default()), active_internal_turns: Arc::new(dashmap::DashMap::new()), suppressed_cancelled_replies: Arc::new(DialogReplySuppressionSet::default()), - retired_deletion_outcomes: Arc::new(DialogReplySuppressionSet::default()), + retired_maintenance_outcomes: Arc::new(DialogReplySuppressionSet::default()), goal_continuation_abort: Arc::new(SessionAbortFlags::default()), outcome_tx, - round_injection_buffer: Arc::new(SessionRoundInjectionBuffer::default()), + round_injection_buffer, + round_injection_source, + pending_background_results, + #[cfg(test)] + background_delivery_before_lock: std::sync::Mutex::new(None), }); let scheduler_for_handler = Arc::clone(&scheduler); @@ -291,7 +417,41 @@ impl DialogScheduler { /// Pass to [`ConversationCoordinator::set_round_injection_source`](super::coordinator::ConversationCoordinator::set_round_injection_source). pub fn round_injection_monitor(&self) -> Arc { - self.round_injection_buffer.clone() + self.round_injection_source.clone() + } + + #[cfg(test)] + fn install_background_delivery_before_lock_signal(&self) -> oneshot::Receiver<()> { + let (tx, rx) = oneshot::channel(); + *self + .background_delivery_before_lock + .lock() + .expect("background delivery test hook") = Some(tx); + rx + } + + #[cfg(test)] + fn signal_background_delivery_before_lock(&self) { + if let Some(tx) = self + .background_delivery_before_lock + .lock() + .expect("background delivery test hook") + .take() + { + let _ = tx.send(()); + } + } + + pub(crate) fn active_turn_requires_tool_confirmation( + &self, + session_id: &str, + turn_id: &str, + ) -> bool { + self.active_turns.user_message_metadata_bool_for_turn( + session_id, + turn_id, + "require_tool_confirmation", + ) == Some(true) } /// Submit a user "steering" message into the currently running dialog turn. @@ -493,55 +653,156 @@ impl DialogScheduler { display_content: Option, user_message_metadata: Option, ) -> Result<(), String> { + #[cfg(test)] + self.signal_background_delivery_before_lock(); + let _operation_guard = self.lock_session_operation(&session_id).await; + let background_task_id = + background_task_id_from_metadata(user_message_metadata.as_ref()).map(str::to_string); + if background_task_id + .as_deref() + .is_some_and(|background_task_id| { + self.coordinator + .take_background_subagent_delivery_suppression(background_task_id) + }) + { + return Ok(()); + } + let display = display_content.unwrap_or_else(|| content.clone()); + let delivery = PendingBackgroundResultDelivery { + session_id: session_id.clone(), + agent_type, + workspace_path, + remote_connection_id, + remote_ssh_host, + content, + display_content: Some(display), + user_message_metadata, + }; let state = self .session_manager .get_session(&session_id) .map(|s| s.state.clone()); match resolve_background_delivery_action(BackgroundDeliveryFacts { - session_state: Self::session_state_fact(state.as_ref()), + session_state: background_result_delivery_state_fact( + &session_id, + state.as_ref(), + delivery.user_message_metadata.as_ref(), + ), }) { BackgroundDeliveryAction::InjectIntoRunningTurn => { - self.round_injection_buffer.push( - &session_id, - resolve_background_delivery_injection( - BackgroundInjectionKind::BackgroundResult, - Uuid::new_v4().to_string(), - content, - Some(display), - SystemTime::now(), - ), + if background_task_id.as_deref().is_some_and(|task_id| { + !self + .coordinator + .background_subagent_control_available_for_injection(task_id) + }) { + if let Some(task_id) = background_task_id.as_deref() { + self.coordinator + .take_background_subagent_delivery_suppression(task_id); + } + return Ok(()); + } + let Some(current_turn_id) = state.as_ref().and_then(|state| match state { + SessionState::Processing { + current_turn_id, .. + } => Some(current_turn_id.clone()), + _ => None, + }) else { + return Err(format!( + "Background result resolved to injection without an active turn: session_id={session_id}" + )); + }; + let injection_id = background_result_injection_id(background_task_id.clone()); + let injection = resolve_background_delivery_injection_for_turn( + BackgroundInjectionKind::BackgroundResult, + injection_id.clone(), + delivery.content.clone(), + delivery.display_content.clone(), + SystemTime::now(), + current_turn_id, ); + if background_task_id.is_some() { + self.pending_background_results + .insert(injection_id, delivery); + } + self.round_injection_buffer.push(&session_id, injection); Ok(()) } BackgroundDeliveryAction::SubmitAgentSessionFollowUp { queue_priority, skip_tool_confirmation, - } => self - .submit( - session_id, - content, - Some(display), - None, - agent_type, - workspace_path, - remote_connection_id, - remote_ssh_host, - DialogSubmissionPolicy::new( - DialogTriggerSource::AgentSession, - queue_priority, - skip_tool_confirmation, - ), - None, - user_message_metadata, - None, + } => { + self.submit_background_result_follow_up_locked( + delivery, + queue_priority, + skip_tool_confirmation, ) .await - .map(|_| ()), + } } } + async fn submit_background_result_follow_up_locked( + &self, + delivery: PendingBackgroundResultDelivery, + queue_priority: DialogQueuePriority, + skip_tool_confirmation: bool, + ) -> Result<(), String> { + let background_task_id = + background_task_id_from_metadata(delivery.user_message_metadata.as_ref()) + .map(str::to_string); + let resolved_turn_id = Uuid::new_v4().to_string(); + let queued_turn = QueuedTurn { + user_input: delivery.content, + original_user_input: delivery.display_content, + prepended_messages: Vec::new(), + turn_id: Some(resolved_turn_id.clone()), + agent_type: delivery.agent_type, + workspace_path: delivery.workspace_path, + remote_connection_id: delivery.remote_connection_id, + remote_ssh_host: delivery.remote_ssh_host, + policy: DialogSubmissionPolicy::new( + DialogTriggerSource::AgentSession, + queue_priority, + skip_tool_confirmation, + ), + reply_route: None, + user_message_metadata: delivery.user_message_metadata, + image_contexts: None, + enqueued_at: SystemTime::now(), + execution: QueuedTurnExecution::Standard, + }; + if background_task_id + .as_deref() + .is_some_and(|background_task_id| { + self.coordinator + .take_background_subagent_delivery_suppression(background_task_id) + }) + { + return Ok(()); + } + let result = self + .submit_queued_turn_locked( + delivery.session_id.clone(), + resolved_turn_id.clone(), + queued_turn, + ) + .await; + if result.is_err() { + if let Some(removed_turn) = + remove_queued_turn_by_id(&self.queues, &delivery.session_id, &resolved_turn_id) + { + self.finish_removed_queued_turn(&delivery.session_id, removed_turn) + .await; + } else if let Some(background_task_id) = background_task_id.as_deref() { + self.coordinator + .finish_background_subagent_delivery(background_task_id); + } + } + result.map(|_| ()) + } + pub async fn submit_init_agents_md( &self, session_id: String, @@ -811,6 +1072,16 @@ impl DialogScheduler { queued_turn: QueuedTurn, ) -> Result { let _operation_guard = self.lock_session_operation(&session_id).await; + self.submit_queued_turn_locked(session_id, resolved_turn_id, queued_turn) + .await + } + + async fn submit_queued_turn_locked( + &self, + session_id: String, + resolved_turn_id: String, + queued_turn: QueuedTurn, + ) -> Result { if let Some(workspace_path) = queued_turn.workspace_path.as_deref() { let requested_storage_path = Self::resolve_session_restore_path( workspace_path, @@ -903,6 +1174,7 @@ impl DialogScheduler { } async fn finish_removed_queued_turn(&self, session_id: &str, removed_turn: QueuedTurn) { + let background_task_id = queued_background_task_id(&removed_turn).map(str::to_string); match removed_turn.execution { QueuedTurnExecution::Standard => { if let Some(turn_id) = removed_turn.turn_id { @@ -926,6 +1198,125 @@ impl DialogScheduler { ))); } } + if let Some(background_task_id) = background_task_id.as_deref() { + self.coordinator + .finish_background_subagent_delivery(background_task_id); + } + } + + fn discard_drained_background_results(&self, drained: Vec) { + for injection in drained { + if !is_background_result_injection(injection.kind) { + continue; + } + if self + .pending_background_results + .remove(&injection.id) + .is_some() + { + self.coordinator + .finish_background_subagent_delivery(&injection.id); + } + } + } + + async fn recover_drained_background_results( + &self, + status: TurnOutcomeStatus, + drained: Vec, + ) { + for injection in drained { + if !is_background_result_injection(injection.kind) { + continue; + } + let Some((background_task_id, delivery)) = + self.pending_background_results.remove(&injection.id) + else { + continue; + }; + if status == TurnOutcomeStatus::Cancelled { + self.coordinator + .finish_background_subagent_delivery(&background_task_id); + continue; + } + let policy = DialogSubmissionPolicy::for_source(DialogTriggerSource::AgentSession); + let _operation_guard = self.lock_session_operation(&delivery.session_id).await; + let result = self + .submit_background_result_follow_up_locked( + delivery.clone(), + policy.queue_priority, + policy.skip_tool_confirmation, + ) + .await; + if let Err(error) = result { + self.coordinator + .finish_background_subagent_delivery(&background_task_id); + warn!( + "Failed to recover an unconsumed background result after turn {}: background_task_id={}, session_id={}, error={}", + status, background_task_id, delivery.session_id, error + ); + } + } + } + + pub(crate) async fn cancel_background_result_delivery( + &self, + session_id: &str, + background_task_id: &str, + ) -> Result { + let _operation_guard = self.lock_session_operation(session_id).await; + let pending_matches_session = self + .pending_background_results + .get(background_task_id) + .is_some_and(|delivery| delivery.session_id == session_id); + if pending_matches_session + && self + .round_injection_buffer + .remove_by_id(session_id, background_task_id) + .is_some() + { + self.pending_background_results.remove(background_task_id); + self.coordinator + .finish_background_subagent_delivery(background_task_id); + return Ok(true); + } + if let Some(removed_turn) = + remove_queued_background_result_by_task_id(&self.queues, session_id, background_task_id) + { + self.finish_removed_queued_turn(session_id, removed_turn) + .await; + return Ok(true); + } + + let Some(turn_id) = self + .active_turns + .turn_id_for_background_subagent_delivery(session_id, background_task_id) + else { + return Ok(false); + }; + self.coordinator + .cancel_dialog_turn(session_id, &turn_id) + .await?; + Ok(true) + } + + pub(crate) async fn cancel_background_result_deliveries( + &self, + session_id: &str, + background_task_ids: &[String], + ) -> Result<(), String> { + let mut first_error = None; + for background_task_id in background_task_ids { + if let Err(error) = self + .cancel_background_result_delivery(session_id, background_task_id) + .await + { + if first_error.is_none() { + first_error = Some(error); + } + } + } + first_error.map_or(Ok(()), Err) } /// Cancel one queued or active turn without allowing it to cross the @@ -1028,10 +1419,10 @@ impl DialogScheduler { .await } - /// Quiesce one session for deletion. Queued turns receive an explicit + /// Quiesce one session for destructive maintenance. Queued turns receive an explicit /// cancelled lifecycle event before active execution is cancelled and /// drained, so no accepted turn disappears silently. - pub(crate) async fn begin_session_deletion( + pub(crate) async fn begin_session_maintenance( &self, session_id: &str, requested_storage_path: &std::path::Path, @@ -1045,31 +1436,52 @@ impl DialogScheduler { self.clear_queue(session_id).await; } abort_thread_goal_continuation_for_session(session_id); + self.coordinator + .cancel_background_subagents_for_parent_session(session_id) + .await?; self.coordinator .cancel_active_turn_for_session(session_id, wait_timeout) .await?; + self.coordinator + .cancel_background_subagents_for_parent_session(session_id) + .await?; self.coordinator .ensure_session_execution_drained(session_id, wait_timeout) .await?; - self.retire_active_turn_for_deletion(session_id); + self.retire_active_turn_for_maintenance(session_id); Ok(SessionMaintenancePermit { _operation_guard: operation_guard, }) } - fn retire_active_turn_for_deletion(&self, session_id: &str) { + pub(crate) async fn begin_session_deletion( + &self, + session_id: &str, + requested_storage_path: &std::path::Path, + wait_timeout: Duration, + ) -> BitFunResult { + self.begin_session_maintenance(session_id, requested_storage_path, wait_timeout) + .await + } + + fn retire_active_turn_for_maintenance(&self, session_id: &str) { let Some(active_turn) = self.active_turns.remove(session_id) else { return; }; + if let Some(background_task_id) = active_turn.background_subagent_task_id() { + self.coordinator + .finish_background_subagent_delivery(background_task_id); + } let turn_id = active_turn.turn_id().to_string(); - self.retired_deletion_outcomes.mark(session_id, &turn_id); + self.retired_maintenance_outcomes.mark(session_id, &turn_id); self.active_internal_turns.remove(session_id); - let _drained = self + let drained = self .round_injection_buffer .drain_for_turn(session_id, &turn_id); + self.discard_drained_background_results(drained); self.take_suppressed_cancelled_reply(session_id, &turn_id); debug!( - "Retired active turn before session deletion: session_id={}, turn_id={}", + "Retired active turn before destructive session maintenance: session_id={}, turn_id={}", session_id, turn_id ); } @@ -1101,6 +1513,7 @@ impl DialogScheduler { let cleared_turns = self.queues.clear(session_id); let count = cleared_turns.len(); for queued_turn in cleared_turns { + let background_task_id = queued_background_task_id(&queued_turn).map(str::to_string); match queued_turn.execution { QueuedTurnExecution::Standard => { if let Some(turn_id) = queued_turn.turn_id { @@ -1131,6 +1544,10 @@ impl DialogScheduler { }); } } + if let Some(background_task_id) = background_task_id.as_deref() { + self.coordinator + .finish_background_subagent_delivery(background_task_id); + } } if count > 0 { info!( @@ -1507,13 +1924,14 @@ impl DialogScheduler { let _operation_guard = self.lock_session_operation(&session_id).await; let Some(active_turn_result) = take_active_turn_for_outcome( &self.active_turns, - &self.retired_deletion_outcomes, + &self.retired_maintenance_outcomes, &session_id, outcome.turn_id(), ) else { - let _drained = self + let drained = self .round_injection_buffer .drain_for_turn(&session_id, outcome.turn_id()); + self.discard_drained_background_results(drained); self.take_suppressed_cancelled_reply(&session_id, outcome.turn_id()); debug!( "Ignoring outcome retired by session deletion: session_id={}, turn_id={}", @@ -1526,9 +1944,10 @@ impl DialogScheduler { ActiveDialogTurnTakeResult::Matched(turn) => Some(turn), ActiveDialogTurnTakeResult::Absent => None, ActiveDialogTurnTakeResult::DifferentTurn => { - let _drained = self + let drained = self .round_injection_buffer .drain_for_turn(&session_id, outcome.turn_id()); + self.discard_drained_background_results(drained); self.take_suppressed_cancelled_reply(&session_id, outcome.turn_id()); debug!( "Ignoring stale turn outcome: session_id={}, turn_id={}", @@ -1538,6 +1957,13 @@ impl DialogScheduler { continue; } }; + if let Some(background_task_id) = active_turn + .as_ref() + .and_then(ActiveDialogTurn::background_subagent_task_id) + { + self.coordinator + .finish_background_subagent_delivery(background_task_id); + } let active_internal_turn = active_turn.as_ref().and_then(|_| { self.active_internal_turns .remove(&session_id) @@ -1563,9 +1989,11 @@ impl DialogScheduler { // outcome is processed (race window between turn finalize and the // next turn starting). Targeting by turn_id keeps those alive. if lifecycle_plan.drain_finished_turn_injections { - let _drained = self + let drained = self .round_injection_buffer .drain_for_turn(&session_id, outcome.turn_id()); + self.recover_drained_background_results(status, drained) + .await; } let suppressed_cancelled_reply = self.take_suppressed_cancelled_reply(&session_id, outcome.turn_id()); @@ -1995,6 +2423,46 @@ fn thread_goal_delivery_messages(reminders: Vec) -> .collect() } +fn background_result_delivery_state_fact( + session_id: &str, + state: Option<&SessionState>, + metadata: Option<&serde_json::Value>, +) -> DialogSessionStateFact { + let Some(SessionState::Processing { + current_turn_id, .. + }) = state + else { + return DialogScheduler::session_state_fact(state); + }; + let Some(metadata) = metadata.and_then(serde_json::Value::as_object) else { + return DialogSessionStateFact::Processing; + }; + let has_exact_parent = + metadata.contains_key("parentSessionId") || metadata.contains_key("parentDialogTurnId"); + if !has_exact_parent { + return DialogSessionStateFact::Processing; + } + + let exact_parent_matches = metadata + .get("parentSessionId") + .and_then(serde_json::Value::as_str) + .zip( + metadata + .get("parentDialogTurnId") + .and_then(serde_json::Value::as_str), + ) + .is_some_and(|(parent_session_id, parent_turn_id)| { + parent_session_id == session_id && parent_turn_id == current_turn_id + }); + if exact_parent_matches { + DialogSessionStateFact::Processing + } else { + // The session is busy, but this result does not belong to the running turn. + // Resolve it as a follow-up; the normal submission path will queue it. + DialogSessionStateFact::Idle + } +} + // ── Global instance ────────────────────────────────────────────────────────── static GLOBAL_SCHEDULER: OnceLock> = OnceLock::new(); @@ -2024,6 +2492,7 @@ pub fn clear_thread_goal_continuation_abort(session_id: &str) { #[cfg(test)] mod tests { use super::*; + use crate::agentic::core::{ProcessingPhase, SessionConfig}; use crate::agentic::events::{EventQueue, EventQueueConfig, EventRouter}; use crate::agentic::execution::{ ExecutionEngine, ExecutionEngineConfig, RoundExecutor, StreamProcessor, @@ -2169,6 +2638,408 @@ mod tests { ); } + #[tokio::test] + async fn background_delivery_rechecks_suppression_after_waiting_for_session_lock() { + let (scheduler, _, event_queue, _root) = test_scheduler(); + let mut events = event_queue.subscribe(); + let operation_guard = scheduler.lock_session_operation("parent-session").await; + let delivery_reached_lock = scheduler.install_background_delivery_before_lock_signal(); + let delivery_scheduler = scheduler.clone(); + let delivery = tokio::spawn(async move { + delivery_scheduler + .deliver_background_result( + "parent-session".to_string(), + "agentic".to_string(), + None, + None, + None, + "background result".to_string(), + None, + Some(serde_json::json!({ + "kind": "background_result", + "sourceKind": "subagent", + "backgroundTaskId": "background-task", + })), + ) + .await + }); + delivery_reached_lock + .await + .expect("delivery should reach the session lock"); + scheduler + .coordinator + .mark_background_subagent_delivery_suppression("background-task".to_string()); + + drop(operation_guard); + delivery + .await + .expect("delivery task") + .expect("suppressed delivery"); + + assert_eq!(scheduler.queue_depth("parent-session"), 0); + assert!( + tokio::time::timeout(Duration::from_millis(20), events.recv()) + .await + .is_err() + ); + assert!(!scheduler + .coordinator + .take_background_subagent_delivery_suppression("background-task")); + } + + #[tokio::test] + async fn background_injection_claims_control_only_when_the_turn_consumes_it() { + let (scheduler, session_manager, _, root) = test_scheduler(); + let session_id = "parent-session"; + let workspace = root.path().join("workspace"); + std::fs::create_dir_all(&workspace).expect("workspace"); + session_manager + .create_session_with_id( + Some(session_id.to_string()), + "Parent".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.to_string_lossy().to_string()), + ..Default::default() + }, + ) + .await + .expect("create parent session"); + session_manager + .update_session_state( + session_id, + SessionState::Processing { + current_turn_id: "peer-turn".to_string(), + phase: ProcessingPhase::Thinking, + }, + ) + .await + .expect("mark peer turn active"); + scheduler + .coordinator + .register_background_subagent_task_for_test( + "background-task", + session_id, + "subagent-session", + ); + + scheduler + .deliver_background_result( + session_id.to_string(), + "agentic".to_string(), + None, + None, + None, + "background result".to_string(), + None, + Some(serde_json::json!({ + "kind": "background_result", + "sourceKind": "subagent", + "backgroundTaskId": "background-task", + "parentSessionId": session_id, + "parentDialogTurnId": "peer-turn", + })), + ) + .await + .expect("buffer background delivery"); + + assert!(scheduler + .coordinator + .has_background_subagent_task_for_test("background-task")); + assert!(scheduler + .pending_background_results + .contains_key("background-task")); + + let source = scheduler.round_injection_monitor(); + let pending = source.take_pending(session_id, "peer-turn"); + assert_eq!(pending.len(), 1); + assert!(!scheduler + .coordinator + .has_background_subagent_task_for_test("background-task")); + assert!(scheduler + .pending_background_results + .contains_key("background-task")); + + source.acknowledge_consumed(session_id, "peer-turn", "background-task", pending[0].kind); + assert!(!scheduler + .pending_background_results + .contains_key("background-task")); + } + + #[tokio::test] + async fn cancellation_removes_a_pending_background_injection_without_recovery() { + let (scheduler, _, _, _root) = test_scheduler(); + let session_id = "parent-session"; + scheduler + .coordinator + .register_background_subagent_task_for_test( + "background-task", + session_id, + "subagent-session", + ); + scheduler.pending_background_results.insert( + "background-task".to_string(), + PendingBackgroundResultDelivery { + session_id: session_id.to_string(), + agent_type: "agentic".to_string(), + workspace_path: None, + remote_connection_id: None, + remote_ssh_host: None, + content: "background result".to_string(), + display_content: Some("background result".to_string()), + user_message_metadata: Some(serde_json::json!({ + "kind": "background_result", + "sourceKind": "subagent", + "backgroundTaskId": "background-task", + })), + }, + ); + let injection = resolve_background_delivery_injection_for_turn( + BackgroundInjectionKind::BackgroundResult, + "background-task".to_string(), + "background result".to_string(), + None, + SystemTime::now(), + "peer-turn".to_string(), + ); + scheduler.round_injection_buffer.push(session_id, injection); + assert!(scheduler + .coordinator + .suppress_background_subagent_task_for_test("background-task")); + + assert!(scheduler + .cancel_background_result_delivery(session_id, "background-task") + .await + .expect("cancel pending injection")); + + assert_eq!( + scheduler.round_injection_buffer.pending_count(session_id), + 0 + ); + assert!(!scheduler + .pending_background_results + .contains_key("background-task")); + assert!(!scheduler + .coordinator + .take_background_subagent_delivery_suppression("background-task")); + } + + #[tokio::test] + async fn finished_turn_recovers_an_unconsumed_background_result_as_a_follow_up() { + let (scheduler, session_manager, _, root) = test_scheduler(); + let session_id = "parent-session"; + let workspace = root.path().join("workspace"); + std::fs::create_dir_all(&workspace).expect("workspace"); + session_manager + .create_session_with_id( + Some(session_id.to_string()), + "Parent".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.to_string_lossy().to_string()), + ..Default::default() + }, + ) + .await + .expect("create parent session"); + session_manager + .update_session_state( + session_id, + SessionState::Processing { + current_turn_id: "peer-turn".to_string(), + phase: ProcessingPhase::Thinking, + }, + ) + .await + .expect("mark peer turn active"); + scheduler + .coordinator + .register_background_subagent_task_for_test( + "background-task", + session_id, + "subagent-session", + ); + scheduler + .deliver_background_result( + session_id.to_string(), + "agentic".to_string(), + None, + None, + None, + "background result".to_string(), + None, + Some(serde_json::json!({ + "kind": "background_result", + "sourceKind": "subagent", + "backgroundTaskId": "background-task", + "parentSessionId": session_id, + "parentDialogTurnId": "peer-turn", + })), + ) + .await + .expect("buffer background delivery"); + + let drained = scheduler + .round_injection_buffer + .drain_for_turn(session_id, "peer-turn"); + session_manager + .update_session_state( + session_id, + SessionState::Processing { + current_turn_id: "next-turn".to_string(), + phase: ProcessingPhase::Thinking, + }, + ) + .await + .expect("mark next turn active"); + scheduler + .recover_drained_background_results(TurnOutcomeStatus::Completed, drained) + .await; + + assert_eq!( + scheduler.round_injection_buffer.pending_count(session_id), + 0 + ); + assert!(!scheduler + .pending_background_results + .contains_key("background-task")); + assert_eq!(scheduler.queue_depth(session_id), 1); + assert!(scheduler + .coordinator + .has_background_subagent_task_for_test("background-task")); + assert!(remove_queued_background_result_by_task_id( + &scheduler.queues, + session_id, + "background-task" + ) + .is_some()); + } + + #[tokio::test] + async fn background_delivery_queues_behind_an_unrelated_running_turn() { + let (scheduler, session_manager, _, root) = test_scheduler(); + let session_id = "parent-session"; + let workspace = root.path().join("workspace"); + std::fs::create_dir_all(&workspace).expect("workspace"); + session_manager + .create_session_with_id( + Some(session_id.to_string()), + "Parent".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.to_string_lossy().to_string()), + ..Default::default() + }, + ) + .await + .expect("create parent session"); + session_manager + .update_session_state( + session_id, + SessionState::Processing { + current_turn_id: "local-turn".to_string(), + phase: ProcessingPhase::Thinking, + }, + ) + .await + .expect("mark unrelated turn active"); + let mut spoofed_user_turn = standard_queued_turn("ordinary-turn"); + spoofed_user_turn.policy = DialogSubmissionPolicy::new( + DialogTriggerSource::DesktopUi, + DialogQueuePriority::High, + false, + ); + spoofed_user_turn.user_message_metadata = Some(serde_json::json!({ + "kind": "background_result", + "sourceKind": "subagent", + "backgroundTaskId": "background-task", + })); + scheduler + .queues + .enqueue(session_id, spoofed_user_turn, DialogQueuePriority::High) + .expect("queue ordinary user turn with colliding metadata"); + + scheduler + .deliver_background_result( + session_id.to_string(), + "agentic".to_string(), + None, + None, + None, + "background result".to_string(), + None, + Some(serde_json::json!({ + "kind": "background_result", + "sourceKind": "subagent", + "backgroundTaskId": "background-task", + "parentSessionId": session_id, + "parentDialogTurnId": "peer-turn", + })), + ) + .await + .expect("queue background delivery"); + + assert_eq!( + scheduler.round_injection_buffer.pending_count(session_id), + 0 + ); + assert_eq!(scheduler.queue_depth(session_id), 2); + assert!(scheduler + .cancel_background_result_delivery(session_id, "background-task") + .await + .expect("cancel exact queued background delivery")); + assert_eq!(scheduler.queue_depth(session_id), 1); + assert!(remove_queued_turn_by_id(&scheduler.queues, session_id, "ordinary-turn").is_some()); + } + + #[tokio::test] + async fn background_delivery_batch_cancellation_attempts_later_ids_after_an_error() { + let (scheduler, _, _, _root) = test_scheduler(); + let session_id = "missing-parent-session"; + scheduler.active_turns.insert( + session_id, + ActiveDialogTurn::new( + "active-background-turn".to_string(), + Some("/workspace".to_string()), + None, + None, + "agentic".to_string(), + "background result".to_string(), + Some(serde_json::json!({ + "kind": "background_result", + "sourceKind": "subagent", + "backgroundTaskId": "background-active", + })), + DialogSubmissionPolicy::for_source(DialogTriggerSource::AgentSession), + None, + ), + ); + let mut queued = standard_queued_turn("queued-background-turn"); + queued.policy = DialogSubmissionPolicy::for_source(DialogTriggerSource::AgentSession); + queued.user_message_metadata = Some(serde_json::json!({ + "kind": "background_result", + "sourceKind": "subagent", + "backgroundTaskId": "background-queued", + })); + scheduler + .queues + .enqueue(session_id, queued, DialogQueuePriority::Normal) + .expect("queue later background delivery"); + + scheduler + .cancel_background_result_deliveries( + session_id, + &[ + "background-active".to_string(), + "background-queued".to_string(), + ], + ) + .await + .expect_err("the active cancellation should fail for a missing session"); + + assert_eq!(scheduler.queue_depth(session_id), 0); + } + #[test] fn queued_submission_without_started_turn_reports_queued() { assert_eq!( @@ -2180,6 +3051,14 @@ mod tests { ); } + #[test] + fn background_result_injection_preserves_the_exact_task_id() { + assert_eq!( + background_result_injection_id(Some("background-task".to_string())), + "background-task" + ); + } + fn desktop_active_turn(turn_id: &str) -> ActiveDialogTurn { ActiveDialogTurn::new( turn_id.to_string(), @@ -2247,7 +3126,7 @@ mod tests { } #[test] - fn retired_delete_outcome_cannot_mutate_a_recreated_session_generation() { + fn retired_maintenance_outcome_cannot_mutate_a_recreated_session_generation() { let active_turns = ActiveDialogTurnStore::default(); let retired = DialogReplySuppressionSet::default(); let session_id = "reused-session"; 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 e6460b0b48..615e10a906 100644 --- a/src/crates/assembly/core/src/agentic/execution/execution_engine.rs +++ b/src/crates/assembly/core/src/agentic/execution/execution_engine.rs @@ -3405,6 +3405,8 @@ impl ExecutionEngine { round_index ); for injection in pending { + let injection_id = injection.id.clone(); + let injection_kind = injection.kind; let wrapped = match injection.kind { RoundInjectionKind::UserSteering => format!( "\nThe user sent a new message while this turn was running. You have just finished the previous atomic action; handle this new user message now as the current direction, while preserving the existing conversation and task context. Do not ignore it or wait for a separate future turn.\n\nNew user message:\n{}\n", @@ -3450,6 +3452,12 @@ impl ExecutionEngine { EventPriority::Normal, ) .await; + source.acknowledge_consumed( + &context.session_id, + &context.dialog_turn_id, + &injection_id, + injection_kind, + ); injection_applied = true; } } diff --git a/src/crates/assembly/core/src/product_runtime.rs b/src/crates/assembly/core/src/product_runtime.rs index 53bd7fbcaf..2f85af27b6 100644 --- a/src/crates/assembly/core/src/product_runtime.rs +++ b/src/crates/assembly/core/src/product_runtime.rs @@ -6,21 +6,29 @@ mod runtime_services; -use std::path::Path; +use std::path::{Path, PathBuf}; use std::sync::Arc; use bitfun_agent_runtime::sdk::AgentRuntime; use bitfun_harness::HarnessRegistry; +use bitfun_runtime_ports::{SessionStoragePathRequest, SessionStorePort, SessionViewRestoreTiming}; use bitfun_runtime_services::RuntimeServices; -use crate::agentic::coordination::{ConversationCoordinator, DialogScheduler}; +use crate::agentic::coordination::{ + ConversationCoordinator, DialogScheduler, SessionMaintenancePermit, +}; use crate::agentic::core::{Message, Session, SessionConfig, SessionState}; +use crate::agentic::keyed_lock::KeyedAsyncLockGuard; use crate::agentic::persistence::session_branch::{SessionBranchRequest, SessionBranchResult}; -use crate::agentic::persistence::PersistenceManager; +use crate::agentic::persistence::{PersistenceManager, SessionMetadataPage}; +use crate::agentic::session::CoreSessionStorePort; use crate::service::session::{DialogTurnData, SessionMetadata}; use crate::service::session_usage::{ generate_session_usage_report, SessionUsageReport, SessionUsageReportRequest, }; +use crate::service::snapshot::{ + get_snapshot_manager_for_workspace, initialize_snapshot_manager_for_workspace, SnapshotManager, +}; use crate::service::token_usage::TokenUsageService; use crate::service_agent_runtime::CoreServiceAgentRuntime; use crate::util::errors::{BitFunError, BitFunResult}; @@ -28,10 +36,38 @@ use crate::util::errors::{BitFunError, BitFunResult}; pub use bitfun_product_capabilities::ProductRuntimeAssembly as CoreProductRuntimeAssembly; pub use runtime_services::CoreRuntimeServicesProvider; +/// Serializes one compatibility mutation with Core's session lifecycle. +pub struct CoreSessionMutationPermit { + _guard: KeyedAsyncLockGuard, + session_id: String, + storage_path: PathBuf, +} + +/// Holds Core's scheduler boundary while a product compatibility operation +/// mutates session state that must not overlap turn dispatch. +pub struct CoreSessionMaintenancePermit { + _permit: SessionMaintenancePermit, +} + fn validate_persisted_session_id(session_id: &str) -> BitFunResult<()> { bitfun_core_types::validate_session_id(session_id).map_err(BitFunError::Validation) } +async fn ensure_snapshot_manager(workspace_path: &Path) -> BitFunResult> { + if let Some(manager) = get_snapshot_manager_for_workspace(workspace_path) { + return Ok(manager); + } + initialize_snapshot_manager_for_workspace(workspace_path.to_path_buf(), None) + .await + .map_err(|error| BitFunError::service(error.to_string()))?; + get_snapshot_manager_for_workspace(workspace_path).ok_or_else(|| { + BitFunError::service(format!( + "Snapshot manager is unavailable for workspace {}", + workspace_path.display() + )) + }) +} + /// Product-assembly entry for the public Agent Runtime SDK. /// /// Concrete coordinator and scheduler ownership remains in Core. Product @@ -63,6 +99,7 @@ impl CoreProductAgentRuntime { #[derive(Clone)] pub struct CoreAgentRuntimeCompatibility { coordinator: Arc, + scheduler: Arc, persistence: Arc, token_usage_service: Arc, } @@ -70,12 +107,14 @@ pub struct CoreAgentRuntimeCompatibility { impl CoreAgentRuntimeCompatibility { pub fn build( coordinator: Arc, + scheduler: Arc, token_usage_service: Arc, ) -> Self { let persistence = coordinator.get_session_manager().persistence_manager(); Self { coordinator, + scheduler, persistence, token_usage_service, } @@ -101,6 +140,108 @@ impl CoreAgentRuntimeCompatibility { .await } + pub async fn create_session_with_workspace( + &self, + session_id: Option, + session_name: String, + agent_type: String, + config: SessionConfig, + workspace_path: String, + ) -> BitFunResult { + self.coordinator + .create_session_with_workspace( + session_id, + session_name, + agent_type, + config, + workspace_path, + ) + .await + } + + pub async fn restore_session_view_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + include_internal: bool, + tail_turn_count: Option, + ) -> BitFunResult<( + Session, + Vec, + usize, + SessionViewRestoreTiming, + )> { + validate_persisted_session_id(session_id)?; + if let Some(tail_turn_count) = tail_turn_count { + let storage_path = self.resolve_persisted_session_storage_path(request).await?; + if include_internal { + self.coordinator + .restore_internal_session_view_from_storage_path_tail_timed( + &storage_path, + session_id, + tail_turn_count, + ) + .await + } else { + self.coordinator + .restore_session_view_from_storage_path_tail_timed( + &storage_path, + session_id, + tail_turn_count, + ) + .await + } + } else { + let (session, turns, timing) = if include_internal { + self.coordinator + .restore_internal_session_view_for_workspace_timed(request, session_id) + .await? + } else { + self.coordinator + .restore_session_view_for_workspace_timed(request, session_id) + .await? + }; + let total_turn_count = turns.len(); + Ok((session, turns, total_turn_count, timing)) + } + } + + pub async fn restore_session_with_turns_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + include_internal: bool, + ) -> BitFunResult<(Session, Vec)> { + validate_persisted_session_id(session_id)?; + if include_internal { + self.coordinator + .restore_internal_session_with_turns_for_workspace(request, session_id) + .await + } else { + self.coordinator + .restore_session_with_turns_for_workspace(request, session_id) + .await + } + } + + pub async fn restore_session_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + include_internal: bool, + ) -> BitFunResult { + validate_persisted_session_id(session_id)?; + if include_internal { + self.coordinator + .restore_internal_session_for_workspace(request, session_id) + .await + } else { + self.coordinator + .restore_session_for_workspace(request, session_id) + .await + } + } + pub async fn restore_session( &self, workspace_path: &Path, @@ -201,6 +342,168 @@ impl CoreAgentRuntimeCompatibility { self.persistence.list_session_metadata(workspace_path).await } + pub async fn list_persisted_sessions_page( + &self, + workspace_path: &Path, + cursor: Option<&str>, + limit: usize, + ) -> BitFunResult { + self.persistence + .list_session_metadata_page(workspace_path, cursor, limit) + .await + } + + pub fn is_session_loaded_in_memory(&self, session_id: &str) -> BitFunResult { + validate_persisted_session_id(session_id)?; + Ok(self + .coordinator + .get_session_manager() + .get_session(session_id) + .is_some()) + } + + pub async fn resolve_persisted_session_storage_path( + &self, + request: SessionStoragePathRequest, + ) -> BitFunResult { + CoreSessionStorePort::with_path_manager(self.persistence.path_manager().clone()) + .resolve_session_storage_path(request) + .await + .map(|resolution| resolution.effective_storage_path) + .map_err(|error| BitFunError::Session(error.to_string())) + } + + pub fn is_session_loaded_from_storage_path( + &self, + storage_path: &Path, + session_id: &str, + ) -> BitFunResult { + validate_persisted_session_id(session_id)?; + self.coordinator + .get_session_manager() + .is_session_loaded_from_storage_path(storage_path, session_id) + } + + pub async fn ensure_session_loaded_from_storage_path( + &self, + storage_path: &Path, + session_id: &str, + include_internal: bool, + ) -> BitFunResult<()> { + if self.is_session_loaded_from_storage_path(storage_path, session_id)? { + return Ok(()); + } + if include_internal { + self.coordinator + .restore_internal_session_from_storage_path(storage_path, session_id) + .await?; + } else { + self.coordinator + .restore_session_from_storage_path(storage_path, session_id) + .await?; + } + Ok(()) + } + + pub async fn begin_persisted_session_mutation( + &self, + storage_path: &Path, + session_id: &str, + ) -> BitFunResult { + validate_persisted_session_id(session_id)?; + let session_manager = self.coordinator.get_session_manager(); + let guard = session_manager.acquire_session_mutation(session_id).await?; + session_manager.validate_session_storage_path_binding(session_id, storage_path)?; + Ok(CoreSessionMutationPermit { + _guard: guard, + session_id: session_id.to_string(), + storage_path: storage_path.to_path_buf(), + }) + } + + pub async fn begin_session_maintenance( + &self, + storage_path: &Path, + session_id: &str, + wait_timeout_ms: u64, + ) -> BitFunResult { + let permit = self + .scheduler + .begin_session_maintenance( + session_id, + storage_path, + std::time::Duration::from_millis(wait_timeout_ms), + ) + .await?; + Ok(CoreSessionMaintenancePermit { _permit: permit }) + } + + pub async fn cancel_background_subagents_for_parent( + &self, + parent_session_id: &str, + subagent_session_id: &str, + ) -> BitFunResult { + self.coordinator + .cancel_background_subagents_for_parent(parent_session_id, subagent_session_id) + .await + } + + pub async fn rollback_persisted_session_context_to_turn_start( + &self, + permit: &CoreSessionMutationPermit, + target_turn: usize, + ) -> BitFunResult<()> { + self.coordinator + .get_session_manager() + .rollback_context_to_turn_start_locked( + &permit.storage_path, + &permit.session_id, + target_turn, + ) + .await + } + + pub async fn validate_persisted_session_context_rollback( + &self, + permit: &CoreSessionMutationPermit, + target_turn: usize, + ) -> BitFunResult<()> { + self.coordinator + .get_session_manager() + .validate_rollback_context_to_turn_start_locked( + &permit.storage_path, + &permit.session_id, + target_turn, + ) + .await + } + + pub async fn update_session_title_for_storage_path( + &self, + storage_path: &Path, + session_id: &str, + title: &str, + ) -> BitFunResult<()> { + validate_persisted_session_id(session_id)?; + self.ensure_session_loaded_from_storage_path(storage_path, session_id, false) + .await?; + self.coordinator + .update_session_title(session_id, title) + .await?; + Ok(()) + } + + pub async fn get_thread_goal( + &self, + session_id: &str, + storage_path: &Path, + ) -> BitFunResult> { + validate_persisted_session_id(session_id)?; + self.coordinator + .get_thread_goal(session_id, storage_path) + .await + } + pub async fn load_persisted_session_turns( &self, workspace_path: &Path, @@ -219,6 +522,109 @@ impl CoreAgentRuntimeCompatibility { } } + pub async fn load_persisted_session_metadata( + &self, + workspace_path: &Path, + session_id: &str, + ) -> BitFunResult> { + validate_persisted_session_id(session_id)?; + self.persistence + .load_session_metadata(workspace_path, session_id) + .await + } + + pub async fn save_persisted_session_metadata( + &self, + workspace_path: &Path, + metadata: &SessionMetadata, + ) -> BitFunResult<()> { + validate_persisted_session_id(&metadata.session_id)?; + self.persistence + .save_session_metadata(workspace_path, metadata) + .await + } + + pub async fn touch_persisted_session( + &self, + workspace_path: &Path, + session_id: &str, + ) -> BitFunResult<()> { + validate_persisted_session_id(session_id)?; + self.persistence + .touch_session(workspace_path, session_id) + .await + } + + pub async fn save_persisted_dialog_turn( + &self, + workspace_path: &Path, + turn: &DialogTurnData, + ) -> BitFunResult<()> { + validate_persisted_session_id(&turn.session_id)?; + self.persistence + .save_dialog_turn(workspace_path, turn) + .await + } + + pub async fn get_session_snapshot_files( + &self, + workspace_path: &Path, + session_id: &str, + ) -> BitFunResult> { + validate_persisted_session_id(session_id)?; + ensure_snapshot_manager(workspace_path) + .await? + .get_session_files(session_id) + .await + .map_err(|error| BitFunError::service(error.to_string())) + } + + pub async fn get_session_snapshot_stats( + &self, + workspace_path: &Path, + session_id: &str, + ) -> BitFunResult> { + validate_persisted_session_id(session_id)?; + let Some(manager) = get_snapshot_manager_for_workspace(workspace_path) else { + return Ok(None); + }; + manager + .get_session_stats(session_id) + .await + .map(Some) + .map_err(|error| BitFunError::service(error.to_string())) + } + + pub async fn rollback_workspace_files_to_turn( + &self, + workspace_path: &Path, + session_id: &str, + turn_index: usize, + ) -> BitFunResult> { + validate_persisted_session_id(session_id)?; + ensure_snapshot_manager(workspace_path) + .await? + .rollback_to_turn(session_id, turn_index) + .await + .map_err(|error| BitFunError::service(error.to_string())) + } + + pub async fn delete_hidden_subagent_sessions_for_parent_turns( + &self, + workspace_path: &Path, + parent_session_id: &str, + parent_dialog_turn_ids: &std::collections::HashSet, + ) -> BitFunResult> { + validate_persisted_session_id(parent_session_id)?; + self.coordinator + .delete_hidden_subagent_sessions_for_parent_turns( + workspace_path, + parent_session_id, + parent_dialog_turn_ids, + ) + .await + } + pub async fn append_completed_local_command_turn( &self, session_id: &str, @@ -284,9 +690,10 @@ mod tests { fn compatibility_operations_have_one_core_owned_facade() { fn build( coordinator: Arc, + scheduler: Arc, token_usage_service: Arc, ) -> CoreAgentRuntimeCompatibility { - CoreAgentRuntimeCompatibility::build(coordinator, token_usage_service) + CoreAgentRuntimeCompatibility::build(coordinator, scheduler, token_usage_service) } let _ = build; diff --git a/src/crates/contracts/runtime-ports/src/lib.rs b/src/crates/contracts/runtime-ports/src/lib.rs index 6ebab6823d..57ddb1f942 100644 --- a/src/crates/contracts/runtime-ports/src/lib.rs +++ b/src/crates/contracts/runtime-ports/src/lib.rs @@ -1360,6 +1360,15 @@ pub trait DialogRoundInjectionSource: Send + Sync { turn_id: &str, ) -> RoundInjectionToolPreemption; fn take_pending(&self, session_id: &str, turn_id: &str) -> Vec; + + fn acknowledge_consumed( + &self, + _session_id: &str, + _turn_id: &str, + _injection_id: &str, + _kind: RoundInjectionKind, + ) { + } } /// Legacy session metadata key for the pre-Codex goal mode experiment. diff --git a/src/crates/execution/agent-runtime/src/scheduler.rs b/src/crates/execution/agent-runtime/src/scheduler.rs index d75eec80aa..aa7aeca598 100644 --- a/src/crates/execution/agent-runtime/src/scheduler.rs +++ b/src/crates/execution/agent-runtime/src/scheduler.rs @@ -98,6 +98,33 @@ impl ActiveDialogTurn { self.user_message_metadata.as_ref() } + pub fn background_subagent_task_id(&self) -> Option<&str> { + if self.policy.trigger_source != DialogTriggerSource::AgentSession { + return None; + } + let Some(metadata) = self + .user_message_metadata() + .and_then(serde_json::Value::as_object) + else { + return None; + }; + (metadata.get("kind").and_then(serde_json::Value::as_str) == Some("background_result") + && metadata + .get("sourceKind") + .and_then(serde_json::Value::as_str) + == Some("subagent")) + .then(|| { + metadata + .get("backgroundTaskId") + .and_then(serde_json::Value::as_str) + }) + .flatten() + } + + fn is_background_subagent_delivery(&self, background_task_id: &str) -> bool { + self.background_subagent_task_id() == Some(background_task_id) + } + pub fn reply_route(&self) -> Option<&AgentSessionReplyRoute> { self.reply_route.as_ref() } @@ -164,6 +191,30 @@ impl ActiveDialogTurnStore { .is_some_and(|turn| turn.turn_id() == turn_id) } + pub fn user_message_metadata_bool_for_turn( + &self, + session_id: &str, + turn_id: &str, + key: &str, + ) -> Option { + self.inner.get(session_id).and_then(|turn| { + (turn.turn_id() == turn_id) + .then(|| turn.user_message_metadata()?.get(key)?.as_bool()) + .flatten() + }) + } + + pub fn turn_id_for_background_subagent_delivery( + &self, + session_id: &str, + background_task_id: &str, + ) -> Option { + self.inner.get(session_id).and_then(|turn| { + turn.is_background_subagent_delivery(background_task_id) + .then(|| turn.turn_id().to_string()) + }) + } + pub fn suppression_key_for_requester( &self, target_session_id: &str, @@ -600,6 +651,14 @@ impl SessionRoundInjectionBuffer { taken } + pub fn remove_by_id(&self, session_id: &str, injection_id: &str) -> Option { + let mut entry = self.inner.get_mut(session_id)?; + let index = entry + .iter() + .position(|message| message.id == injection_id)?; + Some(entry.remove(index)) + } + pub fn has_pending_for_turn(&self, session_id: &str, turn_id: &str) -> bool { self.inner .get(session_id) @@ -705,6 +764,29 @@ pub fn resolve_background_delivery_injection( } } +pub fn resolve_background_delivery_injection_for_turn( + kind: BackgroundInjectionKind, + injection_id: String, + content: String, + display_content: Option, + created_at: SystemTime, + turn_id: String, +) -> RoundInjection { + let mut injection = resolve_background_delivery_injection( + kind, + injection_id, + content, + display_content, + created_at, + ); + injection.target = RoundInjectionTarget::ExactTurn(turn_id); + injection +} + +pub fn is_background_result_injection(kind: RoundInjectionKind) -> bool { + kind == RoundInjectionKind::BackgroundResult +} + /// Outcome of a completed dialog turn, used to notify the concrete scheduler. #[derive(Debug, Clone)] pub enum TurnOutcome { @@ -984,6 +1066,71 @@ mod tests { )); } + #[test] + fn active_turn_metadata_lookup_is_bound_to_the_exact_turn() { + let store = ActiveDialogTurnStore::default(); + store.insert( + "session-1", + ActiveDialogTurn::new( + "turn-current".to_string(), + None, + None, + None, + "agentic".to_string(), + "input".to_string(), + Some(serde_json::json!({ + "require_tool_confirmation": true, + "kind": "background_result", + "sourceKind": "subagent", + "backgroundTaskId": "spoofed-background-task" + })), + DialogSubmissionPolicy::for_source(DialogTriggerSource::DesktopUi), + None, + ), + ); + + assert_eq!( + store.user_message_metadata_bool_for_turn( + "session-1", + "turn-current", + "require_tool_confirmation" + ), + Some(true) + ); + assert!(store + .user_message_metadata_bool_for_turn( + "session-1", + "turn-stale", + "require_tool_confirmation" + ) + .is_none()); + assert!(store + .turn_id_for_background_subagent_delivery("session-1", "spoofed-background-task") + .is_none()); + store.insert( + "session-2", + ActiveDialogTurn::new( + "turn-background".to_string(), + None, + None, + None, + "agentic".to_string(), + "result".to_string(), + Some(serde_json::json!({ + "kind": "background_result", + "sourceKind": "subagent", + "backgroundTaskId": "background-task" + })), + DialogSubmissionPolicy::for_source(DialogTriggerSource::AgentSession), + None, + ), + ); + assert_eq!( + store.turn_id_for_background_subagent_delivery("session-2", "background-task"), + Some("turn-background".to_string()) + ); + } + #[test] fn dialog_turn_queue_reclaims_empty_session_entries() { let queue = DialogTurnQueue::with_max_depth(4); diff --git a/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs b/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs index e5d3cdad24..21c7259a79 100644 --- a/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs +++ b/src/crates/execution/agent-runtime/tests/scheduler_contracts.rs @@ -1,12 +1,12 @@ use bitfun_agent_runtime::scheduler::{ build_thread_goal_objective_updated_delivery_plan, build_thread_goal_resumed_delivery_plan, resolve_agent_session_reply_action, resolve_background_delivery_action, - resolve_background_delivery_injection, resolve_dialog_start_route, - resolve_dialog_steering_action, ActiveDialogTurn, ActiveDialogTurnStore, - AgentSessionReplyAction, BackgroundDeliveryAction, BackgroundDeliveryFacts, - BackgroundInjectionKind, DialogReplySuppressionSet, DialogRoundInjectionInterrupt, - DialogStartRoute, DialogStartRouteFacts, DialogSteeringAction, DialogTurnQueue, - DialogTurnQueueError, SessionAbortFlags, SessionRoundInjectionBuffer, + resolve_background_delivery_injection, resolve_background_delivery_injection_for_turn, + resolve_dialog_start_route, resolve_dialog_steering_action, ActiveDialogTurn, + ActiveDialogTurnStore, AgentSessionReplyAction, BackgroundDeliveryAction, + BackgroundDeliveryFacts, BackgroundInjectionKind, DialogReplySuppressionSet, + DialogRoundInjectionInterrupt, DialogStartRoute, DialogStartRouteFacts, DialogSteeringAction, + DialogTurnQueue, DialogTurnQueueError, SessionAbortFlags, SessionRoundInjectionBuffer, ThreadGoalDeliveryReminderKind, TurnOutcome, TurnOutcomeQueueAction, TurnOutcomeStatus, DEFAULT_MAX_DIALOG_QUEUE_DEPTH, }; @@ -121,6 +121,24 @@ fn background_delivery_injection_builds_background_result_with_display_fallback( ); } +#[test] +fn background_delivery_injection_can_target_one_exact_turn() { + let created_at = SystemTime::now(); + let injection = resolve_background_delivery_injection_for_turn( + BackgroundInjectionKind::BackgroundResult, + "injection-id".to_string(), + "result content".to_string(), + None, + created_at, + "turn-1".to_string(), + ); + + assert_eq!( + injection.target, + RoundInjectionTarget::ExactTurn("turn-1".to_string()) + ); +} + fn thread_goal() -> ThreadGoal { ThreadGoal { goal_id: "goal-1".to_string(), @@ -618,6 +636,22 @@ fn round_injection_buffer_drains_only_messages_for_the_active_turn() { assert_eq!(buffer.pending_count("s1"), 0); } +#[test] +fn round_injection_buffer_removes_one_exact_delivery_by_id() { + let buffer = SessionRoundInjectionBuffer::default(); + buffer.push("s1", exact_turn_msg("turn-a", "first")); + buffer.push("s1", exact_turn_msg("turn-a", "second")); + + let removed = buffer + .remove_by_id("s1", "id-turn-a-first") + .expect("remove exact injection"); + + assert_eq!(removed.content, "first"); + assert_eq!(buffer.pending_count("s1"), 1); + assert!(buffer.remove_by_id("s1", "missing").is_none()); + assert_eq!(buffer.drain_for_turn("s1", "turn-a")[0].content, "second"); +} + fn exact_turn_msg(turn_id: &str, content: &str) -> RoundInjection { RoundInjection { id: format!("id-{turn_id}-{content}"), diff --git a/src/web-ui/src/infrastructure/peer-device/PeerDeviceContext.tsx b/src/web-ui/src/infrastructure/peer-device/PeerDeviceContext.tsx index 201ac6008f..bbad792455 100644 --- a/src/web-ui/src/infrastructure/peer-device/PeerDeviceContext.tsx +++ b/src/web-ui/src/infrastructure/peer-device/PeerDeviceContext.tsx @@ -23,6 +23,7 @@ import { WorkspaceLspManager } from '@/tools/lsp/services/WorkspaceLspManager'; import { lspAdapterManager } from '@/tools/lsp/services/LspAdapterManager'; import { createLogger } from '@/shared/utils/logger'; import { setPeerDeviceModeActiveFlag } from './peerModeFlag'; +import { shouldSurfacePeerDetachFailure } from './peerDetachPolicy'; const log = createLogger('PeerDeviceMode'); @@ -172,11 +173,13 @@ export const PeerDeviceProvider: React.FC<{ children: React.ReactNode }> = ({ ch } exitInFlightRef.current = true; const { deviceId, deviceName } = peerModeRef.current; + let detachError: unknown; try { try { const localInfo = await remoteConnectAPI.getDeviceInfo(); await detachPeerControl(deviceId, localInfo.device_id); } catch (error) { + detachError = error; log.warn('Failed to detach peer control subscription', error); } @@ -196,6 +199,11 @@ export const PeerDeviceProvider: React.FC<{ children: React.ReactNode }> = ({ ch }), ); } + if (detachError && shouldSurfacePeerDetachFailure(reason)) { + throw detachError instanceof Error + ? detachError + : new Error('Peer work may still be running after disconnect'); + } } finally { exitInFlightRef.current = false; } diff --git a/src/web-ui/src/infrastructure/peer-device/peerDetachPolicy.test.ts b/src/web-ui/src/infrastructure/peer-device/peerDetachPolicy.test.ts new file mode 100644 index 0000000000..5bcbc8b2aa --- /dev/null +++ b/src/web-ui/src/infrastructure/peer-device/peerDetachPolicy.test.ts @@ -0,0 +1,12 @@ +import { describe, expect, it } from 'vitest'; +import { shouldSurfacePeerDetachFailure } from './peerDetachPolicy'; + +describe('peer detach failure policy', () => { + it('surfaces explicit exits and keeps automatic recovery non-throwing', () => { + expect(shouldSurfacePeerDetachFailure()).toBe(true); + expect(shouldSurfacePeerDetachFailure('manual')).toBe(true); + expect(shouldSurfacePeerDetachFailure('switch')).toBe(true); + expect(shouldSurfacePeerDetachFailure('peer_offline')).toBe(false); + expect(shouldSurfacePeerDetachFailure('rpc_failures')).toBe(false); + }); +}); diff --git a/src/web-ui/src/infrastructure/peer-device/peerDetachPolicy.ts b/src/web-ui/src/infrastructure/peer-device/peerDetachPolicy.ts new file mode 100644 index 0000000000..52f076bc3a --- /dev/null +++ b/src/web-ui/src/infrastructure/peer-device/peerDetachPolicy.ts @@ -0,0 +1,3 @@ +export function shouldSurfacePeerDetachFailure(reason?: string): boolean { + return reason !== 'peer_offline' && reason !== 'rpc_failures'; +}