diff --git a/docs/architecture/agent-runtime-services-design.md b/docs/architecture/agent-runtime-services-design.md index 9c6a115fa..aeb31af1e 100644 --- a/docs/architecture/agent-runtime-services-design.md +++ b/docs/architecture/agent-runtime-services-design.md @@ -798,10 +798,20 @@ Core 的 Network、Git 和 MCP Catalog 当前仍含兼容 marker，因此该诊 当前 Runtime SDK 已提供会话创建、列出、删除、恢复、模型/模式更新、类型化转录读取、本地分支、用量生成，以及 精确轮次结算。模型与模式更新只接受会话 ID 和被选择的稳定 ID， -不承载目录、提供方配置、选择策略或宿主 UI 语义。`AgentSessionRestoreRequest/Result` 与 +不承载目录、提供方配置、选择策略或宿主 UI 语义。模式 ID 由 Core 对当前有效目录校验；同值更新不刷新活动时间， +有效变更先按会话串行化，再通过与轮次元数据保存、删除共用的规范化物理路径锁持久化，成功后才提交到活动会话。 +该锁只保证单进程内同一会话元数据的读改写顺序，不宣称跨进程事务，也不把 metadata/state 等多文件更新声明为崩溃原子提交。 +恢复标准主会话时，如果持久化模式已从当前目录移除， +Core 会选择仍可执行的内置回退模式并同步修正元数据；内部子会话的专用模式不走这条迁移规则。 +Desktop 的通用会话元数据命令必须显式声明 UI 字段集合，并只在 owner 锁内更新这些字段；Review 状态与未读、关注、标题等 +独立写入意图不得通过完整旧快照互相覆盖。 +Relay 导入的会话元数据以私有 `pending/complete` 标记区分仅有摘要与完整历史；打开本地会话时先检查并补齐该导入， +再让 Core 恢复模型上下文。普通本地会话只做一次元数据读取且不访问账号或网络；部分导入失败保持可重试并停止本次打开， +不能让 UI 与 Core 分别发布不同的截断历史。 +`AgentSessionRestoreRequest/Result` 与 `AgentSessionRestorePort` 归 Agent Runtime SDK，以继续复用 Runtime owner 的完整 `SessionState`；类型化 `SessionTranscript` 归 `runtime-ports`。两者都由 `assembly/core` 注入真实持久化 owner，当前由 CLI/TUI -消费；ACP 的活动会话模式更新也通过窄端口回到同一 Core owner。ACP 为保证模型配置与完整历史来自同一次恢复，继续通过 +消费；TUI 与 ACP 的活动会话模式更新也通过窄端口回到同一 Core owner。ACP 为保证模型配置与完整历史来自同一次恢复，继续通过 Core 兼容门面读取协议回放所需的完整轮次，避免为单一协议扩张通用 transcript，也不绕过附件内容所需的独立授权能力。 会话分支请求显式携带可选远程身份；当前本地 provider 对远程身份返回 `NotAvailable`，不据本地路径 推断远程语义。CLI/TUI 的工具确认、拒绝和用户问题回答，以及 ACP 服务端 / Peer Host 的工具确认与拒绝，通过类型化 @@ -876,6 +886,20 @@ CLI 托管的 ACP 服务端使用 `DeliveryProfile::Acp` 组装一个 Agent Runt 活动会话的模型/模式写入已经通过 SDK 回到 Core owner， 不据此扩张通用 runtime DTO。 +`session/load` 在恢复前占用会话 ID，先完成纯参数校验和临时 MCP 建立，再恢复 Core；只有完整历史通知发送成功后才发布 +活动 ACP 状态并返回成功。同一 ID 的重叠打开或关闭在产生回放和 MCP 副作用前以可重试临时状态拒绝；恢复后的任一步失败 +都会卸载本请求加载的 Core 内存状态并回收临时 MCP，但不删除既有历史。`session/new` 先生成稳定 ID，完成目录校验和 +临时 MCP 建立后再以同一 ID 创建 Core 会话；建立过程失败时尝试回收临时 MCP 和本请求创建的 Core 会话。首次落盘若因 +回滚失败留下目录，会以类型化残余资源结果进入同一补偿路径，不报告“未创建 Core 会话”。补偿失败会携带会话 ID、 +残余资源种类、Core 是否由本请求创建及恢复动作，不伪装成普通输入错误，也不建议 `session/load` 删除既有历史。 +成功的 `session/close` 先阻止新轮次，清空已接收队列、取消后台子会话与活动轮次并确认调度排空，再卸载 Core 临时状态、 +回收临时 MCP 和连接映射；持久化历史及其存储绑定保留，可由后续 `session/load` 重新打开。 +任一步未完成时保留 ACP 会话所有权和持久化历史，返回 `session_close_incomplete`、失败阶段与可重试动作；只有临时 MCP +未回收时才标记具体残余资源，不能沿用 `session/new` 的“本请求创建 Core 会话”语义。 +无效请求与会话不存在分别保持协议可识别的参数错误和资源不存在错误，其他后端故障不泄漏为可重试的客户端输入错误。 +活动会话占用范围是一个 ACP stdio 进程；当前不宣称同一持久化会话可由多个 ACP 进程并发写入。跨进程共享需要先定义 +执行域、权限、冲突和崩溃恢复契约，不在本切片中用临时文件锁提前固化。 + 继续拥有： - ACP protocol。 diff --git a/docs/architecture/cli-product-line-design.md b/docs/architecture/cli-product-line-design.md index b08af076c..b98482c6f 100644 --- a/docs/architecture/cli-product-line-design.md +++ b/docs/architecture/cli-product-line-design.md @@ -97,7 +97,12 @@ BitFun CLI 应成为可独立安装和发布的 Agent 产品，而不是 Desktop `exec --session-id` 和缺失后端会话通过独立固定 ID 方法按原 ID 重建）/列举/删除/恢复、类型化转录、本地分支、 用量生成、轮次提交/取消和精确结算；普通创建 DTO 保持 v1 字段集合，固定 ID 冲突返回 `InvalidRequest`。会话模型更新、工具确认/拒绝和用户问题回答也通过 - SDK 的窄端口回到 Core owner，ACP 的活动会话模式更新也复用同一窄端口；模型/模式目录和提供方配置仍由产品入口解释。TUI 用量卡片持久化、快照及 Peer Host/ACP 维护 + SDK 的窄端口回到 Core owner，TUI 与 ACP 的活动会话模式更新也复用同一窄端口；模型/模式目录和提供方配置仍由产品入口解释。 + TUI 模式切换以异步待提交状态调用 Core，期间终端输入、resize 和重绘保持响应，新的对话提交不会消费用户输入； + 只有 Core 校验并持久化成功后才更新本地显示，失败时保留原模式并给出可重试提示，同值选择不产生持久化写入。 + 等待期间可以切换或新建会话；只有原会话的发送继续等待。首次退出请求在持久化成功后自动退出，失败时留在界面提示重试；再次退出允许立即离开， + 并明确提示下次恢复以 Core 的持久化模式为准。恢复主会话时，已从当前目录移除的持久化模式由 Core 迁移到可执行回退模式； + TUI 对比恢复前后的会话摘要并显示模式变化，如果同时携带启动输入，只预填而不自动执行，须由用户确认后发送。TUI 用量卡片持久化、快照及 Peer Host/ACP 维护 等 SDK v1 缺口由一个 Core 兼容门面转发给原 owner。 - Agentic Event Queue 仍是唯一事件 owner；TUI、`exec` 与 Peer Host 使用独立广播订阅，不互相消费事件。 - 有界旧队列只承担兼容存储；达到容量时不得抑制广播。CLI 保持一个后台 drain，订阅方一旦报告 lag/closed， diff --git a/docs/architecture/product-architecture.md b/docs/architecture/product-architecture.md index 2ba44f3d4..82d026526 100644 --- a/docs/architecture/product-architecture.md +++ b/docs/architecture/product-architecture.md @@ -327,11 +327,11 @@ flowchart LR | 产品形态 | 当前扩展能力 | 入口行为 | |---|---|---| | Desktop / product-full | 生产入口仍依赖 `bitfun-core/product-full` 作为兼容组装层；“外部 AI 应用”设置已消费 OpenCode Command、standalone Tool 和 Subagent 的统一来源快照、审批、冲突选择与诊断；Skills 场景和设置列表显示已发现 Skill 的生态来源与覆盖结果 | 仅本地执行域支持当前三条可执行纵向切片；Skill 仍使用独立 Registry，不据来源展示宣称已并入外部来源协调器。受管 package plugin 仍只有静态预览，不能据设置页接入宣称完整 OpenCode 插件运行时 | -| CLI | 入口仍以 `bitfun-core/product-full` 作为执行兼容 owner；交互式 TUI 已可执行受支持的 Prompt Command，并在通用 `/tools` 与 `/agents` 入口按文字分组消费与 Desktop 相同的审批、冲突和刷新状态；`/agents` 同时承载主 Agent 切换、Subagent 管理和“外部 AI 应用”，不再注册 `/subagents` 或 `external-*` 平行命令。活动 turn 期间仍可查看和管理，只有主 Agent 切换被禁用。Skill 列表显示来源，模式配置按实际选择结果说明覆盖来源 | 已批准的 standalone Tool 进入现有 Tool Runtime；已批准的外部 Subagent 只支持 fresh single-run。CLI/TUI 不解析生态文件、不启动第二套 worker/Agent owner；非交互入口和 Remote 未接入时不得借本机 TUI 路径代执行。本地 Agent 与 Peer Host 路径选择 `DeliveryProfile::Cli` 并消费同一 Runtime Parts/SDK；主会话的恢复、转录、本地分支、用量报告生成与精确结算走 SDK，远程分支明确不支持。TUI 用量卡片持久化、快照及 Peer Host/ACP 持久化维护缺口由单一 Core 兼容门面转发；本切片不改变扩展执行边界 | +| CLI | 入口仍以 `bitfun-core/product-full` 作为执行兼容 owner；交互式 TUI 已可执行受支持的 Prompt Command，并在通用 `/tools` 与 `/agents` 入口按文字分组消费与 Desktop 相同的审批、冲突和刷新状态；`/agents` 同时承载主 Agent 切换、Subagent 管理和“外部 AI 应用”，不再注册 `/subagents` 或 `external-*` 平行命令。活动 turn 期间仍可查看和管理，只有主 Agent 切换被禁用。Skill 列表显示来源，模式配置按实际选择结果说明覆盖来源 | 已批准的 standalone Tool 进入现有 Tool Runtime；已批准的外部 Subagent 只支持 fresh single-run。CLI/TUI 不解析生态文件、不启动第二套 worker/Agent owner；非交互入口和 Remote 未接入时不得借本机 TUI 路径代执行。本地 Agent 与 Peer Host 路径选择 `DeliveryProfile::Cli` 并消费同一 Runtime Parts/SDK；主会话的恢复、转录、本地分支、用量报告生成、模式更新与精确结算走 SDK，远程分支明确不支持。TUI 模式切换异步持久化，期间保持输入和 resize 响应，只在成功后提交本地状态，失败保留原选择；用户可切换会话，退出有等待完成与再次强制退出两级路径。失效的主会话模式由 Core 恢复为可执行模式，TUI 显示迁移并阻止启动输入自动执行。用量卡片持久化、快照及 Peer Host/ACP 持久化维护缺口仍由单一 Core 兼容门面转发。本切片不改变扩展执行边界 | | HarmonyOS PC 原生 CLI/TUI | 未来平台目标，当前未实现 | 目标、问题和风险见平台规约；具体适配另立专题，HAP、手机 Remote App 与远端代执行均不替代 | | HarmonyOS PC GUI | 完整 HarmonyOS PC 支持的另一目标形态，当前未实现 | 与 CLI/TUI 共享稳定能力和 Runtime 语义，但独立设计宿主、界面与发布验证；Web、Remote 或现有 Tauri Desktop 均不能替代 | | HarmonyOS 手机 Remote App | `src/apps/mobile/harmonyos` 是 phone-only ArkTS 远程入口，不持有本地 Rust Agent Runtime | 保持当前能力并按移动端专题独立演进；本轮不提前设计移动 Runtime/TUI/GUI，也不能据此宣称 HarmonyOS PC 本地能力 | -| ACP | CLI 托管的服务端仍以 `bitfun-core/product-full` 作为兼容执行层 | 入口已选择 `DeliveryProfile::Acp` 并消费 Runtime Parts；组装层在入队前原子拒绝忙碌会话，不改变其他产品入口的排队行为；活动会话模型与模式写入走 SDK。会话恢复及完整历史、模型/模式目录与配置读取、MCP、客户端与协议生命周期仍留在现有 Core/ACP 归属，不据此宣称完整解耦 | +| ACP | CLI 托管的服务端仍以 `bitfun-core/product-full` 作为兼容执行层 | 入口已选择 `DeliveryProfile::Acp` 并消费 Runtime Parts；组装层在入队前原子拒绝忙碌会话，不改变其他产品入口的排队行为；活动会话模型与模式写入走 SDK。`session/load` 先校验和建立临时 MCP，再恢复 Core，在历史回放成功后才发布活动状态；失败会卸载本次内存状态而不删除历史。同 ID 的重叠打开/关闭被明确拒绝。成功的 `session/close` 阻止新轮次、排空队列和后台子会话，再卸载临时 Core 状态并回收 MCP 与连接；失败保留会话所有权和历史，返回可重试阶段。持久化历史仍可重新加载。完整历史、模型/模式目录与配置读取仍留在现有 Core/ACP 归属，不据此宣称完整解耦 | | Server / Remote | 当前生产路由没有插件状态消费闭环；Remote 插件执行未实现 | 不在本地替远端项目发现、准备或执行插件；未接入时返回明确不支持 | | Web / Mobile Web | 依赖现有后端入口，不持有插件执行单元 | 对应 profile 当前为空计划或未接入生产，不能据枚举值宣称独立产品能力 | | SDK | 仅有 preview 门面、空 profile 计划和测试替身 | 不牵引 `product-full`、具体服务管理器或插件 host ABI；未满足独立嵌入验证前不宣称可发布 | diff --git a/src/apps/cli/src/agent/runtime_client.rs b/src/apps/cli/src/agent/runtime_client.rs index fb36110c2..a91546cac 100644 --- a/src/apps/cli/src/agent/runtime_client.rs +++ b/src/apps/cli/src/agent/runtime_client.rs @@ -12,10 +12,10 @@ use tokio::sync::Mutex; use bitfun_agent_runtime::sdk::{ AgentDialogTurnRequest, AgentRuntime, AgentSessionCreateRequest, AgentSessionDeleteRequest, AgentSessionForkRequest, AgentSessionForkResult, AgentSessionListRequest, - AgentSessionModelUpdateRequest, AgentSessionRestoreRequest, AgentSessionUsageRequest, - AgentToolConfirmationRequest, AgentToolRejectionRequest, AgentTurnCancellationRequest, - AgentTurnSettlementRequest, AgentUserAnswersRequest, PortErrorKind, RuntimeError, - SessionTranscript, SessionTranscriptRequest, SessionUsageReport, + AgentSessionModeUpdateRequest, AgentSessionModelUpdateRequest, AgentSessionRestoreRequest, + AgentSessionUsageRequest, AgentToolConfirmationRequest, AgentToolRejectionRequest, + AgentTurnCancellationRequest, AgentTurnSettlementRequest, AgentUserAnswersRequest, + PortErrorKind, RuntimeError, SessionTranscript, SessionTranscriptRequest, SessionUsageReport, }; use bitfun_agent_runtime::user_questions::USER_INPUT_AVAILABLE_CONTEXT_KEY; use bitfun_runtime_ports::{AgentSessionSummary, AgentSubmissionSource, DialogSubmissionPolicy}; @@ -41,6 +41,31 @@ fn validated_session_summary( }) } +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) struct SessionModeMigrationNotice { + pub(crate) previous_mode_id: String, + pub(crate) restored_mode_id: String, +} + +impl SessionModeMigrationNotice { + pub(crate) fn user_message(&self) -> String { + format!( + "Session mode \"{}\" is unavailable. This session was restored with \"{}\". Review the mode before continuing.", + self.previous_mode_id, self.restored_mode_id + ) + } +} + +fn session_mode_migration_notice( + previous: &AgentSessionSummary, + restored: &AgentSessionSummary, +) -> Option { + (previous.agent_type != restored.agent_type).then(|| SessionModeMigrationNotice { + previous_mode_id: previous.agent_type.clone(), + restored_mode_id: restored.agent_type.clone(), + }) +} + /// CLI-owned client for the portable Agent Runtime SDK. /// Stateless regarding agent_type; callers pass it per call. pub(crate) struct CliAgentRuntimeClient { @@ -109,14 +134,19 @@ impl CliAgentRuntimeClient { pub(crate) async fn restore_session_in_current_workspace( &self, session_id: &str, - ) -> Result<(AgentSessionSummary, PathBuf)> { + ) -> Result<( + AgentSessionSummary, + PathBuf, + Option, + )> { tracing::info!("Restoring session: {}", session_id); let effective_workspace = self.current_workspace_path(); let sessions = self .list_sessions_in_workspace(&effective_workspace) .await?; - validated_session_summary(&sessions, session_id, &effective_workspace)?; + let previous_summary = + validated_session_summary(&sessions, session_id, &effective_workspace)?; let restored = self .runtime @@ -139,7 +169,8 @@ impl CliAgentRuntimeClient { *session_id_guard = Some(session_id.to_string()); *turn_id_guard = None; - Ok((restored.session, effective_workspace)) + let migration_notice = session_mode_migration_notice(&previous_summary, &restored.session); + Ok((restored.session, effective_workspace, migration_notice)) } pub(crate) async fn delete_session(&self, session_id: &str) -> Result<()> { @@ -178,6 +209,16 @@ impl CliAgentRuntimeClient { .map_err(|error| anyhow::anyhow!(error.into_message())) } + pub(crate) async fn update_session_mode(&self, session_id: &str, mode_id: &str) -> Result<()> { + self.runtime + .update_session_mode(AgentSessionModeUpdateRequest { + session_id: session_id.to_string(), + mode_id: mode_id.to_string(), + }) + .await + .map_err(|error| anyhow::anyhow!(error.into_message())) + } + pub(crate) async fn branch_session_at_latest_turn( &self, source_session_id: &str, @@ -552,7 +593,7 @@ mod tests { use bitfun_runtime_ports::AgentSessionSummary; - use super::validated_session_summary; + use super::{session_mode_migration_notice, validated_session_summary}; #[test] fn model_updates_use_the_runtime_sdk_without_the_core_compatibility_facade() { @@ -565,6 +606,20 @@ mod tests { assert!(!source.contains(&compatibility_update)); } + #[test] + fn mode_updates_use_the_runtime_sdk_without_the_core_compatibility_facade() { + let source = include_str!("runtime_client.rs").replace("\r\n", "\n"); + let runtime_update = ["self.runtime", "\n .update_session_mode"].concat(); + let compatibility_update = [ + "self.compatibility", + "\n .update_session_agent_type", + ] + .concat(); + + assert!(source.contains(&runtime_update)); + assert!(!source.contains(&compatibility_update)); + } + fn session_summary(session_id: &str) -> AgentSessionSummary { AgentSessionSummary { session_id: session_id.to_string(), @@ -605,4 +660,26 @@ mod tests { assert!(message.contains("session-from-another-workspace")); assert!(message.contains("D:/workspace/current")); } + + #[test] + fn restore_reports_a_cli_local_notice_when_core_migrates_the_mode() { + let previous = AgentSessionSummary { + agent_type: "removed-mode".to_string(), + ..session_summary("mode-migration") + }; + let restored = session_summary("mode-migration"); + + let notice = session_mode_migration_notice(&previous, &restored) + .expect("changed mode should be reported to the TUI"); + + assert_eq!(notice.previous_mode_id, "removed-mode"); + assert_eq!(notice.restored_mode_id, "agentic"); + } + + #[test] + fn restore_does_not_report_a_notice_when_the_mode_is_unchanged() { + let summary = session_summary("unchanged-mode"); + + assert!(session_mode_migration_notice(&summary, &summary).is_none()); + } } diff --git a/src/apps/cli/src/modes/chat.rs b/src/apps/cli/src/modes/chat.rs index 564fde17e..827af671c 100644 --- a/src/apps/cli/src/modes/chat.rs +++ b/src/apps/cli/src/modes/chat.rs @@ -139,6 +139,17 @@ enum PendingMcpTask { }, } +struct PendingModeChange { + session_id: String, + mode_id: String, + started_at: Instant, + slow_notice_shown: bool, + exit_warning_shown: bool, + handle: tokio::task::JoinHandle>, +} + +const MODE_CHANGE_SLOW_NOTICE: Duration = Duration::from_secs(15); + #[derive(Default)] struct NonKeyEventOutcome { request_redraw: bool, @@ -171,6 +182,9 @@ pub(crate) struct ChatMode { pending_mcp_op: Option, /// Running MCP tasks (non-blocking, polled in main loop) pending_mcp_tasks: Vec, + /// One durable mode update in flight. The event loop remains responsive + /// while the runtime owner writes session metadata. + pending_mode_change: Option, external_source_snapshot: Option, external_source_conflict_choices: BTreeMap, external_source_conflict_lineage_current_keys: BTreeMap, @@ -215,6 +229,7 @@ impl ChatMode { initial_prompt: None, pending_mcp_op: None, pending_mcp_tasks: Vec::new(), + pending_mode_change: None, external_source_snapshot: None, external_source_conflict_choices: BTreeMap::new(), external_source_conflict_lineage_current_keys: BTreeMap::new(), diff --git a/src/apps/cli/src/modes/chat/commands.rs b/src/apps/cli/src/modes/chat/commands.rs index 283d4ef0a..67b7ae7f1 100644 --- a/src/apps/cli/src/modes/chat/commands.rs +++ b/src/apps/cli/src/modes/chat/commands.rs @@ -1,3 +1,7 @@ +fn mode_change_blocks_typed_submission(pending_for_current_session: bool, input: &str) -> bool { + pending_for_current_session && !input.trim().starts_with('/') +} + impl ChatMode { /// Handle command palette action fn handle_palette_action( @@ -716,8 +720,19 @@ impl ChatMode { return self.handle_action_id(&action_id, chat_view, chat_state, rt_handle); } + let trimmed = chat_view.input_text().trim(); + let pending_for_current_session = self + .pending_mode_change + .as_ref() + .is_some_and(|pending| pending.session_id == chat_state.core_session_id); + if mode_change_blocks_typed_submission(pending_for_current_session, trimmed) { + chat_view.set_status(Some( + "Waiting for the agent mode change to finish before sending.".to_string(), + )); + return Ok(None); + } + if chat_state.is_processing { - let trimmed = chat_view.input_text().trim(); if trimmed.starts_with('/') { if let Some(input) = chat_view.send_input() { return self.handle_command(&input, chat_view, chat_state, rt_handle); diff --git a/src/apps/cli/src/modes/chat/input.rs b/src/apps/cli/src/modes/chat/input.rs index dc7898528..b2969ee92 100644 --- a/src/apps/cli/src/modes/chat/input.rs +++ b/src/apps/cli/src/modes/chat/input.rs @@ -417,6 +417,9 @@ impl ChatMode { } = context; match reason { ChatExitReason::SwitchSession(new_session_id) => { + if let Some(pending) = this.pending_mode_change.as_mut() { + pending.exit_warning_shown = false; + } match this.switch_to_session( &new_session_id, session_id, @@ -432,6 +435,9 @@ impl ChatMode { } } ChatExitReason::NewSession => { + if let Some(pending) = this.pending_mode_change.as_mut() { + pending.exit_warning_shown = false; + } match this.create_new_session(session_id, chat_state, chat_view, rt_handle) { Ok(()) => tracing::info!("Created new session: {}", session_id), Err(e) => { @@ -441,9 +447,19 @@ impl ChatMode { } } } - other => { + ChatExitReason::Quit => { + if let Some(pending) = this.pending_mode_change.as_mut() { + if !pending.exit_warning_shown { + pending.exit_warning_shown = true; + chat_view.set_status(Some( + "Exit requested. Waiting for the agent mode change to finish; exit again to leave now. This mode change may not be saved, and the next restore will use the last successfully persisted mode." + .to_string(), + )); + return; + } + } *should_quit = true; - *exit_reason = other; + *exit_reason = ChatExitReason::Quit; } } } diff --git a/src/apps/cli/src/modes/chat/run.rs b/src/apps/cli/src/modes/chat/run.rs index 07999bd2f..bb1f81912 100644 --- a/src/apps/cli/src/modes/chat/run.rs +++ b/src/apps/cli/src/modes/chat/run.rs @@ -30,8 +30,8 @@ impl ChatMode { // Create or restore core session let rt_handle = tokio::runtime::Handle::current(); - let (mut session_id, mut chat_state) = if let Some(ref restore_id) = self.restore_session_id - { + let (mut session_id, mut chat_state, mode_migration_notice) = + if let Some(ref restore_id) = self.restore_session_id { // Restore existing session tracing::info!("Restoring session: {}", restore_id); let agent = self.agent.clone(); @@ -40,7 +40,7 @@ impl ChatMode { tokio::task::block_in_place(|| { rt_handle.block_on(async { // Restore session in core (loads metadata, messages, managers) - let (summary, effective_workspace_path) = + let (summary, effective_workspace_path, migration_notice) = agent.restore_session_in_current_workspace(&rid).await?; let effective_workspace = Some(effective_workspace_path.to_string_lossy().to_string()); @@ -67,7 +67,7 @@ impl ChatMode { transcript.messages.len() ); - Ok::<_, anyhow::Error>((rid, state)) + Ok::<_, anyhow::Error>((rid, state, migration_notice)) }) })? } else { @@ -83,7 +83,7 @@ impl ChatMode { self.agent_type.clone(), self.workspace.clone(), ); - (session_id, state) + (session_id, state, None) }; // Keep ChatMode workspace in sync with the session's effective workspace @@ -159,13 +159,23 @@ impl ChatMode { let mut event_rx = self.agent.event_source().subscribe(); + if let Some(notice) = &mode_migration_notice { + chat_state.add_system_message(notice.user_message()); + } + // Send initial prompt if provided (from startup page input) if let Some(prompt) = self.initial_prompt.take() { - tracing::info!("Sending initial prompt: {}", prompt); - if prompt.starts_with('/') { + if mode_migration_notice.is_some() { + chat_view.text_input.set_text(&prompt); + chat_view.set_status(Some( + "The restored session uses a fallback mode. Review it, then send the preserved input explicitly." + .to_string(), + )); + } else if prompt.starts_with('/') { // Slash commands will be handled in the main loop chat_view.text_input.set_text(&prompt); } else { + tracing::info!("Sending initial prompt: {}", prompt); let display_name = agent_display_name(&self.agent_type); chat_view.set_status(Some(format!("{} is thinking...", display_name))); @@ -219,6 +229,15 @@ impl ChatMode { if self.poll_mcp_task_completion(&mut chat_view, &mut chat_state, &rt_handle) { needs_redraw = true; } + match self.poll_mode_change_completion(&mut chat_view, &mut chat_state, &rt_handle) { + ModeChangePollOutcome::NoChange => {} + ModeChangePollOutcome::Redraw => needs_redraw = true, + ModeChangePollOutcome::ExitAfterSave => { + should_quit = true; + exit_reason = ChatExitReason::Quit; + continue; + } + } if self.poll_external_tool_mutation(&mut chat_view) { needs_redraw = true; } diff --git a/src/apps/cli/src/modes/chat/selection.rs b/src/apps/cli/src/modes/chat/selection.rs index 5bcc5b67e..364d39024 100644 --- a/src/apps/cli/src/modes/chat/selection.rs +++ b/src/apps/cli/src/modes/chat/selection.rs @@ -5,6 +5,62 @@ enum ModelSelectionApplyOutcome { }, } +enum ModeSelectionApplyOutcome { + SessionUpdateFailed(String), + Applied, +} + +enum ModeChangePollOutcome { + NoChange, + Redraw, + ExitAfterSave, +} + +fn previous_session_mode_change_status( + mode_id: &str, + outcome: &ModeSelectionApplyOutcome, +) -> String { + match outcome { + ModeSelectionApplyOutcome::Applied => format!( + "The previous session mode was changed to {mode_id}; the current session was not modified." + ), + ModeSelectionApplyOutcome::SessionUpdateFailed(error) => format!( + "The previous session mode change to {mode_id} failed: {error}. Return to that session to retry." + ), + } +} + +fn mode_change_completion_should_exit(exit_requested: bool, applied: bool) -> bool { + exit_requested && applied +} + +fn apply_agent_mode_feedback( + current_mode: &mut String, + chat_state: &mut ChatState, + selected_mode: &str, + outcome: ModeSelectionApplyOutcome, +) -> bool { + match outcome { + ModeSelectionApplyOutcome::SessionUpdateFailed(error) => { + tracing::error!( + "Failed to switch agent mode to {}: {}", + selected_mode, + error + ); + chat_state.add_system_message(format!( + "Agent mode was not changed: {error}. Please retry." + )); + false + } + ModeSelectionApplyOutcome::Applied => { + *current_mode = selected_mode.to_string(); + chat_state.agent_type = selected_mode.to_string(); + tracing::info!("Agent mode switched to: {}", selected_mode); + true + } + } +} + fn usage_report_metadata(report: &SessionUsageReport) -> Result { let usage_report = serde_json::to_value(report) .map_err(|error| anyhow!("Failed to serialize usage report: {error}"))?; @@ -263,10 +319,16 @@ impl ChatMode { fn switch_agent_by_offset( &mut self, offset: isize, - _chat_view: &mut ChatView, + chat_view: &mut ChatView, chat_state: &mut ChatState, rt_handle: &tokio::runtime::Handle, ) { + if !agent_mode_switch_allowed(chat_state.is_processing, self.pending_mode_change.is_some()) { + chat_view.set_status(Some(mode_switch_unavailable_message( + chat_state.is_processing, + ))); + return; + } let modes = self.get_mode_agents(rt_handle); if modes.len() <= 1 { return; @@ -281,8 +343,11 @@ impl ChatMode { let next_idx = ((current_idx as isize + offset) % len + len) % len; let next = &modes[next_idx as usize]; - self.agent_type = next.id.clone(); - chat_state.agent_type = next.id.clone(); + let selected = AgentItem { + id: next.id.clone(), + description: next.description.clone(), + }; + self.apply_agent_selection(&selected, chat_view, chat_state, rt_handle); } /// Load current model name from global config for display @@ -473,7 +538,10 @@ impl ChatMode { agent_items, Some(self.agent_type.clone()), true, - agent_mode_switch_allowed(chat_state.is_processing), + agent_mode_switch_allowed( + chat_state.is_processing, + self.pending_mode_change.is_some(), + ), ); } @@ -486,15 +554,17 @@ impl ChatMode { ) { match action { AgentSelectorAction::SwitchMode(selected) => { - if !agent_mode_switch_allowed(chat_state.is_processing) { - chat_view.set_status(Some( - "Agent mode cannot be changed during the current turn. Subagent and external source management remain available." - .to_string(), - )); + if !agent_mode_switch_allowed( + chat_state.is_processing, + self.pending_mode_change.is_some(), + ) { + chat_view.set_status(Some(mode_switch_unavailable_message( + chat_state.is_processing, + ))); return; } chat_view.hide_agent_selector(); - self.apply_agent_selection(&selected, chat_state); + self.apply_agent_selection(&selected, chat_view, chat_state, rt_handle); } AgentSelectorAction::ManageSubagents => { self.show_subagent_selector(chat_view, chat_state, rt_handle); @@ -507,15 +577,107 @@ impl ChatMode { } /// Apply agent selection: switch agent type - fn apply_agent_selection(&mut self, selected: &AgentItem, chat_state: &mut ChatState) { + fn apply_agent_selection( + &mut self, + selected: &AgentItem, + chat_view: &mut ChatView, + chat_state: &mut ChatState, + rt_handle: &tokio::runtime::Handle, + ) { if selected.id == self.agent_type { return; } - self.agent_type = selected.id.clone(); - chat_state.agent_type = selected.id.clone(); - tracing::info!("Switched to agent: {}", selected.id); - if selected.id == "HarmonyOSDev" { + if self.pending_mode_change.is_some() { + chat_view.set_status(Some( + "An agent mode change is already in progress. Please wait.".to_string(), + )); + return; + } + + let session_id = chat_state.core_session_id.clone(); + let mode_id = selected.id.clone(); + let task_mode_id = mode_id.clone(); + let agent = self.agent.clone(); + chat_view.set_status(Some(format!("Switching agent mode to {mode_id}..."))); + let task_session_id = session_id.clone(); + let handle = rt_handle.spawn(async move { + agent + .update_session_mode(&task_session_id, &task_mode_id) + .await + }); + self.pending_mode_change = Some(PendingModeChange { + session_id, + mode_id, + started_at: Instant::now(), + slow_notice_shown: false, + exit_warning_shown: false, + handle, + }); + } + + fn poll_mode_change_completion( + &mut self, + chat_view: &mut ChatView, + chat_state: &mut ChatState, + rt_handle: &tokio::runtime::Handle, + ) -> ModeChangePollOutcome { + let Some(pending) = self.pending_mode_change.as_mut() else { + return ModeChangePollOutcome::NoChange; + }; + if !pending.handle.is_finished() { + if !pending.slow_notice_shown && pending.started_at.elapsed() >= MODE_CHANGE_SLOW_NOTICE + { + pending.slow_notice_shown = true; + if !pending.exit_warning_shown { + chat_view.set_status(Some( + "The agent mode change is still being saved. You can edit or switch sessions; sending in this session waits." + .to_string(), + )); + } + return ModeChangePollOutcome::Redraw; + } + return ModeChangePollOutcome::NoChange; + } + let pending = self + .pending_mode_change + .take() + .expect("finished mode task should remain present"); + let outcome = match tokio::task::block_in_place(|| rt_handle.block_on(pending.handle)) { + Ok(Ok(())) => ModeSelectionApplyOutcome::Applied, + Ok(Err(error)) => ModeSelectionApplyOutcome::SessionUpdateFailed(error.to_string()), + Err(error) => ModeSelectionApplyOutcome::SessionUpdateFailed(format!( + "mode update task failed: {error}" + )), + }; + if chat_state.core_session_id != pending.session_id { + if let ModeSelectionApplyOutcome::SessionUpdateFailed(error) = &outcome { + tracing::error!( + "Failed to switch previous session {} to agent mode {}: {}", + pending.session_id, + pending.mode_id, + error + ); + } + chat_view.set_status(Some(previous_session_mode_change_status( + &pending.mode_id, + &outcome, + ))); + return ModeChangePollOutcome::Redraw; + } + let applied = apply_agent_mode_feedback( + &mut self.agent_type, + chat_state, + &pending.mode_id, + outcome, + ); + if applied { + chat_view.set_status(Some(format!("Agent mode set to {}", pending.mode_id))); + } else { + chat_view.set_status(Some("Agent mode change failed. Please retry.".to_string())); + } + + if applied && pending.mode_id == "HarmonyOSDev" { let deveco_home = std::env::var("DEVECO_HOME").ok(); let missing = deveco_home .as_deref() @@ -528,13 +690,27 @@ impl ChatMode { ); } } + if mode_change_completion_should_exit(pending.exit_warning_shown, applied) { + ModeChangePollOutcome::ExitAfterSave + } else { + ModeChangePollOutcome::Redraw + } } // ============ MCP management ============ } -fn agent_mode_switch_allowed(is_processing: bool) -> bool { - !is_processing +fn agent_mode_switch_allowed(is_processing: bool, mode_change_pending: bool) -> bool { + !is_processing && !mode_change_pending +} + +fn mode_switch_unavailable_message(is_processing: bool) -> String { + if is_processing { + "Agent mode cannot be changed during the current turn. Subagent and external source management remain available." + .to_string() + } else { + "An agent mode change is already in progress. Please wait.".to_string() + } } #[cfg(test)] @@ -543,8 +719,9 @@ mod usage_metadata_tests { #[test] fn mode_switch_is_rechecked_when_an_idle_popup_outlives_turn_start() { - assert!(agent_mode_switch_allowed(false)); - assert!(!agent_mode_switch_allowed(true)); + assert!(agent_mode_switch_allowed(false, false)); + assert!(!agent_mode_switch_allowed(true, false)); + assert!(!agent_mode_switch_allowed(false, true)); } #[test] diff --git a/src/apps/cli/src/modes/chat/sessions.rs b/src/apps/cli/src/modes/chat/sessions.rs index 67456ece7..5603614dc 100644 --- a/src/apps/cli/src/modes/chat/sessions.rs +++ b/src/apps/cli/src/modes/chat/sessions.rs @@ -11,9 +11,9 @@ impl ChatMode { let agent = self.agent.clone(); let sid = new_session_id.to_string(); - let (new_state, restored_agent_type) = tokio::task::block_in_place(|| { + let (new_state, restored_agent_type, migration_notice) = tokio::task::block_in_place(|| { rt_handle.block_on(async { - let (session_summary, effective_workspace_path) = + let (session_summary, effective_workspace_path, migration_notice) = agent.restore_session_in_current_workspace(&sid).await?; let restored_agent_type = session_summary.agent_type.clone(); let effective_workspace = @@ -35,7 +35,7 @@ impl ChatMode { &transcript, ); - Ok::<_, anyhow::Error>((state, restored_agent_type)) + Ok::<_, anyhow::Error>((state, restored_agent_type, migration_notice)) }) })?; @@ -48,6 +48,10 @@ impl ChatMode { // Reload model name self.load_current_model_name(chat_state, rt_handle); + if let Some(notice) = migration_notice { + chat_state.add_system_message(notice.user_message()); + } + // Reset view state chat_view.scroll_to_bottom(); chat_view.set_status(Some(format!("Switched to session: {}", new_session_id))); @@ -102,6 +106,16 @@ impl ChatMode { chat_state: &mut ChatState, rt_handle: &tokio::runtime::Handle, ) { + if self + .pending_mode_change + .as_ref() + .is_some_and(|pending| pending.session_id == chat_state.core_session_id) + { + chat_view.set_status(Some( + "Waiting for the agent mode change to finish before sending.".to_string(), + )); + return; + } if chat_state.is_processing { chat_state.add_system_message("Already processing, please wait.".to_string()); return; diff --git a/src/apps/cli/src/modes/chat/tests.rs b/src/apps/cli/src/modes/chat/tests.rs index ba5d11de0..246ea7856 100644 --- a/src/apps/cli/src/modes/chat/tests.rs +++ b/src/apps/cli/src/modes/chat/tests.rs @@ -4,16 +4,19 @@ mod tests { use super::{ action_opens_extension_management, agent_event_stream_failure, - apply_model_selection_feedback, builtin_command_reconfirmation, command_route, + apply_agent_mode_feedback, apply_model_selection_feedback, builtin_command_reconfirmation, + command_route, external_agent_attention, external_agent_diagnostic_lines, external_agent_pending_notice_key, external_agent_result_is_stale, external_agent_review_text, external_command_projections, external_tool_mutation_result_label, external_tool_pending_notice_key, external_tool_result_is_stale, external_tool_review_text, external_tool_run_location_label, mark_active_turn_failed, merge_external_agent_mutation_snapshot, + mode_change_blocks_typed_submission, mode_change_completion_should_exit, native_command_conflict_key, parse_command_token, parse_external_agent_review_action, parse_external_tool_review_action, CommandQualifier, CommandRoute, ExternalAgentReviewAction, ExternalSourceConflictPreferences, ExternalToolReviewAction, + previous_session_mode_change_status, ModeSelectionApplyOutcome, ModelSelectionApplyOutcome, }; use crate::actions::{action_conflict_behavior_version, ActionState, ResolvedKeymap}; @@ -659,6 +662,86 @@ mod tests { assert!(content.contains("retry")); } + #[test] + fn mode_selection_commits_visible_state_only_after_runtime_success() { + let mut current_mode = "agentic".to_string(); + let mut state = ChatState::new( + "session".to_string(), + "Session".to_string(), + "agentic".to_string(), + Some("D:/workspace/current".to_string()), + ); + + let applied = apply_agent_mode_feedback( + &mut current_mode, + &mut state, + "plan", + ModeSelectionApplyOutcome::Applied, + ); + + assert!(applied); + assert_eq!(current_mode, "plan"); + assert_eq!(state.agent_type, "plan"); + } + + #[test] + fn mode_selection_failure_preserves_visible_state_and_explains_retry() { + let mut current_mode = "agentic".to_string(); + let mut state = ChatState::new( + "session".to_string(), + "Session".to_string(), + "agentic".to_string(), + Some("D:/workspace/current".to_string()), + ); + + let applied = apply_agent_mode_feedback( + &mut current_mode, + &mut state, + "plan", + ModeSelectionApplyOutcome::SessionUpdateFailed( + "session storage unavailable".to_string(), + ), + ); + + assert!(!applied); + assert_eq!(current_mode, "agentic"); + assert_eq!(state.agent_type, "agentic"); + let notice = state.messages.last().expect("failure notice"); + let crate::chat_state::FlowItem::Text { content, .. } = ¬ice.flow_items[0] else { + panic!("failure notice must be text"); + }; + assert!(content.contains("was not changed")); + assert!(content.contains("retry")); + } + + #[test] + fn previous_session_mode_failure_is_not_reported_as_a_success() { + let status = previous_session_mode_change_status( + "Plan", + &ModeSelectionApplyOutcome::SessionUpdateFailed("storage unavailable".to_string()), + ); + + assert!(status.contains("failed")); + assert!(status.contains("storage unavailable")); + assert!(status.contains("retry")); + } + + #[test] + fn pending_mode_change_allows_host_commands_but_blocks_agent_submission() { + assert!(mode_change_blocks_typed_submission(true, "continue")); + assert!(!mode_change_blocks_typed_submission(true, "/new")); + assert!(!mode_change_blocks_typed_submission(true, "/sessions")); + assert!(!mode_change_blocks_typed_submission(true, "/exit")); + assert!(!mode_change_blocks_typed_submission(false, "continue")); + } + + #[test] + fn failed_mode_save_cancels_automatic_exit() { + assert!(mode_change_completion_should_exit(true, true)); + assert!(!mode_change_completion_should_exit(true, false)); + assert!(!mode_change_completion_should_exit(false, true)); + } + #[test] fn shortcut_registry_contract_help_uses_resolved_keymap() { let keymap = ResolvedKeymap::new(&ShortcutsConfig::default()); diff --git a/src/apps/cli/src/peer_host/commands/session.rs b/src/apps/cli/src/peer_host/commands/session.rs index 6bcdb3605..e62e4eb73 100644 --- a/src/apps/cli/src/peer_host/commands/session.rs +++ b/src/apps/cli/src/peer_host/commands/session.rs @@ -7,7 +7,6 @@ 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_runtime_ports::{ AgentSessionDeleteRequest, AgentSessionModelUpdateRequest, SessionStoragePathRequest, }; @@ -300,16 +299,9 @@ pub(crate) async fn archive_session(state: &PeerHostState, args: &Value) -> Resu .begin_persisted_session_mutation(&workspace_path, &session_id) .await .map_err(|error| format!("Failed to lock session archive: {error}"))?; - let mut metadata = state - .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 .compatibility - .save_persisted_session_metadata(&workspace_path, &metadata) + .archive_persisted_session(&workspace_path, &session_id) .await .map_err(|e| format!("Failed to archive session: {e}"))?; Ok(Value::Null) diff --git a/src/apps/cli/tests/acp_stdio_cli.rs b/src/apps/cli/tests/acp_stdio_cli.rs index b40b2c5b0..78d8d1019 100644 --- a/src/apps/cli/tests/acp_stdio_cli.rs +++ b/src/apps/cli/tests/acp_stdio_cli.rs @@ -1,101 +1,492 @@ +mod support; + use std::process::Stdio; use std::time::Duration; -use serde_json::json; +use serde_json::{json, Value}; +use support::{CliTestEnvironment, MockOpenAiServer}; use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; -#[tokio::test] -async fn acp_stdio_initializes_the_production_assembled_runtime() { - let temp = tempfile::tempdir().expect("tempdir"); - let workspace = temp.path().join("workspace"); - let user_root = temp.path().join("user-root"); - let home_root = temp.path().join("home-root"); - let config_root = temp.path().join("host-config"); - std::fs::create_dir_all(&workspace).expect("create workspace"); - - let mut child = tokio::process::Command::new(env!("CARGO_BIN_EXE_bitfun-cli")) - .arg("acp") - .current_dir(&workspace) - .env_remove("BITFUN_USER_ROOT") - .env_remove("BITFUN_HOME") - .env("BITFUN_E2E_STORAGE_GUARD", "1") - .env("BITFUN_E2E_USER_ROOT", &user_root) - .env("BITFUN_E2E_HOME", &home_root) - .env("APPDATA", &config_root) - .env("XDG_CONFIG_HOME", &config_root) - .env("HOME", &home_root) - .stdin(Stdio::piped()) - .stdout(Stdio::piped()) - .stderr(Stdio::piped()) - .kill_on_drop(true) - .spawn() - .expect("start production ACP server"); - - let mut stdin = child.stdin.take().expect("ACP stdin"); - let mut stdout = BufReader::new(child.stdout.take().expect("ACP stdout")); - let mut stderr = child.stderr.take().expect("ACP stderr"); - let stderr_reader = tokio::spawn(async move { - let mut bytes = Vec::new(); - stderr - .read_to_end(&mut bytes) +struct AcpProcess { + child: tokio::process::Child, + stdin: Option, + stdout: Option>, + stderr_reader: tokio::task::JoinHandle, +} + +impl AcpProcess { + async fn spawn(environment: &CliTestEnvironment) -> Self { + let mut command = tokio::process::Command::new(env!("CARGO_BIN_EXE_bitfun-cli")); + command + .arg("acp") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true); + environment.apply_tokio_environment(&mut command); + + let mut child = command.spawn().expect("start production ACP server"); + let stdin = child.stdin.take().expect("ACP stdin"); + let stdout = BufReader::new(child.stdout.take().expect("ACP stdout")); + let mut stderr = child.stderr.take().expect("ACP stderr"); + let stderr_reader = tokio::spawn(async move { + let mut bytes = Vec::new(); + stderr + .read_to_end(&mut bytes) + .await + .expect("read ACP stderr"); + String::from_utf8_lossy(&bytes).into_owned() + }); + + Self { + child, + stdin: Some(stdin), + stdout: Some(stdout), + stderr_reader, + } + } + + async fn request(&mut self, id: i64, method: &str, params: Value) -> (Value, Vec) { + self.send_request(id, method, params).await; + let (mut responses, notifications) = self.read_responses(&[id], method).await; + (responses.remove(0), notifications) + } + + async fn send_request(&mut self, id: i64, method: &str, params: Value) { + let request = json!({ + "jsonrpc": "2.0", + "id": id, + "method": method, + "params": params, + }); + let mut request = serde_json::to_vec(&request).expect("serialize ACP request"); + request.push(b'\n'); + let stdin = self.stdin.as_mut().expect("ACP stdin remains available"); + stdin.write_all(&request).await.expect("write ACP request"); + stdin.flush().await.expect("flush ACP request"); + } + + async fn read_responses( + &mut self, + expected_ids: &[i64], + operation: &str, + ) -> (Vec, Vec) { + let mut responses = Vec::with_capacity(expected_ids.len()); + let mut notifications = Vec::new(); + while responses.len() < expected_ids.len() { + let mut line = String::new(); + let bytes_read = tokio::time::timeout( + Duration::from_secs(60), + self.stdout + .as_mut() + .expect("ACP stdout remains available") + .read_line(&mut line), + ) .await - .expect("read ACP stderr"); - String::from_utf8_lossy(&bytes).into_owned() - }); + .unwrap_or_else(|_| panic!("ACP {operation} request timed out")) + .expect("read ACP stdout"); + assert_ne!( + bytes_read, 0, + "ACP stdout closed while waiting for {operation}" + ); + + let message: Value = serde_json::from_str(&line).unwrap_or_else(|error| { + panic!("ACP stdout contained non-JSON data: {error}: {line}") + }); + if let Some(id) = message.get("id").and_then(Value::as_i64) { + assert!( + expected_ids.contains(&id), + "unexpected ACP response while waiting for {operation}: {message}" + ); + assert!( + !responses.iter().any(|response: &Value| { + response.get("id").and_then(Value::as_i64) == Some(id) + }), + "duplicate ACP response id while waiting for {operation}: {message}" + ); + responses.push(message); + continue; + } - let request = json!({ - "jsonrpc": "2.0", - "id": 1, - "method": "initialize", - "params": { - "protocolVersion": 1, - "clientCapabilities": {} + assert_eq!( + message.get("method"), + Some(&json!("session/update")), + "unexpected ACP message while waiting for {operation}: {message}" + ); + notifications.push(message); } + (responses, notifications) + } + + async fn shutdown(mut self) -> String { + drop(self.stdin.take()); + drop(self.stdout.take()); + match tokio::time::timeout(Duration::from_secs(1), self.child.wait()).await { + Ok(Ok(status)) => assert!(status.success(), "ACP server exited with {status}"), + Ok(Err(error)) => panic!("wait for ACP server: {error}"), + Err(_) => { + self.child.kill().await.expect("stop ACP server"); + let _ = self.child.wait().await.expect("reap ACP server"); + } + } + self.stderr_reader.await.expect("join ACP stderr reader") + } +} + +fn current_config_value<'a>(response: &'a Value, config_id: &str) -> Option<&'a Value> { + response + .pointer("/result/configOptions")? + .as_array()? + .iter() + .find(|option| option.get("id") == Some(&json!(config_id)))? + .get("currentValue") +} + +#[tokio::test] +async fn acp_stdio_preserves_mode_and_history_across_restart_then_closes_active_session() { + let model = MockOpenAiServer::immediate(); + let environment = CliTestEnvironment::new(); + environment.initialize_git_repository(); + environment.configure_mock_model(model.base_url()); + let cwd = environment.workspace().to_string_lossy().to_string(); + + let mut first = AcpProcess::spawn(&environment).await; + let (initialize, _) = first + .request( + 1, + "initialize", + json!({ "protocolVersion": 1, "clientCapabilities": {} }), + ) + .await; + assert_eq!( + initialize.pointer("/result/protocolVersion"), + Some(&json!(1)) + ); + assert_eq!( + initialize.pointer("/result/agentInfo/name"), + Some(&json!("bitfun-acp")) + ); + assert_eq!( + initialize.pointer("/result/agentCapabilities/sessionCapabilities/close"), + Some(&json!({})), + "ACP must advertise session/close before clients can call it" + ); + + let (invalid_new, _) = first + .request( + 2, + "session/new", + json!({ + "cwd": cwd, + "mcpServers": [{ + "name": " ", + "command": "unused", + "args": [], + "env": [] + }] + }), + ) + .await; + assert_eq!(invalid_new.pointer("/error/code"), Some(&json!(-32602))); + let (after_failed_new, _) = first + .request(3, "session/list", json!({ "cwd": cwd })) + .await; + assert_eq!( + after_failed_new.pointer("/result/sessions"), + Some(&json!([])), + "failed session/new must not leave an undisclosed Core session" + ); + + let (created, _) = first + .request(4, "session/new", json!({ "cwd": cwd, "mcpServers": [] })) + .await; + let session_id = created + .pointer("/result/sessionId") + .and_then(Value::as_str) + .expect("new session id") + .to_string(); + + let (configured, mode_updates) = first + .request( + 5, + "session/set_config_option", + json!({ + "sessionId": session_id, + "configId": "mode", + "value": " Plan " + }), + ) + .await; + assert!(configured.get("error").is_none(), "{configured}"); + assert_eq!( + current_config_value(&configured, "mode"), + Some(&json!("Plan")) + ); + assert!( + mode_updates.iter().all(|message| { + message.pointer("/params/update/sessionUpdate") != Some(&json!("current_mode_update")) + }), + "a client-initiated mode change must not be echoed as an autonomous mode update" + ); + + let (prompted, prompt_updates) = first + .request( + 6, + "session/prompt", + json!({ + "sessionId": session_id, + "prompt": [{ "type": "text", "text": "remember this turn" }] + }), + ) + .await; + assert!(prompted.get("error").is_none(), "{prompted}"); + assert!( + !prompt_updates.is_empty(), + "prompt must stream at least one session/update before its response" + ); + let first_stderr = first.shutdown().await; + model.assert_chat_completion_requests(1); + + let restored_model = MockOpenAiServer::immediate(); + environment.configure_mock_model(restored_model.base_url()); + let mut second = AcpProcess::spawn(&environment).await; + let _ = second + .request( + 10, + "initialize", + json!({ "protocolVersion": 1, "clientCapabilities": {} }), + ) + .await; + let (missing_load, _) = second + .request( + 11, + "session/load", + json!({ + "sessionId": "missing-session", + "cwd": environment.workspace().to_string_lossy(), + "mcpServers": [] + }), + ) + .await; + assert_eq!(missing_load.pointer("/error/code"), Some(&json!(-32002))); + assert_eq!( + missing_load.pointer("/error/data/uri"), + Some(&json!("missing-session")) + ); + let (invalid_session_id, _) = second + .request( + 12, + "session/load", + json!({ + "sessionId": "../outside", + "cwd": environment.workspace().to_string_lossy(), + "mcpServers": [{ + "name": "must-not-start", + "command": "bitfun-command-that-must-not-run", + "args": [], + "env": [] + }] + }), + ) + .await; + assert_eq!( + invalid_session_id.pointer("/error/code"), + Some(&json!(-32602)), + "invalid session identity must be rejected before MCP provisioning: {invalid_session_id}" + ); + let (invalid_load, _) = second + .request( + 13, + "session/load", + json!({ + "sessionId": session_id, + "cwd": environment.workspace().to_string_lossy(), + "mcpServers": [{ + "name": " ", + "command": "unused", + "args": [], + "env": [] + }] + }), + ) + .await; + assert_eq!(invalid_load.pointer("/error/code"), Some(&json!(-32602))); + + let load_params = json!({ + "sessionId": session_id, + "cwd": environment.workspace().to_string_lossy(), + "mcpServers": [] }); - let mut request = serde_json::to_vec(&request).expect("serialize initialize request"); - request.push(b'\n'); - stdin - .write_all(&request) - .await - .expect("write ACP initialize request"); - stdin.flush().await.expect("flush ACP initialize request"); - - let mut response_line = String::new(); - tokio::time::timeout( - Duration::from_secs(60), - stdout.read_line(&mut response_line), + second + .send_request(14, "session/load", load_params.clone()) + .await; + second.send_request(15, "session/load", load_params).await; + let (load_responses, replay_updates) = second + .read_responses(&[14, 15], "concurrent session/load") + .await; + let loaded = load_responses + .iter() + .find(|response| response.get("error").is_none()) + .expect("one concurrent session/load must succeed"); + let duplicate_load = load_responses + .iter() + .find(|response| response.get("error").is_some()) + .expect("one concurrent session/load must be rejected"); + assert!(loaded.get("error").is_none(), "{loaded}"); + assert_eq!(current_config_value(&loaded, "mode"), Some(&json!("Plan"))); + assert!( + !replay_updates.is_empty(), + "session/load must replay persisted history before its success response" + ); + assert_eq!(duplicate_load.pointer("/error/code"), Some(&json!(-32603))); + assert_eq!( + duplicate_load.pointer("/error/data/state"), + Some(&json!("session_transition_in_progress")) + ); + assert_eq!( + duplicate_load.pointer("/error/data/retryable"), + Some(&json!(true)) + ); + assert_eq!( + replay_updates + .iter() + .filter(|message| { + message.pointer("/params/update/sessionUpdate") + == Some(&json!("user_message_chunk")) + }) + .count(), + 1, + "overlapping session/load must not replay the persisted user turn twice" + ); + + let (closed, _) = second + .request(16, "session/close", json!({ "sessionId": session_id })) + .await; + assert!(closed.get("error").is_none(), "{closed}"); + + let (reloaded_after_close, replay_after_close) = second + .request( + 17, + "session/load", + json!({ + "sessionId": session_id, + "cwd": environment.workspace().to_string_lossy(), + "mcpServers": [] + }), + ) + .await; + assert!( + reloaded_after_close.get("error").is_none(), + "closing an ACP session must release runtime resources without deleting its history: {reloaded_after_close}" + ); + assert_eq!( + current_config_value(&reloaded_after_close, "mode"), + Some(&json!("Plan")) + ); + assert_eq!( + replay_after_close + .iter() + .filter(|message| { + message.pointer("/params/update/sessionUpdate") + == Some(&json!("user_message_chunk")) + }) + .count(), + 1, + "reloading after close must replay the persisted turn exactly once" + ); + let (closed_after_reload, _) = second + .request(18, "session/close", json!({ "sessionId": session_id })) + .await; + assert!( + closed_after_reload.get("error").is_none(), + "{closed_after_reload}" + ); + + let metadata_path = environment.session_metadata_path(&session_id); + let mut metadata: Value = serde_json::from_slice( + &std::fs::read(&metadata_path).expect("read persisted session metadata"), ) - .await - .expect("production ACP initialize should not hang") - .expect("read ACP initialize response"); - let response: serde_json::Value = - serde_json::from_str(&response_line).expect("valid ACP JSON-RPC response"); - assert_eq!(response.get("id"), Some(&json!(1))); - assert_eq!(response.pointer("/result/protocolVersion"), Some(&json!(1))); + .expect("parse persisted session metadata"); + metadata["agentType"] = json!("RemovedCustomMode"); + std::fs::write( + &metadata_path, + serde_json::to_vec_pretty(&metadata).expect("serialize persisted session metadata"), + ) + .expect("replace persisted session mode"); + + let (loaded_with_fallback, _) = second + .request( + 19, + "session/load", + json!({ + "sessionId": session_id, + "cwd": environment.workspace().to_string_lossy(), + "mcpServers": [] + }), + ) + .await; + assert!( + loaded_with_fallback.get("error").is_none(), + "{loaded_with_fallback}" + ); assert_eq!( - response.pointer("/result/agentInfo/name"), - Some(&json!("bitfun-acp")) + current_config_value(&loaded_with_fallback, "mode"), + Some(&json!("agentic")), + "an unavailable persisted mode must be migrated to an executable fallback" + ); + let (fallback_prompt, _) = second + .request( + 20, + "session/prompt", + json!({ + "sessionId": session_id, + "prompt": [{ "type": "text", "text": "continue after mode fallback" }] + }), + ) + .await; + assert!( + fallback_prompt.get("error").is_none(), + "the restored fallback mode must be executable: {fallback_prompt}" + ); + restored_model.assert_chat_completion_requests(1); + let (closed_after_fallback, _) = second + .request(21, "session/close", json!({ "sessionId": session_id })) + .await; + assert!( + closed_after_fallback.get("error").is_none(), + "{closed_after_fallback}" ); - // ACP is a long-lived subprocess whose host owns termination. Close its - // request stream, then always reap it; the kill fallback covers protocol - // transport tasks that remain alive after the client finishes its check. - drop(stdin); - drop(stdout); - let status = match tokio::time::timeout(Duration::from_secs(1), child.wait()).await { - Ok(result) => Some(result.expect("wait for ACP server")), - Err(_) => { - child.kill().await.expect("stop hung ACP server"); - let _ = child.wait().await.expect("reap ACP server"); - None - } - }; - let stderr = stderr_reader.await.expect("join ACP stderr reader"); - - if let Some(status) = status { - assert!( - status.success(), - "ACP server exited unsuccessfully: {stderr}" - ); - } + let (post_close_prompt, _) = second + .request( + 22, + "session/prompt", + json!({ + "sessionId": session_id, + "prompt": [{ "type": "text", "text": "must be rejected" }] + }), + ) + .await; + assert_eq!( + post_close_prompt.pointer("/error/code"), + Some(&json!(-32002)), + "closed ACP session must no longer accept prompts" + ); + let (post_close_model, _) = second + .request( + 23, + "session/set_model", + json!({ "sessionId": session_id, "modelId": "auto" }), + ) + .await; + assert_eq!( + post_close_model.pointer("/error/code"), + Some(&json!(-32002)) + ); + + let second_stderr = second.shutdown().await; + assert!( + !first_stderr.contains("panicked") && !second_stderr.contains("panicked"), + "ACP process panicked:\nfirst:\n{first_stderr}\nsecond:\n{second_stderr}" + ); } diff --git a/src/apps/cli/tests/support/mod.rs b/src/apps/cli/tests/support/mod.rs index 2ac24abb4..7743b9c50 100644 --- a/src/apps/cli/tests/support/mod.rs +++ b/src/apps/cli/tests/support/mod.rs @@ -122,6 +122,34 @@ impl CliTestEnvironment { &self.workspace } + pub(crate) fn session_metadata_path(&self, session_id: &str) -> PathBuf { + fn find(dir: &Path, session_id: &str) -> Option { + for entry in std::fs::read_dir(dir).ok()? { + let entry = entry.ok()?; + let path = entry.path(); + if path.is_dir() { + if path.file_name().and_then(|name| name.to_str()) == Some(session_id) { + let metadata = path.join("metadata.json"); + if metadata.is_file() { + return Some(metadata); + } + } + if let Some(found) = find(&path, session_id) { + return Some(found); + } + } + } + None + } + + find(self._temp.path(), session_id).unwrap_or_else(|| { + panic!( + "persisted metadata for session {session_id} was not found under {}", + self._temp.path().display() + ) + }) + } + pub(crate) fn configure_mock_model(&self, server_base_url: &str) { let config_dir = self.user_root.join("config"); std::fs::create_dir_all(&config_dir).expect("create model config directory"); @@ -180,6 +208,21 @@ impl CliTestEnvironment { command } + pub(crate) fn apply_tokio_environment(&self, command: &mut tokio::process::Command) { + command + .current_dir(&self.workspace) + .env_remove("BITFUN_USER_ROOT") + .env_remove("BITFUN_HOME") + .env("BITFUN_E2E_STORAGE_GUARD", "1") + .env("BITFUN_E2E_USER_ROOT", &self.user_root) + .env("BITFUN_E2E_HOME", &self.home_root) + .env("APPDATA", &self.config_root) + .env("XDG_CONFIG_HOME", &self.config_root) + .env("HOME", &self.home_root) + .env("USERPROFILE", &self.home_root) + .env("TERM", "xterm-256color"); + } + pub(crate) fn pty_command(&self) -> CommandBuilder { let mut command = CommandBuilder::new(env!("CARGO_BIN_EXE_bitfun-cli")); command.cwd(&self.workspace); diff --git a/src/apps/desktop/src/api/agentic_api.rs b/src/apps/desktop/src/api/agentic_api.rs index 8209bedf4..bf972e443 100644 --- a/src/apps/desktop/src/api/agentic_api.rs +++ b/src/apps/desktop/src/api/agentic_api.rs @@ -791,9 +791,22 @@ pub async fn create_session( repaired = true; } if repaired { + let relationship = request.relationship.clone(); + let deep_review_run_manifest = request.deep_review_run_manifest.clone(); + let review_target_evidence = request.review_target_evidence.clone(); coordinator .get_session_manager() - .save_session_metadata(&effective_path, &metadata) + .update_session_metadata(&effective_path, session_id, |current| { + if current.relationship.is_none() { + current.relationship = relationship; + } + if current.deep_review_run_manifest.is_none() { + current.deep_review_run_manifest = deep_review_run_manifest; + } + if current.review_target_evidence.is_none() { + current.review_target_evidence = review_target_evidence; + } + }) .await .map_err(|error| { format!("Failed to repair Review session metadata: {error}") diff --git a/src/apps/desktop/src/api/remote_connect_api.rs b/src/apps/desktop/src/api/remote_connect_api.rs index 2c980f021..cc05b9777 100644 --- a/src/apps/desktop/src/api/remote_connect_api.rs +++ b/src/apps/desktop/src/api/remote_connect_api.rs @@ -1662,6 +1662,42 @@ pub struct SessionBundle { pub source_device_name: Option, } +const RELAY_TURNS_IMPORT_STATE_KEY: &str = "relayTurnsImportState"; +const RELAY_TURNS_IMPORT_PENDING: &str = "pending"; +const RELAY_TURNS_IMPORT_COMPLETE: &str = "complete"; + +fn relay_turns_import_state(metadata: &SessionMetadata) -> Option<&str> { + metadata + .custom_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .and_then(|custom| custom.get(RELAY_TURNS_IMPORT_STATE_KEY)) + .and_then(serde_json::Value::as_str) +} + +fn relay_turns_import_is_complete(metadata: &SessionMetadata, local_turn_count: usize) -> bool { + metadata.turn_count == local_turn_count + && relay_turns_import_state(metadata) == Some(RELAY_TURNS_IMPORT_COMPLETE) +} + +fn set_relay_turns_import_state(metadata: &mut SessionMetadata, state: &str) { + let mut custom = metadata + .custom_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .cloned() + .unwrap_or_default(); + custom.insert( + RELAY_TURNS_IMPORT_STATE_KEY.to_string(), + serde_json::Value::String(state.to_string()), + ); + metadata.custom_metadata = Some(serde_json::Value::Object(custom)); +} + +fn mark_relay_turns_import_complete(metadata: &mut SessionMetadata) { + set_relay_turns_import_state(metadata, RELAY_TURNS_IMPORT_COMPLETE); +} + /// Export a single local session as an encrypted blob and upload it to the relay. /// Uses the workspace + session_id to load metadata and turns from disk. #[tauri::command] @@ -1834,17 +1870,6 @@ pub async fn account_import_remote_sessions( for fetched in remote_sessions { let session_id = fetched.session_id; let bundle_json = fetched.plaintext; - // Skip if session already exists locally - if manager - .load_session_metadata(&storage_path, &session_id) - .await - .ok() - .flatten() - .is_some() - { - continue; - } - // Deserialize the bundle and write metadata as-is. The source device's // workspace_path is preserved for display (read-only history). Tasks // are always executed on the receiving device's own workspace, so @@ -1852,14 +1877,25 @@ pub async fn account_import_remote_sessions( let bundle: SessionBundle = serde_json::from_str(&bundle_json).map_err(|e| format!("deserialize bundle: {e}"))?; - let metadata: SessionMetadata = serde_json::from_value(bundle.metadata) + let mut metadata: SessionMetadata = serde_json::from_value(bundle.metadata) .map_err(|e| format!("deserialize metadata: {e}"))?; + if metadata.session_id != session_id { + log::warn!( + "Skipping remote session bundle with mismatched metadata identity: expected_session_id={}, metadata_session_id={}", + session_id, + metadata.session_id + ); + continue; + } // Only write metadata — turns are lazy-loaded when the user opens // the session (see `account_fetch_session_turns`). - if manager - .save_session_metadata(&storage_path, &metadata) + set_relay_turns_import_state(&mut metadata, RELAY_TURNS_IMPORT_PENDING); + if !manager + .create_session_metadata_if_absent(&storage_path, &metadata) .await - .is_err() + .map_err(|error| { + format!("persist imported metadata for session {session_id}: {error}") + })? { continue; } @@ -1898,11 +1934,25 @@ pub async fn account_fetch_session_turns( let manager = PersistenceManager::new(path_manager.inner().clone()) .map_err(|e| format!("create persistence manager: {e}"))?; - // If turns already exist locally, no fetch needed. - if let Ok(turns) = manager.load_session_turns(&storage_path, &session_id).await { - if !turns.is_empty() { - return Ok(false); - } + // Ordinary local sessions carry no relay marker and return without an + // account or network lookup. A non-empty turn prefix is not proof that an + // import completed, so pending or inconsistent imports are retried. + let Some(metadata) = manager + .load_session_metadata(&storage_path, &session_id) + .await + .map_err(|error| format!("load imported metadata: {error}"))? + else { + return Ok(false); + }; + if relay_turns_import_state(&metadata).is_none() { + return Ok(false); + } + let local_turns = manager + .load_session_turns(&storage_path, &session_id) + .await + .map_err(|error| format!("load imported turns: {error}"))?; + if relay_turns_import_is_complete(&metadata, local_turns.len()) { + return Ok(false); } // Fetch the full bundle from the relay (which includes turns). @@ -1916,18 +1966,41 @@ pub async fn account_fetch_session_turns( let bundle: SessionBundle = serde_json::from_str(&fetched.plaintext).map_err(|e| format!("deserialize bundle: {e}"))?; - // Write turns first, then metadata (self-healing on crash). - for turn_val in &bundle.turns { - let turn: DialogTurnData = serde_json::from_value(turn_val.clone()) - .map_err(|e| format!("deserialize turn: {e}"))?; - let _ = manager.save_dialog_turn(&storage_path, &turn).await; - } - // Re-save metadata to ensure consistency. - let metadata: SessionMetadata = serde_json::from_value(bundle.metadata) + let metadata: SessionMetadata = serde_json::from_value(bundle.metadata.clone()) .map_err(|e| format!("deserialize metadata: {e}"))?; - let _ = manager - .save_session_metadata(&storage_path, &metadata) - .await; + if metadata.session_id != session_id { + return Err("relay session metadata identity does not match request".to_string()); + } + let turns = bundle + .turns + .iter() + .map(|turn| { + serde_json::from_value::(turn.clone()) + .map_err(|error| format!("deserialize turn: {error}")) + }) + .collect::, _>>()?; + if turns.iter().any(|turn| turn.session_id != session_id) { + return Err("relay session turn identity does not match request".to_string()); + } + + manager + .create_session_metadata_if_absent(&storage_path, &metadata) + .await + .map_err(|e| format!("persist imported metadata: {e}"))?; + + // Each turn save refreshes counts through an owner-side metadata RMW. + for turn in &turns { + manager + .save_dialog_turn(&storage_path, turn) + .await + .map_err(|e| format!("persist imported turn: {e}"))?; + } + manager + .update_session_metadata(&storage_path, &session_id, |metadata| { + mark_relay_turns_import_complete(metadata); + }) + .await + .map_err(|e| format!("mark imported turns complete: {e}"))?; log::info!( "Lazy-loaded {} turns for session {session_id}", @@ -2686,24 +2759,19 @@ async fn import_session_bundle(bundle_json: &str) -> anyhow::Result<()> { dir }); - // Skip if session already exists locally - if manager - .load_session_metadata(&target_dir, &bundle.session_id) - .await - .ok() - .flatten() - .is_some() - { - return Ok(()); + let mut metadata: SessionMetadata = serde_json::from_value(bundle.metadata.clone())?; + if metadata.session_id != bundle.session_id { + return Err(anyhow::anyhow!( + "relay session metadata identity does not match bundle" + )); } - let metadata: SessionMetadata = serde_json::from_value(bundle.metadata.clone())?; - // Only write metadata — turns are lazy-loaded when the user opens the // session. This keeps the import fast and avoids writing potentially // large turn data that may never be read. + set_relay_turns_import_state(&mut metadata, RELAY_TURNS_IMPORT_PENDING); manager - .save_session_metadata(&target_dir, &metadata) + .create_session_metadata_if_absent(&target_dir, &metadata) .await .map_err(|e| anyhow::anyhow!("save metadata: {e}"))?; @@ -2753,4 +2821,28 @@ mod sync_state_tests { state.advance_session_since([9]); assert_eq!(state.last_session_since, 9); } + + #[test] + fn relay_turn_import_requires_an_explicit_complete_marker_and_exact_count() { + let mut metadata = SessionMetadata::new( + "session".to_string(), + "Session".to_string(), + "agentic".to_string(), + "auto".to_string(), + ); + metadata.turn_count = 2; + + assert_eq!(relay_turns_import_state(&metadata), None); + assert!(!relay_turns_import_is_complete(&metadata, 1)); + assert!(!relay_turns_import_is_complete(&metadata, 2)); + set_relay_turns_import_state(&mut metadata, RELAY_TURNS_IMPORT_PENDING); + assert_eq!( + relay_turns_import_state(&metadata), + Some(RELAY_TURNS_IMPORT_PENDING) + ); + assert!(!relay_turns_import_is_complete(&metadata, 2)); + mark_relay_turns_import_complete(&mut metadata); + assert!(!relay_turns_import_is_complete(&metadata, 1)); + assert!(relay_turns_import_is_complete(&metadata, 2)); + } } diff --git a/src/apps/desktop/src/api/session_api.rs b/src/apps/desktop/src/api/session_api.rs index bc592dddc..2dd98b178 100644 --- a/src/apps/desktop/src/api/session_api.rs +++ b/src/apps/desktop/src/api/session_api.rs @@ -19,6 +19,65 @@ use std::sync::Arc; use std::time::Instant; use tauri::State; +const UI_CUSTOM_METADATA_KEYS: [&str; 3] = ["titleSource", "titleKey", "titleParams"]; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub enum UiSessionMetadataField { + SessionName, + Tags, + Todos, + ReviewActionState, + UnreadCompletion, + NeedsUserAttention, + TitleMetadata, +} + +fn merge_ui_owned_session_metadata( + current: &mut SessionMetadata, + incoming: &SessionMetadata, + fields: &[UiSessionMetadataField], +) { + if fields.contains(&UiSessionMetadataField::SessionName) { + current.session_name = incoming.session_name.clone(); + } + if fields.contains(&UiSessionMetadataField::Tags) { + current.tags = incoming.tags.clone(); + } + if fields.contains(&UiSessionMetadataField::Todos) { + current.todos = incoming.todos.clone(); + } + if fields.contains(&UiSessionMetadataField::ReviewActionState) { + current.review_action_state = incoming.review_action_state.clone(); + } + if fields.contains(&UiSessionMetadataField::UnreadCompletion) { + current.unread_completion = incoming.unread_completion.clone(); + } + if fields.contains(&UiSessionMetadataField::NeedsUserAttention) { + current.needs_user_attention = incoming.needs_user_attention.clone(); + } + + if fields.contains(&UiSessionMetadataField::TitleMetadata) { + let mut custom = current + .custom_metadata + .as_ref() + .and_then(serde_json::Value::as_object) + .cloned() + .unwrap_or_default(); + let incoming_custom = incoming + .custom_metadata + .as_ref() + .and_then(serde_json::Value::as_object); + for key in UI_CUSTOM_METADATA_KEYS { + custom.remove(key); + if let Some(value) = incoming_custom.and_then(|metadata| metadata.get(key)) { + custom.insert(key.to_string(), value.clone()); + } + } + current.custom_metadata = (!custom.is_empty()).then(|| serde_json::Value::Object(custom)); + } +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ListPersistedSessionsRequest { pub workspace_path: String, @@ -65,6 +124,7 @@ pub struct SaveSessionTurnRequest { #[derive(Debug, Clone, Serialize, Deserialize)] pub struct SaveSessionMetadataRequest { pub metadata: SessionMetadata, + pub fields: Vec, pub workspace_path: String, #[serde(default, skip_serializing_if = "Option::is_none")] pub remote_connection_id: Option, @@ -359,6 +419,9 @@ pub async fn save_session_metadata( app_state: State<'_, AppState>, path_manager: State<'_, Arc>, ) -> Result<(), String> { + if request.fields.is_empty() { + return Err("At least one session metadata field is required".to_string()); + } let workspace_path = desktop_effective_session_storage_path( &app_state, &request.workspace_path, @@ -369,16 +432,16 @@ pub async fn save_session_metadata( let manager = PersistenceManager::new(path_manager.inner().clone()) .map_err(|e| format!("Failed to create persistence manager: {}", e))?; + let session_id = request.metadata.session_id.clone(); manager - .save_session_metadata(&workspace_path, &request.metadata) + .update_session_metadata(&workspace_path, &session_id, |metadata| { + merge_ui_owned_session_metadata(metadata, &request.metadata, &request.fields); + }) .await .map_err(|e| format!("Failed to save session metadata: {}", e))?; // Notify the auto-sync background task - crate::api::remote_connect_api::notify_session_changed( - &request.metadata.session_id, - &request.workspace_path, - ); + crate::api::remote_connect_api::notify_session_changed(&session_id, &request.workspace_path); Ok(()) } @@ -549,16 +612,10 @@ pub async fn archive_session( let manager = PersistenceManager::new(path_manager.inner().clone()) .map_err(|e| format!("Failed to create persistence manager: {}", e))?; - let mut metadata = manager - .load_session_metadata(&workspace_path, &request.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; - manager - .save_session_metadata(&workspace_path, &metadata) + .update_session_metadata(&workspace_path, &request.session_id, |metadata| { + metadata.status = SessionStatus::Archived; + }) .await .map_err(|e| format!("Failed to save session metadata: {}", e)) } @@ -579,16 +636,10 @@ pub async fn unarchive_session( let manager = PersistenceManager::new(path_manager.inner().clone()) .map_err(|e| format!("Failed to create persistence manager: {}", e))?; - let mut metadata = manager - .load_session_metadata(&workspace_path, &request.session_id) - .await - .map_err(|e| format!("Failed to load session metadata: {}", e))? - .ok_or_else(|| "Session not found".to_string())?; - - metadata.status = SessionStatus::Active; - manager - .save_session_metadata(&workspace_path, &metadata) + .update_session_metadata(&workspace_path, &request.session_id, |metadata| { + metadata.status = SessionStatus::Active; + }) .await .map_err(|e| format!("Failed to save session metadata: {}", e)) } @@ -616,13 +667,16 @@ pub async fn archive_all_sessions( let mut archived_count: u32 = 0; - for mut metadata in sessions { + for metadata in sessions { if metadata.status != SessionStatus::Archived && metadata.session_kind == SessionKind::Standard { - metadata.status = SessionStatus::Archived; manager - .save_session_metadata(&workspace_path, &metadata) + .update_session_metadata(&workspace_path, &metadata.session_id, |current| { + if current.session_kind == SessionKind::Standard { + current.status = SessionStatus::Archived; + } + }) .await .map_err(|e| format!("Failed to save session metadata: {}", e))?; archived_count += 1; @@ -696,3 +750,117 @@ pub async fn delete_all_archived_sessions( Ok(deleted_count) } + +#[cfg(test)] +mod tests { + use super::{merge_ui_owned_session_metadata, UiSessionMetadataField}; + use bitfun_core::service::session::{ + SessionKind, SessionMemoryMode, SessionMetadata, SessionStatus, + }; + use serde_json::json; + + #[test] + fn ui_metadata_merge_preserves_core_authoritative_fields_and_custom_keys() { + let mut current = SessionMetadata::new( + "session".to_string(), + "Current".to_string(), + "plan".to_string(), + "model-a".to_string(), + ); + current.last_submitted_agent_type = Some("plan".to_string()); + current.memory_mode = SessionMemoryMode::Polluted; + current.session_kind = SessionKind::Standard; + current.status = SessionStatus::Archived; + current.turn_count = 7; + current.custom_metadata = Some(json!({ + "threadGoal": { "objective": "preserve" }, + "titleSource": "i18n", + "titleKey": "old" + })); + + let mut incoming = current.clone(); + incoming.session_name = "Renamed".to_string(); + incoming.agent_type = "agentic".to_string(); + incoming.model_name = "stale-model".to_string(); + incoming.memory_mode = SessionMemoryMode::Enabled; + incoming.status = SessionStatus::Active; + incoming.turn_count = 1; + incoming.review_action_state = Some(json!({ "phase": "fixing" })); + incoming.custom_metadata = Some(json!({ + "titleSource": "i18n", + "titleKey": "new", + "untrustedCoreKey": "drop" + })); + + merge_ui_owned_session_metadata( + &mut current, + &incoming, + &[ + UiSessionMetadataField::SessionName, + UiSessionMetadataField::Tags, + UiSessionMetadataField::Todos, + UiSessionMetadataField::ReviewActionState, + UiSessionMetadataField::UnreadCompletion, + UiSessionMetadataField::NeedsUserAttention, + UiSessionMetadataField::TitleMetadata, + ], + ); + + assert_eq!(current.session_name, "Renamed"); + assert_eq!(current.agent_type, "plan"); + assert_eq!(current.model_name, "model-a"); + assert_eq!(current.memory_mode, SessionMemoryMode::Polluted); + assert_eq!(current.status, SessionStatus::Archived); + assert_eq!(current.turn_count, 7); + assert_eq!(current.review_action_state, incoming.review_action_state); + let custom = current.custom_metadata.unwrap(); + assert_eq!(custom["threadGoal"]["objective"], "preserve"); + assert_eq!(custom["titleKey"], "new"); + assert!(custom.get("untrustedCoreKey").is_none()); + } + + #[test] + fn ui_metadata_field_mask_keeps_independent_writers_isolated() { + let mut current = SessionMetadata::new( + "session".to_string(), + "Current".to_string(), + "agentic".to_string(), + "auto".to_string(), + ); + current.review_action_state = Some(json!({ "phase": "review_completed" })); + current.unread_completion = Some("completed".to_string()); + current.needs_user_attention = Some("ask_user".to_string()); + + let mut stale_general_update = current.clone(); + stale_general_update.session_name = "Renamed".to_string(); + stale_general_update.review_action_state = None; + merge_ui_owned_session_metadata( + &mut current, + &stale_general_update, + &[UiSessionMetadataField::SessionName], + ); + assert_eq!(current.session_name, "Renamed"); + assert_eq!( + current.review_action_state, + Some(json!({ "phase": "review_completed" })) + ); + assert_eq!(current.unread_completion.as_deref(), Some("completed")); + assert_eq!(current.needs_user_attention.as_deref(), Some("ask_user")); + + let mut review_update = current.clone(); + review_update.review_action_state = Some(json!({ "phase": "fixing" })); + review_update.unread_completion = None; + review_update.needs_user_attention = None; + merge_ui_owned_session_metadata( + &mut current, + &review_update, + &[UiSessionMetadataField::ReviewActionState], + ); + assert_eq!( + current.review_action_state, + Some(json!({ "phase": "fixing" })) + ); + assert_eq!(current.unread_completion.as_deref(), Some("completed")); + assert_eq!(current.needs_user_attention.as_deref(), Some("ask_user")); + } +} diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index f50f48606..dc6bd52aa 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -1699,7 +1699,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet needs_user_attention: None, }; if let Err(e) = persistence_manager - .save_session_metadata(&workspace_path_buf, &metadata) + .create_session_metadata_if_absent(&workspace_path_buf, &metadata) .await { warn!( @@ -4026,6 +4026,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet loop { let pending = counter.load(Ordering::SeqCst); if pending == 0 { + self.active_turns_per_session + .remove_if(session_id, |_, current| { + Arc::ptr_eq(current, &counter) && current.load(Ordering::SeqCst) == 0 + }); return 0; } if Instant::now() >= deadline { @@ -4090,6 +4094,12 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .try_register_pending(session_id.to_string(), turn_id.to_string()) } + #[cfg(test)] + pub(crate) fn set_active_turn_count_for_test(&self, session_id: &str, count: usize) { + self.active_turns_per_session + .insert(session_id.to_string(), Arc::new(AtomicUsize::new(count))); + } + /// Strict maintenance barrier for callers that must not overlap an older /// turn's tail writes. Unlike normal interactive cancellation, timeout is /// returned as an error instead of being treated as best effort. @@ -4327,7 +4337,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet current_turn_id, wait_timeout.as_millis() ); - break; + return Err(BitFunError::Timeout(format!( + "Active turn cancellation did not drain before timeout: session_id={session_id}, dialog_turn_id={current_turn_id}, timeout_ms={}", + wait_timeout.as_millis() + ))); } sleep(Duration::from_millis(50)).await; } @@ -6992,12 +7005,16 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet pub(crate) async fn cancel_background_subagents_for_parent_session( &self, parent_session_id: &str, - ) -> BitFunResult { + ) -> BitFunResult> { let controls = self.claim_background_subagent_controls(|control| { control.parent_session_id == parent_session_id }); - - self.cancel_background_subagent_controls(controls).await + let subagent_session_ids = controls + .iter() + .map(|(_, control)| control.subagent_session_id.clone()) + .collect::>(); + self.cancel_background_subagent_controls(controls).await?; + Ok(subagent_session_ids) } pub(crate) fn take_background_subagent_delivery_suppression( @@ -7690,6 +7707,30 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .await } + pub async fn update_session_mode(&self, session_id: &str, mode_id: &str) -> BitFunResult<()> { + let mode_id = mode_id.trim(); + if mode_id.is_empty() { + return Err(BitFunError::Validation( + "Session mode must not be empty".to_string(), + )); + } + + let mode_exists = get_agent_registry() + .get_modes_info() + .await + .into_iter() + .any(|mode| mode.id == mode_id); + if !mode_exists { + return Err(BitFunError::Validation(format!( + "Unknown session mode: {mode_id}" + ))); + } + + self.session_manager + .update_session_agent_type(session_id, mode_id) + .await + } + /// Update the session-level prompt-cache guard mode for the latest /// scheduler-accepted user submission. pub async fn update_last_submitted_agent_type( @@ -8132,7 +8173,7 @@ impl bitfun_runtime_ports::AgentSessionModePort for ConversationCoordinator { &self, request: bitfun_runtime_ports::AgentSessionModeUpdateRequest, ) -> bitfun_runtime_ports::PortResult<()> { - self.update_session_agent_type(&request.session_id, &request.mode_id) + self.update_session_mode(&request.session_id, &request.mode_id) .await .map_err(runtime_port_error_preserving_message) } @@ -8679,6 +8720,96 @@ mod tests { assert!(error.message.contains("missing-session")); } + #[tokio::test] + async fn session_mode_port_rejects_blank_mode_for_active_session() { + use bitfun_agent_runtime::sdk::{AgentSessionModePort, AgentSessionModeUpdateRequest}; + + let (coordinator, _) = test_coordinator(); + let workspace_path = std::env::temp_dir().join(format!( + "bitfun-session-mode-validation-test-{}", + uuid::Uuid::new_v4() + )); + std::fs::create_dir_all(&workspace_path).expect("workspace dir should exist"); + let workspace_path_string = workspace_path.to_string_lossy().into_owned(); + let session = TEST_AGENT_MODEL_DEFAULTS + .scope( + AgentModelDefaultsConfig::default(), + coordinator.create_session_with_workspace( + None, + "Runtime mode validation".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace_path_string.clone()), + ..Default::default() + }, + workspace_path_string, + ), + ) + .await + .expect("real Core session should be created"); + + let error = AgentSessionModePort::update_session_mode( + &coordinator, + AgentSessionModeUpdateRequest { + session_id: session.session_id, + mode_id: " ".to_string(), + }, + ) + .await + .expect_err("blank mode must remain a typed invalid request"); + + assert_eq!( + error.kind, + bitfun_runtime_ports::PortErrorKind::InvalidRequest + ); + let _ = std::fs::remove_dir_all(workspace_path); + } + + #[tokio::test] + async fn session_mode_port_rejects_unknown_mode_for_active_session() { + use bitfun_agent_runtime::sdk::{AgentSessionModePort, AgentSessionModeUpdateRequest}; + + let (coordinator, _) = test_coordinator(); + let workspace_path = std::env::temp_dir().join(format!( + "bitfun-session-mode-validation-test-{}", + uuid::Uuid::new_v4() + )); + std::fs::create_dir_all(&workspace_path).expect("workspace dir should exist"); + let workspace_path_string = workspace_path.to_string_lossy().into_owned(); + let session = TEST_AGENT_MODEL_DEFAULTS + .scope( + AgentModelDefaultsConfig::default(), + coordinator.create_session_with_workspace( + None, + "Runtime mode validation".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace_path_string.clone()), + ..Default::default() + }, + workspace_path_string, + ), + ) + .await + .expect("real Core session should be created"); + + let error = AgentSessionModePort::update_session_mode( + &coordinator, + AgentSessionModeUpdateRequest { + session_id: session.session_id, + mode_id: "__missing_runtime_mode__".to_string(), + }, + ) + .await + .expect_err("unknown mode must remain a typed invalid request"); + + assert_eq!( + error.kind, + bitfun_runtime_ports::PortErrorKind::InvalidRequest + ); + let _ = std::fs::remove_dir_all(workspace_path); + } + #[tokio::test] async fn session_mode_runtime_updates_the_real_core_session() { use bitfun_agent_runtime::sdk::{AgentRuntimeBuilder, AgentSessionModeUpdateRequest}; @@ -8716,7 +8847,7 @@ mod tests { runtime .update_session_mode(AgentSessionModeUpdateRequest { session_id: session.session_id.clone(), - mode_id: " plan ".to_string(), + mode_id: " Plan ".to_string(), }) .await .expect("runtime mode port should update the Core owner"); @@ -8726,7 +8857,7 @@ mod tests { .get_session(&session.session_id) .map(|session| session.agent_type.clone()) .as_deref(), - Some("plan") + Some("Plan") ); let _ = std::fs::remove_dir_all(workspace_path); } diff --git a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs index 1f3b7c722..f66be3d21 100644 --- a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs +++ b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs @@ -30,11 +30,12 @@ use crate::agentic::session::SessionManager; use crate::util::errors::{BitFunError, BitFunResult}; use bitfun_runtime_ports::{ThreadGoal, MAX_THREAD_GOAL_AUTO_CONTINUATIONS}; use log::{debug, info, warn}; +use std::collections::HashSet; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering}; use std::sync::Arc; use std::sync::OnceLock; -use std::time::{Duration, SystemTime}; +use std::time::{Duration, Instant, SystemTime}; use tokio::sync::mpsc; use tokio::sync::oneshot; use tokio_util::sync::CancellationToken; @@ -391,6 +392,10 @@ pub struct DialogScheduler { round_injection_buffer: Arc, round_injection_source: Arc, pending_background_results: Arc>, + /// Child sessions already cancelled for a parent maintenance attempt but + /// not yet observed as drained. Retain them across retryable timeouts even + /// after their one-shot cancellation controls have been claimed. + maintenance_background_sessions: Arc>>, #[cfg(test)] background_delivery_before_lock: std::sync::Mutex>>, } @@ -464,6 +469,7 @@ impl DialogScheduler { round_injection_buffer, round_injection_source, pending_background_results, + maintenance_background_sessions: Arc::new(dashmap::DashMap::new()), #[cfg(test)] background_delivery_before_lock: std::sync::Mutex::new(None), }); @@ -1535,18 +1541,51 @@ impl DialogScheduler { self.clear_queue(session_id).await; } abort_thread_goal_continuation_for_session(session_id); - self.coordinator + let deadline = Instant::now() + wait_timeout; + let cancelled_before_parent = self + .coordinator .cancel_background_subagents_for_parent_session(session_id) .await?; + let mut subagent_session_ids = self + .maintenance_background_sessions + .get(session_id) + .map(|sessions| sessions.clone()) + .unwrap_or_default(); + subagent_session_ids.extend(cancelled_before_parent); + if !subagent_session_ids.is_empty() { + self.maintenance_background_sessions + .insert(session_id.to_string(), subagent_session_ids.clone()); + } self.coordinator - .cancel_active_turn_for_session(session_id, wait_timeout) + .cancel_active_turn_for_session( + session_id, + deadline.saturating_duration_since(Instant::now()), + ) .await?; - self.coordinator + let cancelled_during_parent = self + .coordinator .cancel_background_subagents_for_parent_session(session_id) .await?; + subagent_session_ids.extend(cancelled_during_parent); + if !subagent_session_ids.is_empty() { + self.maintenance_background_sessions + .insert(session_id.to_string(), subagent_session_ids.clone()); + } + for subagent_session_id in &subagent_session_ids { + self.coordinator + .ensure_session_execution_drained( + subagent_session_id, + deadline.saturating_duration_since(Instant::now()), + ) + .await?; + } self.coordinator - .ensure_session_execution_drained(session_id, wait_timeout) + .ensure_session_execution_drained( + session_id, + deadline.saturating_duration_since(Instant::now()), + ) .await?; + self.maintenance_background_sessions.remove(session_id); self.retire_active_turn_for_maintenance(session_id); Ok(SessionMaintenancePermit { _operation_guard: operation_guard, @@ -2787,6 +2826,74 @@ mod tests { ); } + #[tokio::test] + async fn maintenance_does_not_release_parent_while_background_child_is_still_running() { + let (scheduler, session_manager, _, root) = test_scheduler(); + let parent_session_id = "parent-session"; + let child_session_id = "background-child-session"; + let workspace = root.path().join("workspace"); + std::fs::create_dir_all(&workspace).expect("workspace"); + session_manager + .create_session_with_id( + Some(parent_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"); + let storage_path = session_manager + .storage_path_binding_for_test(parent_session_id) + .expect("parent storage binding"); + scheduler + .coordinator + .register_background_subagent_task_for_test( + "background-task", + parent_session_id, + child_session_id, + ); + scheduler + .coordinator + .set_active_turn_count_for_test(child_session_id, 1); + + let result = scheduler + .begin_session_maintenance(parent_session_id, &storage_path, Duration::from_millis(40)) + .await; + let error = match result { + Ok(_) => panic!("maintenance must not detach a parent with a running child"), + Err(error) => error, + }; + + assert!(matches!(error, BitFunError::Timeout(_))); + assert!(error.to_string().contains(child_session_id)); + assert!(session_manager.get_session(parent_session_id).is_some()); + + let retry_error = match scheduler + .begin_session_maintenance(parent_session_id, &storage_path, Duration::from_millis(40)) + .await + { + Ok(_) => panic!("retry must retain ownership of the still-running child"), + Err(error) => error, + }; + assert!(matches!(retry_error, BitFunError::Timeout(_))); + assert!(retry_error.to_string().contains(child_session_id)); + + scheduler + .coordinator + .set_active_turn_count_for_test(child_session_id, 0); + let maintenance = scheduler + .begin_session_maintenance(parent_session_id, &storage_path, Duration::from_millis(40)) + .await + .expect("maintenance should succeed after the child drains"); + drop(maintenance); + assert!(!scheduler + .maintenance_background_sessions + .contains_key(parent_session_id)); + } + #[tokio::test] async fn background_delivery_rechecks_suppression_after_waiting_for_session_lock() { let (scheduler, _, event_queue, _root) = test_scheduler(); diff --git a/src/crates/assembly/core/src/agentic/deep_review/report.rs b/src/crates/assembly/core/src/agentic/deep_review/report.rs index 7957e8f33..bbee4fe81 100644 --- a/src/crates/assembly/core/src/agentic/deep_review/report.rs +++ b/src/crates/assembly/core/src/agentic/deep_review/report.rs @@ -152,15 +152,12 @@ pub(crate) async fn persist_deep_review_cache( }; let session_storage_dir = workspace.session_storage_dir(); let session_manager = coordinator.get_session_manager(); - let Some(mut metadata) = session_manager - .load_session_metadata(&session_storage_dir, session_id) - .await? - else { - return Ok(()); - }; - - set_deep_review_cache(&mut metadata, cache_value); session_manager - .save_session_metadata(&session_storage_dir, &metadata) + .persistence_manager() + .update_session_metadata_if_present(&session_storage_dir, session_id, |metadata| { + set_deep_review_cache(metadata, cache_value); + Ok(()) + }) .await + .map(|_| ()) } diff --git a/src/crates/assembly/core/src/agentic/persistence/manager.rs b/src/crates/assembly/core/src/agentic/persistence/manager.rs index f8424826a..272964837 100644 --- a/src/crates/assembly/core/src/agentic/persistence/manager.rs +++ b/src/crates/assembly/core/src/agentic/persistence/manager.rs @@ -43,7 +43,7 @@ use serde::{de::DeserializeOwned, Deserialize, Serialize}; use std::collections::HashMap; use std::io::ErrorKind; use std::path::{Path, PathBuf}; -use std::sync::{Arc, OnceLock}; +use std::sync::{Arc, OnceLock, Weak}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::fs; use tokio::io::AsyncWriteExt; @@ -57,7 +57,7 @@ const COMPRESSION_TRANSCRIPT_CREATE_ATTEMPTS: usize = 32; const TOKEN_ANCHOR_SCHEMA_VERSION: u32 = 1; const SESSION_TURN_READ_CONCURRENCY: usize = 4; -static SESSION_METADATA_UPDATE_LOCKS: OnceLock>>>> = +static SESSION_PERSISTENCE_LOCKS: OnceLock>>>> = OnceLock::new(); static SESSION_BRANCH_ALLOCATION_LOCKS: OnceLock>>>> = OnceLock::new(); @@ -288,6 +288,8 @@ pub struct PersistenceManager { runtime_service: Arc, #[cfg(test)] fail_next_session_state_write: std::sync::Mutex>, + #[cfg(test)] + fail_next_session_metadata_write: std::sync::Mutex>, } impl PersistenceManager { @@ -297,6 +299,8 @@ impl PersistenceManager { path_manager, #[cfg(test)] fail_next_session_state_write: std::sync::Mutex::new(None), + #[cfg(test)] + fail_next_session_metadata_write: std::sync::Mutex::new(None), }) } @@ -321,6 +325,14 @@ impl PersistenceManager { .expect("session state fault lock") = Some(session_id.to_string()); } + #[cfg(test)] + pub(crate) fn fail_next_session_metadata_write_for_test(&self, session_id: &str) { + *self + .fail_next_session_metadata_write + .lock() + .expect("session metadata fault lock") = Some(session_id.to_string()); + } + /// Resolve the on-disk sessions directory for `workspace_path`. /// /// Callers may pass either a logical workspace root or an already-resolved @@ -338,11 +350,6 @@ impl PersistenceManager { CoreSessionStorePort::resolved_sessions_dir_kind(self.path_manager.as_ref(), path).is_some() } - fn metadata_path(&self, workspace_path: &Path, session_id: &str) -> PathBuf { - self.session_layout(workspace_path) - .metadata_path(session_id) - } - fn state_path(&self, workspace_path: &Path, session_id: &str) -> PathBuf { self.session_layout(workspace_path).state_path(session_id) } @@ -522,18 +529,28 @@ impl PersistenceManager { .map_err(Self::json_store_error) } - async fn get_session_metadata_update_lock( + async fn get_session_persistence_lock( &self, workspace_path: &Path, session_id: &str, ) -> Arc> { - let metadata_path = self.metadata_path(workspace_path, session_id); - let registry = SESSION_METADATA_UPDATE_LOCKS.get_or_init(|| Mutex::new(HashMap::new())); + let session_path = self.session_layout(workspace_path).session_dir(session_id); + let session_path = dunce::canonicalize(&session_path).unwrap_or_else(|_| { + session_path + .parent() + .and_then(|parent| dunce::canonicalize(parent).ok()) + .and_then(|parent| session_path.file_name().map(|name| parent.join(name))) + .unwrap_or(session_path) + }); + let registry = SESSION_PERSISTENCE_LOCKS.get_or_init(|| Mutex::new(HashMap::new())); let mut registry_guard = registry.lock().await; - registry_guard - .entry(metadata_path) - .or_insert_with(|| Arc::new(Mutex::new(()))) - .clone() + registry_guard.retain(|_, lock| lock.strong_count() > 0); + if let Some(lock) = registry_guard.get(&session_path).and_then(Weak::upgrade) { + return lock; + } + let lock = Arc::new(Mutex::new(())); + registry_guard.insert(session_path, Arc::downgrade(&lock)); + lock } pub(super) async fn get_session_branch_allocation_lock( @@ -891,15 +908,118 @@ impl PersistenceManager { &self, workspace_path: &Path, metadata: &SessionMetadata, + ) -> BitFunResult<()> { + let persistence_lock = self + .get_session_persistence_lock(workspace_path, &metadata.session_id) + .await; + let _persistence_guard = persistence_lock.lock().await; + self.save_session_metadata_locked(workspace_path, metadata) + .await + } + + async fn save_session_metadata_locked( + &self, + workspace_path: &Path, + metadata: &SessionMetadata, ) -> BitFunResult<()> { Self::validate_session_id(&metadata.session_id)?; self.ensure_runtime_for_write(workspace_path).await?; + #[cfg(test)] + { + let mut fault = self + .fail_next_session_metadata_write + .lock() + .expect("session metadata fault lock"); + if fault.as_deref() == Some(metadata.session_id.as_str()) { + *fault = None; + return Err(BitFunError::io("Injected session metadata write failure")); + } + } self.session_metadata_store(workspace_path) .save_metadata(metadata) .await .map_err(Self::session_metadata_store_error) } + pub async fn create_session_metadata_if_absent( + &self, + workspace_path: &Path, + metadata: &SessionMetadata, + ) -> BitFunResult { + Self::validate_session_id(&metadata.session_id)?; + self.ensure_runtime_for_write(workspace_path).await?; + let persistence_lock = self + .get_session_persistence_lock(workspace_path, &metadata.session_id) + .await; + let _persistence_guard = persistence_lock.lock().await; + if self + .load_session_metadata(workspace_path, &metadata.session_id) + .await? + .is_some() + { + return Ok(false); + } + self.save_session_metadata_locked(workspace_path, metadata) + .await?; + Ok(true) + } + + pub async fn update_session_metadata( + &self, + workspace_path: &Path, + session_id: &str, + update: impl FnOnce(&mut SessionMetadata), + ) -> BitFunResult<()> { + let updated = self + .update_session_metadata_if_present(workspace_path, session_id, |metadata| { + update(metadata); + Ok(()) + }) + .await?; + if updated { + Ok(()) + } else { + Err(BitFunError::NotFound(format!( + "Session metadata not found: {}", + session_id + ))) + } + } + + pub async fn update_session_metadata_if_present( + &self, + workspace_path: &Path, + session_id: &str, + update: impl FnOnce(&mut SessionMetadata) -> BitFunResult<()>, + ) -> BitFunResult { + Self::validate_session_id(session_id)?; + self.ensure_runtime_for_write(workspace_path).await?; + let persistence_lock = self + .get_session_persistence_lock(workspace_path, session_id) + .await; + let _persistence_guard = persistence_lock.lock().await; + self.update_session_metadata_if_present_locked(workspace_path, session_id, update) + .await + } + + async fn update_session_metadata_if_present_locked( + &self, + workspace_path: &Path, + session_id: &str, + update: impl FnOnce(&mut SessionMetadata) -> BitFunResult<()>, + ) -> BitFunResult { + let Some(mut metadata) = self + .load_session_metadata(workspace_path, session_id) + .await? + else { + return Ok(false); + }; + update(&mut metadata)?; + self.save_session_metadata_locked(workspace_path, &metadata) + .await?; + Ok(true) + } + pub async fn set_session_memory_mode( &self, workspace_path: &Path, @@ -907,10 +1027,10 @@ impl PersistenceManager { mode: SessionMemoryMode, ) -> BitFunResult<()> { Self::validate_session_id(session_id)?; - let metadata_update_lock = self - .get_session_metadata_update_lock(workspace_path, session_id) + let persistence_lock = self + .get_session_persistence_lock(workspace_path, session_id) .await; - let _metadata_update_guard = metadata_update_lock.lock().await; + let _persistence_guard = persistence_lock.lock().await; let mut metadata = self .load_session_metadata(workspace_path, session_id) .await? @@ -918,7 +1038,8 @@ impl PersistenceManager { BitFunError::NotFound(format!("Session metadata not found: {}", session_id)) })?; metadata.memory_mode = mode; - self.save_session_metadata(workspace_path, &metadata).await + self.save_session_metadata_locked(workspace_path, &metadata) + .await } pub async fn mark_session_memory_mode_polluted( @@ -927,10 +1048,10 @@ impl PersistenceManager { session_id: &str, ) -> BitFunResult<()> { Self::validate_session_id(session_id)?; - let metadata_update_lock = self - .get_session_metadata_update_lock(workspace_path, session_id) + let persistence_lock = self + .get_session_persistence_lock(workspace_path, session_id) .await; - let _metadata_update_guard = metadata_update_lock.lock().await; + let _persistence_guard = persistence_lock.lock().await; let mut metadata = self .load_session_metadata(workspace_path, session_id) .await? @@ -943,7 +1064,7 @@ impl PersistenceManager { ); if metadata.memory_mode == SessionMemoryMode::Enabled { metadata.memory_mode = SessionMemoryMode::Polluted; - self.save_session_metadata(workspace_path, &metadata) + self.save_session_metadata_locked(workspace_path, &metadata) .await?; } if should_enqueue_phase2 { @@ -1407,6 +1528,10 @@ impl PersistenceManager { error )) })?; + let persistence_lock = self + .get_session_persistence_lock(workspace_path, &session.session_id) + .await; + let _persistence_guard = persistence_lock.lock().await; let session_dir = self .session_layout(workspace_path) .session_dir(&session.session_id); @@ -1427,7 +1552,10 @@ impl PersistenceManager { } } - if let Err(error) = self.save_session_files(workspace_path, session).await { + if let Err(error) = self + .save_session_files_locked(workspace_path, session) + .await + { if let Err(cleanup_error) = self .session_metadata_store(workspace_path) .delete_session_dir_and_index(&session.session_id) @@ -1437,6 +1565,11 @@ impl PersistenceManager { "Failed to clean up partial session persistence: session_id={}, error={}", session.session_id, cleanup_error ); + return Err(BitFunError::SessionCreateCleanupRequired { + session_id: session.session_id.clone(), + error: error.to_string(), + cleanup_error: cleanup_error.to_string(), + }); } return Err(error); } @@ -1448,12 +1581,17 @@ impl PersistenceManager { pub async fn save_session(&self, workspace_path: &Path, session: &Session) -> BitFunResult<()> { Self::validate_session_id(&session.session_id)?; self.ensure_runtime_for_write(workspace_path).await?; + let persistence_lock = self + .get_session_persistence_lock(workspace_path, &session.session_id) + .await; + let _persistence_guard = persistence_lock.lock().await; self.ensure_session_dir(workspace_path, &session.session_id) .await?; - self.save_session_files(workspace_path, session).await + self.save_session_files_locked(workspace_path, session) + .await } - async fn save_session_files( + async fn save_session_files_locked( &self, workspace_path: &Path, session: &Session, @@ -1464,7 +1602,7 @@ impl PersistenceManager { let metadata = self .build_session_metadata(workspace_path, session, existing_metadata.as_ref()) .await; - self.save_session_metadata(workspace_path, &metadata) + self.save_session_metadata_locked(workspace_path, &metadata) .await?; let state = StoredSessionStateFile { @@ -1811,6 +1949,10 @@ impl PersistenceManager { ) -> BitFunResult<()> { Self::validate_session_id(session_id)?; self.ensure_runtime_for_write(workspace_path).await?; + let persistence_lock = self + .get_session_persistence_lock(workspace_path, session_id) + .await; + let _persistence_guard = persistence_lock.lock().await; let mut stored_state = self .load_stored_session_state(workspace_path, session_id) .await? @@ -1839,6 +1981,10 @@ impl PersistenceManager { session_id: &str, ) -> BitFunResult<()> { Self::validate_session_id(session_id)?; + let persistence_lock = self + .get_session_persistence_lock(workspace_path, session_id) + .await; + let _persistence_guard = persistence_lock.lock().await; self.session_metadata_store(workspace_path) .delete_session_dir_and_index(session_id) .await @@ -1886,10 +2032,10 @@ impl PersistenceManager { Self::validate_session_id(&turn.session_id)?; let save_started_at = Instant::now(); self.ensure_runtime_for_write(workspace_path).await?; - let metadata_update_lock = self - .get_session_metadata_update_lock(workspace_path, &turn.session_id) + let persistence_lock = self + .get_session_persistence_lock(workspace_path, &turn.session_id) .await; - let _metadata_update_guard = metadata_update_lock.lock().await; + let _persistence_guard = persistence_lock.lock().await; let mut metadata = self .load_session_metadata(workspace_path, &turn.session_id) .await? @@ -1968,7 +2114,7 @@ impl PersistenceManager { } let metadata_started_at = Instant::now(); - self.save_session_metadata(workspace_path, &metadata) + self.save_session_metadata_locked(workspace_path, &metadata) .await?; if should_enqueue_phase2_for_pollution { self.enqueue_phase2_if_session_selected(&turn.session_id, current_unix_secs()) @@ -2225,6 +2371,10 @@ impl PersistenceManager { turn_index: usize, ) -> BitFunResult<()> { Self::validate_session_id(session_id)?; + let persistence_lock = self + .get_session_persistence_lock(workspace_path, session_id) + .await; + let _persistence_guard = persistence_lock.lock().await; if !self.turns_dir(workspace_path, session_id).exists() { return Ok(()); } @@ -2234,20 +2384,27 @@ impl PersistenceManager { .await .map_err(|e| BitFunError::io(format!("Failed to delete dialog turn files: {}", e)))?; - if let Some(mut metadata) = self + if self .load_session_metadata(workspace_path, session_id) .await? + .is_some() { let turns = self.load_session_turns(workspace_path, session_id).await?; let workspace_path_text = workspace_path.to_string_lossy(); - refresh_session_metadata_from_turns( - &mut metadata, - workspace_path_text.as_ref(), - &turns, - Self::system_time_to_unix_ms(SystemTime::now()), - ); - self.save_session_metadata(workspace_path, &metadata) - .await?; + self.update_session_metadata_if_present_locked( + workspace_path, + session_id, + |metadata| { + refresh_session_metadata_from_turns( + metadata, + workspace_path_text.as_ref(), + &turns, + Self::system_time_to_unix_ms(SystemTime::now()), + ); + Ok(()) + }, + ) + .await?; } Ok(()) @@ -2649,6 +2806,10 @@ impl PersistenceManager { turn_index: usize, ) -> BitFunResult { Self::validate_session_id(session_id)?; + let persistence_lock = self + .get_session_persistence_lock(workspace_path, session_id) + .await; + let _persistence_guard = persistence_lock.lock().await; let turns = self.load_session_turns(workspace_path, session_id).await?; let mut deleted = 0usize; @@ -2665,20 +2826,27 @@ impl PersistenceManager { } } - if let Some(mut metadata) = self + if self .load_session_metadata(workspace_path, session_id) .await? + .is_some() { let remaining_turns = self.load_session_turns(workspace_path, session_id).await?; let workspace_path_text = workspace_path.to_string_lossy(); - refresh_session_metadata_from_turns( - &mut metadata, - workspace_path_text.as_ref(), - &remaining_turns, - Self::system_time_to_unix_ms(SystemTime::now()), - ); - self.save_session_metadata(workspace_path, &metadata) - .await?; + self.update_session_metadata_if_present_locked( + workspace_path, + session_id, + |metadata| { + refresh_session_metadata_from_turns( + metadata, + workspace_path_text.as_ref(), + &remaining_turns, + Self::system_time_to_unix_ms(SystemTime::now()), + ); + Ok(()) + }, + ) + .await?; } Ok(deleted) @@ -2691,6 +2859,10 @@ impl PersistenceManager { turn_index: usize, ) -> BitFunResult { Self::validate_session_id(session_id)?; + let persistence_lock = self + .get_session_persistence_lock(workspace_path, session_id) + .await; + let _persistence_guard = persistence_lock.lock().await; let turns = self.load_session_turns(workspace_path, session_id).await?; let mut deleted = 0usize; @@ -2707,36 +2879,39 @@ impl PersistenceManager { } } - if let Some(mut metadata) = self + if self .load_session_metadata(workspace_path, session_id) .await? + .is_some() { let remaining_turns = self.load_session_turns(workspace_path, session_id).await?; let workspace_path_text = workspace_path.to_string_lossy(); - refresh_session_metadata_from_turns( - &mut metadata, - workspace_path_text.as_ref(), - &remaining_turns, - Self::system_time_to_unix_ms(SystemTime::now()), - ); - self.save_session_metadata(workspace_path, &metadata) - .await?; + self.update_session_metadata_if_present_locked( + workspace_path, + session_id, + |metadata| { + refresh_session_metadata_from_turns( + metadata, + workspace_path_text.as_ref(), + &remaining_turns, + Self::system_time_to_unix_ms(SystemTime::now()), + ); + Ok(()) + }, + ) + .await?; } Ok(deleted) } pub async fn touch_session(&self, workspace_path: &Path, session_id: &str) -> BitFunResult<()> { - Self::validate_session_id(session_id)?; - if let Some(mut metadata) = self - .load_session_metadata(workspace_path, session_id) - .await? - { + self.update_session_metadata_if_present(workspace_path, session_id, |metadata| { metadata.touch(); - self.save_session_metadata(workspace_path, &metadata) - .await?; - } - Ok(()) + Ok(()) + }) + .await + .map(|_| ()) } } @@ -3437,6 +3612,82 @@ mod tests { .exists()); } + #[tokio::test] + async fn metadata_patch_and_turn_save_share_one_read_modify_write_lock() { + let workspace = TestWorkspace::new(); + let manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let session_id = Uuid::new_v4().to_string(); + let session = Session::new_with_id( + session_id.clone(), + "Concurrent metadata".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ); + manager + .save_session(workspace.path(), &session) + .await + .expect("session should save"); + + let metadata_lock = manager + .get_session_persistence_lock(workspace.path(), &session_id) + .await; + let metadata_guard = metadata_lock.lock().await; + let workspace_path = workspace.path().to_path_buf(); + + let patch_task = tokio::spawn({ + let manager = manager.clone(); + let workspace_path = workspace_path.clone(); + let session_id = session_id.clone(); + async move { + manager + .update_session_metadata(&workspace_path, &session_id, |metadata| { + metadata.agent_type = "Plan".to_string(); + }) + .await + } + }); + + let mut turn = DialogTurnData::new( + "turn-0".to_string(), + 0, + session_id.clone(), + user_message("concurrent turn"), + ); + turn.mark_completed(); + let turn_task = tokio::spawn({ + let manager = manager.clone(); + let workspace_path = workspace_path.clone(); + async move { manager.save_dialog_turn(&workspace_path, &turn).await } + }); + + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + assert!(!patch_task.is_finished()); + assert!(!turn_task.is_finished()); + drop(metadata_guard); + + patch_task + .await + .expect("metadata patch task should join") + .expect("metadata patch should save"); + turn_task + .await + .expect("turn save task should join") + .expect("turn should save"); + + let metadata = manager + .load_session_metadata(&workspace_path, &session_id) + .await + .expect("metadata should load") + .expect("metadata should exist"); + assert_eq!(metadata.agent_type, "Plan"); + assert_eq!(metadata.turn_count, 1); + } + #[tokio::test] async fn save_dialog_turn_updates_metadata_without_scanning_unrelated_turn_files() { let workspace = TestWorkspace::new(); @@ -3513,6 +3764,158 @@ mod tests { assert_eq!(metadata.message_count, 5); } + #[tokio::test] + async fn turn_deletion_waits_for_the_session_metadata_transaction() { + let workspace = TestWorkspace::new(); + let manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let session_id = Uuid::new_v4().to_string(); + let session = Session::new_with_id( + session_id.clone(), + "Transactional deletion".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ); + manager + .save_session(workspace.path(), &session) + .await + .expect("session should save"); + let mut turn = DialogTurnData::new( + "turn-0".to_string(), + 0, + session_id.clone(), + user_message("turn to delete"), + ); + turn.mark_completed(); + manager + .save_dialog_turn(workspace.path(), &turn) + .await + .expect("turn should save"); + + let metadata_lock = manager + .get_session_persistence_lock(workspace.path(), &session_id) + .await; + let metadata_guard = metadata_lock.lock().await; + let turn_path = manager.turn_path(workspace.path(), &session_id, 0); + let delete_task = tokio::spawn({ + let manager = manager.clone(); + let workspace_path = workspace.path().to_path_buf(); + let session_id = session_id.clone(); + async move { + manager + .delete_turns_from(&workspace_path, &session_id, 0) + .await + } + }); + + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + assert!( + turn_path.exists(), + "turn files must not change before the metadata transaction is acquired" + ); + assert!(!delete_task.is_finished()); + drop(metadata_guard); + + assert_eq!( + delete_task + .await + .expect("delete task should join") + .expect("delete should succeed"), + 1 + ); + assert!(!turn_path.exists()); + let metadata = manager + .load_session_metadata(workspace.path(), &session_id) + .await + .expect("metadata should load") + .expect("metadata should exist"); + assert_eq!(metadata.turn_count, 0); + } + + #[tokio::test] + async fn whole_session_deletion_waits_for_the_persistence_transaction() { + let workspace = TestWorkspace::new(); + let manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let session_id = Uuid::new_v4().to_string(); + let session = Session::new_with_id( + session_id.clone(), + "Transactional session deletion".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ); + manager + .save_session(workspace.path(), &session) + .await + .expect("session should save"); + + let persistence_lock = manager + .get_session_persistence_lock(workspace.path(), &session_id) + .await; + let persistence_guard = persistence_lock.lock().await; + let session_dir = manager + .session_layout(workspace.path()) + .session_dir(&session_id); + let delete_task = tokio::spawn({ + let manager = manager.clone(); + let workspace_path = workspace.path().to_path_buf(); + let session_id = session_id.clone(); + async move { manager.delete_session(&workspace_path, &session_id).await } + }); + + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + assert!(session_dir.exists()); + assert!(!delete_task.is_finished()); + drop(persistence_guard); + + delete_task + .await + .expect("delete task should join") + .expect("session delete should succeed"); + assert!(!session_dir.exists()); + } + + #[tokio::test] + async fn metadata_lock_identity_normalizes_workspace_path_aliases() { + let workspace = TestWorkspace::new(); + let manager = + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"); + let session_id = Uuid::new_v4().to_string(); + let session = Session::new_with_id( + session_id.clone(), + "Canonical metadata lock".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ); + manager + .save_session(workspace.path(), &session) + .await + .expect("session should save"); + + std::fs::create_dir_all(workspace.path().join("alias-component")) + .expect("alias component should exist"); + let alias = workspace.path().join("alias-component").join(".."); + let canonical_lock = manager + .get_session_persistence_lock(workspace.path(), &session_id) + .await; + let alias_lock = manager + .get_session_persistence_lock(&alias, &session_id) + .await; + + assert!(Arc::ptr_eq(&canonical_lock, &alias_lock)); + } + #[tokio::test] async fn save_dialog_turn_persists_last_finished_at() { let workspace = TestWorkspace::new(); diff --git a/src/crates/assembly/core/src/agentic/session/context_store.rs b/src/crates/assembly/core/src/agentic/session/context_store.rs index 3e486e97b..50ad6d3b7 100644 --- a/src/crates/assembly/core/src/agentic/session/context_store.rs +++ b/src/crates/assembly/core/src/agentic/session/context_store.rs @@ -56,4 +56,9 @@ impl SessionContextStore { self.session_contexts.remove(session_id); debug!("Deleted session context cache: session_id={}", session_id); } + + #[cfg(test)] + pub(crate) fn has_session(&self, session_id: &str) -> bool { + self.session_contexts.contains_key(session_id) + } } diff --git a/src/crates/assembly/core/src/agentic/session/session_manager.rs b/src/crates/assembly/core/src/agentic/session/session_manager.rs index 9f2c9db50..1446a564d 100644 --- a/src/crates/assembly/core/src/agentic/session/session_manager.rs +++ b/src/crates/assembly/core/src/agentic/session/session_manager.rs @@ -2,6 +2,7 @@ //! //! Responsible for session CRUD, lifecycle management, and resource association +use crate::agentic::agents::get_agent_registry; use crate::agentic::core::{ new_turn_id, CompressionContract, CompressionState, InternalReminderKind, Message, MessageContent, MessageRole, MessageSemanticKind, ProcessingPhase, Session, SessionConfig, @@ -56,6 +57,7 @@ use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Instant; use std::time::{Duration, SystemTime}; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use tokio::time; /// Session manager configuration @@ -138,6 +140,11 @@ pub struct SessionManager { /// Active sessions in memory sessions: Arc>, + /// Exact admission accounting for loaded sessions. A permit is acquired + /// before create/restore publishes runtime state and released on unload/delete/eviction. + active_session_capacity: Arc, + active_session_permits: Arc>, + /// Runtime cache of session_id -> effective session storage path. /// Populated on session create/restore and used to restore evicted sessions /// or resolve workspace-bound operations that only receive a session_id. @@ -167,6 +174,25 @@ pub struct SessionManager { config: SessionManagerConfig, } +fn clear_session_runtime_stores( + session_id: &str, + context_store: &SessionContextStore, + prompt_cache_store: &SessionPromptCacheStore, + token_anchor_store: &TokenAnchorStore, + turn_skill_agent_snapshot_store: &TurnSkillAgentSnapshotStore, + skill_agent_baseline_override_snapshot_store: &DashMap, + file_read_state_store: &FileReadStateStore, + evidence_ledger: &SessionEvidenceLedger, +) { + context_store.delete_session(session_id); + prompt_cache_store.delete_session(session_id); + token_anchor_store.delete_session(session_id); + turn_skill_agent_snapshot_store.delete_session(session_id); + skill_agent_baseline_override_snapshot_store.remove(session_id); + file_read_state_store.delete_session(session_id); + evidence_ledger.delete_session(session_id); +} + #[derive(Clone)] struct SessionAutoSaveSnapshot { session_id: String, @@ -202,6 +228,42 @@ impl SessionManager { Ok(self.lock_session_mutation(session_id).await) } + fn reserve_active_session(&self) -> BitFunResult { + self.active_session_capacity + .clone() + .try_acquire_owned() + .map_err(|_| { + BitFunError::Validation(format!( + "Exceeded maximum session limit: {}", + self.config.max_active_sessions + )) + }) + } + + fn commit_active_session_reservation(&self, session_id: &str, permit: OwnedSemaphorePermit) { + let previous = self + .active_session_permits + .insert(session_id.to_string(), permit); + debug_assert!(previous.is_none(), "active session permit already existed"); + } + + fn release_active_session_reservation(&self, session_id: &str) { + self.active_session_permits.remove(session_id); + } + + #[cfg(test)] + fn evict_loaded_session_for_test(&self, session_id: &str) { + self.sessions.remove(session_id); + self.release_active_session_reservation(session_id); + } + + #[cfg(test)] + pub(crate) fn storage_path_binding_for_test(&self, session_id: &str) -> Option { + self.session_storage_path_index + .get(session_id) + .map(|binding| binding.path.clone()) + } + fn normalize_session_storage_path(path: &Path) -> PathBuf { dunce::canonicalize(path).unwrap_or_else(|_| path.to_path_buf()) } @@ -1461,6 +1523,8 @@ impl SessionManager { let manager = Self { sessions: Arc::new(DashMap::new()), + active_session_capacity: Arc::new(Semaphore::new(config.max_active_sessions)), + active_session_permits: Arc::new(DashMap::new()), session_storage_path_index: Arc::new(DashMap::new()), session_mutation_locks: KeyedAsyncLock::default(), context_store, @@ -1664,6 +1728,8 @@ impl SessionManager { fn spawn_model_reconciliation_listener(&self) { let sessions = self.sessions.clone(); + let active_session_capacity = self.active_session_capacity.clone(); + let active_session_permits = self.active_session_permits.clone(); let session_storage_path_index = self.session_storage_path_index.clone(); let session_mutation_locks = self.session_mutation_locks.clone(); let context_store = self.context_store.clone(); @@ -1691,6 +1757,8 @@ impl SessionManager { // surface area we need from the cloned shared fields above. let manager = Self { sessions, + active_session_capacity, + active_session_permits, session_storage_path_index, session_mutation_locks, context_store, @@ -1815,14 +1883,6 @@ impl SessionManager { BitFunError::Validation("Session workspace_path is required".to_string()) })?; - // Check session count limit - if self.sessions.len() >= self.config.max_active_sessions { - return Err(BitFunError::Validation(format!( - "Exceeded maximum session limit: {}", - self.config.max_active_sessions - ))); - } - let mut session = if let Some(id) = session_id { Session::new_with_id(id, session_name, agent_type.clone(), config) } else { @@ -1851,6 +1911,7 @@ impl SessionManager { "Persisted session ID already exists: {session_id}" ))); } + let active_session_permit = self.reserve_active_session()?; let storage_claim = self.claim_session_storage_path(&session_id, &session_storage_path, true)?; @@ -1893,6 +1954,7 @@ impl SessionManager { self.turn_skill_agent_snapshot_store .delete_session(&session_id); self.file_read_state_store.delete_session(&session_id); + self.evidence_ledger.delete_session(&session_id); self.release_failed_session_storage_path_claim( &session_id, &session_storage_path, @@ -1902,6 +1964,7 @@ impl SessionManager { } } self.commit_session_storage_path_claim(&session_id, &session_storage_path, storage_claim); + self.commit_active_session_reservation(&session_id, active_session_permit); info!("Session created: session_name={}", session.session_name); @@ -2741,10 +2804,38 @@ impl SessionManager { session_id: &str, agent_type: &str, ) -> BitFunResult<()> { - if let Some(mut session) = self.sessions.get_mut(session_id) { - session.agent_type = agent_type.to_string(); - session.updated_at = SystemTime::now(); - session.last_activity_at = SystemTime::now(); + let _mutation_guard = self.acquire_session_mutation(session_id).await?; + let mut session = self + .sessions + .get(session_id) + .map(|session| session.clone()) + .ok_or_else(|| BitFunError::NotFound(format!("Session not found: {}", session_id)))?; + + if session.agent_type == agent_type { + return Ok(()); + } + + let now = SystemTime::now(); + session.agent_type = agent_type.to_string(); + session.updated_at = now; + session.last_activity_at = now; + + if self.should_persist_session_id(session_id) { + let last_active_at = now + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64; + self.update_persisted_session_metadata(session_id, |metadata| { + metadata.agent_type = agent_type.to_string(); + metadata.last_active_at = last_active_at; + }) + .await?; + } + + if let Some(mut active_session) = self.sessions.get_mut(session_id) { + active_session.agent_type = session.agent_type.clone(); + active_session.updated_at = now; + active_session.last_activity_at = now; } else { return Err(BitFunError::NotFound(format!( "Session not found: {}", @@ -2752,17 +2843,6 @@ impl SessionManager { ))); } - if self.should_persist_session_id(session_id) { - let effective_path = self.effective_session_storage_path(session_id).await; - let session_snapshot = self.sessions.get(session_id).map(|s| s.clone()); - // Ref guard released -- DashMap shard lock is free. - if let (Some(workspace_path), Some(session)) = (effective_path, session_snapshot) { - self.persistence_manager - .save_session(&workspace_path, &session) - .await?; - } - } - debug!( "Session agent type updated: session_id={}, agent_type={}", session_id, agent_type @@ -2781,6 +2861,7 @@ impl SessionManager { session_id: &str, agent_type: &str, ) -> BitFunResult<()> { + let _mutation_guard = self.acquire_session_mutation(session_id).await?; if let Some(mut session) = self.sessions.get_mut(session_id) { session.last_submitted_agent_type = Some(agent_type.to_string()); session.updated_at = SystemTime::now(); @@ -3079,6 +3160,54 @@ impl SessionManager { .await } + /// Release one loaded session and its transient runtime stores while keeping + /// persisted history and the storage-path binding available for a later restore. + /// + /// Callers must quiesce scheduler execution before unloading. A processing + /// session is rejected so close/failure compensation cannot detach live work. + pub(crate) async fn unload_session_from_memory(&self, session_id: &str) -> BitFunResult { + bitfun_core_types::validate_session_id(session_id).map_err(BitFunError::Validation)?; + let _mutation_guard = self.lock_session_mutation(session_id).await; + let Some(session) = self.get_session(session_id) else { + return Ok(false); + }; + if matches!(session.state, SessionState::Processing { .. }) { + return Err(BitFunError::Validation(format!( + "Cannot unload a processing session: {session_id}" + ))); + } + + if self.config.enable_persistence && Self::should_persist_session(&session) { + let storage_path = self + .effective_session_storage_path(session_id) + .await + .ok_or_else(|| { + BitFunError::NotFound(format!( + "Session storage path is unavailable: {session_id}" + )) + })?; + self.persistence_manager + .save_session(&storage_path, &session) + .await?; + } + + if self.sessions.remove(session_id).is_none() { + return Ok(false); + } + self.release_active_session_reservation(session_id); + clear_session_runtime_stores( + session_id, + self.context_store.as_ref(), + self.prompt_cache_store.as_ref(), + self.token_anchor_store.as_ref(), + self.turn_skill_agent_snapshot_store.as_ref(), + self.skill_agent_baseline_override_snapshot_store.as_ref(), + self.file_read_state_store.as_ref(), + self.evidence_ledger.as_ref(), + ); + Ok(true) + } + async fn delete_session_from_paths_locked( &self, cleanup_workspace_path: &Path, @@ -3143,14 +3272,16 @@ impl SessionManager { "Session deletion stage starting: session_id={}, stage=context_store_delete", session_id ); - self.context_store.delete_session(session_id); - self.prompt_cache_store.delete_session(session_id); - self.token_anchor_store.delete_session(session_id); - self.turn_skill_agent_snapshot_store - .delete_session(session_id); - self.skill_agent_baseline_override_snapshot_store - .remove(session_id); - self.file_read_state_store.delete_session(session_id); + clear_session_runtime_stores( + session_id, + self.context_store.as_ref(), + self.prompt_cache_store.as_ref(), + self.token_anchor_store.as_ref(), + self.turn_skill_agent_snapshot_store.as_ref(), + self.skill_agent_baseline_override_snapshot_store.as_ref(), + self.file_read_state_store.as_ref(), + self.evidence_ledger.as_ref(), + ); debug!( "Session deletion stage completed: session_id={}, stage=context_store_delete, duration_ms={}", session_id, @@ -3219,6 +3350,7 @@ impl SessionManager { session_id ); self.sessions.remove(session_id); + self.release_active_session_reservation(session_id); debug!( "Session deletion stage completed: session_id={}, stage=in_memory_remove, duration_ms={}", session_id, @@ -3810,6 +3942,11 @@ impl SessionManager { let restore_started_at = Instant::now(); // Check if session is already in memory let session_already_in_memory = self.sessions.contains_key(session_id); + let active_session_permit = if session_already_in_memory { + None + } else { + Some(self.reserve_active_session()?) + }; debug!( "Session restore phase completed: session_id={}, phase=use_storage_path, duration_ms=0", @@ -3853,6 +3990,33 @@ impl SessionManager { let ai_config_for_restore = Self::load_ai_config_for_model_resolution().await; let mut should_persist_restored_session = false; + let mut auto_migrated_model_id = None; + + if !include_internal { + let available_modes = get_agent_registry().get_modes_info().await; + if !available_modes + .iter() + .any(|mode| mode.id == session.agent_type) + { + let fallback_mode = available_modes + .iter() + .find(|mode| mode.id == "agentic") + .or_else(|| available_modes.first()) + .map(|mode| mode.id.clone()) + .ok_or_else(|| { + BitFunError::Validation( + "No executable main agent mode is available for session restore" + .to_string(), + ) + })?; + warn!( + "Persisted session mode is unavailable; applying executable fallback: session_id={}, persisted_mode={}, fallback_mode={}", + session.session_id, session.agent_type, fallback_mode + ); + session.agent_type = fallback_mode; + should_persist_restored_session = true; + } + } // Lazy migration: if the persisted model_id is no longer usable // (model deleted or disabled while the session was on disk), repoint @@ -3878,17 +4042,7 @@ impl SessionManager { let previous_model_id = trimmed.to_string(); session.config.model_id = Some("auto".to_string()); should_persist_restored_session = true; - - if let Some(coordinator) = crate::agentic::coordination::get_global_coordinator() { - coordinator - .emit_session_model_auto_migrated( - session_id, - &previous_model_id, - "auto", - "model_unavailable_on_restore", - ) - .await; - } + auto_migrated_model_id = Some(previous_model_id); } } @@ -3984,29 +4138,6 @@ impl SessionManager { ); } - // 3. Restore the in-memory context cache from the recovered messages. - // If session already exists, delete old one first then create (ensure clean state) - if session_already_in_memory { - self.context_store.delete_session(session_id); - self.prompt_cache_store.delete_session(session_id); - self.token_anchor_store.delete_session(session_id); - self.turn_skill_agent_snapshot_store - .delete_session(session_id); - self.skill_agent_baseline_override_snapshot_store - .remove(session_id); - self.file_read_state_store.delete_session(session_id); - } - - let context_replace_started_at = Instant::now(); - self.context_store - .replace_context(session_id, messages.clone()); - debug!( - "Session restore phase completed: session_id={}, phase=replace_context, message_count={}, duration_ms={}", - session_id, - messages.len(), - elapsed_ms_u64(context_replace_started_at) - ); - let recoverable_turn_count = latest_turn_index .map(|turn_index| turn_index + 1) .unwrap_or(0) @@ -4020,6 +4151,7 @@ impl SessionManager { persisted_turns.len() ); session.dialog_turn_ids = persisted_turn_ids; + should_persist_restored_session = true; } else if session.dialog_turn_ids.len() > recoverable_turn_count { warn!( "Session metadata exceeds recoverable history, truncating: session_id={}, session_turn_count={}, recoverable_turn_count={}", @@ -4028,6 +4160,7 @@ impl SessionManager { recoverable_turn_count ); session.dialog_turn_ids.truncate(recoverable_turn_count); + should_persist_restored_session = true; } else if persisted_turns.len() == session.dialog_turn_ids.len() && session.dialog_turn_ids != persisted_turn_ids { @@ -4036,6 +4169,7 @@ impl SessionManager { session_id ); session.dialog_turn_ids = persisted_turn_ids; + should_persist_restored_session = true; } if recoverable_turn_count == 0 && !session.dialog_turn_ids.is_empty() && messages.is_empty() @@ -4045,8 +4179,41 @@ impl SessionManager { session_id ); session.dialog_turn_ids.clear(); + should_persist_restored_session = true; + } + + // Complete all fallible restore migrations before publishing any runtime state. + // A failed write keeps the session unloaded; restore-time recovery handles any + // partial metadata/state update left by the existing multi-file persistence format. + if should_persist_restored_session && self.should_persist_session_id(session_id) { + self.persistence_manager + .save_session(session_storage_path, &session) + .await?; + } + + // 3. Publish the recovered runtime context only after migrations are durable. + if session_already_in_memory { + clear_session_runtime_stores( + session_id, + self.context_store.as_ref(), + self.prompt_cache_store.as_ref(), + self.token_anchor_store.as_ref(), + self.turn_skill_agent_snapshot_store.as_ref(), + self.skill_agent_baseline_override_snapshot_store.as_ref(), + self.file_read_state_store.as_ref(), + self.evidence_ledger.as_ref(), + ); } + let context_replace_started_at = Instant::now(); + self.context_store + .replace_context(session_id, messages.clone()); + debug!( + "Session restore phase completed: session_id={}, phase=replace_context, message_count={}, duration_ms={}", + session_id, + messages.len(), + elapsed_ms_u64(context_replace_started_at) + ); let context_msg_count = self.context_store.get_context_messages(session_id).len(); debug!( @@ -4064,17 +4231,27 @@ impl SessionManager { // as completed would surface misleading unread indicators after an upgrade. // Unread completion is now written only by runtime completion/persist paths. - if should_persist_restored_session && self.should_persist_session_id(session_id) { - self.persistence_manager - .save_session(session_storage_path, &session) - .await?; - } - // 4. Add to memory (will overwrite if already exists) self.sessions .insert(session_id.to_string(), session.clone()); + if let Some(permit) = active_session_permit { + self.commit_active_session_reservation(session_id, permit); + } self.bind_session_storage_path_committed(session_id, session_storage_path.to_path_buf()); + if let Some(previous_model_id) = auto_migrated_model_id { + if let Some(coordinator) = crate::agentic::coordination::get_global_coordinator() { + coordinator + .emit_session_model_auto_migrated( + session_id, + &previous_model_id, + "auto", + "model_unavailable_on_restore", + ) + .await; + } + } + Ok((session, persisted_turns)) } @@ -4313,6 +4490,18 @@ impl SessionManager { .await } + pub async fn update_session_metadata( + &self, + workspace_path: &Path, + session_id: &str, + update: impl FnOnce(&mut SessionMetadata), + ) -> BitFunResult<()> { + self.persistence_manager + .update_session_metadata(workspace_path, session_id, update) + .await + } + + #[cfg(test)] pub async fn save_session_metadata( &self, workspace_path: &Path, @@ -4388,36 +4577,28 @@ impl SessionManager { }) } - async fn load_or_persist_session_metadata( + async fn ensure_session_metadata_persisted( &self, workspace_path: &Path, session_id: &str, - ) -> BitFunResult { - match self + ) -> BitFunResult<()> { + if self .persistence_manager .load_session_metadata(workspace_path, session_id) .await? + .is_some() { - Some(metadata) => Ok(metadata), - None => { - let session = self - .sessions - .get(session_id) - .map(|value| value.clone()) - .ok_or_else(|| { - BitFunError::NotFound(format!("Session not found: {}", session_id)) - })?; - self.persistence_manager - .save_session(workspace_path, &session) - .await?; - self.persistence_manager - .load_session_metadata(workspace_path, session_id) - .await? - .ok_or_else(|| { - BitFunError::NotFound(format!("Session not found: {}", session_id)) - }) - } + return Ok(()); } + + let session = self + .sessions + .get(session_id) + .map(|value| value.clone()) + .ok_or_else(|| BitFunError::NotFound(format!("Session not found: {}", session_id)))?; + self.persistence_manager + .save_session(workspace_path, &session) + .await } async fn update_session_metadata_at_workspace( @@ -4426,12 +4607,10 @@ impl SessionManager { session_id: &str, update: impl FnOnce(&mut SessionMetadata), ) -> BitFunResult<()> { - let mut metadata = self - .load_or_persist_session_metadata(workspace_path, session_id) + self.ensure_session_metadata_persisted(workspace_path, session_id) .await?; - update(&mut metadata); self.persistence_manager - .save_session_metadata(workspace_path, &metadata) + .update_session_metadata(workspace_path, session_id, update) .await } @@ -4539,6 +4718,7 @@ impl SessionManager { processing_phase: ProcessingPhase, user_message_metadata: Option, ) -> BitFunResult { + let _mutation_guard = self.acquire_session_mutation(session_id).await?; let session = self .get_session(session_id) .ok_or_else(|| BitFunError::NotFound(format!("Session not found: {}", session_id)))?; @@ -5563,6 +5743,7 @@ impl SessionManager { session_id: &str, compression_state: CompressionState, ) -> BitFunResult<()> { + let _mutation_guard = self.acquire_session_mutation(session_id).await?; let effective_path = self.effective_session_storage_path(session_id).await; // IMPORTANT: keep the DashMap guard scope short -- do NOT hold it across .await. @@ -5778,6 +5959,7 @@ impl SessionManager { /// Start cleanup task for expired sessions fn spawn_cleanup_task(&self) { let sessions = self.sessions.clone(); + let active_session_permits = self.active_session_permits.clone(); let timeout = self.config.session_idle_timeout; let persistence = self.persistence_manager.clone(); let enable_persistence = self.config.enable_persistence; @@ -5789,6 +5971,7 @@ impl SessionManager { let skill_agent_baseline_override_snapshot_store = self.skill_agent_baseline_override_snapshot_store.clone(); let file_read_state_store = self.file_read_state_store.clone(); + let evidence_ledger = self.evidence_ledger.clone(); tokio::spawn(async move { let mut ticker = time::interval(Duration::from_secs(60)); @@ -5849,12 +6032,17 @@ impl SessionManager { }) .is_some() { - context_store.delete_session(&candidate.session_id); - prompt_cache_store.delete_session(&candidate.session_id); - token_anchor_store.delete_session(&candidate.session_id); - turn_skill_agent_snapshot_store.delete_session(&candidate.session_id); - skill_agent_baseline_override_snapshot_store.remove(&candidate.session_id); - file_read_state_store.delete_session(&candidate.session_id); + active_session_permits.remove(&candidate.session_id); + clear_session_runtime_stores( + &candidate.session_id, + context_store.as_ref(), + prompt_cache_store.as_ref(), + token_anchor_store.as_ref(), + turn_skill_agent_snapshot_store.as_ref(), + skill_agent_baseline_override_snapshot_store.as_ref(), + file_read_state_store.as_ref(), + evidence_ledger.as_ref(), + ); } } } @@ -5871,8 +6059,8 @@ mod tests { SessionManagerConfig, }; use crate::agentic::core::{ - Message, MessageContent, MessageRole, ProcessingPhase, Session, SessionConfig, - SessionModelBindingPolicy, SessionState, ToolCall, ToolResult, + CompressionState, Message, MessageContent, MessageRole, ProcessingPhase, Session, + SessionConfig, SessionModelBindingPolicy, SessionState, ToolCall, ToolResult, }; use crate::agentic::persistence::PersistenceManager; use crate::agentic::session::{ @@ -6059,50 +6247,299 @@ mod tests { } #[tokio::test] - async fn session_model_update_is_restored_from_persistence() { + async fn unloading_a_session_releases_capacity_without_deleting_persistence() { let workspace = TestWorkspace::new(); let persistence_manager = Arc::new( PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), ); - let manager = test_manager(persistence_manager); - let session = manager + let manager = test_manager_with_config( + persistence_manager.clone(), + SessionManagerConfig { + max_active_sessions: 1, + session_idle_timeout: Duration::from_secs(3600), + auto_save_interval: Duration::from_secs(300), + enable_persistence: true, + prompt_cache_policy: PromptCachePolicy::default(), + }, + ); + let config = SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }; + let first = manager .create_session( - "Persisted model update".to_string(), + "First loaded session".to_string(), "agentic".to_string(), - SessionConfig { - workspace_path: Some(workspace.path().to_string_lossy().into_owned()), - model_id: Some("primary".to_string()), - ..Default::default() - }, + config.clone(), ) .await - .expect("session should create"); + .expect("first session should be created"); - manager - .update_session_model_id(&session.session_id, "auto") + assert!(manager + .unload_session_from_memory(&first.session_id) .await - .expect("model update should persist"); - manager.sessions.remove(&session.session_id); + .expect("session should unload")); + assert!(manager.get_session(&first.session_id).is_none()); + assert!( + persistence_manager + .load_session_metadata(workspace.path(), &first.session_id) + .await + .expect("metadata should load") + .is_some(), + "unload must preserve persisted history" + ); - let restored = manager - .restore_session(workspace.path(), &session.session_id) + let second = manager + .create_session( + "Second loaded session".to_string(), + "agentic".to_string(), + config, + ) .await - .expect("session should restore from persistence"); - assert_eq!(restored.config.model_id.as_deref(), Some("auto")); + .expect("unload should release the active-session slot"); + assert_ne!(first.session_id, second.session_id); } #[tokio::test] - async fn session_storage_identity_rejects_same_id_in_another_workspace() { + async fn restores_share_the_same_exact_active_session_capacity_as_creates() { let workspace = TestWorkspace::new(); - let other_workspace = TestWorkspace::new(); let persistence_manager = Arc::new( PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), ); - let manager = test_manager(persistence_manager); - let session_id = "shared-session-id"; + let config = SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }; + let first = Session::new( + "First persisted".to_string(), + "agentic".to_string(), + config.clone(), + ); + let second = Session::new( + "Second persisted".to_string(), + "agentic".to_string(), + config, + ); + persistence_manager + .save_session(workspace.path(), &first) + .await + .expect("first fixture should persist"); + persistence_manager + .save_session(workspace.path(), &second) + .await + .expect("second fixture should persist"); + let manager = test_manager_with_config( + persistence_manager, + SessionManagerConfig { + max_active_sessions: 1, + enable_persistence: true, + ..Default::default() + }, + ); - assert!(manager - .claim_session_storage_path(session_id, workspace.path(), true) + manager + .restore_session(workspace.path(), &first.session_id) + .await + .expect("first restore should reserve the only slot"); + let error = manager + .restore_session(workspace.path(), &second.session_id) + .await + .expect_err("second restore must respect active-session capacity"); + assert!(error.to_string().contains("maximum session limit")); + + manager + .unload_session_from_memory(&first.session_id) + .await + .expect("first session should unload"); + manager + .restore_session(workspace.path(), &second.session_id) + .await + .expect("unload should release capacity for a later restore"); + } + + #[tokio::test] + async fn concurrent_creates_cannot_overbook_active_session_capacity() { + let workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let manager = Arc::new(test_manager_with_config( + persistence_manager, + SessionManagerConfig { + max_active_sessions: 1, + enable_persistence: true, + ..Default::default() + }, + )); + let config = SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }; + + let first = { + let manager = manager.clone(); + let config = config.clone(); + tokio::spawn(async move { + manager + .create_session_with_id( + Some("capacity-first".to_string()), + "First".to_string(), + "agentic".to_string(), + config, + ) + .await + }) + }; + let second = { + let manager = manager.clone(); + tokio::spawn(async move { + manager + .create_session_with_id( + Some("capacity-second".to_string()), + "Second".to_string(), + "agentic".to_string(), + config, + ) + .await + }) + }; + let first = first.await.expect("first create task should join"); + let second = second.await.expect("second create task should join"); + + assert_eq!(usize::from(first.is_ok()) + usize::from(second.is_ok()), 1); + assert_eq!(manager.sessions.len(), 1); + assert_eq!(manager.active_session_permits.len(), 1); + } + + #[tokio::test] + async fn failed_unavailable_mode_migration_does_not_publish_the_restored_session() { + let workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let session_id = uuid::Uuid::new_v4().to_string(); + let session = Session::new_with_id( + session_id.clone(), + "Unavailable mode".to_string(), + "removed-mode-that-cannot-exist".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ); + persistence_manager + .save_session(workspace.path(), &session) + .await + .expect("invalid historical mode fixture should persist"); + persistence_manager.fail_next_session_metadata_write_for_test(&session_id); + let manager = test_manager_with_config( + persistence_manager, + SessionManagerConfig { + enable_persistence: true, + ..Default::default() + }, + ); + + let error = manager + .restore_session(workspace.path(), &session_id) + .await + .expect_err("mode migration write failure must fail restore"); + + assert!(error.to_string().contains("Injected session metadata")); + assert!( + manager.get_session(&session_id).is_none(), + "failed migration must not consume active-session capacity" + ); + assert!(manager.active_session_permits.is_empty()); + assert!(manager + .session_storage_path_index + .get(&session_id) + .is_none()); + } + + #[tokio::test] + async fn failed_restore_state_write_does_not_publish_context_or_capacity() { + let workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let session_id = uuid::Uuid::new_v4().to_string(); + let session = Session::new_with_id( + session_id.clone(), + "Unavailable mode".to_string(), + "removed-mode-that-cannot-exist".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ); + persistence_manager + .save_session(workspace.path(), &session) + .await + .expect("historical session fixture should persist"); + persistence_manager.fail_next_session_state_write_for_test(&session_id); + let manager = test_manager(persistence_manager); + + let error = manager + .restore_session(workspace.path(), &session_id) + .await + .expect_err("state migration write failure must fail restore"); + + assert!(error.to_string().contains("Injected session state")); + assert!(manager.get_session(&session_id).is_none()); + assert!(manager.active_session_permits.is_empty()); + assert!(manager + .session_storage_path_index + .get(&session_id) + .is_none()); + assert!(!manager.context_store.has_session(&session_id)); + } + + #[tokio::test] + async fn session_model_update_is_restored_from_persistence() { + let workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let manager = test_manager(persistence_manager); + let session = manager + .create_session( + "Persisted model update".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().into_owned()), + model_id: Some("primary".to_string()), + ..Default::default() + }, + ) + .await + .expect("session should create"); + + manager + .update_session_model_id(&session.session_id, "auto") + .await + .expect("model update should persist"); + manager.evict_loaded_session_for_test(&session.session_id); + + let restored = manager + .restore_session(workspace.path(), &session.session_id) + .await + .expect("session should restore from persistence"); + assert_eq!(restored.config.model_id.as_deref(), Some("auto")); + } + + #[tokio::test] + async fn session_storage_identity_rejects_same_id_in_another_workspace() { + let workspace = TestWorkspace::new(); + let other_workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let manager = test_manager(persistence_manager); + let session_id = "shared-session-id"; + + assert!(manager + .claim_session_storage_path(session_id, workspace.path(), true) .expect("first workspace claim")); let error = manager .claim_session_storage_path(session_id, other_workspace.path(), true) @@ -6176,7 +6613,7 @@ mod tests { .expect_err("an active duplicate must fail"); assert!(duplicate.to_string().contains("already exists")); - manager.sessions.remove(session_id); + manager.evict_loaded_session_for_test(session_id); manager .create_session_with_id_and_details( Some(session_id.to_string()), @@ -6219,7 +6656,7 @@ mod tests { std::fs::create_dir_all(&turns_dir).expect("turns directory"); let sentinel = turns_dir.join("existing-turn.json"); std::fs::write(&sentinel, b"existing history").expect("persisted turn sentinel"); - manager.sessions.remove(session_id); + manager.evict_loaded_session_for_test(session_id); let error = manager .create_session_with_id( @@ -6353,7 +6790,7 @@ mod tests { .delete_session(&sessions_dir, &session_id) .await .expect("persistence deletion"); - manager.sessions.remove(&session_id); + manager.evict_loaded_session_for_test(&session_id); manager.session_storage_path_index.remove(&session_id); drop(deletion_guard); @@ -6494,7 +6931,7 @@ mod tests { ) .await .expect("session should create"); - manager.sessions.remove(&session.session_id); + manager.evict_loaded_session_for_test(&session.session_id); let guard = manager.lock_session_mutation(&session.session_id).await; let manager_for_restore = manager.clone(); @@ -6515,6 +6952,239 @@ mod tests { .expect("restore should continue after the permit is released"); } + #[tokio::test] + async fn session_mode_update_waits_for_the_same_session_mutation_permit() { + let workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let manager = Arc::new(test_manager(persistence_manager)); + let session = manager + .create_session( + "Serialized mode update".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ) + .await + .expect("session should create"); + + let guard = manager.lock_session_mutation(&session.session_id).await; + let manager_for_update = manager.clone(); + let session_id = session.session_id.clone(); + let update_task = tokio::spawn(async move { + manager_for_update + .update_session_agent_type(&session_id, "Plan") + .await + }); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!update_task.is_finished()); + + drop(guard); + update_task + .await + .expect("update task should join") + .expect("mode update should continue after the permit is released"); + } + + #[tokio::test] + async fn compression_update_waits_for_the_same_session_mutation_permit() { + let workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let manager = Arc::new(test_manager(persistence_manager)); + let session = manager + .create_session( + "Serialized compression update".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ) + .await + .expect("session should create"); + + let guard = manager.lock_session_mutation(&session.session_id).await; + let manager_for_update = manager.clone(); + let session_id = session.session_id.clone(); + let update_task = tokio::spawn(async move { + manager_for_update + .update_compression_state( + &session_id, + CompressionState { + last_compression_at: None, + compression_count: 1, + }, + ) + .await + }); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!update_task.is_finished()); + + drop(guard); + update_task + .await + .expect("update task should join") + .expect("compression update should continue after the permit is released"); + } + + #[tokio::test] + async fn turn_start_waits_for_the_same_session_mutation_permit() { + let workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let manager = Arc::new(test_manager(persistence_manager)); + let session = manager + .create_session( + "Serialized turn start".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ) + .await + .expect("session should create"); + + let guard = manager.lock_session_mutation(&session.session_id).await; + let manager_for_turn = manager.clone(); + let session_id = session.session_id.clone(); + let turn_task = tokio::spawn(async move { + manager_for_turn + .start_dialog_turn( + &session_id, + "agentic".to_string(), + "hello".to_string(), + Some("serialized-turn".to_string()), + None, + None, + ) + .await + }); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!turn_task.is_finished()); + + drop(guard); + turn_task + .await + .expect("turn task should join") + .expect("turn start should continue after the permit is released"); + } + + #[tokio::test] + async fn same_session_mode_is_a_timestamp_preserving_noop() { + let workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let manager = test_manager(persistence_manager); + let session = manager + .create_session( + "Idempotent mode update".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ) + .await + .expect("session should create"); + let before = manager + .get_session(&session.session_id) + .expect("active session before update"); + let before_updated_at = before.updated_at; + let before_last_activity_at = before.last_activity_at; + drop(before); + tokio::time::sleep(Duration::from_millis(20)).await; + + manager + .update_session_agent_type(&session.session_id, "agentic") + .await + .expect("same mode should succeed"); + + let after = manager + .get_session(&session.session_id) + .expect("active session after update"); + assert_eq!(after.updated_at, before_updated_at); + assert_eq!(after.last_activity_at, before_last_activity_at); + } + + #[tokio::test] + async fn session_mode_persists_without_a_turn_and_survives_restore() { + let workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let manager = test_manager(persistence_manager.clone()); + let session = manager + .create_session( + "Durable mode update".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ) + .await + .expect("session should create"); + + manager + .update_session_agent_type(&session.session_id, "Plan") + .await + .expect("mode update should persist without a turn"); + let metadata = persistence_manager + .load_session_metadata(workspace.path(), &session.session_id) + .await + .expect("metadata should load") + .expect("metadata should exist"); + assert_eq!(metadata.agent_type, "Plan"); + + manager.evict_loaded_session_for_test(&session.session_id); + let restored = manager + .restore_session(workspace.path(), &session.session_id) + .await + .expect("session should restore"); + assert_eq!(restored.agent_type, "Plan"); + } + + #[tokio::test] + async fn session_mode_update_does_not_rewrite_the_runtime_state_file() { + let workspace = TestWorkspace::new(); + let persistence_manager = Arc::new( + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"), + ); + let manager = test_manager(persistence_manager.clone()); + let session = manager + .create_session( + "Metadata-only mode update".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ) + .await + .expect("session should create"); + persistence_manager.fail_next_session_state_write_for_test(&session.session_id); + + manager + .update_session_agent_type(&session.session_id, "Plan") + .await + .expect("mode updates must not depend on rewriting runtime state"); + manager.evict_loaded_session_for_test(&session.session_id); + + let restored = manager + .restore_session(workspace.path(), &session.session_id) + .await + .expect("metadata-only mode update should remain restorable"); + assert_eq!(restored.agent_type, "Plan"); + } + #[tokio::test] async fn persistence_manager_accessor_reuses_runtime_owner() { let persistence_manager = @@ -7797,7 +8467,7 @@ mod tests { .expect("snapshot load should succeed") .is_none()); - manager.sessions.remove(&session.session_id); + manager.evict_loaded_session_for_test(&session.session_id); let restored = manager .restore_session(workspace.path(), &session.session_id) .await @@ -8543,7 +9213,7 @@ mod tests { ) .await .expect("session should create"); - manager.sessions.remove(&session.session_id); + manager.evict_loaded_session_for_test(&session.session_id); let cleanup_workspace_path = manager .resolve_session_cleanup_workspace_path( diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/thread_goal_tools.rs b/src/crates/assembly/core/src/agentic/tools/implementations/thread_goal_tools.rs index eb29319f8..400f18303 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/thread_goal_tools.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/thread_goal_tools.rs @@ -57,7 +57,9 @@ fn thread_goal_port_error(port_error: PortError) -> BitFunError { PortErrorKind::NotAvailable => { user_facing_thread_goal_error(BitFunError::NotImplemented(port_error.message)) } - PortErrorKind::PermissionDenied | PortErrorKind::Backend => { + PortErrorKind::PermissionDenied + | PortErrorKind::CleanupRequired + | PortErrorKind::Backend => { user_facing_thread_goal_error(BitFunError::Tool(port_error.message)) } } diff --git a/src/crates/assembly/core/src/product_runtime.rs b/src/crates/assembly/core/src/product_runtime.rs index e31d1f9fe..e22b06719 100644 --- a/src/crates/assembly/core/src/product_runtime.rs +++ b/src/crates/assembly/core/src/product_runtime.rs @@ -30,7 +30,7 @@ use crate::agentic::keyed_lock::KeyedAsyncLockGuard; use crate::agentic::persistence::session_branch::{SessionBranchRequest, SessionBranchResult}; use crate::agentic::persistence::{PersistenceManager, SessionMetadataPage}; use crate::agentic::session::CoreSessionStorePort; -use crate::service::session::{DialogTurnData, SessionMetadata}; +use crate::service::session::{DialogTurnData, SessionMetadata, SessionStatus}; use crate::service::session_usage::{ generate_session_usage_report, SessionUsageReport, SessionUsageReportRequest, }; @@ -432,6 +432,17 @@ impl CoreAgentRuntimeCompatibility { Ok(CoreSessionMaintenancePermit { _permit: permit }) } + /// Compatibility-only lifecycle operation for ACP setup compensation and + /// session/close. It releases loaded Core state but preserves persistence + /// and the storage binding so the same session can be restored later. + pub async fn unload_persisted_session(&self, session_id: &str) -> BitFunResult { + validate_persisted_session_id(session_id)?; + self.coordinator + .get_session_manager() + .unload_session_from_memory(session_id) + .await + } + pub async fn cancel_background_subagents_for_parent( &self, parent_session_id: &str, @@ -516,25 +527,16 @@ impl CoreAgentRuntimeCompatibility { } } - pub async fn load_persisted_session_metadata( + pub async fn archive_persisted_session( &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)?; + validate_persisted_session_id(session_id)?; self.persistence - .save_session_metadata(workspace_path, metadata) + .update_session_metadata(workspace_path, session_id, |metadata| { + metadata.status = SessionStatus::Archived; + }) .await } @@ -669,6 +671,7 @@ fn runtime_port_error(error: BitFunError) -> PortError { BitFunError::NotFound(_) => PortErrorKind::NotFound, BitFunError::Timeout(_) => PortErrorKind::Timeout, BitFunError::Cancelled(_) => PortErrorKind::Cancelled, + BitFunError::SessionCreateCleanupRequired { .. } => PortErrorKind::CleanupRequired, _ => PortErrorKind::Backend, }; PortError::new(kind, error.to_string()) @@ -752,11 +755,12 @@ mod tests { use bitfun_runtime_services::RuntimeServices; use super::{ - validate_local_session_fork_request, validate_persisted_session_id, + runtime_port_error, validate_local_session_fork_request, validate_persisted_session_id, CoreAgentRuntimeCompatibility, CoreProductAgentRuntime, }; use crate::agentic::coordination::{ConversationCoordinator, DialogScheduler}; use crate::service::token_usage::TokenUsageService; + use crate::util::errors::BitFunError; use bitfun_agent_runtime::sdk::{AgentSessionForkRequest, PortErrorKind}; #[test] @@ -799,6 +803,7 @@ mod tests { let _ = CoreAgentRuntimeCompatibility::list_persisted_sessions; let _ = CoreAgentRuntimeCompatibility::load_persisted_session_turns; let _ = CoreAgentRuntimeCompatibility::update_session_agent_type; + let _ = CoreAgentRuntimeCompatibility::unload_persisted_session; } #[test] @@ -809,6 +814,18 @@ mod tests { assert!(error.to_string().contains("session_id"), "{error}"); } + #[test] + fn session_create_rollback_residual_remains_typed_across_the_runtime_port() { + let error = runtime_port_error(BitFunError::SessionCreateCleanupRequired { + session_id: "session-1".to_string(), + error: "metadata write failed".to_string(), + cleanup_error: "session directory is locked".to_string(), + }); + + assert_eq!(error.kind, PortErrorKind::CleanupRequired); + assert!(error.message.contains("session-1"), "{error}"); + } + #[test] fn local_session_fork_rejects_remote_identity() { let error = validate_local_session_fork_request(&AgentSessionForkRequest { diff --git a/src/crates/assembly/core/src/service/mcp/server/manager/lifecycle.rs b/src/crates/assembly/core/src/service/mcp/server/manager/lifecycle.rs index a6c09822f..a88a3c300 100644 --- a/src/crates/assembly/core/src/service/mcp/server/manager/lifecycle.rs +++ b/src/crates/assembly/core/src/service/mcp/server/manager/lifecycle.rs @@ -428,7 +428,10 @@ impl MCPServerManager { let server_id = config.id.clone(); if self.runtime.contains(&server_id).await { - let _ = self.remove_ephemeral_server(&server_id).await; + return Err(BitFunError::Configuration(format!( + "MCP server already exists: {}", + server_id + ))); } self.runtime.insert_runtime_config(config.clone()).await?; @@ -448,25 +451,30 @@ impl MCPServerManager { pub async fn remove_ephemeral_server(&self, server_id: &str) -> BitFunResult<()> { info!("Removing ephemeral MCP server: id={}", server_id); - let _ = self.stop_server(server_id).await; - self.stop_connection_event_listener(server_id).await; - - match self.runtime.unregister(server_id).await { - Ok(_) => { - info!("Unregistered ephemeral MCP server: id={}", server_id); - } - Err(e) => { - warn!( - "Ephemeral MCP server was not registered, skipping unregister: id={} error={}", - server_id, e - ); - } + if !self.runtime.contains(server_id).await { + self.runtime.remove_runtime_config(server_id).await; + self.clear_reconnect_state(server_id).await; + self.runtime.remove_catalog(server_id).await; + Self::unregister_mcp_tools(server_id).await; + return Ok(()); } - self.runtime.remove_runtime_config(server_id).await; + let stop_result = self.stop_server(server_id).await; + self.stop_connection_event_listener(server_id).await; self.clear_reconnect_state(server_id).await; self.runtime.remove_catalog(server_id).await; + if let Err(error) = stop_result { + warn!( + "Failed to stop ephemeral MCP server; retaining runtime ownership for retry: id={} error={}", + server_id, error + ); + return Err(error); + } + + self.runtime.unregister(server_id).await?; + self.runtime.remove_runtime_config(server_id).await; + info!("Unregistered ephemeral MCP server: id={}", server_id); Ok(()) } diff --git a/src/crates/assembly/core/src/util/errors.rs b/src/crates/assembly/core/src/util/errors.rs index 9622216e6..3b91aff8f 100644 --- a/src/crates/assembly/core/src/util/errors.rs +++ b/src/crates/assembly/core/src/util/errors.rs @@ -26,6 +26,15 @@ pub enum BitFunError { #[error("Session error: {0}")] Session(String), + #[error( + "Session creation persistence failed and rollback did not complete: session_id={session_id}, error={error}, cleanup_error={cleanup_error}" + )] + SessionCreateCleanupRequired { + session_id: String, + error: String, + cleanup_error: String, + }, + #[error("Workspace error: {0}")] Workspace(String), diff --git a/src/crates/contracts/runtime-ports/src/lib.rs b/src/crates/contracts/runtime-ports/src/lib.rs index a99059928..dd151debd 100644 --- a/src/crates/contracts/runtime-ports/src/lib.rs +++ b/src/crates/contracts/runtime-ports/src/lib.rs @@ -46,6 +46,7 @@ pub enum PortErrorKind { PermissionDenied, Cancelled, Timeout, + CleanupRequired, Backend, } diff --git a/src/crates/execution/agent-runtime/src/evidence_ledger.rs b/src/crates/execution/agent-runtime/src/evidence_ledger.rs index 117fcc5fa..55035035b 100644 --- a/src/crates/execution/agent-runtime/src/evidence_ledger.rs +++ b/src/crates/execution/agent-runtime/src/evidence_ledger.rs @@ -193,6 +193,10 @@ impl SessionEvidenceLedger { Self::default() } + pub fn delete_session(&self, session_id: &str) -> bool { + self.events_by_session.remove(session_id).is_some() + } + pub fn append(&self, event: EvidenceLedgerEvent) -> EvidenceLedgerEvent { self.events_by_session .entry(event.session_id.clone()) @@ -401,6 +405,24 @@ mod tests { assert!(ledger.events_for_turn("other-session", "turn-a").is_empty()); } + #[test] + fn deleting_a_session_releases_its_evidence_events() { + let ledger = SessionEvidenceLedger::new(); + ledger.append(EvidenceLedgerEvent::new( + "session-a", + "turn-a", + "Task", + EvidenceLedgerTargetKind::Subagent, + "Review", + EvidenceLedgerEventStatus::Created, + "Started", + )); + + assert!(ledger.delete_session("session-a")); + assert!(ledger.events_for_turn("session-a", "turn-a").is_empty()); + assert!(!ledger.delete_session("session-a")); + } + #[test] fn checkpoint_created_event_preserves_recovery_boundary_metadata() { let checkpoint = EvidenceLedgerCheckpoint { diff --git a/src/crates/interfaces/acp/Cargo.toml b/src/crates/interfaces/acp/Cargo.toml index a4eee5850..519fe4d5b 100644 --- a/src/crates/interfaces/acp/Cargo.toml +++ b/src/crates/interfaces/acp/Cargo.toml @@ -13,6 +13,7 @@ bitfun-core = { path = "../../assembly/core", default-features = false, features bitfun-agent-runtime = { path = "../../execution/agent-runtime" } bitfun-agent-tools = { path = "../../execution/tool-contracts" } bitfun-events = { path = "../../contracts/events" } +bitfun-core-types = { path = "../../contracts/core-types" } agent-client-protocol = { workspace = true } tokio = { workspace = true } @@ -25,6 +26,7 @@ chrono = { workspace = true } dashmap = { workspace = true } log = { workspace = true } uuid = { workspace = true } +sha2 = { workspace = true } [lints] workspace = true diff --git a/src/crates/interfaces/acp/src/client/session_persistence.rs b/src/crates/interfaces/acp/src/client/session_persistence.rs index 6b53e5fc4..6855be1bd 100644 --- a/src/crates/interfaces/acp/src/client/session_persistence.rs +++ b/src/crates/interfaces/acp/src/client/session_persistence.rs @@ -63,9 +63,16 @@ impl AcpSessionPersistence { CUSTOM_METADATA_LAST_RESUME_ERROR_KEY: null, })); - self.manager - .save_session_metadata(session_storage_path, &metadata) - .await?; + if !self + .manager + .create_session_metadata_if_absent(session_storage_path, &metadata) + .await? + { + return Err(BitFunError::Validation(format!( + "ACP flow session ID already exists: {}", + session_id + ))); + } Ok(CreateAcpFlowSessionRecordResponse { session_id, @@ -150,18 +157,10 @@ impl AcpSessionPersistence { bitfun_session_id: &str, update: impl FnOnce(&mut SessionMetadata) -> BitFunResult<()>, ) -> BitFunResult<()> { - let Some(mut metadata) = self - .manager - .load_session_metadata(session_storage_path, bitfun_session_id) - .await? - else { - return Ok(()); - }; - - update(&mut metadata)?; self.manager - .save_session_metadata(session_storage_path, &metadata) + .update_session_metadata_if_present(session_storage_path, bitfun_session_id, update) .await + .map(|_| ()) } } diff --git a/src/crates/interfaces/acp/src/runtime.rs b/src/crates/interfaces/acp/src/runtime.rs index c1893d0fc..4fd400e7a 100644 --- a/src/crates/interfaces/acp/src/runtime.rs +++ b/src/crates/interfaces/acp/src/runtime.rs @@ -1,17 +1,20 @@ use std::sync::Arc; use agent_client_protocol::schema::{ - AgentCapabilities, CancelNotification, Implementation, InitializeRequest, InitializeResponse, - ListSessionsRequest, ListSessionsResponse, LoadSessionRequest, LoadSessionResponse, - McpCapabilities, NewSessionRequest, NewSessionResponse, PromptCapabilities, PromptRequest, - PromptResponse, ProtocolVersion, SessionCapabilities, SessionListCapabilities, + AgentCapabilities, CancelNotification, CloseSessionRequest, CloseSessionResponse, + Implementation, InitializeRequest, InitializeResponse, ListSessionsRequest, + ListSessionsResponse, LoadSessionRequest, LoadSessionResponse, McpCapabilities, + NewSessionRequest, NewSessionResponse, PromptCapabilities, PromptRequest, PromptResponse, + ProtocolVersion, SessionCapabilities, SessionCloseCapabilities, SessionListCapabilities, SetSessionConfigOptionRequest, SetSessionConfigOptionResponse, SetSessionModeRequest, SetSessionModeResponse, SetSessionModelRequest, SetSessionModelResponse, }; use agent_client_protocol::{Client, ConnectionTo, Error, Result}; use async_trait::async_trait; -use bitfun_agent_runtime::sdk::{AgentRuntime, RuntimeError}; +use bitfun_agent_runtime::sdk::{AgentRuntime, PortErrorKind, RuntimeError}; use bitfun_core::product_runtime::CoreAgentRuntimeCompatibility; +use bitfun_core::util::errors::BitFunError; +use dashmap::mapref::entry::Entry; use dashmap::DashMap; use crate::server::{AcpRuntime, AcpServer}; @@ -30,6 +33,7 @@ pub struct BitfunAcpRuntime { pub(crate) compatibility: CoreAgentRuntimeCompatibility, pub(crate) sessions: DashMap, pub(crate) connections: DashMap>, + session_transitions: DashMap, } #[derive(Clone)] @@ -39,8 +43,39 @@ pub(crate) struct AcpSessionState { pub(crate) cwd: String, pub(crate) mode_id: String, pub(crate) model_id: String, - #[allow(dead_code)] pub(crate) mcp_server_ids: Vec, + pub(crate) lifecycle: Arc>, +} + +pub(crate) struct AcpSessionTransition<'a> { + session_id: String, + transitions: &'a DashMap, +} + +impl<'a> AcpSessionTransition<'a> { + fn claim(transitions: &'a DashMap, session_id: &str) -> Result { + match transitions.entry(session_id.to_string()) { + Entry::Vacant(entry) => { + entry.insert(()); + Ok(Self { + session_id: session_id.to_string(), + transitions, + }) + } + Entry::Occupied(_) => Err(Error::internal_error().data(serde_json::json!({ + "state": "session_transition_in_progress", + "sessionId": session_id, + "retryable": true, + "recoveryAction": "Wait for the active session open or close request to finish, then retry the same request" + }))), + } + } +} + +impl Drop for AcpSessionTransition<'_> { + fn drop(&mut self) { + self.transitions.remove(&self.session_id); + } } impl BitfunAcpRuntime { @@ -50,6 +85,7 @@ impl BitfunAcpRuntime { compatibility, sessions: DashMap::new(), connections: DashMap::new(), + session_transitions: DashMap::new(), } } @@ -66,8 +102,97 @@ impl BitfunAcpRuntime { Error::internal_error().data(serde_json::json!(error.to_string())) } + pub(crate) fn cleanup_required_error( + session_id: &str, + stage: &str, + cleanup_kinds: &[&str], + core_session_created: bool, + recovery_action: &str, + ) -> Error { + Error::internal_error().data(serde_json::json!({ + "message": "ACP session lifecycle operation failed and automatic cleanup did not complete", + "sessionId": session_id, + "stage": stage, + "cleanupRequired": true, + "cleanupKinds": cleanup_kinds, + "coreSessionCreated": core_session_created, + "recoveryAction": recovery_action + })) + } + + pub(crate) fn session_close_incomplete_error( + session_id: &str, + stage: &str, + cause: impl std::fmt::Display, + cleanup_kinds: &[&str], + ) -> Error { + Error::internal_error().data(serde_json::json!({ + "message": "ACP session close did not complete", + "sessionId": session_id, + "state": "session_close_incomplete", + "stage": stage, + "cause": cause.to_string(), + "retryable": true, + "acpSessionRetained": true, + "persistedSessionPreserved": true, + "cleanupRequired": !cleanup_kinds.is_empty(), + "cleanupKinds": cleanup_kinds, + "recoveryAction": "Retry session/close for the same sessionId; restart the ACP process only if the same failure continues" + })) + } + pub(crate) fn runtime_error(error: RuntimeError) -> Error { - Self::internal_error(error.into_message()) + match error { + RuntimeError::Port(error) => match error.kind { + PortErrorKind::InvalidRequest => Error::invalid_params().data(error.message), + PortErrorKind::NotFound => Error::resource_not_found(None), + _ => Self::internal_error(error.message), + }, + other => Self::internal_error(other.into_message()), + } + } + + pub(crate) fn session_runtime_error(session_id: &str, error: RuntimeError) -> Error { + match error { + RuntimeError::Port(error) if error.kind == PortErrorKind::NotFound => { + Error::resource_not_found(Some(session_id.to_string())) + } + other => Self::runtime_error(other), + } + } + + pub(crate) fn session_core_error(session_id: &str, error: BitFunError) -> Error { + match error { + BitFunError::NotFound(_) => Error::resource_not_found(Some(session_id.to_string())), + BitFunError::Validation(message) => Error::invalid_params().data(message), + other => Self::internal_error(other), + } + } + + pub(crate) async fn lock_active_session( + &self, + session_id: &str, + ) -> Result<(AcpSessionState, tokio::sync::OwnedMutexGuard<()>)> { + let session = self + .sessions + .get(session_id) + .ok_or_else(|| Error::resource_not_found(Some(session_id.to_string())))? + .clone(); + let lifecycle_guard = session.lifecycle.clone().lock_owned().await; + let active_session = self + .sessions + .get(session_id) + .filter(|active| Arc::ptr_eq(&active.lifecycle, &session.lifecycle)) + .ok_or_else(|| Error::resource_not_found(Some(session_id.to_string())))? + .clone(); + Ok((active_session, lifecycle_guard)) + } + + pub(crate) fn claim_session_transition( + &self, + session_id: &str, + ) -> Result> { + AcpSessionTransition::claim(&self.session_transitions, session_id) } } @@ -83,7 +208,9 @@ impl AcpRuntime for BitfunAcpRuntime { ) .mcp_capabilities(McpCapabilities::new().http(true)) .session_capabilities( - SessionCapabilities::new().list(SessionListCapabilities::new()), + SessionCapabilities::new() + .list(SessionListCapabilities::new()) + .close(SessionCloseCapabilities::new()), ), ) .agent_info( @@ -119,6 +246,10 @@ impl AcpRuntime for BitfunAcpRuntime { self.cancel_prompt(notification).await } + async fn close_session(&self, request: CloseSessionRequest) -> Result { + self.close_active_session(request).await + } + async fn set_session_mode( &self, request: SetSessionModeRequest, @@ -140,3 +271,107 @@ impl AcpRuntime for BitfunAcpRuntime { self.update_session_model(request).await } } + +#[cfg(test)] +mod tests { + use agent_client_protocol::schema::ErrorCode; + use bitfun_agent_runtime::sdk::{PortError, PortErrorKind, RuntimeError}; + use dashmap::DashMap; + + use super::{AcpSessionTransition, BitfunAcpRuntime}; + + #[test] + fn invalid_runtime_request_remains_invalid_params_at_the_protocol_boundary() { + let error = BitfunAcpRuntime::runtime_error(RuntimeError::Port(PortError::new( + PortErrorKind::InvalidRequest, + "unknown session mode", + ))); + + assert_eq!(error.code, ErrorCode::InvalidParams); + assert_eq!(error.data, Some(serde_json::json!("unknown session mode"))); + } + + #[test] + fn missing_runtime_session_remains_resource_not_found_at_the_protocol_boundary() { + let error = BitfunAcpRuntime::runtime_error(RuntimeError::Port(PortError::new( + PortErrorKind::NotFound, + "Session not found: session-404", + ))); + + assert_eq!(error.code, ErrorCode::ResourceNotFound); + assert_eq!(error.data, None); + } + + #[test] + fn session_runtime_error_uses_the_requested_resource_id() { + let error = BitfunAcpRuntime::session_runtime_error( + "session-404", + RuntimeError::Port(PortError::new( + PortErrorKind::NotFound, + "Session not found: session-404", + )), + ); + + assert_eq!(error.code, ErrorCode::ResourceNotFound); + assert_eq!( + error.data, + Some(serde_json::json!({ "uri": "session-404" })) + ); + } + + #[test] + fn backend_runtime_failure_stays_internal_at_the_protocol_boundary() { + let error = BitfunAcpRuntime::runtime_error(RuntimeError::Port(PortError::new( + PortErrorKind::Backend, + "storage unavailable", + ))); + + assert_eq!(error.code, ErrorCode::InternalError); + assert_eq!(error.data, Some(serde_json::json!("storage unavailable"))); + } + + #[test] + fn session_transition_claim_rejects_overlap_and_releases_on_drop() { + let transitions = DashMap::new(); + let first = AcpSessionTransition::claim(&transitions, "session-1") + .expect("first transition should claim the session"); + + let overlapping = match AcpSessionTransition::claim(&transitions, "session-1") { + Ok(_) => panic!("overlapping load or close must fail before side effects"), + Err(error) => error, + }; + assert_eq!(overlapping.code, ErrorCode::InternalError); + assert_eq!( + overlapping.data.as_ref().and_then(|data| data.get("state")), + Some(&serde_json::json!("session_transition_in_progress")) + ); + assert_eq!( + overlapping + .data + .as_ref() + .and_then(|data| data.get("retryable")), + Some(&serde_json::json!(true)) + ); + + drop(first); + AcpSessionTransition::claim(&transitions, "session-1") + .expect("the claim should be reusable after the transition ends"); + } + + #[test] + fn incomplete_close_error_exposes_retry_and_ownership_state() { + let error = BitfunAcpRuntime::session_close_incomplete_error( + "session-1", + "ephemeral MCP cleanup", + "server stop failed", + &["ephemeralMcp"], + ); + let data = error.data.expect("close error should carry recovery data"); + + assert_eq!(data["state"], "session_close_incomplete"); + assert_eq!(data["retryable"], true); + assert_eq!(data["acpSessionRetained"], true); + assert_eq!(data["persistedSessionPreserved"], true); + assert_eq!(data["cleanupKinds"], serde_json::json!(["ephemeralMcp"])); + } +} diff --git a/src/crates/interfaces/acp/src/runtime/mcp.rs b/src/crates/interfaces/acp/src/runtime/mcp.rs index a69810b63..1d9df56bf 100644 --- a/src/crates/interfaces/acp/src/runtime/mcp.rs +++ b/src/crates/interfaces/acp/src/runtime/mcp.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use agent_client_protocol::schema::{McpServer, McpServerSse, McpServerStdio}; @@ -8,38 +8,104 @@ use bitfun_core::service::mcp::{ get_global_mcp_service, set_global_mcp_service, ConfigLocation, MCPServerConfig, MCPServerManager, MCPServerTransport, MCPServerType, MCPService, }; +use sha2::{Digest, Sha256}; use super::BitfunAcpRuntime; impl BitfunAcpRuntime { + pub(super) fn validate_mcp_servers(&self, servers: &[McpServer]) -> Result<()> { + let configs = acp_mcp_server_configs("validation", servers.iter().cloned())?; + ensure_unique_server_ids(&configs) + } + pub(super) async fn provision_mcp_servers( &self, acp_session_id: &str, servers: Vec, + cleanup_recovery_action: &'static str, ) -> Result> { if servers.is_empty() { return Ok(Vec::new()); } let manager = mcp_server_manager().await?; - let mut server_ids: Vec = Vec::with_capacity(servers.len()); + let configs = acp_mcp_server_configs(acp_session_id, servers)?; + ensure_unique_server_ids(&configs)?; + let mut server_ids: Vec = Vec::with_capacity(configs.len()); - for server in servers { - let config = acp_mcp_server_config(acp_session_id, server)?; + for config in configs { let server_id = config.id.clone(); + // Claim cleanup responsibility before startup. Registration can + // succeed before handshake/start later fails, so the current ID + // must participate in compensation as well as earlier servers. + server_ids.push(server_id.clone()); if let Err(error) = manager.add_ephemeral_server(config).await { - for provisioned_id in &server_ids { - let _ = manager.remove_ephemeral_server(provisioned_id).await; + if let Err(cleanup_error) = self.release_mcp_servers(&server_ids).await { + log::warn!( + "Failed to clean up ACP MCP servers after provisioning error: session_id={}, error={}", + acp_session_id, + cleanup_error + ); + return Err(Self::cleanup_required_error( + acp_session_id, + "MCP provisioning", + &["ephemeralMcp"], + false, + cleanup_recovery_action, + )); } return Err(Self::internal_error(error)); } - - server_ids.push(server_id); } Ok(server_ids) } + + pub(super) async fn release_mcp_servers(&self, server_ids: &[String]) -> Result<()> { + if server_ids.is_empty() { + return Ok(()); + } + + let manager = mcp_server_manager().await?; + let mut first_error = None; + for server_id in server_ids { + if let Err(error) = manager.remove_ephemeral_server(server_id).await { + log::warn!( + "Failed to remove ephemeral ACP MCP server: server_id={}, error={}", + server_id, + error + ); + if first_error.is_none() { + first_error = Some(error); + } + } + } + + match first_error { + Some(error) => Err(Self::internal_error(error)), + None => Ok(()), + } + } +} + +fn acp_mcp_server_configs( + acp_session_id: &str, + servers: impl IntoIterator, +) -> Result> { + servers + .into_iter() + .map(|server| acp_mcp_server_config(acp_session_id, server)) + .collect() +} + +fn ensure_unique_server_ids(configs: &[MCPServerConfig]) -> Result<()> { + let mut ids = HashSet::with_capacity(configs.len()); + if configs.iter().all(|config| ids.insert(config.id.clone())) { + Ok(()) + } else { + Err(Error::invalid_params().data("MCP server names must be unique within a session")) + } } async fn mcp_server_manager() -> Result> { @@ -151,23 +217,27 @@ fn clean_server_name(name: &str) -> Result { } fn ephemeral_server_id(acp_session_id: &str, server_name: &str) -> String { - format!( - "acp-{}-{}", - sanitize_id_part(acp_session_id), - sanitize_id_part(server_name) - ) + let mut digest = Sha256::new(); + for value in [acp_session_id, server_name] { + digest.update((value.len() as u64).to_le_bytes()); + digest.update(value.as_bytes()); + } + format!("acp-{:x}", digest.finalize()) } -fn sanitize_id_part(value: &str) -> String { - let sanitized = value - .chars() - .map(|ch| { - if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' { - ch - } else { - '-' - } - }) - .collect::(); - sanitized.trim_matches('-').to_string() +#[cfg(test)] +mod tests { + use super::ephemeral_server_id; + + #[test] + fn ephemeral_server_ids_do_not_collapse_distinct_session_or_server_names() { + assert_ne!( + ephemeral_server_id("foo:bar", "tools"), + ephemeral_server_id("foo bar", "tools") + ); + assert_ne!( + ephemeral_server_id("session", "tools:read"), + ephemeral_server_id("session", "tools read") + ); + } } diff --git a/src/crates/interfaces/acp/src/runtime/model.rs b/src/crates/interfaces/acp/src/runtime/model.rs index 4776cb88a..da9657001 100644 --- a/src/crates/interfaces/acp/src/runtime/model.rs +++ b/src/crates/interfaces/acp/src/runtime/model.rs @@ -22,7 +22,9 @@ impl BitfunAcpRuntime { ) -> Result { let session_id = request.session_id.to_string(); let model_id = request.model_id.to_string(); - self.set_session_model_id(&session_id, &model_id).await?; + let (session, _lifecycle_guard) = self.lock_active_session(&session_id).await?; + self.set_session_model_id_for_active(&session, &model_id) + .await?; Ok(SetSessionModelResponse::new()) } @@ -37,13 +39,16 @@ impl BitfunAcpRuntime { .as_value_id() .ok_or_else(|| Error::invalid_params().data("config option value must be a string"))? .to_string(); + let (session, _lifecycle_guard) = self.lock_active_session(&session_id).await?; match config_id.as_str() { MODEL_CONFIG_ID => { - self.set_session_model_id(&session_id, &value).await?; + self.set_session_model_id_for_active(&session, &value) + .await?; } MODE_CONFIG_ID => { - self.update_session_mode_inner(&session_id, &value).await?; + self.update_session_mode_for_active(&session, &value) + .await?; } _ => { return Err(Error::invalid_params() @@ -64,27 +69,26 @@ impl BitfunAcpRuntime { )) } - async fn set_session_model_id(&self, session_id: &str, model_id: &str) -> Result<()> { - let acp_session = self - .sessions - .get(session_id) - .ok_or_else(|| Error::resource_not_found(Some(session_id.to_string())))?; - let bitfun_session_id = acp_session.bitfun_session_id.clone(); - drop(acp_session); - + async fn set_session_model_id_for_active( + &self, + session: &super::AcpSessionState, + model_id: &str, + ) -> Result<()> { let normalized_model_id = normalize_model_selection(model_id).await?; self.agent_runtime .update_session_model(AgentSessionModelUpdateRequest { - session_id: bitfun_session_id, + session_id: session.bitfun_session_id.clone(), model_id: normalized_model_id.clone(), }) .await - .map_err(|error| Self::internal_error(error.into_message()))?; + .map_err(|error| Self::session_runtime_error(&session.acp_session_id, error))?; - if let Some(mut state) = self.sessions.get_mut(session_id) { - state.model_id = normalized_model_id; - } + let mut state = self + .sessions + .get_mut(&session.acp_session_id) + .ok_or_else(|| Error::resource_not_found(Some(session.acp_session_id.clone())))?; + state.model_id = normalized_model_id; Ok(()) } diff --git a/src/crates/interfaces/acp/src/runtime/prompt.rs b/src/crates/interfaces/acp/src/runtime/prompt.rs index 680ab4b5d..52bdb3c60 100644 --- a/src/crates/interfaces/acp/src/runtime/prompt.rs +++ b/src/crates/interfaces/acp/src/runtime/prompt.rs @@ -26,11 +26,7 @@ use super::{AcpSessionState, BitfunAcpRuntime}; impl BitfunAcpRuntime { pub(super) async fn run_prompt(&self, request: PromptRequest) -> Result { let session_id = request.session_id.to_string(); - let acp_session = self - .sessions - .get(&session_id) - .ok_or_else(|| Error::resource_not_found(Some(session_id.clone())))?; - let acp_session = acp_session.clone(); + let (acp_session, lifecycle_guard) = self.lock_active_session(&session_id).await?; let connection = self .connections .get(&session_id) @@ -46,12 +42,12 @@ impl BitfunAcpRuntime { let mut event_rx = self .agent_runtime .subscribe_session_events(&acp_session.bitfun_session_id) - .map_err(Self::runtime_error)?; + .map_err(|error| Self::session_runtime_error(&session_id, error))?; let outcome = self .agent_runtime .submit_dialog_turn(dialog_turn_request(&acp_session, parsed_prompt)) .await - .map_err(Self::runtime_error)?; + .map_err(|error| Self::session_runtime_error(&session_id, error))?; let turn_id = match resolve_started_prompt_turn(outcome) { Ok(turn_id) => turn_id, Err(queued_turn_id) => { @@ -62,11 +58,12 @@ impl BitfunAcpRuntime { "acp_busy_rejected", )) .await - .map_err(Self::runtime_error)?; + .map_err(|error| Self::session_runtime_error(&session_id, error))?; return Err(Error::internal_error() .data("Session state does not allow starting new dialog: Processing")); } }; + drop(lifecycle_guard); let stop_reason = wait_for_prompt_completion( self, @@ -83,11 +80,7 @@ impl BitfunAcpRuntime { pub(super) async fn cancel_prompt(&self, notification: CancelNotification) -> Result<()> { let session_id = notification.session_id.to_string(); - let acp_session = self - .sessions - .get(&session_id) - .ok_or_else(|| Error::resource_not_found(Some(session_id.clone())))?; - let acp_session = acp_session.clone(); + let (acp_session, _lifecycle_guard) = self.lock_active_session(&session_id).await?; self.agent_runtime .cancel_turn(turn_cancellation_request( @@ -96,7 +89,7 @@ impl BitfunAcpRuntime { "acp_client_cancelled", )) .await - .map_err(Self::runtime_error)?; + .map_err(|error| Self::session_runtime_error(&session_id, error))?; Ok(()) } @@ -129,7 +122,7 @@ fn resolve_started_prompt_turn( } } -fn turn_cancellation_request( +pub(super) fn turn_cancellation_request( session_id: &str, turn_id: Option<&str>, reason: &str, @@ -420,6 +413,7 @@ mod tests { mode_id: "agentic".to_string(), model_id: "auto".to_string(), mcp_server_ids: Vec::new(), + lifecycle: std::sync::Arc::new(tokio::sync::Mutex::new(())), } } diff --git a/src/crates/interfaces/acp/src/runtime/session.rs b/src/crates/interfaces/acp/src/runtime/session.rs index 3f181efe0..310f20522 100644 --- a/src/crates/interfaces/acp/src/runtime/session.rs +++ b/src/crates/interfaces/acp/src/runtime/session.rs @@ -1,19 +1,19 @@ use std::path::Path; use agent_client_protocol::schema::{ - CurrentModeUpdate, ListSessionsRequest, ListSessionsResponse, LoadSessionRequest, - LoadSessionResponse, NewSessionRequest, NewSessionResponse, SessionId, SessionInfo, - SessionMode, SessionModeState, SessionUpdate, SetSessionModeRequest, SetSessionModeResponse, + CloseSessionRequest, CloseSessionResponse, ListSessionsRequest, ListSessionsResponse, + LoadSessionRequest, LoadSessionResponse, NewSessionRequest, NewSessionResponse, SessionId, + SessionInfo, SessionMode, SessionModeState, SetSessionModeRequest, SetSessionModeResponse, }; use agent_client_protocol::{Client, ConnectionTo, Error, Result}; use bitfun_agent_runtime::sdk::{ - AgentSessionCreateRequest, AgentSessionListRequest, AgentSessionModeUpdateRequest, - SessionStoragePathRequest, + AgentSessionCreateRequest, AgentSessionDeleteRequest, AgentSessionListRequest, + AgentSessionModeUpdateRequest, SessionStoragePathRequest, }; use bitfun_core::agentic::agents::get_agent_registry; use chrono::{DateTime, Utc}; +use dashmap::mapref::entry::Entry; -use super::events::send_update; use super::model::{ build_session_config_options, build_session_model_state, normalize_session_model_id, }; @@ -21,6 +21,15 @@ use super::replay::replay_session_history; use super::{AcpSessionState, BitfunAcpRuntime}; impl BitfunAcpRuntime { + fn validate_session_target(session_id: &str, cwd: &Path) -> Result<()> { + bitfun_core_types::validate_session_id(session_id) + .map_err(|message| Error::invalid_params().data(message))?; + if !cwd.is_absolute() { + return Err(Error::invalid_params().data("cwd must be an absolute path")); + } + Ok(()) + } + pub(super) async fn create_session( &self, request: NewSessionRequest, @@ -28,48 +37,110 @@ impl BitfunAcpRuntime { ) -> Result { let cwd = request.cwd.to_string_lossy().to_string(); let mcp_servers = request.mcp_servers; - let session = self + self.validate_mcp_servers(&mcp_servers)?; + let modes = build_session_modes(Some("agentic")).await; + let models = build_session_model_state(None).await?; + let config_options = build_session_config_options(None, Some("agentic")).await?; + let session_id = uuid::Uuid::new_v4().to_string(); + Self::validate_session_target(&session_id, Path::new(&cwd))?; + let _session_transition = self.claim_session_transition(&session_id)?; + let mcp_server_ids = self + .provision_mcp_servers( + &session_id, + mcp_servers, + "Restart the ACP process before retrying session/new; no persisted Core session was created", + ) + .await?; + let create_request = AgentSessionCreateRequest { + session_name: format!( + "ACP Session - {}", + chrono::Local::now().format("%Y-%m-%d %H:%M:%S") + ), + agent_type: "agentic".to_string(), + workspace_path: Some(cwd.clone()), + remote_connection_id: None, + remote_ssh_host: None, + metadata: serde_json::Map::new(), + }; + let session = match self .agent_runtime - .create_session(AgentSessionCreateRequest { - session_name: format!( - "ACP Session - {}", - chrono::Local::now().format("%Y-%m-%d %H:%M:%S") - ), - agent_type: "agentic".to_string(), - workspace_path: Some(cwd.clone()), - remote_connection_id: None, - remote_ssh_host: None, - metadata: serde_json::Map::new(), - }) + .create_session_with_id(session_id.clone(), create_request) .await - .map_err(Self::runtime_error)?; - + { + Ok(session) => session, + Err(error) => { + let core_cleanup_required = matches!( + &error, + bitfun_agent_runtime::sdk::RuntimeError::Port(port_error) + if port_error.kind == bitfun_agent_runtime::sdk::PortErrorKind::CleanupRequired + ); + let mcp_cleaned = self.release_mcp_servers(&mcp_server_ids).await.is_ok(); + let core_cleaned = if core_cleanup_required { + self.delete_failed_new_core_session( + &session_id, + &cwd, + "Core session creation rollback", + ) + .await + } else { + true + }; + if !mcp_cleaned || !core_cleaned { + let mut cleanup_kinds = Vec::with_capacity(2); + if !mcp_cleaned { + cleanup_kinds.push("ephemeralMcp"); + } + if !core_cleaned { + cleanup_kinds.push("coreSession"); + } + return Err(Self::cleanup_required_error( + &session_id, + "Core session creation", + &cleanup_kinds, + core_cleanup_required, + "Restart the ACP process, inspect session/list for the returned sessionId, and remove only that failed session through a supported session manager before retrying", + )); + } + return Err(Self::runtime_error(error)); + } + }; let acp_session = AcpSessionState { acp_session_id: session.session_id.clone(), bitfun_session_id: session.session_id.clone(), cwd, mode_id: session.agent_type.clone(), model_id: normalize_session_model_id(None), - mcp_server_ids: self - .provision_mcp_servers(&session.session_id, mcp_servers) - .await?, + mcp_server_ids, + lifecycle: std::sync::Arc::new(tokio::sync::Mutex::new(())), }; - self.sessions - .insert(acp_session.acp_session_id.clone(), acp_session.clone()); - self.connections - .insert(acp_session.acp_session_id.clone(), connection); - - let modes = build_session_modes(Some(session.agent_type.as_str())).await; - let models = build_session_model_state(Some(&acp_session.model_id)).await?; - let config_options = - build_session_config_options(Some(&acp_session.model_id), Some(&acp_session.mode_id)) - .await?; - Ok( - NewSessionResponse::new(SessionId::new(acp_session.acp_session_id)) - .modes(modes) - .models(models) - .config_options(config_options), - ) + let response = NewSessionResponse::new(SessionId::new(acp_session.acp_session_id.clone())) + .modes(modes) + .models(models) + .config_options(config_options); + + if let Err(error) = self.publish_active_session(&acp_session, connection).await { + let (mcp_cleaned, core_cleaned) = self + .cleanup_failed_new_session_setup(&acp_session, "session publication") + .await; + if !mcp_cleaned || !core_cleaned { + let mut cleanup_kinds = Vec::with_capacity(2); + if !mcp_cleaned { + cleanup_kinds.push("ephemeralMcp"); + } + if !core_cleaned { + cleanup_kinds.push("coreSession"); + } + return Err(Self::cleanup_required_error( + &acp_session.acp_session_id, + "session publication", + &cleanup_kinds, + true, + "Restart the ACP process, inspect session/list for the returned sessionId, and remove only that newly created session through a supported session manager before retrying", + )); + } + return Err(error); + } + Ok(response) } pub(super) async fn restore_session( @@ -80,10 +151,23 @@ impl BitfunAcpRuntime { let cwd = request.cwd.to_string_lossy().to_string(); let session_id = request.session_id.to_string(); let mcp_servers = request.mcp_servers; + self.validate_mcp_servers(&mcp_servers)?; + Self::validate_session_target(&session_id, Path::new(&cwd))?; + let _session_transition = self.claim_session_transition(&session_id)?; + if self.sessions.contains_key(&session_id) { + return Err(Error::invalid_params().data("session is already active")); + } + let mcp_server_ids = self + .provision_mcp_servers( + &session_id, + mcp_servers, + "Restart the ACP process, then retry session/load with the same sessionId and cwd; preserve the existing persisted Core session", + ) + .await?; // ACP history replay and model selection must come from one persisted // snapshot. Keep this compatibility path until the runtime contract can // return the rich turn data ACP actually projects. - let (session, turns) = self + let restore = self .compatibility .restore_session_with_turns_for_workspace( SessionStoragePathRequest { @@ -94,46 +178,251 @@ impl BitfunAcpRuntime { &session_id, false, ) - .await - .map_err(Self::internal_error)?; - + .await; + let (session, turns) = match restore { + Ok(restored) => restored, + Err(error) => { + if self.release_mcp_servers(&mcp_server_ids).await.is_err() { + return Err(Self::cleanup_required_error( + &session_id, + "Core session restore", + &["ephemeralMcp"], + false, + "Restart the ACP process, then retry session/load with the same sessionId and cwd; preserve the existing persisted Core session", + )); + } + return Err(Self::session_core_error(&session_id, error)); + } + }; let acp_session = AcpSessionState { acp_session_id: session.session_id.clone(), bitfun_session_id: session.session_id.clone(), cwd, mode_id: session.agent_type.clone(), model_id: normalize_session_model_id(session.config.model_id.as_deref()), - mcp_server_ids: self - .provision_mcp_servers(&session.session_id, mcp_servers) - .await?, + mcp_server_ids, + lifecycle: std::sync::Arc::new(tokio::sync::Mutex::new(())), + }; + let response = async { + let modes = build_session_modes(Some(session.agent_type.as_str())).await; + let models = build_session_model_state(Some(&acp_session.model_id)).await?; + let config_options = build_session_config_options( + Some(&acp_session.model_id), + Some(&acp_session.mode_id), + ) + .await?; + Ok(LoadSessionResponse::new() + .modes(modes) + .models(models) + .config_options(config_options)) + } + .await; + let response = match response { + Ok(response) => response, + Err(error) => { + if let Some(cleanup_error) = self + .cleanup_failed_loaded_session_setup(&acp_session, "load session response") + .await + { + return Err(cleanup_error); + } + return Err(error); + } }; - self.sessions - .insert(acp_session.acp_session_id.clone(), acp_session.clone()); - self.connections - .insert(acp_session.acp_session_id.clone(), connection.clone()); - // Replay the restored transcript so the client can rebuild its UI. - // Failures to send individual notifications shouldn't abort the load; - // the runtime context is already restored and the client can still - // issue prompts against the session. if let Err(error) = replay_session_history(&connection, &acp_session.acp_session_id, &turns) { + if let Some(cleanup_error) = self + .cleanup_failed_loaded_session_setup(&acp_session, "history replay") + .await + { + return Err(cleanup_error); + } + return Err(error); + } + + if let Err(error) = self.publish_active_session(&acp_session, connection).await { + if let Some(cleanup_error) = self + .cleanup_failed_loaded_session_setup(&acp_session, "session publication") + .await + { + return Err(cleanup_error); + } + return Err(error); + } + Ok(response) + } + + async fn publish_active_session( + &self, + session: &AcpSessionState, + connection: ConnectionTo, + ) -> Result<()> { + let _lifecycle_guard = session.lifecycle.lock().await; + match self.sessions.entry(session.acp_session_id.clone()) { + Entry::Vacant(entry) => { + entry.insert(session.clone()); + } + Entry::Occupied(_) => { + return Err(Error::invalid_params().data("session is already active")); + } + } + self.connections + .insert(session.acp_session_id.clone(), connection); + Ok(()) + } + + async fn cleanup_failed_session_setup(&self, session: &AcpSessionState, stage: &str) -> bool { + if let Err(error) = self.release_mcp_servers(&session.mcp_server_ids).await { log::warn!( - "ACP session history replay interrupted: session_id={}, error={}", - acp_session.acp_session_id, + "Failed to clean up ACP session setup: session_id={}, stage={}, error={}", + session.acp_session_id, + stage, error ); + return false; } + true + } - let modes = build_session_modes(Some(session.agent_type.as_str())).await; - let models = build_session_model_state(Some(&acp_session.model_id)).await?; - let config_options = - build_session_config_options(Some(&acp_session.model_id), Some(&acp_session.mode_id)) - .await?; - Ok(LoadSessionResponse::new() - .modes(modes) - .models(models) - .config_options(config_options)) + async fn cleanup_failed_loaded_session_setup( + &self, + session: &AcpSessionState, + stage: &str, + ) -> Option { + let mcp_cleaned = self.cleanup_failed_session_setup(session, stage).await; + let core_unloaded = match self + .compatibility + .unload_persisted_session(&session.bitfun_session_id) + .await + { + Ok(_) => true, + Err(error) => { + log::warn!( + "Failed to unload Core session after ACP load error: session_id={}, stage={}, error={}", + session.bitfun_session_id, + stage, + error + ); + false + } + }; + if mcp_cleaned && core_unloaded { + return None; + } + + let mut cleanup_kinds = Vec::with_capacity(2); + if !mcp_cleaned { + cleanup_kinds.push("ephemeralMcp"); + } + if !core_unloaded { + cleanup_kinds.push("coreRuntime"); + } + Some(Self::cleanup_required_error( + &session.acp_session_id, + stage, + &cleanup_kinds, + false, + "Restart the ACP process, then retry session/load with the same sessionId and cwd; preserve the existing persisted Core session", + )) + } + + async fn cleanup_failed_new_session_setup( + &self, + session: &AcpSessionState, + stage: &str, + ) -> (bool, bool) { + let mcp_cleaned = self.cleanup_failed_session_setup(session, stage).await; + let core_cleaned = self + .delete_failed_new_core_session(&session.bitfun_session_id, &session.cwd, stage) + .await; + (mcp_cleaned, core_cleaned) + } + + async fn delete_failed_new_core_session( + &self, + session_id: &str, + cwd: &str, + stage: &str, + ) -> bool { + if let Err(error) = self + .agent_runtime + .delete_session(AgentSessionDeleteRequest { + workspace_path: cwd.to_string(), + session_id: session_id.to_string(), + remote_connection_id: None, + remote_ssh_host: None, + }) + .await + { + log::warn!( + "Failed to delete Core session after ACP setup error: session_id={}, stage={}, error={}", + session_id, + stage, + error + ); + return false; + } + true + } + + pub(super) async fn close_active_session( + &self, + request: CloseSessionRequest, + ) -> Result { + let session_id = request.session_id.to_string(); + let _session_transition = self.claim_session_transition(&session_id)?; + let (active_session, _lifecycle_guard) = self.lock_active_session(&session_id).await?; + let storage_path = self + .compatibility + .resolve_persisted_session_storage_path(SessionStoragePathRequest { + workspace_path: Path::new(&active_session.cwd).to_path_buf(), + remote_connection_id: None, + remote_ssh_host: None, + }) + .await + .map_err(|error| { + Self::session_close_incomplete_error( + &session_id, + "storage path resolution", + error, + &[], + ) + })?; + let _maintenance = self + .compatibility + .begin_session_maintenance(&storage_path, &active_session.bitfun_session_id, 5_000) + .await + .map_err(|error| { + Self::session_close_incomplete_error(&session_id, "active work drain", error, &[]) + })?; + self.compatibility + .unload_persisted_session(&active_session.bitfun_session_id) + .await + .map_err(|error| { + Self::session_close_incomplete_error(&session_id, "Core runtime unload", error, &[]) + })?; + + if let Err(error) = self + .release_mcp_servers(&active_session.mcp_server_ids) + .await + { + log::warn!( + "Failed to release ACP MCP servers after Core session close; retaining ACP ownership for retry: session_id={}, error={}", + session_id, + error + ); + return Err(Self::session_close_incomplete_error( + &session_id, + "ephemeral MCP cleanup", + error, + &["ephemeralMcp"], + )); + } + + self.sessions.remove(&session_id); + self.connections.remove(&session_id); + Ok(CloseSessionResponse::new()) } pub(super) async fn list_sessions_for_cwd( @@ -198,46 +487,35 @@ impl BitfunAcpRuntime { &self, request: SetSessionModeRequest, ) -> Result { + let session_id = request.session_id.to_string(); let mode_id = request.mode_id.to_string(); - self.update_session_mode_inner(&request.session_id.to_string(), &mode_id) + let (session, _lifecycle_guard) = self.lock_active_session(&session_id).await?; + self.update_session_mode_for_active(&session, &mode_id) .await?; Ok(SetSessionModeResponse::new()) } - pub(super) async fn update_session_mode_inner( + pub(super) async fn update_session_mode_for_active( &self, - session_id: &str, + session: &AcpSessionState, mode_id: &str, ) -> Result<()> { - let acp_session = self - .sessions - .get(session_id) - .ok_or_else(|| Error::resource_not_found(Some(session_id.to_string())))?; - let bitfun_session_id = acp_session.bitfun_session_id.clone(); - drop(acp_session); - - validate_mode_id(mode_id).await?; - + let mode_id = mode_id.trim(); self.agent_runtime .update_session_mode(AgentSessionModeUpdateRequest { - session_id: bitfun_session_id, + session_id: session.bitfun_session_id.clone(), mode_id: mode_id.to_string(), }) .await - .map_err(Self::runtime_error)?; + .map_err(|error| Self::session_runtime_error(&session.acp_session_id, error))?; - if let Some(mut state) = self.sessions.get_mut(session_id) { - state.mode_id = mode_id.to_string(); - } - - if let Some(connection) = self.connections.get(session_id) { - send_update( - &connection, - session_id, - SessionUpdate::CurrentModeUpdate(CurrentModeUpdate::new(mode_id.to_string())), - )?; - } + let mut state = self + .sessions + .get_mut(&session.acp_session_id) + .ok_or_else(|| Error::resource_not_found(Some(session.acp_session_id.clone())))?; + state.mode_id = mode_id.to_string(); + drop(state); Ok(()) } @@ -270,20 +548,6 @@ async fn build_session_modes(preferred_mode_id: Option<&str>) -> SessionModeStat SessionModeState::new(current_mode_id, available_modes) } -async fn validate_mode_id(mode_id: &str) -> Result<()> { - let mode_exists = get_agent_registry() - .get_modes_info() - .await - .into_iter() - .any(|info| info.id == mode_id); - - if mode_exists { - Ok(()) - } else { - Err(Error::invalid_params().data(format!("unknown session mode: {}", mode_id))) - } -} - fn unix_ms_to_rfc3339(time_ms: u64) -> String { let time_ms = i64::try_from(time_ms).unwrap_or(i64::MAX); DateTime::::from_timestamp_millis(time_ms) @@ -293,7 +557,7 @@ fn unix_ms_to_rfc3339(time_ms: u64) -> String { #[cfg(test)] mod tests { - use super::{unix_ms_to_rfc3339, validate_mode_id}; + use super::unix_ms_to_rfc3339; #[test] fn session_timestamps_remain_rfc3339_after_runtime_projection() { @@ -303,9 +567,4 @@ mod tests { "2023-11-14T22:13:20+00:00" ); } - - #[tokio::test] - async fn unknown_session_mode_remains_an_invalid_request() { - assert!(validate_mode_id("__missing_acp_mode__").await.is_err()); - } } diff --git a/src/crates/interfaces/acp/src/server.rs b/src/crates/interfaces/acp/src/server.rs index 23a16a28b..b5b8d059a 100644 --- a/src/crates/interfaces/acp/src/server.rs +++ b/src/crates/interfaces/acp/src/server.rs @@ -1,11 +1,12 @@ use std::sync::Arc; use agent_client_protocol::schema::{ - AuthenticateRequest, AuthenticateResponse, CancelNotification, InitializeRequest, - InitializeResponse, ListSessionsRequest, ListSessionsResponse, LoadSessionRequest, - LoadSessionResponse, NewSessionRequest, NewSessionResponse, PromptRequest, PromptResponse, - SetSessionConfigOptionRequest, SetSessionConfigOptionResponse, SetSessionModeRequest, - SetSessionModeResponse, SetSessionModelRequest, SetSessionModelResponse, + AuthenticateRequest, AuthenticateResponse, CancelNotification, CloseSessionRequest, + CloseSessionResponse, InitializeRequest, InitializeResponse, ListSessionsRequest, + ListSessionsResponse, LoadSessionRequest, LoadSessionResponse, NewSessionRequest, + NewSessionResponse, PromptRequest, PromptResponse, SetSessionConfigOptionRequest, + SetSessionConfigOptionResponse, SetSessionModeRequest, SetSessionModeResponse, + SetSessionModelRequest, SetSessionModelResponse, }; use agent_client_protocol::{ Agent, ByteStreams, Client, ConnectTo, ConnectionTo, Dispatch, Error, Result, @@ -42,6 +43,10 @@ pub trait AcpRuntime: Send + Sync + 'static { async fn cancel(&self, notification: CancelNotification) -> Result<()>; + async fn close_session(&self, _request: CloseSessionRequest) -> Result { + Err(Error::method_not_found().data("session/close is not implemented")) + } + async fn set_session_mode( &self, _request: SetSessionModeRequest, @@ -188,6 +193,21 @@ where }, agent_client_protocol::on_receive_notification!(), ) + .on_receive_request( + { + let runtime = runtime.clone(); + async move |request: CloseSessionRequest, + responder, + cx: ConnectionTo| { + let runtime = runtime.clone(); + cx.spawn(async move { + responder.respond_with_result(runtime.close_session(request).await) + })?; + Ok(()) + } + }, + agent_client_protocol::on_receive_request!(), + ) .on_receive_request( { let runtime = runtime.clone(); diff --git a/src/crates/services/services-integrations/src/mcp/server/process.rs b/src/crates/services/services-integrations/src/mcp/server/process.rs index c5b0e0c01..7ca867c91 100644 --- a/src/crates/services/services-integrations/src/mcp/server/process.rs +++ b/src/crates/services/services-integrations/src/mcp/server/process.rs @@ -31,6 +31,8 @@ pub struct MCPServerProcess { last_ping_time: Arc>>, last_error_message: Arc>>, message_rx: Option>, + #[cfg(test)] + fail_next_stop: bool, } impl MCPServerProcess { @@ -51,9 +53,16 @@ impl MCPServerProcess { last_ping_time: Arc::new(RwLock::new(None)), last_error_message: Arc::new(RwLock::new(None)), message_rx: None, + #[cfg(test)] + fail_next_stop: false, } } + #[cfg(test)] + pub(crate) fn fail_next_stop_for_test(&mut self) { + self.fail_next_stop = true; + } + /// Starts the server process. pub async fn start( &mut self, @@ -263,13 +272,24 @@ impl MCPServerProcess { info!("Stopping MCP server: name={} id={}", self.name, self.id); self.set_status(MCPServerStatus::Stopping).await; - if let Some(mut child) = self.child.take() { - if let Err(e) = child.kill().await { - warn!( + #[cfg(test)] + if self.fail_next_stop { + self.fail_next_stop = false; + return Err(MCPRuntimeError::process("Injected MCP stop failure")); + } + + if let Some(child) = self.child.as_mut() { + if let Err(error) = child.kill().await { + let message = format!( "Failed to kill MCP server process: name={} id={} error={}", - self.name, self.id, e + self.name, self.id, error ); + warn!("{}", message); + self.set_status_with_error(MCPServerStatus::Failed, Some(message.clone())) + .await; + return Err(MCPRuntimeError::process(message)); } + self.child = None; } self.connection = None; diff --git a/src/crates/services/services-integrations/src/mcp/server/registry.rs b/src/crates/services/services-integrations/src/mcp/server/registry.rs index 981842641..51a57995c 100644 --- a/src/crates/services/services-integrations/src/mcp/server/registry.rs +++ b/src/crates/services/services-integrations/src/mcp/server/registry.rs @@ -83,14 +83,13 @@ impl MCPServerRegistry { /// Unregisters a server. pub async fn unregister(&self, server_id: &str) -> MCPRuntimeResult<()> { let _lifecycle_guard = self.lifecycle_lock.lock().await; - let process = { - let mut servers = self.servers.write().await; - servers.remove(server_id) - }; + let process = self.servers.read().await.get(server_id).cloned(); if let Some(process) = process { let mut proc = process.write().await; proc.stop().await?; + drop(proc); + self.servers.write().await.remove(server_id); info!("Unregistered MCP server: id={}", server_id); Ok(()) } else { @@ -151,17 +150,32 @@ impl MCPServerRegistry { /// Clears the registry. pub async fn clear(&self) -> MCPRuntimeResult<()> { let _lifecycle_guard = self.lifecycle_lock.lock().await; - let processes = { - let mut servers = self.servers.write().await; - servers - .drain() - .map(|(_, process)| process) - .collect::>() - }; + let processes = self + .servers + .read() + .await + .iter() + .map(|(server_id, process)| (server_id.clone(), process.clone())) + .collect::>(); + let mut first_error = None; - for process in processes { + for (server_id, process) in processes { let mut proc = process.write().await; - let _ = proc.stop().await; + match proc.stop().await { + Ok(()) => { + drop(proc); + self.servers.write().await.remove(&server_id); + } + Err(error) => { + if first_error.is_none() { + first_error = Some(error); + } + } + } + } + + if let Some(error) = first_error { + return Err(error); } info!("Cleared MCP server registry"); @@ -221,6 +235,61 @@ mod tests { assert!(!registry.contains("test").await); } + #[tokio::test] + async fn failed_unregister_retains_process_ownership_for_retry() { + let registry = MCPServerRegistry::new(); + let config = local_config("retryable-stop"); + registry.register(&config).await.unwrap(); + let process = registry + .get_process(&config.id) + .await + .expect("registered process should exist"); + process.write().await.fail_next_stop_for_test(); + + let error = registry + .unregister(&config.id) + .await + .expect_err("injected stop failure must propagate"); + + assert_eq!(error.kind(), MCPRuntimeErrorKind::Process); + assert!(registry.contains(&config.id).await); + registry + .unregister(&config.id) + .await + .expect("retained process should be retryable"); + assert!(!registry.contains(&config.id).await); + } + + #[tokio::test] + async fn failed_clear_retains_only_processes_that_still_need_cleanup() { + let registry = MCPServerRegistry::new(); + let retryable = local_config("retryable-clear"); + let stoppable = local_config("stoppable-clear"); + registry.register(&retryable).await.unwrap(); + registry.register(&stoppable).await.unwrap(); + registry + .get_process(&retryable.id) + .await + .expect("retryable process") + .write() + .await + .fail_next_stop_for_test(); + + let error = registry + .clear() + .await + .expect_err("one failed stop must fail registry clear"); + + assert_eq!(error.kind(), MCPRuntimeErrorKind::Process); + assert!(registry.contains(&retryable.id).await); + assert!(!registry.contains(&stoppable.id).await); + registry + .clear() + .await + .expect("retained process should retry"); + assert!(registry.get_all_server_ids().await.is_empty()); + } + #[tokio::test] async fn registry_reports_validation_and_missing_errors_as_runtime_errors() { let registry = MCPServerRegistry::new(); diff --git a/src/web-ui/src/flow_chat/services/ReviewActionBarPersistenceService.test.ts b/src/web-ui/src/flow_chat/services/ReviewActionBarPersistenceService.test.ts index 92ecd489a..4f875348b 100644 --- a/src/web-ui/src/flow_chat/services/ReviewActionBarPersistenceService.test.ts +++ b/src/web-ui/src/flow_chat/services/ReviewActionBarPersistenceService.test.ts @@ -151,7 +151,7 @@ describe('ReviewActionBarPersistenceService', () => { } as any); expect(sessionAPI.saveSessionMetadata).toHaveBeenCalledTimes(1); - const [metadata, workspacePath] = (sessionAPI.saveSessionMetadata as any).mock.calls[0]; + const [metadata, workspacePath, fields] = (sessionAPI.saveSessionMetadata as any).mock.calls[0]; expect(metadata.sessionId).toBe('session-1'); expect(metadata.sessionName).toBe('Existing Session'); expect(metadata.agentType).toBe('agentic'); @@ -171,6 +171,7 @@ describe('ReviewActionBarPersistenceService', () => { persistedAt: expect.any(Number), }); expect(workspacePath).toBe('/workspace/project'); + expect(fields).toEqual(['reviewActionState']); }); it('builds complete metadata when no existing metadata is available', async () => { @@ -253,7 +254,8 @@ describe('ReviewActionBarPersistenceService', () => { } as any); expect(sessionAPI.saveSessionMetadata).toHaveBeenCalledTimes(1); - const [, , remoteConnectionId, remoteSshHost] = (sessionAPI.saveSessionMetadata as any).mock.calls[0]; + const [, , fields, remoteConnectionId, remoteSshHost] = (sessionAPI.saveSessionMetadata as any).mock.calls[0]; + expect(fields).toEqual(['reviewActionState']); expect(remoteConnectionId).toBe('remote-1'); expect(remoteSshHost).toBe('ssh-host-1'); }); @@ -275,10 +277,11 @@ describe('ReviewActionBarPersistenceService', () => { await clearPersistedReviewState('session-1', '/workspace/project'); expect(sessionAPI.saveSessionMetadata).toHaveBeenCalledTimes(1); - const [metadata] = (sessionAPI.saveSessionMetadata as any).mock.calls[0]; + const [metadata, , fields] = (sessionAPI.saveSessionMetadata as any).mock.calls[0]; expect(metadata.sessionId).toBe('session-1'); expect(metadata.sessionName).toBe('Test Session'); expect(metadata.reviewActionState).toBeUndefined(); + expect(fields).toEqual(['reviewActionState']); }); }); diff --git a/src/web-ui/src/flow_chat/services/ReviewActionBarPersistenceService.ts b/src/web-ui/src/flow_chat/services/ReviewActionBarPersistenceService.ts index e39b9d41b..a946a1db9 100644 --- a/src/web-ui/src/flow_chat/services/ReviewActionBarPersistenceService.ts +++ b/src/web-ui/src/flow_chat/services/ReviewActionBarPersistenceService.ts @@ -81,6 +81,7 @@ export async function persistReviewActionState(state: ReviewActionBarState): Pro await sessionAPI.saveSessionMetadata( metadata, session.workspacePath, + ['reviewActionState'], session.remoteConnectionId, session.remoteSshHost ); @@ -100,7 +101,8 @@ export async function clearPersistedReviewState(sessionId: string, workspacePath await sessionAPI.saveSessionMetadata( metadata, - workspacePath + workspacePath, + ['reviewActionState'] ); } catch (error) { log.warn('Failed to clear persisted review action state', { sessionId, error }); diff --git a/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.ts b/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.ts index fbc5e11ca..415b8275b 100644 --- a/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.ts +++ b/src/web-ui/src/flow_chat/services/flow-chat-manager/PersistenceModule.ts @@ -534,6 +534,14 @@ export async function updateSessionMetadata( await sessionAPI.saveSessionMetadata( metadata, workspacePath, + [ + 'sessionName', + 'tags', + 'todos', + 'unreadCompletion', + 'needsUserAttention', + 'titleMetadata', + ], session.remoteConnectionId, session.remoteSshHost ); diff --git a/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts b/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts index e8339af71..fca03ba78 100644 --- a/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts +++ b/src/web-ui/src/flow_chat/store/FlowChatStore.test.ts @@ -13,6 +13,7 @@ const apiMocks = vi.hoisted(() => ({ restoreSession: vi.fn(), restoreSessionView: vi.fn(), restoreSessionWithTurns: vi.fn(), + accountFetchSessionTurns: vi.fn(), })); const configManagerMock = vi.hoisted(() => { @@ -68,6 +69,12 @@ vi.mock('@/infrastructure/api/service-api/AgentAPI', () => ({ }, })); +vi.mock('@/infrastructure/api/service-api/RemoteConnectAPI', () => ({ + remoteConnectAPI: { + accountFetchSessionTurns: apiMocks.accountFetchSessionTurns, + }, +})); + vi.mock('@/infrastructure/config/services/ConfigManager', () => ({ configManager: configManagerMock, })); @@ -693,6 +700,7 @@ describe('FlowChatStore session model selection', () => { describe('FlowChatStore historical session hydration state', () => { beforeEach(() => { + apiMocks.accountFetchSessionTurns.mockResolvedValue(false); vi.stubGlobal('CustomEvent', class { type: string; detail: unknown; @@ -739,6 +747,64 @@ describe('FlowChatStore historical session hydration state', () => { }); }); + it('checks relay history completeness before restoring Core context', async () => { + const order: string[] = []; + apiMocks.accountFetchSessionTurns.mockImplementationOnce(async () => { + order.push('relay'); + return true; + }); + apiMocks.restoreSessionView.mockImplementationOnce(async () => { + order.push('restore'); + return { + session: { + sessionId: 'history-1', + sessionName: 'History 1', + agentType: 'agentic', + state: 'Idle', + turnCount: 0, + createdAt: 1, + }, + turns: [], + contextRestoreState: 'ready', + }; + }); + flowChatStore.setState(() => ({ + sessions: new Map([ + ['history-1', createSession({ + sessionId: 'history-1', + isHistorical: true, + historyState: 'metadata-only', + })], + ]), + activeSessionId: 'history-1', + })); + + await flowChatStore.loadSessionHistory('history-1', 'D:/workspace/BitFun'); + + expect(order).toEqual(['relay', 'restore']); + }); + + it('fails closed before Core restore when relay history is incomplete', async () => { + apiMocks.accountFetchSessionTurns.mockRejectedValueOnce(new Error('relay unavailable')); + flowChatStore.setState(() => ({ + sessions: new Map([ + ['history-1', createSession({ + sessionId: 'history-1', + isHistorical: true, + historyState: 'metadata-only', + })], + ]), + activeSessionId: 'history-1', + })); + + await expect( + flowChatStore.loadSessionHistory('history-1', 'D:/workspace/BitFun') + ).rejects.toThrow('relay unavailable'); + + expect(apiMocks.restoreSessionView).not.toHaveBeenCalled(); + expect(flowChatStore.getState().sessions.get('history-1')?.historyState).toBe('failed'); + }); + it('loads model config once while processing multiple persisted sessions', async () => { configManagerMock.getConfig.mockImplementation(async (path: string) => { if (path === 'ai.models') return [{ id: 'primary-model', context_window: 256000 }]; diff --git a/src/web-ui/src/flow_chat/store/FlowChatStore.ts b/src/web-ui/src/flow_chat/store/FlowChatStore.ts index d34b0f2d4..971424bd7 100644 --- a/src/web-ui/src/flow_chat/store/FlowChatStore.ts +++ b/src/web-ui/src/flow_chat/store/FlowChatStore.ts @@ -3985,6 +3985,48 @@ export class FlowChatStore { let restoredLoadedTurnCount: number | undefined; let restoredTotalTurnCount: number | undefined; let restoredTiming: SessionViewRestoreTiming | undefined; + + // Finish or resume relay history import before Core restores its model + // context. Ordinary local sessions return after one metadata read, while + // an incomplete relay import fails closed instead of publishing a + // truncated UI/Core history pair. + if (!remote && workspacePath) { + const relayImportStartedAt = nowMs(); + startupTrace.markPhase('historical_session_relay_import_start', { + remote, + sessionId, + sessionTraceId, + }); + try { + const { remoteConnectAPI } = await import( + '@/infrastructure/api/service-api/RemoteConnectAPI' + ); + const fetched = await remoteConnectAPI.accountFetchSessionTurns( + sessionId, + workspacePath + ); + startupTrace.markPhase('historical_session_relay_import_end', { + remote, + sessionId, + sessionTraceId, + fetched, + durationMs: elapsedMs(relayImportStartedAt), + }); + } catch (fetchErr) { + startupTrace.markPhase('historical_session_relay_import_failed', { + remote, + sessionId, + sessionTraceId, + durationMs: elapsedMs(relayImportStartedAt), + }); + log.warn('Relay session history is incomplete; retry opening the session', { + sessionId, + error: fetchErr, + }); + throw fetchErr; + } + } + const stateMachineManagerPromise = import('../state-machine'); if (!isAcpSession) { const restoreStartedAt = nowMs(); @@ -4138,36 +4180,6 @@ export class FlowChatStore { remoteConnectionId, remoteSshHost ); - // Cloud-imported sessions may only have metadata locally; lazy-fetch turns. - if ( - !remote && - (!Array.isArray(turns) || turns.length === 0) && - workspacePath - ) { - try { - const { remoteConnectAPI } = await import( - '@/infrastructure/api/service-api/RemoteConnectAPI' - ); - const fetched = await remoteConnectAPI.accountFetchSessionTurns( - sessionId, - workspacePath - ); - if (fetched) { - turns = await sessionAPI.loadSessionTurns( - sessionId, - workspacePath, - limit, - remoteConnectionId, - remoteSshHost - ); - } - } catch (fetchErr) { - log.warn('accountFetchSessionTurns failed during hydrate', { - sessionId, - error: fetchErr, - }); - } - } startupTrace.markPhase('historical_session_turns_load_end', { remote, sessionId, diff --git a/src/web-ui/src/infrastructure/api/service-api/RemoteConnectAPI.ts b/src/web-ui/src/infrastructure/api/service-api/RemoteConnectAPI.ts index 973b283c0..b7ecac21b 100644 --- a/src/web-ui/src/infrastructure/api/service-api/RemoteConnectAPI.ts +++ b/src/web-ui/src/infrastructure/api/service-api/RemoteConnectAPI.ts @@ -463,7 +463,7 @@ class RemoteConnectAPIService { } } - /** Lazy-load a session's turns from the relay on first open. */ + /** Complete or resume a relay-imported session's lazy turn import. */ async accountFetchSessionTurns(sessionId: string, workspacePath: string): Promise { try { return await this.adapter.request('account_fetch_session_turns', { @@ -471,8 +471,8 @@ class RemoteConnectAPIService { workspacePath, }); } catch (e) { - log.warn('accountFetchSessionTurns failed', e); - return false; + log.error('accountFetchSessionTurns failed', e); + throw e; } } diff --git a/src/web-ui/src/infrastructure/api/service-api/SessionAPI.ts b/src/web-ui/src/infrastructure/api/service-api/SessionAPI.ts index c4a44ce95..ee56e3280 100644 --- a/src/web-ui/src/infrastructure/api/service-api/SessionAPI.ts +++ b/src/web-ui/src/infrastructure/api/service-api/SessionAPI.ts @@ -3,6 +3,15 @@ import { api } from './ApiClient'; import { createTauriCommandError } from '../errors/TauriCommandError'; import type { SessionMetadata, DialogTurnData } from '@/shared/types/session-history'; +export type UiSessionMetadataField = + | 'sessionName' + | 'tags' + | 'todos' + | 'reviewActionState' + | 'unreadCompletion' + | 'needsUserAttention' + | 'titleMetadata'; + export interface SessionMetadataPageRequest { workspacePath: string; limit: number; @@ -295,6 +304,7 @@ export class SessionAPI { async saveSessionMetadata( metadata: SessionMetadata, workspacePath: string, + fields: UiSessionMetadataField[], remoteConnectionId?: string, remoteSshHost?: string ): Promise { @@ -302,6 +312,7 @@ export class SessionAPI { await api.invoke('save_session_metadata', { request: { metadata, + fields, workspace_path: workspacePath, ...remoteSessionFields(remoteConnectionId, remoteSshHost), } diff --git a/src/web-ui/src/locales/en-US/flow-chat.json b/src/web-ui/src/locales/en-US/flow-chat.json index a780c46d5..f3ec7541b 100644 --- a/src/web-ui/src/locales/en-US/flow-chat.json +++ b/src/web-ui/src/locales/en-US/flow-chat.json @@ -43,7 +43,7 @@ "loadingTitle": "Loading saved session", "loadingDescription": "Preparing the conversation history.", "failedTitle": "Session history did not load", - "failedDescription": "Retry loading the saved conversation.", + "failedDescription": "Retry loading the saved conversation. If it was imported from another device, check your account and connection first.", "preparingOlderHistory": "Preparing older history...", "olderHistoryNotReady": "Older history is not ready yet.", "remoteSessionMissing": "The backend session data was not found. If you just reconnected an SSH remote workspace, close and reopen the remote project, or create a new session and try again.", diff --git a/src/web-ui/src/locales/zh-CN/flow-chat.json b/src/web-ui/src/locales/zh-CN/flow-chat.json index 0a9871f04..0fa0822c7 100644 --- a/src/web-ui/src/locales/zh-CN/flow-chat.json +++ b/src/web-ui/src/locales/zh-CN/flow-chat.json @@ -43,7 +43,7 @@ "loadingTitle": "正在加载历史会话", "loadingDescription": "正在准备该会话的历史内容。", "failedTitle": "历史会话加载失败", - "failedDescription": "可以重新加载已保存的会话内容。", + "failedDescription": "可以重新加载已保存的会话内容；若会话来自其他设备，请先检查账号登录与网络连接。", "preparingOlderHistory": "正在准备更早的历史记录...", "olderHistoryNotReady": "更早的历史记录暂未准备好。", "remoteSessionMissing": "在后端找不到该会话数据。若刚重新连接过 SSH 远程工作区，请关闭并重新打开该远程项目，或新建会话后再试。", diff --git a/src/web-ui/src/locales/zh-TW/flow-chat.json b/src/web-ui/src/locales/zh-TW/flow-chat.json index 9cfde38c7..834e82e5a 100644 --- a/src/web-ui/src/locales/zh-TW/flow-chat.json +++ b/src/web-ui/src/locales/zh-TW/flow-chat.json @@ -43,7 +43,7 @@ "loadingTitle": "正在載入歷史會話", "loadingDescription": "正在準備該會話的歷史內容。", "failedTitle": "歷史會話載入失敗", - "failedDescription": "可以重新載入已儲存的會話內容。", + "failedDescription": "可以重新載入已儲存的會話內容；若會話來自其他裝置，請先檢查帳號登入與網路連線。", "preparingOlderHistory": "正在準備更早的歷史記錄...", "olderHistoryNotReady": "更早的歷史記錄暫未準備好。", "remoteSessionMissing": "在後端找不到該會話資料。若剛重新連接過 SSH 遠端工作區，請關閉並重新開啟該遠端項目，或新增會話後再試。",