From 292060e9039a8941239b481e362302cf1170dfef Mon Sep 17 00:00:00 2001 From: limityan Date: Fri, 31 Jul 2026 12:09:35 +0800 Subject: [PATCH] feat(cli): add OpenCode-compatible session undo and redo --- .../agent-runtime-services-design.md | 2 +- docs/architecture/cli-product-line-design.md | 4 +- .../rules/source/forbidden-rules.mjs | 2 +- scripts/core-boundaries/self-test.mjs | 5 +- src/apps/cli/src/account_sync.rs | 30 +- src/apps/cli/src/actions.rs | 63 + src/apps/cli/src/agent/runtime_client.rs | 77 +- src/apps/cli/src/chat_state.rs | 120 ++ src/apps/cli/src/modes/chat.rs | 3 +- src/apps/cli/src/modes/chat/commands.rs | 94 +- src/apps/cli/src/modes/chat/run.rs | 14 +- src/apps/cli/src/modes/chat/tests.rs | 28 + .../cli/src/peer_host/commands/session.rs | 18 +- .../cli/src/peer_host/commands/snapshot.rs | 100 +- src/apps/cli/src/shared_runtime.rs | 14 + src/apps/cli/src/ui/startup.rs | 2 + .../desktop/src/api/remote_connect_api.rs | 161 ++- src/apps/desktop/src/api/session_api.rs | 160 +-- src/apps/desktop/src/api/snapshot_service.rs | 586 +++++--- src/apps/desktop/src/runtime/mod.rs | 32 +- .../src/runtime/session_application.rs | 116 +- .../adapters/agent-runtime-ipc/AGENTS-CN.md | 2 +- .../adapters/agent-runtime-ipc/AGENTS.md | 3 +- .../agent-runtime-ipc/src/operation.rs | 18 +- .../adapters/agent-runtime-ipc/src/server.rs | 50 +- .../src/tests/protocol_contracts.rs | 31 +- .../src/tests/shared_controller.rs | 67 +- .../src/agentic/coordination/coordinator.rs | 1268 +++++++++++++++-- .../src/agentic/coordination/scheduler.rs | 163 ++- .../core/src/agentic/insights/collector.rs | 31 +- .../core/src/agentic/memories/service.rs | 15 +- .../core/src/agentic/persistence/manager.rs | 300 +++- .../assembly/core/src/agentic/session/mod.rs | 1 + .../core/src/agentic/session/revert.rs | 261 ++++ .../src/agentic/session/session_manager.rs | 705 ++++++++- .../implementations/session_history_tool.rs | 12 +- .../assembly/core/src/product_runtime.rs | 633 +++++++- .../remote_connect/bot/command_router.rs | 10 +- .../core/src/service/session_usage/mod.rs | 3 +- .../core/src/service/session_usage/service.rs | 121 +- .../core/src/service/snapshot/manager.rs | 156 +- .../core/src/service/snapshot/service.rs | 146 +- .../src/service/snapshot/snapshot_core.rs | 711 ++++++++- .../src/service/snapshot/snapshot_system.rs | 61 +- .../core/src/service_agent_runtime.rs | 167 ++- src/crates/contracts/events/src/agentic.rs | 33 + .../events/src/frontend_projection.rs | 4 + src/crates/contracts/runtime-ports/src/lib.rs | 80 ++ .../src/local_workspace_snapshot.rs | 4 + .../execution/agent-runtime/src/runtime.rs | 99 +- src/crates/execution/agent-runtime/src/sdk.rs | 58 +- .../src/remote_connect.rs | 163 ++- .../src/remote_connect/account.rs | 96 ++ 53 files changed, 6261 insertions(+), 842 deletions(-) create mode 100644 src/crates/assembly/core/src/agentic/session/revert.rs diff --git a/docs/architecture/agent-runtime-services-design.md b/docs/architecture/agent-runtime-services-design.md index 645755e77..a9e910b5c 100644 --- a/docs/architecture/agent-runtime-services-design.md +++ b/docs/architecture/agent-runtime-services-design.md @@ -1129,7 +1129,7 @@ Product 测试： turn、处理基础会话控制、更新会话模型并处理工具确认/拒绝；本地工作区快照准备、文件清单、统计和文件回滚通过独立 owner port 复用 Core 实现， 富历史和其余持久化维护缺口仍通过单一 Core 兼容接口处理，不再构造独立调度器、持久化 manager 或事件队列； wire schema、Relay ACK/重放和重连协议未在该切换中扩张。 -- CLI 主会话客户端通过 Rust Runtime SDK 处理 session、transcript、fork、usage report、用量卡片完成态本地命令轮次、turn、cancel 与 settlement；其他 preview 缺口仍通过一个 Core 兼容接口处理； +- CLI 主会话客户端通过 Rust Runtime SDK 处理 session、transcript、fork、本地 Session undo/redo、usage report、用量卡片完成态本地命令轮次、turn、cancel 与 settlement；undo/redo 使用独立窄 port，由 Core 统一暂存 transcript、模型上下文与工作区边界，不扩展 `RuntimeServices` 为 service locator；其他 preview 缺口仍通过一个 Core 兼容接口处理； 该接口复用现有归属模块，不建立第二套状态或事件格式。 - CLI 托管的 ACP 服务端已以 `DeliveryProfile::Acp` 构造真实 Runtime Parts；会话创建/列举、轮次、取消、会话模型更新、工具确认/拒绝和 Agent 事件订阅复用同一 Agent Runtime API 语义，ACP stdio、连接与协议转换保持不变。Agentic Event Queue 仍是唯一事件归属模块； diff --git a/docs/architecture/cli-product-line-design.md b/docs/architecture/cli-product-line-design.md index 71cda26de..ac54547d5 100644 --- a/docs/architecture/cli-product-line-design.md +++ b/docs/architecture/cli-product-line-design.md @@ -256,13 +256,15 @@ Headless CLI 和公开 Agent SDK 都调用同一 Agent Runtime API，但交付 | 形态 | 默认部署 | 当前 Shared 范围 | |---|---|---| -| 交互式 TUI | Embedded | 显式 `--shared` 后支持 Session list/create/restore/delete/fork、transcript、当前 Session rename/Agent mode/model、声明式上下文 reload、当前 Session 手动 context compaction、Turn submit/cancel、Permission 和 UserInput | +| 交互式 TUI | Embedded | 显式 `--shared` 后支持 Session list/create/restore/delete/fork、transcript、当前 Session rename/Agent mode/model、声明式上下文 reload、当前 Session 手动 context compaction、Session undo/redo、Turn submit/cancel、Permission 和 UserInput | | `bitfun exec` / CI | Embedded | 不接受 Shared；保持独立进程、stdout/stderr 和退出码语义 | | ACP / SDK Host / GUI / Remote / Peer | 各自既有部署 | 不消费 TUI IPC，也不因本开关改变生命周期 | Shared TUI 不提供 Session archive、模型目录/默认值、Agent/Subagent 管理、MCP/扩展、账号同步、用量、observer、replay 或通用 controller transfer；对应入口给出明确的 Embedded 恢复建议，不在 Client 进程初始化第二套 Core owner。 Shared 模式的斜杠命令、快捷键帮助和底部提示使用同一能力投影：OpenCode 对齐的 `/fork` 以 `Full session` 或历史用户提示词选择分支边界；选择提示词时 fork 只复制该 Turn 之前的历史，并把提示词放回 composer 而不自动发送。`/rename ` 修改当前 Session 名称；`/agent`、Tab 和 Shift+Tab 只切换当前 Session 的 Agent mode；`/models` 只切换当前 Session 的 model；`/reload [skills|instructions]` 刷新下一条消息使用的声明式上下文；OpenCode 对齐的 `/compact` 及其 `/summarize` alias 以一个可取消的 maintenance Turn 压缩当前 Session 上下文，不增加自创命令或快捷键。该 Turn 与普通对话共用 Session 原子准入，取得所有权后再读取待压缩上下文；权威 transcript 保留完整 tool payload，但重建模型上下文时排除该 maintenance Turn。Embedded 与 Shared 的 `/help` 都从 Action Registry 展示这些入口；在 slash menu 中选择 rename 只预填命令并等待用户输入名称。若外部来源使用相同命令名，用户明确选择的 BitFun 命令可完成这一次参数提交，即使偏好保存失败也不会重新弹出来源选择。它们不进入管理页面，也不修改未来 Session 的默认值。其他不支持动作不显示为可执行入口。Session 切换或 fork 失败保留原控制权；Shared fork 只有在新 Session 与 transcript 的响应可编码后才原子转移 controller。单个连接已有活动 Turn 时拒绝重复提交、fork、manual compaction 以及 Session rename/mode/model update，但允许 reload 只影响下一条消息；事件订阅失效后当前视图立即失效并要求重启 Shared TUI。 +OpenCode 对齐的 `/undo` 与 `/redo` 只提供这两个命令名，不增加 alias、快捷键或通用 checkpoint API。Core 以持久化的 staged boundary 同步裁剪可见 transcript、模型上下文和该边界后的受跟踪工作区改动；`/undo` 把被撤销的用户提示词放回 composer，连续调用可继续向前，`/redo` 逐级恢复直到清除 boundary，usage 与 fork 也只读取该 boundary 前的可见事实。活动 Turn 必须先由 scheduler 完成取消和 drain，失败时不写入 boundary；Shared 仍要求 current controller，并以 Runtime 返回的权威 transcript 更新 TUI。暂存期间提交新提示词或执行既有本地 snapshot accept/reject/rollback mutation，会先进入可恢复的 committing 阶段，永久删除隐藏后缀再建立新分支；同一阶段的 snapshot record 被拒绝，避免追加到隐藏历史。崩溃恢复会按持久化阶段重放 workspace/context 对齐或完成提交；跨文件系统与 Session 存储不宣称原子事务，部分失败返回 `outcome_unknown` 并要求恢复 Session 后检查。该能力当前只支持本地工作区，Remote 在任何回退写入前返回 `NotAvailable`。 + 部署差异由 CLI Runtime client 封装。Embedded 以 Rust 类型直接调用 `AgentRuntime`，不初始化 IPC 或执行 JSON 编解码；Shared 将同一业务请求映射为一个有界本机 frame，Client/Server 各自只编码一次，再交给同一 Runtime owner。多 TUI 复用一个 Runtime 进程，连接和队列保持有界，不按 TUI 数量复制 Session owner。详细的 4+1 视图、帧上限和并发边界见 [`agent-runtime-deployment-design.md`](agent-runtime-deployment-design.md)。 diff --git a/scripts/core-boundaries/rules/source/forbidden-rules.mjs b/scripts/core-boundaries/rules/source/forbidden-rules.mjs index 773bb3991..939a12688 100644 --- a/scripts/core-boundaries/rules/source/forbidden-rules.mjs +++ b/scripts/core-boundaries/rules/source/forbidden-rules.mjs @@ -6,7 +6,7 @@ export const forbiddenContentRules = [ reason: 'agent-runtime-ipc operation scope is frozen to the reviewed Shared TUI slice', patterns: [ { - regex: /^\s+(?!(?:Health|ListSessions|CreateSession|RestoreSession|DeleteSession|ForkSession|RenameSession|UpdateSessionMode|UpdateSessionModel|ReloadSessionContext|CompactSession|SubmitTurn|CancelTurn|PendingPermissions|RespondPermission|SubmitUserAnswers|Unit|Sessions|SessionCreated|SessionRestored|SessionForked|TurnAccepted|TurnCancelled|None|CurrentController|AttachExisting|UncontrolledTarget|Self|RuntimeIpcSessionRequirement|RuntimeIpcOperationRules|RuntimeSessionForkRequest|AgentContextReloadRequest|AgentDialogTurnRequest|AgentSessionCompactionRequest|AgentSessionCreateRequest|AgentSessionCreateResult|AgentSessionListRequest|AgentSessionModeUpdateRequest|AgentSessionModelUpdateRequest|AgentSessionSummary|AgentTurnCancellationRequest|AgentTurnCancellationResult|SessionTranscript)\b)[A-Z][A-Za-z0-9_]*\b/, + regex: /^\s+(?!(?:Health|ListSessions|CreateSession|RestoreSession|DeleteSession|ForkSession|RenameSession|UpdateSessionMode|UpdateSessionModel|ReloadSessionContext|CompactSession|UndoSession|RedoSession|SubmitTurn|CancelTurn|PendingPermissions|RespondPermission|SubmitUserAnswers|Unit|Sessions|SessionCreated|SessionRestored|SessionForked|SessionReverted|TurnAccepted|TurnCancelled|None|CurrentController|AttachExisting|UncontrolledTarget|Self|RuntimeIpcSessionRequirement|RuntimeIpcOperationRules|RuntimeSessionForkRequest|AgentContextReloadRequest|AgentDialogTurnRequest|AgentSessionCompactionRequest|AgentSessionCreateRequest|AgentSessionCreateResult|AgentSessionListRequest|AgentSessionModeUpdateRequest|AgentSessionModelUpdateRequest|AgentSessionRevertRequest|AgentSessionRevertResult|AgentSessionSummary|AgentTurnCancellationRequest|AgentTurnCancellationResult|SessionTranscript)\b)[A-Z][A-Za-z0-9_]*\b/, message: 'agent-runtime-ipc may not add archive, replay, observer, general controller-transfer, or other operations beyond the reviewed Shared TUI slice', }, diff --git a/scripts/core-boundaries/self-test.mjs b/scripts/core-boundaries/self-test.mjs index 17de9ff7f..2ccd32057 100644 --- a/scripts/core-boundaries/self-test.mjs +++ b/scripts/core-boundaries/self-test.mjs @@ -4863,8 +4863,11 @@ export function runManifestParserSelfTest({ runtimeIpcOperationPattern.test(' RenameSession {') || runtimeIpcOperationPattern.test(' UpdateSessionMode {') || runtimeIpcOperationPattern.test(' UpdateSessionModel {') || + runtimeIpcOperationPattern.test(' UndoSession {') || + runtimeIpcOperationPattern.test(' RedoSession {') || runtimeIpcOperationPattern.test(' SubmitTurn {') || - runtimeIpcOperationPattern.test(' SessionForked {') + runtimeIpcOperationPattern.test(' SessionForked {') || + runtimeIpcOperationPattern.test(' SessionReverted {') ) { throw new Error('agent-runtime-ipc operation guard must preserve the Shared TUI operation budget'); } diff --git a/src/apps/cli/src/account_sync.rs b/src/apps/cli/src/account_sync.rs index d27b4bc3f..01119db9e 100644 --- a/src/apps/cli/src/account_sync.rs +++ b/src/apps/cli/src/account_sync.rs @@ -12,6 +12,9 @@ use tokio::sync::RwLock; use bitfun_core::product_runtime::CoreAgentRuntimeCompatibility; use bitfun_core::service::config::get_global_config_service; +use bitfun_core::service::remote_connect::account::{ + ensure_relay_session_history_exportable, relay_session_export_metadata, +}; use bitfun_core::service::remote_connect::settings_sync; use bitfun_core::service::remote_connect::{sync_state, AccountClient}; @@ -95,8 +98,7 @@ pub(crate) async fn push_settings_after_local_change() { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -#[derive(Default)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub(crate) enum SyncStatus { #[default] Idle, @@ -105,7 +107,6 @@ pub(crate) enum SyncStatus { Failed, } - #[derive(Debug, Clone)] pub(crate) struct SyncProgress { pub status: SyncStatus, @@ -317,12 +318,17 @@ pub(crate) async fn run_auto_sync( if !account_context_is_current(generation) { return Err(anyhow!("account sync cancelled")); } + if let Err(error) = ensure_relay_session_history_exportable(meta) { + tracing::debug!("Skipping CLI account session export: {error}"); + continue; + } let turns = compatibility .load_persisted_session_turns(&storage_path, &meta.session_id, None) .await .map_err(|e| anyhow!("load turns: {e}"))?; + let metadata = relay_session_export_metadata(meta, turns.len()); let metadata_json = - serde_json::to_value(meta).map_err(|e| anyhow!("serialize metadata: {e}"))?; + serde_json::to_value(metadata).map_err(|e| anyhow!("serialize metadata: {e}"))?; let turns_json: Vec = turns .iter() .map(|t| serde_json::to_value(t).unwrap_or(serde_json::Value::Null)) @@ -482,4 +488,20 @@ mod tests { .to_string() .contains("HTTP 507")); } + + #[test] + fn cli_session_backup_uses_the_shared_import_guard_and_visible_count() { + let source = include_str!("account_sync.rs").replace("\r\n", "\n"); + let export_loop = source + .split_once("for meta in local_sessions.iter()") + .expect("CLI account Session export loop") + .1 + .split_once("let upload_total = pending_uploads.len()") + .expect("CLI account Session export loop boundary") + .0; + + assert!(export_loop.contains("ensure_relay_session_history_exportable(meta)")); + assert!(export_loop.contains("relay_session_export_metadata(meta, turns.len())")); + assert!(export_loop.contains("pending_uploads.push")); + } } diff --git a/src/apps/cli/src/actions.rs b/src/apps/cli/src/actions.rs index 26c101dad..701d92710 100644 --- a/src/apps/cli/src/actions.rs +++ b/src/apps/cli/src/actions.rs @@ -75,6 +75,8 @@ pub(crate) enum ActionHandler { NewSession, Sessions, ForkSession, + UndoSession, + RedoSession, RenameSession, Skills, Reload, @@ -128,6 +130,8 @@ impl ActionHandler { | Self::NewSession | Self::Sessions | Self::ForkSession + | Self::UndoSession + | Self::RedoSession | Self::RenameSession | Self::AcpHelp | Self::Init @@ -391,6 +395,36 @@ static ACTION_SPECS: &[ActionSpec] = &[ shortcut_label: None, slash_on_startup: false, }, + ActionSpec { + id: "undo_session", + name: "Undo session", + aliases: &["/undo"], + description: "Undo the latest user prompt and restore it to the composer", + contexts: CHAT, + availability: ActionAvailability::Always, + handler: ActionHandler::UndoSession, + default_bindings: &[], + fallback_bindings: &[], + shortcut_field: None, + palette: None, + shortcut_label: None, + slash_on_startup: false, + }, + ActionSpec { + id: "redo_session", + name: "Redo session", + aliases: &["/redo"], + description: "Redo the most recently undone session history", + contexts: CHAT, + availability: ActionAvailability::Always, + handler: ActionHandler::RedoSession, + default_bindings: &[], + fallback_bindings: &[], + shortcut_field: None, + palette: None, + shortcut_label: None, + slash_on_startup: false, + }, ActionSpec { id: "skills", name: "Skills", @@ -1910,6 +1944,35 @@ mod tests { assert!(action_for_alias("/branch", ActionContext::Chat).is_none()); } + #[test] + fn undo_and_redo_use_only_the_opencode_commands_in_both_deployments() { + let undo = action_by_id("undo_session", ActionContext::Chat) + .expect("OpenCode-compatible session undo action"); + let redo = action_by_id("redo_session", ActionContext::Chat) + .expect("OpenCode-compatible session redo action"); + + assert_eq!(undo.aliases, &["/undo"]); + assert_eq!(undo.handler, ActionHandler::UndoSession); + assert_eq!(undo.availability, ActionAvailability::Always); + assert!(undo.default_bindings.is_empty()); + assert!(undo.fallback_bindings.is_empty()); + assert!(undo.available(ActionState::chat(false, false).for_shared_tui())); + assert!(undo.available(ActionState::chat(true, false).for_shared_tui())); + + assert_eq!(redo.aliases, &["/redo"]); + assert_eq!(redo.handler, ActionHandler::RedoSession); + assert_eq!(redo.availability, ActionAvailability::Always); + assert!(redo.default_bindings.is_empty()); + assert!(redo.fallback_bindings.is_empty()); + assert!(redo.available(ActionState::chat(false, false).for_shared_tui())); + assert!(redo.available(ActionState::chat(true, false).for_shared_tui())); + + assert!(action_for_alias("/rewind", ActionContext::Chat).is_none()); + assert!(action_for_alias("/revert", ActionContext::Chat).is_none()); + assert!(action_by_id("undo_session", ActionContext::Startup).is_none()); + assert!(action_by_id("redo_session", ActionContext::Startup).is_none()); + } + #[test] fn compact_uses_opencode_commands_without_an_invented_shortcut() { let action = action_by_id("compact_session", ActionContext::Chat) diff --git a/src/apps/cli/src/agent/runtime_client.rs b/src/apps/cli/src/agent/runtime_client.rs index d72b7e2b6..742889a18 100644 --- a/src/apps/cli/src/agent/runtime_client.rs +++ b/src/apps/cli/src/agent/runtime_client.rs @@ -16,10 +16,11 @@ use bitfun_agent_runtime::sdk::{ AgentSessionCompactionRequest, AgentSessionCreateRequest, AgentSessionDeleteRequest, AgentSessionForkBeforeTurnRequest, AgentSessionForkRequest, AgentSessionForkResult, AgentSessionListRequest, AgentSessionModeUpdateRequest, AgentSessionModelUpdateRequest, - AgentSessionRenameRequest, AgentSessionRestoreRequest, AgentSessionUsageRequest, - AgentTurnCancellationRequest, AgentTurnSettlementRequest, AgentUserAnswersRequest, - PermissionReply, PermissionRequest, PermissionRequestEventReceiver, PortError, PortErrorKind, - RuntimeError, SessionTranscript, SessionTranscriptRequest, SessionUsageReport, + AgentSessionRenameRequest, AgentSessionRestoreRequest, AgentSessionRevertRequest, + AgentSessionRevertResult, AgentSessionUsageRequest, AgentTurnCancellationRequest, + AgentTurnSettlementRequest, AgentUserAnswersRequest, PermissionReply, PermissionRequest, + PermissionRequestEventReceiver, PortError, PortErrorKind, RuntimeError, SessionTranscript, + SessionTranscriptRequest, SessionUsageReport, }; use bitfun_agent_runtime_ipc::{ RuntimeIpcClient, RuntimeIpcClientError, RuntimeIpcClientEvent, RuntimeIpcErrorCode, @@ -815,6 +816,55 @@ impl CliAgentRuntimeClient { Ok((session, binding, transcript)) } + pub(crate) async fn revert_current_session( + &self, + undo: bool, + ) -> Result { + let session_id = self.require_session_id().await?; + let request = AgentSessionRevertRequest { + workspace_path: self.project_workspace_path_string(), + session_id: session_id.clone(), + remote_connection_id: None, + remote_ssh_host: None, + }; + let locally_active_turn_id = self.current_turn_id.lock().await.clone(); + let mut reverted = match &self.backend { + CliAgentRuntimeBackend::Embedded(runtime) => if undo { + runtime.undo_session(request).await + } else { + runtime.redo_session(request).await + } + .map_err(|error| anyhow::anyhow!(error.into_message()))?, + CliAgentRuntimeBackend::Shared(client) => { + let operation = if undo { + RuntimeIpcOperation::UndoSession { request } + } else { + RuntimeIpcOperation::RedoSession { request } + }; + match client.request(operation).await? { + RuntimeIpcOperationResult::SessionReverted { revert } + if revert.session_id == session_id => + { + revert + } + _ => return Err(unexpected_shared_result("revert_session")), + } + } + }; + if reverted.session_id != session_id { + return Err(anyhow::anyhow!( + "Runtime reverted an unexpected session identity" + )); + } + if let Some(turn_id) = locally_active_turn_id { + if !reverted.retired_turn_ids.contains(&turn_id) { + reverted.retired_turn_ids.push(turn_id); + } + } + *self.current_turn_id.lock().await = None; + Ok(reverted) + } + pub(crate) async fn generate_session_usage_report( &self, request: AgentSessionUsageRequest, @@ -1656,6 +1706,25 @@ mod tests { assert!(fork.contains("RuntimeIpcOperationResult::SessionForked")); } + #[test] + fn session_revert_uses_the_same_authoritative_result_in_both_deployments() { + let source = include_str!("runtime_client.rs").replace("\r\n", "\n"); + let revert = source + .split_once("pub(crate) async fn revert_current_session(") + .expect("session revert method") + .1 + .split_once("pub(crate) async fn generate_session_usage_report(") + .expect("session revert method boundary") + .0; + + assert!(revert.contains("CliAgentRuntimeBackend::Embedded(runtime)")); + assert!(revert.contains("runtime.undo_session(request)")); + assert!(revert.contains("runtime.redo_session(request)")); + assert!(revert.contains("RuntimeIpcOperation::UndoSession")); + assert!(revert.contains("RuntimeIpcOperation::RedoSession")); + assert!(revert.contains("RuntimeIpcOperationResult::SessionReverted")); + } + #[test] fn session_delete_uses_direct_runtime_or_private_shared_ipc() { let source = include_str!("runtime_client.rs").replace("\r\n", "\n"); diff --git a/src/apps/cli/src/chat_state.rs b/src/apps/cli/src/chat_state.rs index 3d4df90fd..bf93e806e 100644 --- a/src/apps/cli/src/chat_state.rs +++ b/src/apps/cli/src/chat_state.rs @@ -358,6 +358,9 @@ pub(crate) struct ChatState { // -- Streaming state (transient, not persisted) -- /// Current turn ID being processed current_turn_id: Option, + /// Turns retired by authoritative Session operations. Their already queued + /// events are fenced from the rebuilt visible transcript. + ignored_turn_ids: HashSet, /// Ordered flow items for the current streaming message. /// Text, thinking, and tool blocks are interleaved in chronological order, /// matching the actual conversation flow (inspired by opencode's Part model). @@ -414,6 +417,7 @@ impl ChatState { metadata: ChatMetadata::default(), last_primary_model_usage: None, current_turn_id: None, + ignored_turn_ids: HashSet::new(), current_flow_items: Vec::new(), tool_index: HashMap::new(), is_processing: false, @@ -743,6 +747,37 @@ impl ChatState { state } + /// Replace the render projection from the Runtime-owned transcript while + /// retaining Session/workspace/product settings owned by the TUI shell. + pub(crate) fn replace_from_authoritative_transcript( + &mut self, + transcript: &SessionTranscript, + retired_turn_ids: &[String], + ) { + debug_assert_eq!(transcript.session_id, self.core_session_id); + let projected = Self::from_session_transcript( + self.core_session_id.clone(), + self.session_name.clone(), + self.agent_type.clone(), + self.workspace.clone(), + transcript, + ); + self.messages = projected.messages; + self.metadata = projected.metadata; + self.last_primary_model_usage = None; + if let Some(turn_id) = self.current_turn_id.take() { + self.ignored_turn_ids.insert(turn_id); + } + self.ignored_turn_ids + .extend(retired_turn_ids.iter().cloned()); + self.current_flow_items.clear(); + self.tool_index.clear(); + self.is_processing = false; + self.permission_prompt = None; + self.permission_queue.clear(); + self.question_prompt = None; + } + // ============ Event Handlers ============ /// Handle the start of a new dialog turn @@ -1201,6 +1236,17 @@ impl ChatState { self.question_prompt = None; } + pub(crate) fn should_apply_turn_cancelled(&mut self, turn_id: &str) -> bool { + if self.should_ignore_turn_event(turn_id) { + return false; + } + self.current_turn_id.is_none() || self.current_turn_id.as_deref() == Some(turn_id) + } + + pub(crate) fn should_ignore_turn_event(&self, turn_id: &str) -> bool { + self.ignored_turn_ids.contains(turn_id) + } + /// Record the latest primary-model request observed by this TUI. pub(crate) fn handle_primary_model_usage(&mut self, usage: ModelTokenUsageSnapshot) { self.last_primary_model_usage = Some(usage); @@ -1792,6 +1838,80 @@ mod tests { ); } + #[test] + fn authoritative_transcript_replaces_only_session_projection_and_clears_active_turn() { + let mut state = ChatState::new( + "session-1".to_string(), + "Session".to_string(), + "agentic".to_string(), + Some("D:/workspace/project".to_string()), + ); + state.current_model_id = Some("model-1".to_string()); + state.handle_turn_started("turn-2", "Hidden prompt"); + let transcript = SessionTranscript { + session_id: "session-1".to_string(), + messages: vec![TranscriptMessage { + id: Some("user-1".to_string()), + role: "user".to_string(), + turn_id: Some("turn-1".to_string()), + timestamp_ms: Some(1_000), + content: TranscriptContent::Text("Visible prompt".to_string()), + }], + }; + + state.replace_from_authoritative_transcript(&transcript, &[]); + + assert_eq!(state.messages.len(), 1); + assert_eq!(state.messages[0].turn_id.as_deref(), Some("turn-1")); + assert_eq!(state.current_model_id.as_deref(), Some("model-1")); + assert_eq!(state.workspace.as_deref(), Some("D:/workspace/project")); + assert!(!state.is_processing); + assert!(state.current_turn_id().is_none()); + assert!(!state.should_apply_turn_cancelled("turn-2")); + } + + #[test] + fn authoritative_transcript_fences_a_turn_start_that_arrives_after_revert() { + let mut state = ChatState::new( + "session-1".to_string(), + "Session".to_string(), + "agentic".to_string(), + None, + ); + let transcript = SessionTranscript { + session_id: "session-1".to_string(), + messages: vec![TranscriptMessage { + id: Some("assistant-visible".to_string()), + role: "assistant".to_string(), + turn_id: Some("turn-visible".to_string()), + timestamp_ms: Some(1_000), + content: TranscriptContent::Text("Visible answer".to_string()), + }], + }; + + state.replace_from_authoritative_transcript( + &transcript, + &["retired-turn".to_string(), "queued-turn".to_string()], + ); + + assert!(state.should_ignore_turn_event("retired-turn")); + assert!(state.should_ignore_turn_event("queued-turn")); + if !state.should_ignore_turn_event("retired-turn") { + state.handle_turn_started("retired-turn", "late prompt"); + } + assert_eq!(state.messages.len(), 1); + assert!(!state.is_processing); + assert!(!state.should_apply_turn_cancelled("retired-turn")); + assert!(!state.should_apply_turn_cancelled("queued-turn")); + if state.should_apply_turn_cancelled("queued-turn") { + state.handle_turn_cancelled(); + } + assert!(matches!( + state.messages[0].flow_items.as_slice(), + [FlowItem::Text { content, .. }] if content == "Visible answer" + )); + } + #[test] fn session_fork_points_keep_stable_turn_ids_and_newest_prompt_first() { let transcript = SessionTranscript { diff --git a/src/apps/cli/src/modes/chat.rs b/src/apps/cli/src/modes/chat.rs index 418db4987..eb533374f 100644 --- a/src/apps/cli/src/modes/chat.rs +++ b/src/apps/cli/src/modes/chat.rs @@ -18,7 +18,8 @@ use std::time::{Duration, Instant}; use tokio::sync::broadcast::error::TryRecvError; use bitfun_agent_runtime::sdk::{ - AgentLocalCommandTurnRecordRequest, AgentSessionUsageRequest, SessionUsageReport, + AgentLocalCommandTurnRecordRequest, AgentSessionComposerUpdate, AgentSessionUsageRequest, + SessionUsageReport, }; use bitfun_events::{AgenticEvent, ToolEventData, ToolEventIdentity}; use resize::ResizeRedrawState; diff --git a/src/apps/cli/src/modes/chat/commands.rs b/src/apps/cli/src/modes/chat/commands.rs index 777a8c695..862b7ae3e 100644 --- a/src/apps/cli/src/modes/chat/commands.rs +++ b/src/apps/cli/src/modes/chat/commands.rs @@ -37,16 +37,19 @@ fn pending_session_operation_blocks_runtime_action( pending_for_current_session: bool, handler: ActionHandler, ) -> bool { - shared_tui - && pending_for_current_session - && matches!( + pending_for_current_session + && (matches!( handler, - ActionHandler::Sessions - | ActionHandler::ForkSession - | ActionHandler::RenameSession - | ActionHandler::CompactSession - | ActionHandler::Init - ) + ActionHandler::UndoSession | ActionHandler::RedoSession + ) || shared_tui + && matches!( + handler, + ActionHandler::Sessions + | ActionHandler::ForkSession + | ActionHandler::RenameSession + | ActionHandler::CompactSession + | ActionHandler::Init + )) } fn requested_session_name(arguments: &str) -> Option { @@ -94,6 +97,8 @@ fn builtin_arguments_error( match handler { ActionHandler::CompactSession => Some("Usage: /compact"), ActionHandler::ForkSession => Some("Usage: /fork"), + ActionHandler::UndoSession => Some("Usage: /undo"), + ActionHandler::RedoSession => Some("Usage: /redo"), _ => None, } } @@ -151,9 +156,13 @@ fn session_command_help_note() -> String { .expect("current session rename action must remain registered"); let fork = action_for_alias("/fork", ActionContext::Chat) .expect("current session fork action must remain registered"); + let undo = action_for_alias("/undo", ActionContext::Chat) + .expect("current session undo action must remain registered"); + let redo = action_for_alias("/redo", ActionContext::Chat) + .expect("current session redo action must remain registered"); format!( - "Session Commands\n {}\n {}", - fork.description, rename.description + "Session Commands\n /fork - {}\n /rename - {}\n /undo - {}\n /redo - {}", + fork.description, rename.description, undo.description, redo.description ) } @@ -917,6 +926,12 @@ impl ChatMode { ActionHandler::ForkSession => { self.show_fork_selector(chat_view, chat_state); } + ActionHandler::UndoSession => { + self.revert_session(true, chat_view, chat_state, rt_handle); + } + ActionHandler::RedoSession => { + self.revert_session(false, chat_view, chat_state, rt_handle); + } ActionHandler::RenameSession => { return self.start_session_rename("", chat_view, chat_state, rt_handle); } @@ -1070,6 +1085,63 @@ impl ChatMode { } } + fn revert_session( + &mut self, + undo: bool, + chat_view: &mut ChatView, + chat_state: &mut ChatState, + rt_handle: &tokio::runtime::Handle, + ) { + let operation = if undo { "Undo" } else { "Redo" }; + chat_view.set_status(Some(format!("{operation}ing session..."))); + let agent = self.agent.clone(); + let result = tokio::task::block_in_place(|| { + rt_handle.block_on(async move { agent.revert_current_session(undo).await }) + }); + let reverted = match result { + Ok(reverted) => reverted, + Err(error) => { + chat_view.set_status(Some(format!( + "Could not {} session: {error}", + operation.to_ascii_lowercase() + ))); + return; + } + }; + + chat_state.replace_from_authoritative_transcript( + &reverted.transcript, + &reverted.retired_turn_ids, + ); + match reverted.composer { + AgentSessionComposerUpdate::Preserve => {} + AgentSessionComposerUpdate::Replace { text } => chat_view.set_input(&text), + AgentSessionComposerUpdate::Clear => chat_view.clear_input(), + } + self.selected_native_command_once = None; + chat_view.clear_screen(); + chat_view.scroll_to_bottom(); + chat_view.set_status(Some(if reverted.changed { + if undo { + format!( + "Undid the latest prompt; {} persisted turn(s) are hidden.", + reverted.hidden_turn_count + ) + } else if reverted.hidden_turn_count == 0 { + "Restored the full session history.".to_string() + } else { + format!( + "Redid session history; {} persisted turn(s) remain hidden.", + reverted.hidden_turn_count + ) + } + } else if undo { + "Nothing to undo.".to_string() + } else { + "Nothing to redo.".to_string() + })); + } + fn start_session_rename( &mut self, arguments: &str, diff --git a/src/apps/cli/src/modes/chat/run.rs b/src/apps/cli/src/modes/chat/run.rs index a2a481426..cf31673e3 100644 --- a/src/apps/cli/src/modes/chat/run.rs +++ b/src/apps/cli/src/modes/chat/run.rs @@ -712,6 +712,17 @@ impl ChatMode { tracing::debug!("Processing core event: {:?}", event); + if event + .turn_id() + .is_some_and(|turn_id| chat_state.should_ignore_turn_event(turn_id)) + { + tracing::debug!( + "Ignoring event for a Turn retired by authoritative Session state: turn_id={:?}", + event.turn_id() + ); + continue; + } + match event { AgenticEvent::DialogTurnStarted { turn_id, @@ -825,8 +836,7 @@ impl ChatMode { } AgenticEvent::DialogTurnCancelled { turn_id, .. } => { - let active_turn_id = chat_state.current_turn_id(); - if active_turn_id.is_none() || active_turn_id == Some(turn_id.as_str()) { + if chat_state.should_apply_turn_cancelled(turn_id) { chat_state.handle_turn_cancelled(); self.refresh_workspace_git_status(&mut chat_state, &rt_handle); chat_view.invalidate_lines_cache(); diff --git a/src/apps/cli/src/modes/chat/tests.rs b/src/apps/cli/src/modes/chat/tests.rs index 4558e1c2e..12377c4a0 100644 --- a/src/apps/cli/src/modes/chat/tests.rs +++ b/src/apps/cli/src/modes/chat/tests.rs @@ -1197,6 +1197,22 @@ mod tests { ), Some("Usage: /fork") ); + assert_eq!( + builtin_arguments_error( + CommandRoute::Builtin, + ActionHandler::UndoSession, + "unexpected" + ), + Some("Usage: /undo") + ); + assert_eq!( + builtin_arguments_error( + CommandRoute::Builtin, + ActionHandler::RedoSession, + "unexpected" + ), + Some("Usage: /redo") + ); } #[test] @@ -1878,6 +1894,16 @@ mod tests { true, ActionHandler::Sessions, )); + assert!(pending_session_operation_blocks_runtime_action( + false, + true, + ActionHandler::UndoSession, + )); + assert!(pending_session_operation_blocks_runtime_action( + true, + true, + ActionHandler::RedoSession, + )); assert!(!pending_session_operation_blocks_runtime_action( true, false, @@ -1967,6 +1993,8 @@ mod tests { assert!(help.contains("Session Commands")); assert!(help.contains(rename.description)); assert!(help.contains("/rename ")); + assert!(help.contains("/undo")); + assert!(help.contains("/redo")); } #[test] diff --git a/src/apps/cli/src/peer_host/commands/session.rs b/src/apps/cli/src/peer_host/commands/session.rs index d30fcf74c..b7f936022 100644 --- a/src/apps/cli/src/peer_host/commands/session.rs +++ b/src/apps/cli/src/peer_host/commands/session.rs @@ -546,10 +546,19 @@ pub(crate) async fn get_session_stats( .map_err(session_stats_validation_error)?; require_local_snapshot_workspace(request, &workspace_path).await?; + let scope = ensure_session_workspace_runtime_ownership(state, request)?; + let storage_path = resolved_session_storage_scope(state, scope).await?; + let read = state + .compatibility + .begin_persisted_session_read(&storage_path, &session_id) + .await + .map_err(|error| format!("Failed to open a consistent snapshot view: {error}"))?; + let stats = local_snapshot_session_stats( state.local_workspace_snapshot.as_ref(), PathBuf::from(&workspace_path), session_id, + read.visible_turn_end(), ) .await?; @@ -587,7 +596,12 @@ pub(crate) async fn save_session_turn( return Err("turn_data session_id does not match request session_id".to_string()); } } - let _mutation = state + state + .compatibility + .ensure_session_loaded_from_storage_path(&workspace_path, &turn.session_id, false) + .await + .map_err(|error| format!("Failed to load session before saving a Turn: {error}"))?; + let mutation = state .compatibility .begin_persisted_session_mutation(&workspace_path, &turn.session_id) .await @@ -595,7 +609,7 @@ pub(crate) async fn save_session_turn( state .compatibility - .save_persisted_dialog_turn(&workspace_path, &turn) + .save_persisted_dialog_turn(&mutation, &turn) .await .map_err(|e| format!("Failed to save session turn: {e}"))?; Ok(Value::Null) diff --git a/src/apps/cli/src/peer_host/commands/snapshot.rs b/src/apps/cli/src/peer_host/commands/snapshot.rs index 94ccfdf4f..f6a25c2b1 100644 --- a/src/apps/cli/src/peer_host/commands/snapshot.rs +++ b/src/apps/cli/src/peer_host/commands/snapshot.rs @@ -61,10 +61,12 @@ pub(super) async fn local_snapshot_session_files( port: &dyn LocalWorkspaceSnapshotPort, workspace_path: PathBuf, session_id: String, + max_turn_exclusive: Option, ) -> Result, String> { port.get_session_files(LocalWorkspaceSnapshotSessionRequest { workspace_path, session_id, + max_turn_exclusive, }) .await .map_err(|error| { @@ -79,10 +81,12 @@ pub(super) async fn local_snapshot_session_stats( port: &dyn LocalWorkspaceSnapshotPort, workspace_path: PathBuf, session_id: String, + max_turn_exclusive: Option, ) -> Result { port.get_session_stats(LocalWorkspaceSnapshotSessionRequest { workspace_path, session_id, + max_turn_exclusive, }) .await .map_err(|error| { @@ -160,10 +164,18 @@ pub(crate) async fn get_session_files( bitfun_agent_runtime::session_control::validate_session_id(&session_id)?; require_local_snapshot_workspace(request, &workspace_path).await?; + let scope = ensure_session_workspace_runtime_ownership(state, request)?; + let storage_path = resolved_session_storage_scope(state, scope).await?; + let read = state + .compatibility + .begin_persisted_session_read(&storage_path, &session_id) + .await + .map_err(|error| format!("Failed to open a consistent snapshot view: {error}"))?; let files = local_snapshot_session_files( state.local_workspace_snapshot.as_ref(), PathBuf::from(&workspace_path), session_id, + read.visible_turn_end(), ) .await?; @@ -185,13 +197,11 @@ pub(crate) async fn rollback_to_turn(state: &PeerHostState, args: &Value) -> Res let workspace = PathBuf::from(&workspace_path); let scope = ensure_session_workspace_runtime_ownership(state, request)?; let session_storage_path = resolved_session_storage_scope(state, scope).await?; - if delete_turns { - state - .compatibility - .ensure_session_loaded_from_storage_path(&session_storage_path, &session_id, false) - .await - .map_err(|error| format!("Failed to load session before rollback: {error}"))?; - } + state + .compatibility + .ensure_session_loaded_from_storage_path(&session_storage_path, &session_id, false) + .await + .map_err(|error| format!("Failed to load session before rollback: {error}"))?; let maintenance = state .compatibility .begin_session_maintenance(&session_storage_path, &session_id, 2_000) @@ -207,32 +217,28 @@ pub(crate) async fn rollback_to_turn(state: &PeerHostState, args: &Value) -> Res .map_err(|error| format!("Failed to cancel Peer descendants before rollback: {error}"))?; state.turns.drain_session_turns(&session_id); - let mutation = if delete_turns { - Some( - state - .compatibility - .begin_persisted_session_mutation(&session_storage_path, &session_id) - .await - .map_err(|error| format!("Failed to lock session rollback: {error}"))?, - ) - } else { - None - }; + let mutation = state + .compatibility + .begin_persisted_session_mutation(&session_storage_path, &session_id) + .await + .map_err(|error| format!("Failed to lock session rollback: {error}"))?; + state + .compatibility + .commit_session_revert_before_snapshot_mutation(&mutation) + .await + .map_err(|error| { + format!("Failed to commit the staged Session undo before rollback: {error}") + })?; let rolled_back_parent_turn_ids = if delete_turns { let turns = state .compatibility - .load_persisted_session_turns(&session_storage_path, &session_id, None) + .load_persisted_session_turns_for_mutation(&mutation, None) .await .map_err(|error| format!("Failed to load turns before rollback: {error}"))?; state .compatibility - .validate_persisted_session_context_rollback( - mutation - .as_ref() - .expect("mutation exists when deleting turns"), - turn_index, - ) + .validate_persisted_session_context_rollback(&mutation, turn_index) .await .map_err(|error| format!("Failed to validate session rollback: {error}"))?; turns @@ -261,12 +267,7 @@ pub(crate) async fn rollback_to_turn(state: &PeerHostState, args: &Value) -> Res if delete_turns { if let Err(error) = state .compatibility - .rollback_persisted_session_context_to_turn_start( - mutation - .as_ref() - .expect("mutation exists when deleting turns"), - turn_index, - ) + .rollback_persisted_session_context_to_turn_start(&mutation, turn_index) .await { return Err(history_rollback_partial_failure(error)); @@ -414,6 +415,9 @@ mod tests { let cancellation = rollback_source .find("cancel_peer_turns") .expect("descendant cancellation must remain in the rollback flow"); + let revert_commit = rollback_source + .find("commit_session_revert_before_snapshot_mutation") + .expect("staged Session undo must be committed before legacy rollback"); let file_rollback = rollback_source .find("rollback_local_workspace_files(") .expect("workspace-file rollback must remain in the rollback flow"); @@ -426,7 +430,8 @@ mod tests { assert!( remote_guard < maintenance && maintenance < cancellation - && cancellation < file_rollback + && cancellation < revert_commit + && revert_commit < file_rollback && file_rollback < history_rollback && history_rollback < event_projection, "rollback must preserve remote guard, maintenance, cancellation, files, history, and event order" @@ -438,12 +443,22 @@ mod tests { let port = RecordingSnapshotPort::default(); let workspace = PathBuf::from("workspace"); - let files = local_snapshot_session_files(&port, workspace.clone(), "session-1".to_string()) - .await - .expect("file projection should succeed"); - let stats = local_snapshot_session_stats(&port, workspace.clone(), "session-1".to_string()) - .await - .expect("stats projection should succeed"); + let files = local_snapshot_session_files( + &port, + workspace.clone(), + "session-1".to_string(), + Some(2), + ) + .await + .expect("file projection should succeed"); + let stats = local_snapshot_session_stats( + &port, + workspace.clone(), + "session-1".to_string(), + Some(2), + ) + .await + .expect("stats projection should succeed"); let restored = rollback_local_workspace_files(&port, workspace.clone(), "session-1".to_string(), 4) .await @@ -464,6 +479,15 @@ mod tests { .workspace_path, workspace ); + assert_eq!( + port.file_request + .lock() + .expect("file request lock") + .as_ref() + .expect("file request") + .max_turn_exclusive, + Some(2) + ); assert_eq!( port.rollback_request .lock() diff --git a/src/apps/cli/src/shared_runtime.rs b/src/apps/cli/src/shared_runtime.rs index 5963b876c..7e5b5e53c 100644 --- a/src/apps/cli/src/shared_runtime.rs +++ b/src/apps/cli/src/shared_runtime.rs @@ -354,6 +354,18 @@ impl RuntimeIpcRequestHandler for SharedRuntimeHandler { turn_id: result.turn_id, }) .map_err(runtime_ipc_error), + RuntimeIpcOperation::UndoSession { request } => self + .runtime + .undo_session(request) + .await + .map(|revert| RuntimeIpcOperationResult::SessionReverted { revert }) + .map_err(runtime_ipc_error), + RuntimeIpcOperation::RedoSession { request } => self + .runtime + .redo_session(request) + .await + .map(|revert| RuntimeIpcOperationResult::SessionReverted { revert }) + .map_err(runtime_ipc_error), RuntimeIpcOperation::SubmitTurn { request } => { let outcome = self .runtime @@ -599,6 +611,8 @@ impl SharedRuntimeHandler { .as_deref() .ok_or_else(workspace_mismatch_error)?, ), + RuntimeIpcOperation::UndoSession { request } + | RuntimeIpcOperation::RedoSession { request } => Some(request.workspace_path.as_str()), _ => None, }; let Some(requested) = requested else { diff --git a/src/apps/cli/src/ui/startup.rs b/src/apps/cli/src/ui/startup.rs index 35d8da9d5..2b69ab9b5 100644 --- a/src/apps/cli/src/ui/startup.rs +++ b/src/apps/cli/src/ui/startup.rs @@ -1068,6 +1068,8 @@ impl StartupPage { ActionHandler::NavigateBack => self.navigate_back(), ActionHandler::RenameSession | ActionHandler::ForkSession + | ActionHandler::UndoSession + | ActionHandler::RedoSession | ActionHandler::Reload | ActionHandler::Tools | ActionHandler::Extensions diff --git a/src/apps/desktop/src/api/remote_connect_api.rs b/src/apps/desktop/src/api/remote_connect_api.rs index b4ea03dc8..099bf2ad5 100644 --- a/src/apps/desktop/src/api/remote_connect_api.rs +++ b/src/apps/desktop/src/api/remote_connect_api.rs @@ -2,11 +2,14 @@ use crate::api::session_storage_path::desktop_effective_session_storage_path; use crate::embedded_relay_host::DesktopEmbeddedRelayHost; -use bitfun_core::agentic::coordination::{get_global_coordinator, ConversationCoordinator}; +use bitfun_core::agentic::coordination::{ + get_global_coordinator, get_global_scheduler, ConversationCoordinator, +}; use bitfun_core::agentic::persistence::PersistenceManager; use bitfun_core::agentic::tools::account_login_capability::set_account_login_available; use bitfun_core::agentic::tools::page_deploy_host::set_page_deploy_handler; use bitfun_core::agentic::tools::page_publish_host::set_page_publish_handler; +use bitfun_core::product_runtime::CoreAgentRuntimeCompatibility; use bitfun_core::service::remote_connect::session_store::{ clear_credential_hint, load_credential_hint, save_credential_hint, AccountHint, }; @@ -20,7 +23,10 @@ use bitfun_core::service::session::{DialogTurnData, SessionMetadata}; use bitfun_core::service::workspace::{get_global_workspace_service, WorkspaceKind}; use bitfun_core::service::workspace_runtime::WorkspaceRuntimeService; use bitfun_services_integrations::remote_connect::account::{ - error_indicates_expired_token, validate_relay_base_url, + ensure_relay_session_history_exportable, error_indicates_expired_token, + mark_relay_session_history_import_complete, mark_relay_session_history_import_pending, + relay_session_export_metadata, relay_session_history_import_is_complete, + relay_session_history_import_state, validate_relay_base_url, }; use bitfun_services_integrations::remote_connect::{ deploy_page_version_on_relay, join_relay_url, list_pages_from_relay, @@ -2981,40 +2987,16 @@ 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); +async fn load_account_visible_session_turns( + storage_path: &std::path::Path, + session_id: &str, +) -> Result, String> { + let coordinator = get_global_coordinator() + .ok_or_else(|| "Core coordinator is not initialized for session sync".to_string())?; + coordinator + .load_visible_persisted_session_turns(storage_path, session_id) + .await + .map_err(|error| format!("load visible session history: {error}")) } /// Export a single local session as an encrypted blob and upload it to the relay. @@ -3042,12 +3024,11 @@ pub async fn account_export_local_session( .await .map_err(|e| format!("load metadata: {e}"))? .ok_or_else(|| format!("session not found: {session_id}"))?; + ensure_relay_session_history_exportable(&metadata)?; // Load all turns - let turns = manager - .load_session_turns(&storage_path, &session_id) - .await - .map_err(|e| format!("load turns: {e}"))?; + let turns = load_account_visible_session_turns(&storage_path, &session_id).await?; + let metadata = relay_session_export_metadata(&metadata, turns.len()); // Serialize to bundle let metadata_json = @@ -3106,13 +3087,17 @@ pub async fn account_export_all_sessions( let mut state = sync_state::load(&acct_session.user_id); let mut pending: Vec<(String, String, String)> = Vec::new(); for meta in &sessions { - let turns = manager - .load_session_turns(&storage_path, &meta.session_id) + if let Err(error) = ensure_relay_session_history_exportable(meta) { + log::debug!("Skipping account session export: {error}"); + continue; + } + let turns = load_account_visible_session_turns(&storage_path, &meta.session_id) .await .map_err(|e| format!("load turns for {}: {e}", meta.session_id))?; + let metadata = relay_session_export_metadata(meta, turns.len()); let metadata_json = - serde_json::to_value(meta).map_err(|e| format!("serialize metadata: {e}"))?; + serde_json::to_value(metadata).map_err(|e| format!("serialize metadata: {e}"))?; let turns_json: Vec = turns .iter() .map(|t| serde_json::to_value(t).unwrap_or(serde_json::Value::Null)) @@ -3219,7 +3204,7 @@ pub async fn account_import_remote_sessions( } // Only write metadata — turns are lazy-loaded when the user opens // the session (see `account_fetch_session_turns`). - set_relay_turns_import_state(&mut metadata, RELAY_TURNS_IMPORT_PENDING); + mark_relay_session_history_import_pending(&mut metadata); if !manager .create_session_metadata_if_absent(&storage_path, &metadata) .await @@ -3273,8 +3258,8 @@ pub async fn account_fetch_session_turns( .map_err(|e| format!("create persistence manager: {e}"))?; // 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. + // account or network lookup. Only the durable complete marker proves that + // the imported turn batch finished; a partial prefix remains pending. let Some(metadata) = manager .load_session_metadata(&storage_path, &session_id) .await @@ -3282,15 +3267,11 @@ pub async fn account_fetch_session_turns( else { return Ok(false); }; - if relay_turns_import_state(&metadata).is_none() { + if relay_session_history_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()) { + if relay_session_history_import_is_complete(&metadata) { return Ok(false); } @@ -3325,6 +3306,16 @@ pub async fn account_fetch_session_turns( return Err("relay session turn identity does not match request".to_string()); } + let coordinator = get_global_coordinator() + .ok_or_else(|| "Core coordinator is not initialized for session import".to_string())?; + let scheduler = get_global_scheduler() + .ok_or_else(|| "Core scheduler is not initialized for session import".to_string())?; + let compatibility = CoreAgentRuntimeCompatibility::build(coordinator, scheduler); + let _history_write = compatibility + .begin_external_persisted_history_write(&storage_path, &session_id) + .await + .map_err(|error| format!("session import is unavailable during undo or redo: {error}"))?; + manager .create_session_metadata_if_absent(&storage_path, &metadata) .await @@ -3339,7 +3330,7 @@ pub async fn account_fetch_session_turns( } manager .update_session_metadata(&storage_path, &session_id, |metadata| { - mark_relay_turns_import_complete(metadata); + mark_relay_session_history_import_complete(metadata); }) .await .map_err(|e| format!("mark imported turns complete: {e}"))?; @@ -3749,12 +3740,16 @@ async fn account_auto_sync_inner( let mut pending_uploads: Vec<(String, String, String)> = Vec::new(); for meta in local_sessions.iter() { ensure_account_auto_sync_current(sync_operation_id)?; - let turns = manager - .load_session_turns(&storage_path, &meta.session_id) + if let Err(error) = ensure_relay_session_history_exportable(meta) { + log::debug!("Skipping account auto-sync export: {error}"); + continue; + } + let turns = load_account_visible_session_turns(&storage_path, &meta.session_id) .await .map_err(|e| format!("load turns: {e}"))?; + let metadata = relay_session_export_metadata(meta, turns.len()); let metadata_json = - serde_json::to_value(meta).map_err(|e| format!("serialize metadata: {e}"))?; + serde_json::to_value(metadata).map_err(|e| format!("serialize metadata: {e}"))?; let turns_json: Vec = turns .iter() .map(|t| serde_json::to_value(t).unwrap_or(serde_json::Value::Null)) @@ -4187,10 +4182,12 @@ async fn export_and_upload_session( .load_session_metadata(&storage_path, session_id) .await? .ok_or_else(|| anyhow::anyhow!("session not found: {session_id}"))?; + ensure_relay_session_history_exportable(&metadata).map_err(anyhow::Error::msg)?; - let turns = manager - .load_session_turns(&storage_path, session_id) - .await?; + let turns = load_account_visible_session_turns(&storage_path, session_id) + .await + .map_err(anyhow::Error::msg)?; + let metadata = relay_session_export_metadata(&metadata, turns.len()); let metadata_json = serde_json::to_value(&metadata)?; let turns_json: Vec = turns @@ -4378,7 +4375,7 @@ async fn import_session_bundle(bundle_json: &str, account_generation: u64) -> an // 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); + mark_relay_session_history_import_pending(&mut metadata); manager .create_session_metadata_if_absent(&target_dir, &metadata) .await @@ -4734,7 +4731,7 @@ mod sync_state_tests { } #[test] - fn relay_turn_import_requires_an_explicit_complete_marker_and_exact_count() { + fn pending_relay_turn_imports_are_never_exportable() { let mut metadata = SessionMetadata::new( "session".to_string(), "Session".to_string(), @@ -4743,18 +4740,40 @@ mod sync_state_tests { ); 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_session_history_import_state(&metadata), None); + assert!(!relay_session_history_import_is_complete(&metadata)); + assert!(ensure_relay_session_history_exportable(&metadata).is_ok()); + mark_relay_session_history_import_pending(&mut metadata); assert_eq!( - relay_turns_import_state(&metadata), - Some(RELAY_TURNS_IMPORT_PENDING) + relay_session_history_import_state(&metadata), + Some("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)); + assert!(!relay_session_history_import_is_complete(&metadata)); + assert!(ensure_relay_session_history_exportable(&metadata).is_err()); + mark_relay_session_history_import_complete(&mut metadata); + assert!(relay_session_history_import_is_complete(&metadata)); + assert!(ensure_relay_session_history_exportable(&metadata).is_ok()); + + metadata.custom_metadata = Some(serde_json::json!({ + "relayTurnsImportState": "unknown" + })); + assert!(ensure_relay_session_history_exportable(&metadata).is_err()); + } + + #[test] + fn account_export_metadata_matches_the_visible_history_projection() { + let mut metadata = SessionMetadata::new( + "session".to_string(), + "Session".to_string(), + "agentic".to_string(), + "auto".to_string(), + ); + metadata.turn_count = 3; + + let exported = relay_session_export_metadata(&metadata, 2); + + assert_eq!(metadata.turn_count, 3); + assert_eq!(exported.turn_count, 2); } } diff --git a/src/apps/desktop/src/api/session_api.rs b/src/apps/desktop/src/api/session_api.rs index 018edc08e..41e42d584 100644 --- a/src/apps/desktop/src/api/session_api.rs +++ b/src/apps/desktop/src/api/session_api.rs @@ -1,17 +1,13 @@ //! Session persistence API use crate::api::app_state::AppState; -use crate::api::session_storage_path::desktop_effective_session_storage_path; use crate::runtime::{ DesktopRuntimeContext, DesktopSessionApplicationError, DesktopSessionScopeRequest, UiSessionMetadataField, }; use crate::startup_trace::DesktopStartupTrace; use bitfun_core::agentic::coordination::get_global_scheduler; -use bitfun_core::agentic::persistence::{ - PersistenceManager, SessionBranchResult, SessionMetadataPage, -}; -use bitfun_core::infrastructure::PathManager; +use bitfun_core::agentic::persistence::{SessionBranchResult, SessionMetadataPage}; use bitfun_core::service::remote_ssh::normalize_remote_workspace_path; use bitfun_core::service::session::{ DialogTurnData, SessionKind, SessionMetadata, SessionStatus, SessionTranscriptExport, @@ -20,7 +16,6 @@ use bitfun_core::service::session::{ use bitfun_core::service::session_usage::SessionUsageReport; use bitfun_core::service::workspace::WorkspaceKind; use serde::{Deserialize, Serialize}; -use std::sync::Arc; use std::time::Instant; use tauri::State; @@ -430,33 +425,20 @@ pub async fn get_session_usage_report( #[tauri::command] pub async fn save_session_turn( request: SaveSessionTurnRequest, - app_state: State<'_, AppState>, - path_manager: State<'_, Arc>, runtime: State<'_, DesktopRuntimeContext>, ) -> Result<(), String> { runtime .session_application() - .ensure_workspace_runtime_ownership(desktop_session_scope( - request.workspace_path.clone(), - request.remote_connection_id.clone(), - request.remote_ssh_host.clone(), - )) - .await - .map_err(desktop_session_error)?; - let workspace_path = desktop_effective_session_storage_path( - &app_state, - &request.workspace_path, - request.remote_connection_id.as_deref(), - request.remote_ssh_host.as_deref(), - ) - .await; - let manager = PersistenceManager::new(path_manager.inner().clone()) - .map_err(|e| format!("Failed to create persistence manager: {}", e))?; - - manager - .save_dialog_turn(&workspace_path, &request.turn_data) + .save_session_turn( + desktop_session_scope( + request.workspace_path.clone(), + request.remote_connection_id, + request.remote_ssh_host, + ), + &request.turn_data, + ) .await - .map_err(|e| format!("Failed to save session turn: {}", e))?; + .map_err(|error| format!("Failed to save session turn: {error}"))?; // Notify the auto-sync background task (debounced upload to relay) crate::api::remote_connect_api::notify_session_changed( @@ -495,22 +477,16 @@ pub async fn save_session_metadata( #[tauri::command] pub async fn export_session_transcript( request: ExportSessionTranscriptRequest, - app_state: State<'_, AppState>, - path_manager: State<'_, Arc>, + runtime: State<'_, DesktopRuntimeContext>, ) -> Result { - let workspace_path = desktop_effective_session_storage_path( - &app_state, - &request.workspace_path, - request.remote_connection_id.as_deref(), - request.remote_ssh_host.as_deref(), - ) - .await; - let manager = PersistenceManager::new(path_manager.inner().clone()) - .map_err(|e| format!("Failed to create persistence manager: {}", e))?; - - manager + runtime + .session_application() .export_session_transcript( - &workspace_path, + desktop_session_scope( + request.workspace_path, + request.remote_connection_id, + request.remote_ssh_host, + ), &request.session_id, &SessionTranscriptExportOptions { tools: request.tools, @@ -520,7 +496,12 @@ pub async fn export_session_transcript( }, ) .await - .map_err(|e| format!("Failed to export session transcript: {}", e)) + .map_err(|error| { + format!( + "Failed to export session transcript: {}", + desktop_session_error(error) + ) + }) } #[tauri::command] @@ -682,33 +663,18 @@ pub async fn unarchive_session( #[tauri::command] pub async fn archive_all_sessions( request: ArchiveAllSessionsRequest, - app_state: State<'_, AppState>, - path_manager: State<'_, Arc>, runtime: State<'_, DesktopRuntimeContext>, ) -> Result { - runtime + let scope = desktop_session_scope( + request.workspace_path, + request.remote_connection_id, + request.remote_ssh_host, + ); + let sessions = runtime .session_application() - .ensure_workspace_runtime_ownership(desktop_session_scope( - request.workspace_path.clone(), - request.remote_connection_id.clone(), - request.remote_ssh_host.clone(), - )) - .await - .map_err(desktop_session_error)?; - let workspace_path = desktop_effective_session_storage_path( - &app_state, - &request.workspace_path, - request.remote_connection_id.as_deref(), - request.remote_ssh_host.as_deref(), - ) - .await; - let manager = PersistenceManager::new(path_manager.inner().clone()) - .map_err(|e| format!("Failed to create persistence manager: {}", e))?; - - let sessions = manager - .list_session_metadata(&workspace_path) + .list_persisted_sessions(scope.clone()) .await - .map_err(|e| format!("Failed to list sessions: {}", e))?; + .map_err(|error| format!("Failed to list sessions: {}", desktop_session_error(error)))?; let mut archived_count: u32 = 0; @@ -716,14 +682,16 @@ pub async fn archive_all_sessions( if metadata.status != SessionStatus::Archived && metadata.session_kind == SessionKind::Standard { - manager - .update_session_metadata(&workspace_path, &metadata.session_id, |current| { - if current.session_kind == SessionKind::Standard { - current.status = SessionStatus::Archived; - } - }) + runtime + .session_application() + .set_session_archived(scope.clone(), metadata.session_id, true) .await - .map_err(|e| format!("Failed to save session metadata: {}", e))?; + .map_err(|error| { + format!( + "Failed to save session metadata: {}", + desktop_session_error(error) + ) + })?; archived_count += 1; } } @@ -750,44 +718,30 @@ pub async fn list_archived_sessions( #[tauri::command] pub async fn delete_all_archived_sessions( request: DeleteAllArchivedSessionsRequest, - app_state: State<'_, AppState>, - path_manager: State<'_, Arc>, runtime: State<'_, DesktopRuntimeContext>, ) -> Result { - runtime + let scope = desktop_session_scope( + request.workspace_path, + request.remote_connection_id, + request.remote_ssh_host, + ); + let sessions = runtime .session_application() - .ensure_workspace_runtime_ownership(desktop_session_scope( - request.workspace_path.clone(), - request.remote_connection_id.clone(), - request.remote_ssh_host.clone(), - )) + .list_archived_sessions(scope.clone()) .await - .map_err(desktop_session_error)?; - let workspace_path = desktop_effective_session_storage_path( - &app_state, - &request.workspace_path, - request.remote_connection_id.as_deref(), - request.remote_ssh_host.as_deref(), - ) - .await; - let manager = PersistenceManager::new(path_manager.inner().clone()) - .map_err(|e| format!("Failed to create persistence manager: {}", e))?; - - let sessions = manager - .list_session_metadata(&workspace_path) - .await - .map_err(|e| format!("Failed to list sessions: {}", e))?; + .map_err(|error| format!("Failed to list sessions: {}", desktop_session_error(error)))?; let mut deleted_count: u32 = 0; for metadata in sessions { - if metadata.status == SessionStatus::Archived { - manager - .delete_session(&workspace_path, &metadata.session_id) - .await - .map_err(|e| format!("Failed to delete session: {}", e))?; - deleted_count += 1; - } + runtime + .session_application() + .delete_session(scope.clone(), metadata.session_id) + .await + .map_err(|error| { + format!("Failed to delete session: {}", desktop_session_error(error)) + })?; + deleted_count += 1; } Ok(deleted_count) diff --git a/src/apps/desktop/src/api/snapshot_service.rs b/src/apps/desktop/src/api/snapshot_service.rs index d66d259ad..6daa965c1 100644 --- a/src/apps/desktop/src/api/snapshot_service.rs +++ b/src/apps/desktop/src/api/snapshot_service.rs @@ -1,20 +1,24 @@ //! Snapshot Service API -use bitfun_core::infrastructure::try_get_path_manager_arc; +use bitfun_core::product_runtime::{ + CoreSessionMaintenancePermit, CoreSessionMutationPermit, CoreSessionReadPermit, +}; use bitfun_core::service::remote_ssh::workspace_state::is_remote_path; use bitfun_core::service::snapshot::{ ensure_snapshot_manager_for_workspace, get_snapshot_manager_for_workspace, - initialize_snapshot_manager_for_workspace, open_snapshot_manager_for_view, OperationType, - SnapshotConfig, SnapshotManager, + initialize_snapshot_manager_for_workspace, open_snapshot_manager_for_view, FileChangeEntry, + OperationType, SnapshotConfig, SnapshotManager, }; +#[cfg(test)] +use bitfun_runtime_ports::LocalWorkspaceSnapshotSessionRequest; use bitfun_runtime_ports::{ - LocalWorkspaceSnapshotPort, LocalWorkspaceSnapshotSessionRequest, - LocalWorkspaceSnapshotTurnRequest, PortError, PortErrorKind, + LocalWorkspaceSnapshotPort, LocalWorkspaceSnapshotTurnRequest, PortError, PortErrorKind, + SessionStoragePathRequest, }; use log::{info, warn}; use serde::{Deserialize, Serialize}; -use std::collections::HashSet; -use std::{path::PathBuf, sync::Arc, time::Duration}; +use std::collections::{BTreeSet, HashSet}; +use std::{path::PathBuf, sync::Arc}; use tauri::{AppHandle, Emitter, State}; use crate::runtime::{DesktopRuntimeContext, DesktopSessionScopeRequest}; @@ -346,6 +350,7 @@ async fn rollback_local_workspace_files( .collect()) } +#[cfg(test)] async fn local_snapshot_session_files( port: &dyn LocalWorkspaceSnapshotPort, workspace_path: PathBuf, @@ -356,6 +361,7 @@ async fn local_snapshot_session_files( .get_session_files(LocalWorkspaceSnapshotSessionRequest { workspace_path, session_id, + max_turn_exclusive: None, }) .await .map_err(|error| { @@ -367,6 +373,7 @@ async fn local_snapshot_session_files( .collect()) } +#[cfg(test)] async fn local_snapshot_session_stats( port: &dyn LocalWorkspaceSnapshotPort, workspace_path: PathBuf, @@ -377,6 +384,7 @@ async fn local_snapshot_session_stats( .get_session_stats(LocalWorkspaceSnapshotSessionRequest { workspace_path, session_id, + max_turn_exclusive: None, }) .await .map_err(|error| { @@ -486,6 +494,102 @@ async fn snapshot_manager_for_view( .map_err(|error| format!("Failed to open snapshot view: {error}")) } +struct SnapshotHistoryMutation { + _maintenance: CoreSessionMaintenancePermit, + mutation: CoreSessionMutationPermit, + storage_path: PathBuf, +} + +async fn begin_snapshot_history_read( + runtime: &DesktopRuntimeContext, + workspace_path: &str, + session_id: &str, +) -> Result { + let compatibility = runtime.session_application().compatibility(); + let storage_path = compatibility + .resolve_persisted_session_storage_path(SessionStoragePathRequest { + workspace_path: PathBuf::from(workspace_path), + remote_connection_id: None, + remote_ssh_host: None, + }) + .await + .map_err(|error| { + format!("Failed to resolve session storage before snapshot view: {error}") + })?; + compatibility + .begin_persisted_session_read(&storage_path, session_id) + .await + .map_err(|error| format!("Failed to open a consistent snapshot view: {error}")) +} + +async fn begin_snapshot_history_mutation( + runtime: &DesktopRuntimeContext, + workspace_path: &str, + session_id: &str, +) -> Result { + let compatibility = runtime.session_application().compatibility(); + let storage_path = compatibility + .resolve_persisted_session_storage_path(SessionStoragePathRequest { + workspace_path: PathBuf::from(workspace_path), + remote_connection_id: None, + remote_ssh_host: None, + }) + .await + .map_err(|error| { + format!("Failed to resolve session storage before snapshot mutation: {error}") + })?; + compatibility + .ensure_session_loaded_from_storage_path(&storage_path, session_id, false) + .await + .map_err(|error| format!("Failed to load session before snapshot mutation: {error}"))?; + let maintenance = compatibility + .begin_session_maintenance(&storage_path, session_id, 2_000) + .await + .map_err(|error| format!("Failed to quiesce session before snapshot mutation: {error}"))?; + let mutation = compatibility + .begin_persisted_session_mutation(&storage_path, session_id) + .await + .map_err(|error| format!("Failed to lock snapshot mutation: {error}"))?; + compatibility + .commit_session_revert_before_snapshot_mutation(&mutation) + .await + .map_err(|error| { + format!("Failed to commit the staged Session undo before snapshot mutation: {error}") + })?; + Ok(SnapshotHistoryMutation { + _maintenance: maintenance, + mutation, + storage_path, + }) +} + +async fn begin_snapshot_record_mutation( + runtime: &DesktopRuntimeContext, + workspace_path: &str, + session_id: &str, +) -> Result { + let compatibility = runtime.session_application().compatibility(); + let storage_path = compatibility + .resolve_persisted_session_storage_path(SessionStoragePathRequest { + workspace_path: PathBuf::from(workspace_path), + remote_connection_id: None, + remote_ssh_host: None, + }) + .await + .map_err(|error| { + format!("Failed to resolve session storage before snapshot recording: {error}") + })?; + let mutation = compatibility + .begin_persisted_session_mutation(&storage_path, session_id) + .await + .map_err(|error| format!("Failed to lock snapshot recording: {error}"))?; + compatibility + .ensure_snapshot_record_allowed(&mutation) + .await + .map_err(|error| format!("Failed to admit snapshot recording: {error}"))?; + Ok(mutation) +} + #[tauri::command] pub async fn record_file_change( app_handle: AppHandle, @@ -510,6 +614,13 @@ pub async fn record_file_change( } }; + let _record_mutation = begin_snapshot_record_mutation( + runtime.inner(), + &request.workspace_path, + &request.session_id, + ) + .await?; + let snapshot_id = manager .record_file_change( &request.session_id, @@ -542,6 +653,12 @@ pub async fn rollback_session( ) -> Result, String> { ensure_complete_rollback_supported(&request.workspace_path, &request.remote_scope).await?; ensure_local_runtime_ownership(runtime.inner(), &request.workspace_path).await?; + let _history_mutation = begin_snapshot_history_mutation( + runtime.inner(), + &request.workspace_path, + &request.session_id, + ) + .await?; let manager = ensure_snapshot_manager_ready_for(&request.workspace_path, "rollback_session").await?; @@ -577,22 +694,33 @@ pub async fn rollback_to_turn( ensure_complete_rollback_supported(&request.workspace_path, &request.remote_scope).await?; ensure_local_runtime_ownership(runtime.inner(), &request.workspace_path).await?; let workspace_path = resolve_workspace_dir(&request.workspace_path).await?; + let compatibility = runtime.session_application().compatibility(); + let history_mutation = begin_snapshot_history_mutation( + runtime.inner(), + &request.workspace_path, + &request.session_id, + ) + .await?; - { - use bitfun_core::agentic::coordination::get_global_coordinator; - - if let Some(coordinator) = get_global_coordinator() { - if let Err(e) = coordinator - .cancel_active_turn_for_session(&request.session_id, Duration::from_secs(2)) - .await - { - warn!( - "Failed to cancel active turn before rollback: session_id={}, turn_index={}, error={}", - request.session_id, request.turn_index, e - ); - } - } - } + let rolled_back_parent_turn_ids = if request.delete_turns { + compatibility + .validate_persisted_session_context_rollback( + &history_mutation.mutation, + request.turn_index, + ) + .await + .map_err(|error| format!("Failed to validate session rollback: {error}"))?; + compatibility + .load_persisted_session_turns_for_mutation(&history_mutation.mutation, None) + .await + .map_err(|error| format!("Failed to load turns before rollback: {error}"))? + .into_iter() + .filter(|turn| turn.turn_index >= request.turn_index) + .map(|turn| turn.turn_id) + .collect::>() + } else { + HashSet::new() + }; let restored_files_str = rollback_local_workspace_files( runtime.local_workspace_snapshot(), @@ -603,107 +731,32 @@ pub async fn rollback_to_turn( ) .await?; - let mut deleted_turns_count = 0; + let deleted_turns_count = rolled_back_parent_turn_ids.len(); if request.delete_turns { - let workspace_path = PathBuf::from(&request.workspace_path); - let mut rolled_back_parent_turn_ids = HashSet::new(); - - use bitfun_core::agentic::persistence::PersistenceManager; - - match try_get_path_manager_arc() { - Ok(path_manager) => match PersistenceManager::new(path_manager) { - Ok(persistence_manager) => { - match persistence_manager - .load_session_turns(&workspace_path, &request.session_id) - .await - { - Ok(turns) => { - rolled_back_parent_turn_ids = turns - .into_iter() - .filter(|turn| turn.turn_index >= request.turn_index) - .map(|turn| turn.turn_id) - .collect(); - } - Err(e) => { - warn!( - "Failed to load parent turns before rollback cleanup: session_id={}, turn_index={}, error={}", - request.session_id, request.turn_index, e - ); - } - } - } - Err(e) => { - warn!("Failed to create PersistenceManager: error={}", e); - } - }, - Err(e) => { - warn!("Failed to create PathManager: error={}", e); - } - } - - { - use bitfun_core::agentic::coordination::get_global_coordinator; - - if let Some(coordinator) = get_global_coordinator() { - if !rolled_back_parent_turn_ids.is_empty() { - if let Err(e) = coordinator - .delete_hidden_subagent_sessions_for_parent_turns( - &workspace_path, - &request.session_id, - &rolled_back_parent_turn_ids, - ) - .await - { - warn!( - "Failed to delete hidden subagent sessions during rollback: session_id={}, turn_index={}, error={}", - request.session_id, request.turn_index, e - ); - } - } - - if let Err(e) = coordinator - .get_session_manager() - .rollback_context_to_turn_start( - &workspace_path, - &request.session_id, - request.turn_index, - ) - .await - { - warn!( - "Rollback agentic context failed: session_id={}, turn_index={}, error={}", - request.session_id, request.turn_index, e - ); - } - } else { - warn!("Global coordinator not initialized, skipping agentic context rollback"); - } - } - - match try_get_path_manager_arc() { - Ok(path_manager) => match PersistenceManager::new(path_manager) { - Ok(persistence_manager) => { - match persistence_manager - .delete_turns_from(&workspace_path, &request.session_id, request.turn_index) - .await - { - Ok(count) => { - deleted_turns_count = count; - } - Err(e) => { - warn!( - "Failed to delete conversation turns: session_id={}, turn_index={}, error={}", - request.session_id, request.turn_index, e - ); - } - } - } - Err(e) => { - warn!("Failed to create PersistenceManager: error={}", e); - } - }, - Err(e) => { - warn!("Failed to create PathManager: error={}", e); + compatibility + .rollback_persisted_session_context_to_turn_start( + &history_mutation.mutation, + request.turn_index, + ) + .await + .map_err(|error| { + format!( + "Workspace files were rolled back, but session history rollback failed. Reload the session before retrying: {error}" + ) + })?; + if !rolled_back_parent_turn_ids.is_empty() { + if let Err(error) = compatibility + .delete_hidden_subagent_sessions_for_parent_turns( + &history_mutation.storage_path, + &request.session_id, + &rolled_back_parent_turn_ids, + ) + .await + { + warn!( + "Failed to delete hidden subagent sessions during rollback: session_id={}, turn_index={}, error={}", + request.session_id, request.turn_index, error + ); } } @@ -717,6 +770,8 @@ pub async fn rollback_to_turn( ); } + drop(history_mutation); + let _ = app_handle.emit( "turn_rolled_back", serde_json::json!({ @@ -739,6 +794,12 @@ pub async fn accept_session( ) -> Result { ensure_local_snapshot_mutation_path(&request.workspace_path, &request.remote_scope).await?; ensure_local_runtime_ownership(runtime.inner(), &request.workspace_path).await?; + let _history_mutation = begin_snapshot_history_mutation( + runtime.inner(), + &request.workspace_path, + &request.session_id, + ) + .await?; let manager = ensure_snapshot_manager_ready_for(&request.workspace_path, "accept_session").await?; @@ -768,6 +829,12 @@ pub async fn accept_file( ) -> Result { ensure_local_snapshot_mutation_path(&request.workspace_path, &request.remote_scope).await?; ensure_local_runtime_ownership(runtime.inner(), &request.workspace_path).await?; + let _history_mutation = begin_snapshot_history_mutation( + runtime.inner(), + &request.workspace_path, + &request.session_id, + ) + .await?; let manager = ensure_snapshot_manager_ready_for(&request.workspace_path, "accept_file").await?; manager @@ -797,6 +864,12 @@ pub async fn reject_file( ) -> Result { ensure_local_snapshot_mutation_path(&request.workspace_path, &request.remote_scope).await?; ensure_local_runtime_ownership(runtime.inner(), &request.workspace_path).await?; + let _history_mutation = begin_snapshot_history_mutation( + runtime.inner(), + &request.workspace_path, + &request.session_id, + ) + .await?; let manager = ensure_snapshot_manager_ready_for(&request.workspace_path, "reject_file").await?; let restored_files = manager @@ -832,77 +905,67 @@ pub async fn get_session_files( if request.remote_scope.declares_remote() || is_remote_path(&request.workspace_path).await { return Ok(vec![]); } - let workspace_path = resolve_workspace_dir(&request.workspace_path).await?; - - local_snapshot_session_files( - runtime.local_workspace_snapshot(), - workspace_path, + let read = begin_snapshot_history_read( + runtime.inner(), &request.workspace_path, - request.session_id, + &request.session_id, ) - .await + .await?; + let manager = snapshot_manager_for_view(&request.workspace_path, &request.remote_scope).await?; + manager + .get_session_files_before(&request.session_id, read.visible_turn_end()) + .await + .map(|files| { + files + .into_iter() + .map(|path| path.to_string_lossy().into_owned()) + .collect() + }) + .map_err(|error| format!("Failed to get session files: {error}")) } #[tauri::command] pub async fn get_session_turns( - _app_handle: AppHandle, + runtime: State<'_, DesktopRuntimeContext>, request: GetSessionTurnsRequest, ) -> Result, String> { - use bitfun_core::agentic::persistence::PersistenceManager; - if request.remote_scope.declares_remote() || is_remote_path(&request.workspace_path).await { return Ok(vec![]); } - - let workspace_path = PathBuf::from(&request.workspace_path); - if let Ok(path_manager) = try_get_path_manager_arc() { - match PersistenceManager::new(path_manager) { - Ok(persistence_manager) => { - match persistence_manager - .load_session_metadata(&workspace_path, &request.session_id) - .await - { - Ok(Some(metadata)) => { - let turns: Vec = (0..metadata.turn_count).collect(); - return Ok(turns); - } - Ok(None) => {} - Err(e) => { - warn!( - "Failed to load conversation metadata: session_id={}, error={}, falling back to snapshot", - request.session_id, e - ); - } - } - } - Err(e) => { - warn!( - "Failed to create PersistenceManager: error={}, falling back to snapshot", - e - ); - } - } - } - + let read = begin_snapshot_history_read( + runtime.inner(), + &request.workspace_path, + &request.session_id, + ) + .await?; let manager = snapshot_manager_for_view(&request.workspace_path, &request.remote_scope).await?; - - let turns = manager - .get_session_turns(&request.session_id) + manager + .get_session_turns_before(&request.session_id, read.visible_turn_end()) .await - .map_err(|e| format!("Failed to get session turns: {}", e))?; - - Ok(turns) + .map_err(|error| format!("Failed to get session turns: {error}")) } #[tauri::command] -pub async fn get_turn_files(request: GetTurnFilesRequest) -> Result, String> { +pub async fn get_turn_files( + runtime: State<'_, DesktopRuntimeContext>, + request: GetTurnFilesRequest, +) -> Result, String> { if request.remote_scope.declares_remote() || is_remote_path(&request.workspace_path).await { return Ok(vec![]); } + let read = begin_snapshot_history_read( + runtime.inner(), + &request.workspace_path, + &request.session_id, + ) + .await?; let manager = snapshot_manager_for_view(&request.workspace_path, &request.remote_scope).await?; - let files = manager - .get_turn_files(&request.session_id, request.turn_index) + .get_turn_files_before( + &request.session_id, + request.turn_index, + read.visible_turn_end(), + ) .await .map_err(|e| format!("Failed to get turn files: {}", e))?; @@ -913,14 +976,25 @@ pub async fn get_turn_files(request: GetTurnFilesRequest) -> Result, } #[tauri::command] -pub async fn get_file_diff(request: GetFileDiffRequest) -> Result { +pub async fn get_file_diff( + runtime: State<'_, DesktopRuntimeContext>, + request: GetFileDiffRequest, +) -> Result { + ensure_local_snapshot_mutation_path(&request.workspace_path, &request.remote_scope).await?; + let read = begin_snapshot_history_read( + runtime.inner(), + &request.workspace_path, + &request.session_id, + ) + .await?; let manager = snapshot_manager_for_view(&request.workspace_path, &request.remote_scope).await?; let diff = manager - .get_file_diff( + .get_file_diff_before( &request.session_id, &request.file_path, request.operation_id.as_deref(), + read.visible_turn_end(), ) .await .map_err(|e| format!("Failed to get file diff: {}", e))?; @@ -930,15 +1004,21 @@ pub async fn get_file_diff(request: GetFileDiffRequest) -> Result, request: GetOperationDiffRequest, ) -> Result { + ensure_local_snapshot_mutation_path(&request.workspace_path, &request.remote_scope).await?; + let read = + begin_snapshot_history_read(runtime.inner(), &request.workspace_path, &request.sessionId) + .await?; let manager = snapshot_manager_for_view(&request.workspace_path, &request.remote_scope).await?; let diff = manager - .get_file_diff( + .get_file_diff_before( &request.sessionId, &request.filePath, request.operationId.as_deref(), + read.visible_turn_end(), ) .await .map_err(|e| format!("Failed to get file diff: {}", e))?; @@ -972,12 +1052,21 @@ pub async fn get_operation_diff( #[tauri::command] pub async fn get_session_file_diff_stats( + runtime: State<'_, DesktopRuntimeContext>, request: GetSessionFileDiffStatsRequest, ) -> Result { + ensure_local_snapshot_mutation_path(&request.workspace_path, &request.remote_scope).await?; + let read = + begin_snapshot_history_read(runtime.inner(), &request.workspace_path, &request.sessionId) + .await?; let manager = snapshot_manager_for_view(&request.workspace_path, &request.remote_scope).await?; let stats = manager - .get_session_file_diff_stats(&request.sessionId, &request.filePath) + .get_session_file_diff_stats_before( + &request.sessionId, + &request.filePath, + read.visible_turn_end(), + ) .await .map_err(|e| format!("Failed to get session file diff stats: {}", e))?; @@ -986,12 +1075,21 @@ pub async fn get_session_file_diff_stats( #[tauri::command] pub async fn get_operation_summary( + runtime: State<'_, DesktopRuntimeContext>, request: GetOperationSummaryRequest, ) -> Result { + ensure_local_snapshot_mutation_path(&request.workspace_path, &request.remote_scope).await?; + let read = + begin_snapshot_history_read(runtime.inner(), &request.workspace_path, &request.sessionId) + .await?; let manager = snapshot_manager_for_view(&request.workspace_path, &request.remote_scope).await?; let summary = manager - .get_operation_summary(&request.sessionId, &request.operationId) + .get_operation_summary_before( + &request.sessionId, + &request.operationId, + read.visible_turn_end(), + ) .await .map_err(|e| format!("Failed to get operation summary: {}", e))?; @@ -1010,15 +1108,21 @@ pub async fn get_operation_summary( #[tauri::command] pub async fn get_session_operations( + runtime: State<'_, DesktopRuntimeContext>, request: GetSessionFilesRequest, ) -> Result { if request.remote_scope.declares_remote() || is_remote_path(&request.workspace_path).await { return Ok(serde_json::Value::Array(Vec::new())); } + let read = begin_snapshot_history_read( + runtime.inner(), + &request.workspace_path, + &request.session_id, + ) + .await?; let manager = snapshot_manager_for_view(&request.workspace_path, &request.remote_scope).await?; - let session = manager - .get_session(&request.session_id) + .get_session_before(&request.session_id, read.visible_turn_end()) .await .map_err(|e| format!("Failed to get session operations: {}", e))?; @@ -1063,6 +1167,12 @@ pub async fn accept_operation( ) -> Result { ensure_local_snapshot_mutation_path(&request.workspace_path, &request.remote_scope).await?; ensure_local_runtime_ownership(runtime.inner(), &request.workspace_path).await?; + let _history_mutation = begin_snapshot_history_mutation( + runtime.inner(), + &request.workspace_path, + &request.sessionId, + ) + .await?; let manager = ensure_snapshot_manager_ready_for(&request.workspace_path, "accept_operation").await?; @@ -1103,6 +1213,12 @@ pub async fn reject_operation( ) -> Result { ensure_local_snapshot_mutation_path(&request.workspace_path, &request.remote_scope).await?; ensure_local_runtime_ownership(runtime.inner(), &request.workspace_path).await?; + let _history_mutation = begin_snapshot_history_mutation( + runtime.inner(), + &request.workspace_path, + &request.sessionId, + ) + .await?; let manager = ensure_snapshot_manager_ready_for(&request.workspace_path, "reject_operation").await?; @@ -1154,15 +1270,17 @@ pub async fn get_session_stats( "total_changes": 0 })); } - let workspace_path = resolve_workspace_dir(&request.workspace_path).await?; - - local_snapshot_session_stats( - runtime.local_workspace_snapshot(), - workspace_path, + let read = begin_snapshot_history_read( + runtime.inner(), &request.workspace_path, - request.session_id, + &request.session_id, ) - .await + .await?; + let manager = snapshot_manager_for_view(&request.workspace_path, &request.remote_scope).await?; + manager + .get_session_stats_before(&request.session_id, read.visible_turn_end()) + .await + .map_err(|error| format!("Failed to get session stats: {error}")) } #[tauri::command] @@ -1213,6 +1331,7 @@ pub async fn check_git_isolation( #[tauri::command] pub async fn get_file_change_history( + runtime: State<'_, DesktopRuntimeContext>, request: GetFileChangeHistoryRequest, ) -> Result { if request.remote_scope.declares_remote() || is_remote_path(&request.workspace_path).await { @@ -1221,16 +1340,45 @@ pub async fn get_file_change_history( let manager = snapshot_manager_for_view(&request.workspace_path, &request.remote_scope).await?; let file_path = PathBuf::from(&request.file_path); - let changes = manager - .get_file_change_history(&file_path) + let session_ids = manager + .list_sessions() .await - .map_err(|e| format!("Failed to get file change history: {}", e))?; + .map_err(|error| format!("Failed to list snapshot sessions: {error}"))?; + let mut changes = Vec::new(); + for session_id in session_ids { + let read = + begin_snapshot_history_read(runtime.inner(), &request.workspace_path, &session_id) + .await?; + let session = manager + .get_session_before(&session_id, read.visible_turn_end()) + .await + .map_err(|error| format!("Failed to get file change history: {error}"))?; + changes.extend( + session + .operations + .into_iter() + .filter(|operation| operation.file_path == file_path) + .map(|operation| FileChangeEntry { + session_id: operation.session_id, + turn_index: operation.turn_index, + snapshot_id: operation + .before_snapshot_id + .unwrap_or_else(|| format!("empty_snapshot_{}", operation.operation_id)), + timestamp: operation.timestamp, + operation_type: operation.operation_type, + tool_name: operation.tool_context.tool_name, + }), + ); + drop(read); + } + changes.sort_by_key(|entry| (entry.session_id.clone(), entry.turn_index, entry.timestamp)); serde_json::to_value(changes).map_err(|e| format!("Serialization failed: {}", e)) } #[tauri::command] pub async fn get_all_modified_files( + runtime: State<'_, DesktopRuntimeContext>, request: GetAllModifiedFilesRequest, ) -> Result, String> { if request.remote_scope.declares_remote() || is_remote_path(&request.workspace_path).await { @@ -1238,14 +1386,27 @@ pub async fn get_all_modified_files( } let manager = snapshot_manager_for_view(&request.workspace_path, &request.remote_scope).await?; - let files = manager - .get_all_modified_files() + let session_ids = manager + .list_sessions() .await - .map_err(|e| format!("Failed to get modified files: {}", e))?; + .map_err(|error| format!("Failed to list snapshot sessions: {error}"))?; + let mut files = BTreeSet::new(); + for session_id in session_ids { + let read = + begin_snapshot_history_read(runtime.inner(), &request.workspace_path, &session_id) + .await?; + files.extend( + manager + .get_session_files_before(&session_id, read.visible_turn_end()) + .await + .map_err(|error| format!("Failed to get modified files: {error}"))?, + ); + drop(read); + } Ok(files - .iter() - .map(|p| p.to_string_lossy().to_string()) + .into_iter() + .map(|path| path.to_string_lossy().into_owned()) .collect()) } @@ -1413,7 +1574,42 @@ mod tests { assert_remote_guard_precedes(rollback_session, "ensure_local_runtime_ownership"); assert_remote_guard_precedes(rollback_session, "ensure_snapshot_manager_ready_for"); assert_remote_guard_precedes(rollback_to_turn, "ensure_local_runtime_ownership"); - assert_remote_guard_precedes(rollback_to_turn, "cancel_active_turn_for_session"); + assert_remote_guard_precedes(rollback_to_turn, "begin_snapshot_history_mutation"); + } + + #[test] + fn snapshot_mutators_share_session_revert_admission() { + let source = include_str!("snapshot_service.rs"); + for (command, next_command) in [ + ("rollback_session", "rollback_to_turn"), + ("rollback_to_turn", "accept_session"), + ("accept_session", "accept_file"), + ("accept_file", "reject_file"), + ("reject_file", "get_session_files"), + ("accept_operation", "reject_operation"), + ("reject_operation", "get_session_stats"), + ] { + let body = source + .split_once(&format!("pub async fn {command}")) + .unwrap_or_else(|| panic!("{command} remains present")) + .1 + .split_once(&format!("pub async fn {next_command}")) + .unwrap_or_else(|| panic!("{next_command} remains present")) + .0; + assert!( + body.contains("begin_snapshot_history_mutation"), + "{command} must commit staged Session undo under the shared mutation owner" + ); + } + + let record = source + .split_once("pub async fn record_file_change") + .expect("record_file_change remains present") + .1 + .split_once("pub async fn rollback_session") + .expect("rollback_session remains present") + .0; + assert!(record.contains("begin_snapshot_record_mutation")); } #[tokio::test] diff --git a/src/apps/desktop/src/runtime/mod.rs b/src/apps/desktop/src/runtime/mod.rs index e999d245f..d5d9189e5 100644 --- a/src/apps/desktop/src/runtime/mod.rs +++ b/src/apps/desktop/src/runtime/mod.rs @@ -166,18 +166,23 @@ mod tests { let session_commands = include_str!("../api/session_api.rs"); assert_eq!( session_commands.matches("PersistenceManager::new").count(), - 4, - "only raw turn save, transcript export, and the two excluded bulk operations keep direct persistence" + 0, + "Tauri Session commands must delegate persistence ownership to the application boundary" ); + let session_application = include_str!("session_application.rs"); + assert!(session_application.contains("save_persisted_dialog_turn")); + assert!(session_application.contains("export_persisted_session_transcript")); let snapshot_commands = include_str!("../api/snapshot_service.rs"); assert_eq!( snapshot_commands .matches(".local_workspace_snapshot()") .count(), - 3, - "only file listing, typed stats, and workspace rollback use the local owner port" + 1, + "only workspace rollback uses the mutation port; reads use Core's bounded snapshot view" ); + assert!(snapshot_commands.contains("begin_snapshot_history_read")); + assert!(snapshot_commands.contains("open_snapshot_manager_for_view")); assert!(snapshot_commands.contains("ensure_local_snapshot_mutation_path")); let rollback_source = &snapshot_commands[snapshot_commands @@ -186,15 +191,16 @@ mod tests { let remote_guard = rollback_source .find("ensure_complete_rollback_supported") .expect("remote rollback guard must remain host-owned"); - let cancellation = rollback_source - .find("cancel_active_turn_for_session") - .expect("active-turn cancellation must precede rollback"); + let maintenance = rollback_source + .find("begin_snapshot_history_mutation") + .expect("scheduler maintenance must precede rollback"); let file_rollback = rollback_source .find("rollback_local_workspace_files(") .expect("workspace files must be restored through the port adapter"); - let history_cleanup = rollback_source - .find("if request.delete_turns") - .expect("history cleanup must remain host-owned"); + let history_cleanup = file_rollback + + rollback_source[file_rollback..] + .find("if request.delete_turns") + .expect("history cleanup must remain host-owned"); let history_event = rollback_source .find("conversation_turns_deleted") .expect("history event must remain host-projected"); @@ -202,12 +208,12 @@ mod tests { .find("turn_rolled_back") .expect("rollback event must remain host-projected"); assert!( - remote_guard < cancellation - && cancellation < file_rollback + remote_guard < maintenance + && maintenance < file_rollback && file_rollback < history_cleanup && history_cleanup < history_event && history_event < rollback_event, - "Desktop rollback must preserve remote, cancellation, files, history, and event order" + "Desktop rollback must preserve remote, maintenance, files, history, and event order" ); let sdk_source = include_str!("../../../../crates/execution/agent-runtime/src/sdk.rs"); diff --git a/src/apps/desktop/src/runtime/session_application.rs b/src/apps/desktop/src/runtime/session_application.rs index eb1b7ca53..7f71e049b 100644 --- a/src/apps/desktop/src/runtime/session_application.rs +++ b/src/apps/desktop/src/runtime/session_application.rs @@ -10,9 +10,9 @@ use std::time::Instant; use async_trait::async_trait; use bitfun_agent_runtime::sdk::{ - AgentRuntime, AgentSessionArchiveStateRequest, AgentSessionDeleteRequest, - AgentSessionForkAtTurnRequest, AgentSessionRenameRequest, AgentSessionUsageRequest, - PortErrorKind, RuntimeError, + AgentLocalCommandTurnRecordRequest, AgentRuntime, AgentSessionArchiveStateRequest, + AgentSessionDeleteRequest, AgentSessionForkAtTurnRequest, AgentSessionRenameRequest, + AgentSessionUsageRequest, PortErrorKind, RuntimeError, }; use bitfun_core::agentic::coordination::{ConversationCoordinator, DialogScheduler}; use bitfun_core::agentic::core::Session; @@ -21,7 +21,10 @@ use bitfun_core::agentic::session::SessionViewRestoreTiming; use bitfun_core::product_runtime::{CoreAgentRuntimeCompatibility, CoreProductAgentRuntime}; use bitfun_core::service::remote_ssh::workspace_state::get_effective_session_path; use bitfun_core::service::remote_ssh::SSHConnectionManager; -use bitfun_core::service::session::{DialogTurnData, SessionMetadata, SessionStatus}; +use bitfun_core::service::session::{ + DialogTurnData, DialogTurnKind, SessionMetadata, SessionStatus, SessionTranscriptExport, + SessionTranscriptExportOptions, +}; use bitfun_core::service::session_usage::SessionUsageReport; use bitfun_core::service::token_usage::TokenUsageService; use bitfun_core::service::workspace::WorkspaceService; @@ -90,6 +93,30 @@ fn desktop_runtime_session_error(error: RuntimeError) -> DesktopSessionApplicati } } +fn local_command_turn_record_request( + turn: &DialogTurnData, +) -> DesktopSessionApplicationResult> { + if turn.kind != DialogTurnKind::LocalCommand { + return Ok(None); + } + let metadata = match turn.user_message.metadata.clone() { + None => serde_json::Map::new(), + Some(serde_json::Value::Object(metadata)) => metadata, + Some(_) => { + return Err(DesktopSessionApplicationError::Validation( + "Local command Turn metadata must be an object".to_string(), + )); + } + }; + Ok(Some(AgentLocalCommandTurnRecordRequest { + session_id: turn.session_id.clone(), + content: turn.user_message.content.clone(), + turn_id: Some(turn.turn_id.clone()), + timestamp_ms: Some(turn.timestamp), + metadata, + })) +} + #[derive(Debug)] pub(crate) struct DesktopSessionViewRestore { pub session: Session, @@ -254,6 +281,10 @@ impl DesktopSessionApplication { &self.agent_runtime } + pub(crate) fn compatibility(&self) -> &CoreAgentRuntimeCompatibility { + &self.compatibility + } + pub(crate) async fn reload_session_context( &self, request: AgentContextReloadRequest, @@ -344,6 +375,20 @@ impl DesktopSessionApplication { .collect()) } + pub(crate) async fn export_session_transcript( + &self, + request: DesktopSessionScopeRequest, + session_id: &str, + options: &SessionTranscriptExportOptions, + ) -> DesktopSessionApplicationResult { + let scope = self.resolved_scope(request).await; + self.ensure_runtime_ownership(&scope)?; + self.compatibility + .export_persisted_session_transcript(&self.storage_path(&scope), session_id, options) + .await + .map_err(desktop_core_session_error) + } + pub(crate) async fn load_session_turns( &self, request: DesktopSessionScopeRequest, @@ -371,6 +416,36 @@ impl DesktopSessionApplication { .map_err(|error| DesktopSessionApplicationError::Core(error.to_string())) } + pub(crate) async fn save_session_turn( + &self, + request: DesktopSessionScopeRequest, + turn: &DialogTurnData, + ) -> DesktopSessionApplicationResult<()> { + let scope = self.resolved_scope(request).await; + self.ensure_runtime_ownership(&scope)?; + let storage_path = self.storage_path(&scope); + self.compatibility + .ensure_session_loaded_from_storage_path(&storage_path, &turn.session_id, false) + .await + .map_err(desktop_core_session_error)?; + if let Some(local_command) = local_command_turn_record_request(turn)? { + return self + .agent_runtime + .record_completed_local_command_turn(local_command) + .await + .map_err(desktop_runtime_session_error); + } + let mutation = self + .compatibility + .begin_persisted_session_mutation(&storage_path, &turn.session_id) + .await + .map_err(desktop_core_session_error)?; + self.compatibility + .save_persisted_dialog_turn(&mutation, turn) + .await + .map_err(desktop_core_session_error) + } + pub(crate) async fn touch_session( &self, request: DesktopSessionScopeRequest, @@ -791,6 +866,39 @@ mod tests { ); } + #[test] + fn desktop_local_usage_turn_maps_to_the_fixed_runtime_append_contract() { + let mut turn = DialogTurnData::new_with_kind( + DialogTurnKind::LocalCommand, + "local-usage-report-1".to_string(), + 2, + "session-1".to_string(), + None, + bitfun_core::service::session::UserMessageData { + id: "local-usage-user-report-1".to_string(), + content: "# Usage".to_string(), + timestamp: 42, + metadata: Some(json!({ + "localCommandKind": "usage_report", + "reportId": "report-1" + })), + }, + ); + turn.timestamp = 42; + + let request = local_command_turn_record_request(&turn) + .expect("valid local command") + .expect("local command request"); + assert_eq!(request.session_id, "session-1"); + assert_eq!(request.turn_id.as_deref(), Some("local-usage-report-1")); + assert_eq!(request.content, "# Usage"); + assert_eq!(request.timestamp_ms, Some(42)); + assert_eq!(request.metadata["localCommandKind"], "usage_report"); + + let source = include_str!("session_application.rs"); + assert!(source.contains("record_completed_local_command_turn(local_command)")); + } + struct RecordingDeletePort { events: Arc>>, workspace_path: Arc>>, diff --git a/src/crates/adapters/agent-runtime-ipc/AGENTS-CN.md b/src/crates/adapters/agent-runtime-ipc/AGENTS-CN.md index 3a680b345..80e4324e1 100644 --- a/src/crates/adapters/agent-runtime-ipc/AGENTS-CN.md +++ b/src/crates/adapters/agent-runtime-ipc/AGENTS-CN.md @@ -15,7 +15,7 @@ ## 边界 - 只导出 CLI adapter 实际使用的 workspace-private API，且 crate 不得发布，也不得把 wire 作为 SDK 合同。 -- 封闭 operation 范围为 Health、Session list/create/restore/delete/fork（restore/fork 结果包含 transcript）、当前 Session rename、Agent mode/model update 和手动 context compaction、声明式上下文 reload、Turn submit/cancel、pending/respond Permission 和 UserInput answers。delete 只允许作用于未被任何 Client 控制的空闲 Session。fork 要求当前 controller 且 Session 空闲：可以复制到最新持久化 Turn，也可以停在显式选中 Turn 之前；只有包含新 Session 与 transcript 的成功结果完成编码后，Server 才能把连接 lease 从源 Session 原子切换到 fork。手动 compaction 要求当前 controller 且 Session 空闲；Client 在准入前提供精确 Turn ID，使超时或断连 cleanup 可以取消同一个 owned task；Core 开始原子 context commit 后，晚到取消不能暴露错误的空闲状态。上下文 reload 可在活动 Turn 中执行，不改写该 Turn，并通过缓存保护保证下一条消息重新读取已失效的 instructions。断连 cleanup 属于内部生命周期，不是 detach operation。模型目录和默认值仍是 wire 之外的产品配置；禁止顺带加入 archive、replay、observer、通用 controller transfer、Tool/MCP/Hook 管理或其他产品配置。 +- 封闭 operation 范围为 Health、Session list/create/restore/delete/fork（restore/fork 结果包含 transcript）、当前 Session rename、Agent mode/model update、手动 context compaction、Session undo/redo、声明式上下文 reload、Turn submit/cancel、pending/respond Permission 和 UserInput answers。delete 只允许作用于未被任何 Client 控制的空闲 Session。fork 要求当前 controller 且 Session 空闲：可以复制到最新持久化 Turn，也可以停在显式选中 Turn 之前；只有包含新 Session 与 transcript 的成功结果完成编码后，Server 才能把连接 lease 从源 Session 原子切换到 fork。手动 compaction 要求当前 controller 且 Session 空闲；Client 在准入前提供精确 Turn ID，使超时或断连 cleanup 可以取消同一个 owned task；Core 开始原子 context commit 后，晚到取消不能暴露错误的空闲状态。undo/redo 要求当前 controller，但可在活动 Turn 中进入，因为取消、drain 与回退写入顺序由 Core 统一负责；成功结果携带权威 transcript，并清除连接侧活动 Turn 投影。该能力只支持本地工作区，不暴露通用 checkpoint 协议。上下文 reload 可在活动 Turn 中执行，不改写该 Turn，并通过缓存保护保证下一条消息重新读取已失效的 instructions。断连 cleanup 属于内部生命周期，不是 detach operation。模型目录和默认值仍是 wire 之外的产品配置；禁止顺带加入 archive、replay、observer、通用 controller transfer、Tool/MCP/Hook 管理或其他产品配置。 - 可以复用稳定 Event、Product Domain 和 Runtime Port DTO。禁止依赖 `bitfun-core`、Agent Runtime 实现、SDK Host、services、Tauri、terminal、tool runtime 或远程 transport。 - 只使用 Windows Named Pipe 或 Unix Domain Socket；禁止 TCP、HTTP、WebSocket、浏览器访问或远程 fallback。 - 这是本机同用户隔离，不是沙箱。未来产品 composition 必须提供当前用户私有 runtime 目录。 diff --git a/src/crates/adapters/agent-runtime-ipc/AGENTS.md b/src/crates/adapters/agent-runtime-ipc/AGENTS.md index 9bc72bf7b..973b69973 100644 --- a/src/crates/adapters/agent-runtime-ipc/AGENTS.md +++ b/src/crates/adapters/agent-runtime-ipc/AGENTS.md @@ -22,11 +22,12 @@ session controller leases, event delivery, connection bounds, and cleanup. It is - Export only the exact workspace-private API needed by the CLI adapter. Do not publish this crate or expose its wire as an SDK contract. -- The closed operation budget is Health, Session list/create/restore/delete/fork (including transcript on restore/fork), current-Session rename, Agent mode/model update and manual context compaction, +- The closed operation budget is Health, Session list/create/restore/delete/fork (including transcript on restore/fork), current-Session rename, Agent mode/model update, manual context compaction, and Session undo/redo, declarative context reload, Turn submit/cancel, pending/respond Permission, and UserInput answers. Delete is limited to an idle Session not controlled by any client. Fork is a current-controller, idle-only operation. It either copies through the latest persisted Turn or stops immediately before an explicitly selected Turn. The encoded success result carries the authoritative new Session and transcript; only then may the server atomically switch the connection lease from the source Session to the fork. Manual compaction is a current-controller, idle-only Turn operation. The client supplies its exact Turn ID before admission so timeout or disconnect cleanup can cancel the same owned task; once Core begins the atomic context commit, a late cancellation does not expose a false idle state. Context reload may run during an active Turn, does not rewrite that Turn, and guards the cache so the next message reads invalidated instructions. + Undo/redo is a current-controller operation that may enter during an active Turn because Core owns cancel-and-drain before mutation. Its success response carries the authoritative transcript and clears the connection's active-Turn projection. It is local-workspace only and does not expose a generic checkpoint protocol. Disconnect cleanup is internal lifecycle, not a detach operation. Model catalogs and defaults remain product configuration outside this wire. Do not add archive, replay, observer, general controller transfer, Tool/MCP/Hook management, or other product configuration incidentally. diff --git a/src/crates/adapters/agent-runtime-ipc/src/operation.rs b/src/crates/adapters/agent-runtime-ipc/src/operation.rs index 1a347a518..2405fc271 100644 --- a/src/crates/adapters/agent-runtime-ipc/src/operation.rs +++ b/src/crates/adapters/agent-runtime-ipc/src/operation.rs @@ -2,8 +2,9 @@ use bitfun_product_domains::tool_permissions::{PermissionReply, PermissionReques use bitfun_runtime_ports::{ AgentContextReloadRequest, AgentDialogTurnRequest, AgentSessionCompactionRequest, AgentSessionCreateRequest, AgentSessionCreateResult, AgentSessionListRequest, - AgentSessionModeUpdateRequest, AgentSessionModelUpdateRequest, AgentSessionSummary, - AgentTurnCancellationRequest, AgentTurnCancellationResult, SessionTranscript, + AgentSessionModeUpdateRequest, AgentSessionModelUpdateRequest, AgentSessionRevertRequest, + AgentSessionRevertResult, AgentSessionSummary, AgentTurnCancellationRequest, + AgentTurnCancellationResult, SessionTranscript, }; use serde::{Deserialize, Serialize}; @@ -78,6 +79,12 @@ pub enum RuntimeIpcOperation { CompactSession { request: AgentSessionCompactionRequest, }, + UndoSession { + request: AgentSessionRevertRequest, + }, + RedoSession { + request: AgentSessionRevertRequest, + }, SubmitTurn { request: AgentDialogTurnRequest, }, @@ -108,6 +115,8 @@ impl RuntimeIpcOperation { Self::ForkSession { request } => Some(&request.session_id), Self::ReloadSessionContext { request } => Some(&request.session_id), Self::CompactSession { request } => Some(&request.session_id), + Self::UndoSession { request } => Some(&request.session_id), + Self::RedoSession { request } => Some(&request.session_id), Self::SubmitTurn { request } => Some(&request.session_id), Self::CancelTurn { request } => Some(&request.session_id), Self::PendingPermissions { session_id } @@ -144,6 +153,8 @@ impl RuntimeIpcOperation { RuntimeIpcOperationRules::new(CurrentController, true, true, true) } Self::ReloadSessionContext { .. } + | Self::UndoSession { .. } + | Self::RedoSession { .. } | Self::CancelTurn { .. } | Self::RespondPermission { .. } | Self::SubmitUserAnswers { .. } => { @@ -216,6 +227,9 @@ pub enum RuntimeIpcOperationResult { session: AgentSessionSummary, transcript: SessionTranscript, }, + SessionReverted { + revert: AgentSessionRevertResult, + }, TurnAccepted { session_id: String, turn_id: String, diff --git a/src/crates/adapters/agent-runtime-ipc/src/server.rs b/src/crates/adapters/agent-runtime-ipc/src/server.rs index 432bf04aa..25bc0f357 100644 --- a/src/crates/adapters/agent-runtime-ipc/src/server.rs +++ b/src/crates/adapters/agent-runtime-ipc/src/server.rs @@ -526,24 +526,50 @@ async fn run_initialized_connection( continue; } }; + if matches!(result, RuntimeIpcOperationResult::SessionReverted { .. }) { + // Core has cancelled and drained the owned turn before applying + // the revert, so the connection must drop its provisional view. + *active_turn_id = None; + } let response = RuntimeIpcFrame::Response { request_id, result }; - let response_bytes = - match serialize_frame_with_limit(&response, MAX_RESPONSE_FRAME_BYTES) { - Err(RuntimeIpcIoError::FrameTooLarge { .. }) => { - config.leases.rollback(connection_id, lease_transition); - send_error( + let revert_response = matches!( + &response, + RuntimeIpcFrame::Response { + result: RuntimeIpcOperationResult::SessionReverted { .. }, + .. + } + ); + let response_bytes = match serialize_frame_with_limit( + &response, + MAX_RESPONSE_FRAME_BYTES, + ) { + Err(RuntimeIpcIoError::FrameTooLarge { .. }) if revert_response => { + config.leases.rollback(connection_id, lease_transition); + send_error( stream, config.request_timeout, Some(request_id), - RuntimeIpcErrorCode::FrameTooLarge, - "runtime IPC response exceeds the supported frame size", + RuntimeIpcErrorCode::OutcomeUnknown, + "session revert completed but its authoritative transcript exceeds the Shared TUI response limit; reconnect and restore the Session", ) .await?; - continue; - } - Err(error) => return Err(RuntimeIpcServerError::Io(error)), - Ok(bytes) => bytes, - }; + return Err(RuntimeIpcServerError::Disconnected); + } + Err(RuntimeIpcIoError::FrameTooLarge { .. }) => { + config.leases.rollback(connection_id, lease_transition); + send_error( + stream, + config.request_timeout, + Some(request_id), + RuntimeIpcErrorCode::FrameTooLarge, + "runtime IPC response exceeds the supported frame size", + ) + .await?; + continue; + } + Err(error) => return Err(RuntimeIpcServerError::Io(error)), + Ok(bytes) => bytes, + }; let RuntimeIpcFrame::Response { result, .. } = &response else { unreachable!("response frame was just constructed") diff --git a/src/crates/adapters/agent-runtime-ipc/src/tests/protocol_contracts.rs b/src/crates/adapters/agent-runtime-ipc/src/tests/protocol_contracts.rs index e34900b27..ba35e4d6e 100644 --- a/src/crates/adapters/agent-runtime-ipc/src/tests/protocol_contracts.rs +++ b/src/crates/adapters/agent-runtime-ipc/src/tests/protocol_contracts.rs @@ -9,7 +9,7 @@ use bitfun_product_domains::tool_permissions::PermissionReply; use bitfun_runtime_ports::{ AgentContextReloadRequest, AgentContextReloadTarget, AgentDialogTurnRequest, AgentSessionCompactionRequest, AgentSessionModeUpdateRequest, AgentSessionModelUpdateRequest, - AgentSubmissionSource, DialogSubmissionPolicy, + AgentSessionRevertRequest, AgentSubmissionSource, DialogSubmissionPolicy, }; use serde_json::{json, Map}; @@ -219,6 +219,35 @@ fn protocol_round_trips_manual_compaction_as_an_idle_controller_turn() { assert!(rules.side_effecting); } +#[test] +fn protocol_round_trips_undo_as_an_active_controller_operation() { + let operation = RuntimeIpcOperation::UndoSession { + request: AgentSessionRevertRequest { + workspace_path: "D:/workspace/project".to_string(), + session_id: "session-1".to_string(), + remote_connection_id: None, + remote_ssh_host: None, + }, + }; + + let encoded = serde_json::to_value(&operation).expect("serialize session undo"); + assert_eq!(encoded["operation"], "undo_session"); + assert_eq!(encoded["request"]["sessionId"], "session-1"); + let decoded: RuntimeIpcOperation = + serde_json::from_value(encoded).expect("deserialize session undo"); + + assert_eq!(decoded, operation); + assert_eq!(decoded.session_id(), Some("session-1")); + let rules = decoded.rules(); + assert_eq!( + rules.session_requirement, + RuntimeIpcSessionRequirement::CurrentController + ); + assert!(!rules.requires_idle); + assert!(!rules.serializes_session_selection); + assert!(rules.side_effecting); +} + #[test] fn protocol_round_trips_session_delete_and_not_found() { let operation = RuntimeIpcOperation::DeleteSession { diff --git a/src/crates/adapters/agent-runtime-ipc/src/tests/shared_controller.rs b/src/crates/adapters/agent-runtime-ipc/src/tests/shared_controller.rs index c44b2bec1..d7c1d005f 100644 --- a/src/crates/adapters/agent-runtime-ipc/src/tests/shared_controller.rs +++ b/src/crates/adapters/agent-runtime-ipc/src/tests/shared_controller.rs @@ -8,8 +8,9 @@ use crate::{ use async_trait::async_trait; use bitfun_events::{AgenticEvent, AgenticEventEnvelope, AgenticEventPriority}; use bitfun_runtime_ports::{ - AgentDialogTurnRequest, AgentSessionCompactionRequest, AgentSessionCreateRequest, - AgentSessionCreateResult, AgentSessionModeUpdateRequest, AgentSessionModelUpdateRequest, + AgentDialogTurnRequest, AgentSessionCompactionRequest, AgentSessionComposerUpdate, + AgentSessionCreateRequest, AgentSessionCreateResult, AgentSessionModeUpdateRequest, + AgentSessionModelUpdateRequest, AgentSessionRevertRequest, AgentSessionRevertResult, AgentSessionSummary, AgentSubmissionSource, DialogSubmissionPolicy, SessionTranscript, }; use serde_json::Map; @@ -227,6 +228,22 @@ impl RuntimeIpcRequestHandler for FakeHandler { turn_id: request.turn_id, }) } + RuntimeIpcOperation::UndoSession { request } + | RuntimeIpcOperation::RedoSession { request } => { + Ok(RuntimeIpcOperationResult::SessionReverted { + revert: AgentSessionRevertResult { + transcript: SessionTranscript { + session_id: request.session_id.clone(), + messages: Vec::new(), + }, + session_id: request.session_id, + composer: AgentSessionComposerUpdate::Preserve, + retired_turn_ids: Vec::new(), + changed: true, + hidden_turn_count: 1, + }, + }) + } RuntimeIpcOperation::CancelTurn { request } => { if self.settle_cancel { let _ = self.events.send(test_agent_event( @@ -537,6 +554,17 @@ fn compact_operation(session_id: &str, turn_id: &str) -> RuntimeIpcOperation { } } +fn undo_operation(workspace: &Path, session_id: &str) -> RuntimeIpcOperation { + RuntimeIpcOperation::UndoSession { + request: AgentSessionRevertRequest { + workspace_path: workspace.to_string_lossy().to_string(), + session_id: session_id.to_string(), + remote_connection_id: None, + remote_ssh_host: None, + }, + } +} + fn update_mode_operation(session_id: &str, mode_id: &str) -> RuntimeIpcOperation { RuntimeIpcOperation::UpdateSessionMode { request: AgentSessionModeUpdateRequest { @@ -1138,6 +1166,41 @@ async fn rename_requires_the_controlled_idle_session() { server.finish().await; } +#[tokio::test] +async fn undo_can_cancel_the_controlled_active_turn_and_clears_its_projection() { + let handler = Arc::new(FakeHandler::default()); + let server = TestServer::start(server_config(), handler.clone()).await; + let mut client = server.connect("undo-controller").await; + + expect_response( + &mut client, + 2, + restore_operation(server.workspace.path(), "session-a"), + ) + .await; + expect_response( + &mut client, + 3, + submit_operation(server.workspace.path(), "session-a", "turn-a"), + ) + .await; + expect_response( + &mut client, + 4, + undo_operation(server.workspace.path(), "session-a"), + ) + .await; + expect_response(&mut client, 5, rename_operation("session-a", "After undo")).await; + + let calls = handler.calls.lock().expect("calls"); + assert!(calls + .iter() + .any(|operation| matches!(operation, RuntimeIpcOperation::UndoSession { .. }))); + drop(calls); + drop(client); + server.finish().await; +} + #[tokio::test] async fn delete_requires_an_uncontrolled_target_and_an_idle_connection() { let handler = Arc::new(FakeHandler::default()); diff --git a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs index 8455331b4..db7fe8fe7 100644 --- a/src/crates/assembly/core/src/agentic/coordination/coordinator.rs +++ b/src/crates/assembly/core/src/agentic/coordination/coordinator.rs @@ -36,6 +36,9 @@ use crate::agentic::goal_mode::{ use crate::agentic::image_analysis::ImageContextData; use crate::agentic::memories::{start_memory_startup_task, MemoryStartupRequest}; use crate::agentic::round_preempt::DialogRoundInjectionSource; +use crate::agentic::session::revert::{ + resolve_redo, resolve_undo, SessionRevertPhase, SessionRevertTransition, +}; use crate::agentic::session::session_store_port::CoreSessionStorePort; use crate::agentic::session::{SessionManager, SessionReferenceLocator}; use crate::agentic::side_question::build_btw_user_input; @@ -82,10 +85,11 @@ use bitfun_agent_runtime::remote_file_delivery::{ use bitfun_agent_runtime::sdk::PermissionReply; use bitfun_agent_runtime::user_questions::USER_INPUT_AVAILABLE_CONTEXT_KEY; use bitfun_runtime_ports::{ - AgentSessionWorkspaceBinding, AgentThreadGoalDeliveryKind, AgentThreadGoalDeliveryRequest, - DelegationPolicy, PermissionDelegationContext, PermissionRuntimeCeiling, RemoteExecPort, - SessionStoragePathRequest, SessionStoragePathResolution, SessionStorePort, SubagentContextMode, - TerminalPort, ThreadGoal, ThreadGoalContinuationPlan, ThreadGoalStatus, + AgentSessionComposerUpdate, AgentSessionWorkspaceBinding, AgentThreadGoalDeliveryKind, + AgentThreadGoalDeliveryRequest, DelegationPolicy, PermissionDelegationContext, + PermissionRuntimeCeiling, RemoteExecPort, SessionStoragePathRequest, + SessionStoragePathResolution, SessionStorePort, SubagentContextMode, TerminalPort, ThreadGoal, + ThreadGoalContinuationPlan, ThreadGoalStatus, }; use dashmap::DashMap; use log::{debug, error, info, warn}; @@ -2651,7 +2655,6 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet match self.restore_path_for_existing_session(session_id).await { Ok(restore_path) => { match self - .session_manager .restore_session_from_storage_path(&restore_path, session_id) .await { @@ -2807,11 +2810,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet let session = match self.session_manager.get_session(&session_id) { Some(session) => session, - None => { - self.session_manager - .restore_session(&workspace_root, &session_id) - .await? - } + None => self.restore_session(&workspace_root, &session_id).await?, }; let turn_count = self.session_manager.get_turn_count(&session_id); @@ -3630,6 +3629,10 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet "Manual compaction turn_id must not be empty".to_string(), )); } + let mutation_guard = self + .session_manager + .acquire_session_mutation(&session_id) + .await?; let initial_session = self .session_manager .get_session(&session_id) @@ -3663,16 +3666,19 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet ))); } + self.commit_session_revert_before_persisted_turn_locked(&session_id, "Manual compaction") + .await?; let user_message_metadata = Some(Self::manual_compaction_metadata()); let turn_id = self .session_manager - .start_maintenance_turn( + .start_maintenance_turn_locked( &session_id, MANUAL_COMPACTION_COMMAND.to_string(), requested_turn_id, user_message_metadata.clone(), ) .await?; + drop(mutation_guard); // Once the maintenance turn owns Processing, competing dialog turns // can no longer mutate context. Capture the authoritative context only // after that atomic admission so a just-completed turn cannot be lost. @@ -4050,8 +4056,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet if !restore.is_remote_storage() { self.ensure_runtime_ownership(&restore.requested_workspace_path, None, None)?; } - self.session_manager - .restore_session_from_storage_path(&restore.effective_storage_path, &session_id) + self.restore_session_from_storage_path(&restore.effective_storage_path, &session_id) .await? } }; @@ -4251,7 +4256,6 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet ) .await?; match self - .session_manager .restore_session_from_storage_path(&restore_path, &session_id) .await { @@ -5233,6 +5237,32 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet workspace_path: &Path, session_id: &str, ) -> BitFunResult<()> { + let session_storage_path = self + .session_manager + .resolve_storage_path_for_workspace_path(workspace_path) + .await; + let has_revert_state = self + .session_manager + .persistence_manager() + .load_session_revert_state(&session_storage_path, session_id) + .await? + .is_some(); + if has_revert_state + && !self + .session_manager + .is_session_loaded_from_storage_path(&session_storage_path, session_id)? + { + self.restore_internal_session_from_storage_path(&session_storage_path, session_id) + .await?; + } + let _mutation_guard = self + .session_manager + .acquire_session_mutation(session_id) + .await?; + self.session_manager + .validate_session_storage_path_binding(session_id, &session_storage_path)?; + self.reconcile_session_revert_locked(&session_storage_path, session_id) + .await?; // SessionEnd hooks observe the session before its state is gone. // Their timeout is capped tightly so deletion cannot hang. let session_hook_facts = match self.session_manager.get_session(session_id) { @@ -5259,7 +5289,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet native_hooks::clear_session_hook_state(session_id); } self.session_manager - .delete_session(workspace_path, session_id) + .delete_session_locked(workspace_path, session_id) .await?; self.background_subagent_outcomes .delete_session_references(session_id) @@ -5387,122 +5417,591 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet session_id: &str, ) -> BitFunResult { self.ensure_runtime_ownership(workspace_path, None, None)?; - self.session_manager + let session = self + .session_manager .restore_session(workspace_path, session_id) - .await + .await?; + self.reconcile_restored_session(session_id, session).await } - pub async fn restore_session_from_storage_path( - &self, - session_storage_path: &Path, - session_id: &str, - ) -> BitFunResult { - self.session_manager - .restore_session_from_storage_path(session_storage_path, session_id) - .await + pub(crate) fn local_revert_workspace(&self, session_id: &str) -> BitFunResult { + let session = self + .session_manager + .get_session(session_id) + .ok_or_else(|| BitFunError::NotFound(format!("Session not found: {session_id}")))?; + if session.config.remote_connection_id.is_some() || session.config.remote_ssh_host.is_some() + { + return Err(BitFunError::Validation( + "Session undo and redo are unavailable for remote workspaces".to_string(), + )); + } + let workspace_path = session.config.workspace_path.as_deref().ok_or_else(|| { + BitFunError::Validation(format!("Session workspace_path is missing: {session_id}")) + })?; + let workspace_path = PathBuf::from(workspace_path); + if !workspace_path.is_dir() { + return Err(BitFunError::Validation(format!( + "Session workspace directory does not exist: {}", + workspace_path.display() + ))); + } + Ok(workspace_path) } - pub async fn restore_internal_session_from_storage_path( + pub(crate) async fn apply_session_revert_locked( &self, session_storage_path: &Path, session_id: &str, - ) -> BitFunResult { - self.session_manager - .restore_internal_session_from_storage_path(session_storage_path, session_id) - .await - } + undo: bool, + ) -> BitFunResult<(AgentSessionComposerUpdate, bool, usize)> { + let workspace_path = self.local_revert_workspace(session_id)?; + let snapshot_manager = + crate::service::snapshot::get_or_create_snapshot_manager(workspace_path.clone(), None) + .await + .map_err(|error| BitFunError::service(error.to_string()))?; + let persistence = self.session_manager.persistence_manager(); + let mut current = persistence + .load_session_revert_state(session_storage_path, session_id) + .await?; + if current + .as_ref() + .is_some_and(|state| state.phase != SessionRevertPhase::Staged) + { + self.reconcile_session_revert_locked(session_storage_path, session_id) + .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "Session revert could not finish a pending transition: session_id={session_id}, error={error}" + )) + })?; + current = persistence + .load_session_revert_state(session_storage_path, session_id) + .await?; + } + let turns = persistence + .load_session_turns(session_storage_path, session_id) + .await?; + let transition = if undo { + resolve_undo(&turns, current.as_ref()) + } else { + resolve_redo(&turns, current.as_ref()) + }; + let Some(transition) = transition else { + return Ok((AgentSessionComposerUpdate::Preserve, false, 0)); + }; - pub async fn restore_session_for_workspace( - &self, - request: SessionStoragePathRequest, - session_id: &str, - ) -> BitFunResult { - self.ensure_runtime_ownership( - &request.workspace_path, - request.remote_connection_id.as_deref(), - request.remote_ssh_host.as_deref(), - )?; - self.session_manager - .restore_session_for_workspace(request, session_id) - .await + match transition { + SessionRevertTransition::Stage { + mut state, + replacement_prompt, + hidden_turn_count, + } => { + state.phase = SessionRevertPhase::Applying; + snapshot_manager + .prepare_workspace_revert(session_id, &mut state) + .await + .map_err(|error| BitFunError::service(error.to_string()))?; + persistence + .save_session_revert_state(session_storage_path, session_id, &state) + .await?; + snapshot_manager + .apply_workspace_revert(session_id, &state) + .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "Staged Session boundary was persisted but workspace reconciliation failed: session_id={session_id}, error={error}" + )) + })?; + self.session_manager + .apply_staged_revert_context_locked( + session_storage_path, + session_id, + state.boundary_turn, + ) + .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "Staged Session boundary and workspace were updated but runtime context reconciliation failed: session_id={session_id}, error={error}" + )) + })?; + state.phase = SessionRevertPhase::Staged; + persistence + .save_session_revert_state(session_storage_path, session_id, &state) + .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "Session boundary was applied but its stable phase could not be persisted: session_id={session_id}, error={error}" + )) + })?; + let composer = replacement_prompt + .map(|text| AgentSessionComposerUpdate::Replace { text }) + .unwrap_or(AgentSessionComposerUpdate::Preserve); + Ok((composer, true, hidden_turn_count)) + } + SessionRevertTransition::Clear { mut previous_state } => { + previous_state.boundary_turn = previous_state.original_turn_end; + previous_state.phase = SessionRevertPhase::Clearing; + persistence + .save_session_revert_state(session_storage_path, session_id, &previous_state) + .await?; + snapshot_manager + .apply_workspace_revert(session_id, &previous_state) + .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "Session redo may have partially restored the workspace: session_id={session_id}, error={error}" + )) + })?; + self.session_manager + .apply_staged_revert_context_locked( + session_storage_path, + session_id, + previous_state.original_turn_end, + ) + .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "Session redo restored the workspace but could not reconcile runtime context: session_id={session_id}, error={error}" + )) + })?; + persistence + .delete_session_revert_state(session_storage_path, session_id) + .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "Session redo restored history but could not clear its staged marker: session_id={session_id}, error={error}" + )) + })?; + if let Err(error) = snapshot_manager + .delete_workspace_revert_checkpoint(&previous_state) + .await + { + warn!( + "Failed to delete cleared Session revert checkpoint: session_id={}, error={}", + session_id, error + ); + } + Ok((AgentSessionComposerUpdate::Clear, true, 0)) + } + } } - pub async fn restore_internal_session_for_workspace( + pub(crate) async fn reconcile_session_revert_locked( &self, - request: SessionStoragePathRequest, + session_storage_path: &Path, session_id: &str, - ) -> BitFunResult { - self.ensure_runtime_ownership( - &request.workspace_path, - request.remote_connection_id.as_deref(), - request.remote_ssh_host.as_deref(), - )?; - self.session_manager - .restore_internal_session_for_workspace(request, session_id) - .await + ) -> BitFunResult<()> { + let persistence = self.session_manager.persistence_manager(); + let Some(state) = persistence + .load_session_revert_state(session_storage_path, session_id) + .await? + else { + return Ok(()); + }; + match state.phase { + SessionRevertPhase::Committing => self + .commit_session_revert_locked(session_storage_path, session_id) + .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "Session restore could not finish a pending revert commit: session_id={session_id}, error={error}" + )) + }), + SessionRevertPhase::Staged => self + .session_manager + .apply_staged_revert_context_locked( + session_storage_path, + session_id, + state.boundary_turn, + ) + .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "Session restore could not reconcile staged runtime context: session_id={session_id}, error={error}" + )) + }), + SessionRevertPhase::Applying | SessionRevertPhase::Clearing => self + .reconcile_session_revert_application_locked( + session_storage_path, + session_id, + state, + ) + .await, + } } - pub async fn restore_internal_session( + async fn prepare_persisted_session_read_locked( &self, - workspace_path: &Path, + session_storage_path: &Path, session_id: &str, - ) -> BitFunResult { - self.ensure_runtime_ownership(workspace_path, None, None)?; + ) -> BitFunResult<()> { + bitfun_core_types::validate_session_id(session_id).map_err(BitFunError::Validation)?; self.session_manager - .restore_internal_session(workspace_path, session_id) - .await + .validate_session_storage_path_binding(session_id, session_storage_path)?; + if let Some(state) = self + .session_manager + .persistence_manager() + .load_session_revert_state(session_storage_path, session_id) + .await? + { + if state.phase != SessionRevertPhase::Staged { + if self.session_manager.get_session(session_id).is_none() { + return Err(BitFunError::OutcomeUnknown(format!( + "Session history is unavailable until the unfinished undo transition is restored: session_id={session_id}" + ))); + } + self.reconcile_session_revert_locked(session_storage_path, session_id) + .await?; + } + } + Ok(()) } - /// Restore session and return the persisted turns read during restore. - pub async fn restore_session_with_turns( + /// Read the product-visible persisted Turn history through Core's + /// per-Session mutation owner. Persistence supplies cross-process + /// exclusion; this keyed guard supplies the missing in-process ordering + /// against undo, redo, commit, and external history imports. + pub async fn load_visible_persisted_session_turns( &self, - workspace_path: &Path, + session_storage_path: &Path, session_id: &str, - ) -> BitFunResult<(Session, Vec)> { - self.ensure_runtime_ownership(workspace_path, None, None)?; + ) -> BitFunResult> { + let _mutation = self + .session_manager + .acquire_session_mutation(session_id) + .await?; + self.prepare_persisted_session_read_locked(session_storage_path, session_id) + .await?; self.session_manager - .restore_session_with_turns(workspace_path, session_id) + .persistence_manager() + .load_visible_session_turns(session_storage_path, session_id) .await } - pub async fn restore_session_with_turns_from_storage_path( + /// Export a transcript while retaining the same Session history boundary + /// from marker admission through artifact generation. + pub async fn export_visible_persisted_session_transcript( &self, session_storage_path: &Path, session_id: &str, - ) -> BitFunResult<(Session, Vec)> { + options: &crate::service::session::SessionTranscriptExportOptions, + ) -> BitFunResult { + let _mutation = self + .session_manager + .acquire_session_mutation(session_id) + .await?; + self.prepare_persisted_session_read_locked(session_storage_path, session_id) + .await?; self.session_manager - .restore_session_with_turns_from_storage_path(session_storage_path, session_id) + .persistence_manager() + .export_session_transcript(session_storage_path, session_id, options) .await } - pub async fn restore_internal_session_with_turns_from_storage_path( + async fn reconcile_session_revert_application_locked( &self, session_storage_path: &Path, session_id: &str, - ) -> BitFunResult<(Session, Vec)> { + state: crate::agentic::session::revert::SessionRevertState, + ) -> BitFunResult<()> { + let persistence = self.session_manager.persistence_manager(); + let workspace_path = self.local_revert_workspace(session_id)?; + let snapshot_manager = + crate::service::snapshot::get_or_create_snapshot_manager(workspace_path.clone(), None) + .await + .map_err(|error| BitFunError::service(error.to_string()))?; + snapshot_manager + .apply_workspace_revert(session_id, &state) + .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "Session restore could not reconcile an applying workspace boundary: session_id={session_id}, error={error}" + )) + })?; self.session_manager - .restore_internal_session_with_turns_from_storage_path(session_storage_path, session_id) + .apply_staged_revert_context_locked( + session_storage_path, + session_id, + state.boundary_turn, + ) .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "Session restore reconciled the workspace but not runtime context: session_id={session_id}, error={error}" + )) + })?; + if state.phase == SessionRevertPhase::Clearing { + persistence + .delete_session_revert_state(session_storage_path, session_id) + .await?; + if let Err(error) = snapshot_manager + .delete_workspace_revert_checkpoint(&state) + .await + { + warn!( + "Failed to delete recovered Session revert checkpoint: session_id={}, error={}", + session_id, error + ); + } + } else { + let mut staged = state; + staged.phase = SessionRevertPhase::Staged; + persistence + .save_session_revert_state(session_storage_path, session_id, &staged) + .await?; + } + Ok(()) } - pub async fn restore_session_with_turns_for_workspace( + pub(crate) async fn commit_session_revert_locked( &self, - request: SessionStoragePathRequest, + session_storage_path: &Path, session_id: &str, - ) -> BitFunResult<(Session, Vec)> { - self.ensure_runtime_ownership( - &request.workspace_path, - request.remote_connection_id.as_deref(), - request.remote_ssh_host.as_deref(), - )?; + ) -> BitFunResult<()> { + let persistence = self.session_manager.persistence_manager(); + let Some(mut state) = persistence + .load_session_revert_state(session_storage_path, session_id) + .await? + else { + return Ok(()); + }; + if matches!( + state.phase, + SessionRevertPhase::Applying | SessionRevertPhase::Clearing + ) { + self.reconcile_session_revert_application_locked( + session_storage_path, + session_id, + state.clone(), + ) + .await?; + let Some(reconciled) = persistence + .load_session_revert_state(session_storage_path, session_id) + .await? + else { + return Ok(()); + }; + state = reconciled; + } + let workspace_path = self.local_revert_workspace(session_id)?; + let snapshot_manager = + crate::service::snapshot::get_or_create_snapshot_manager(workspace_path.clone(), None) + .await + .map_err(|error| BitFunError::service(error.to_string()))?; + if state.phase == SessionRevertPhase::Staged { + state.phase = SessionRevertPhase::Committing; + persistence + .save_session_revert_state(session_storage_path, session_id, &state) + .await?; + } + let discarded_parent_turn_ids = persistence + .load_session_turns(session_storage_path, session_id) + .await? + .into_iter() + .filter(|turn| turn.turn_index >= state.boundary_turn) + .map(|turn| turn.turn_id) + .collect::>(); + if !discarded_parent_turn_ids.is_empty() { + Box::pin(self.delete_hidden_subagent_sessions_for_parent_turns( + &workspace_path, + session_id, + &discarded_parent_turn_ids, + )) + .await?; + } self.session_manager - .restore_session_with_turns_for_workspace(request, session_id) + .commit_staged_revert_context_locked( + session_storage_path, + session_id, + state.boundary_turn, + ) + .await?; + snapshot_manager + .commit_workspace_revert(session_id, &state) + .await + .map_err(|error| BitFunError::service(error.to_string()))?; + persistence + .delete_session_revert_state(session_storage_path, session_id) .await } - pub async fn restore_internal_session_with_turns_for_workspace( + async fn commit_session_revert_before_persisted_turn_locked( &self, - request: SessionStoragePathRequest, + session_id: &str, + operation: &str, + ) -> BitFunResult<()> { + let Some(session_storage_path) = self + .session_manager + .effective_session_storage_path(session_id) + .await + else { + return Ok(()); + }; + if self + .session_manager + .persistence_manager() + .load_session_revert_state(&session_storage_path, session_id) + .await? + .is_none() + { + return Ok(()); + } + self.commit_session_revert_locked(&session_storage_path, session_id) + .await + .map_err(|error| { + BitFunError::OutcomeUnknown(format!( + "{operation} was not admitted because the staged Session suffix could not be committed safely: session_id={session_id}, error={error}" + )) + }) + } + + /// Commit an existing staged boundary before the scheduler admits a new + /// user Turn. The scheduler's per-Session operation lock must already be held. + pub(crate) async fn commit_session_revert_before_submission( + &self, + session_id: &str, + ) -> BitFunResult<()> { + let _mutation_guard = self + .session_manager + .acquire_session_mutation(session_id) + .await?; + self.commit_session_revert_before_persisted_turn_locked(session_id, "A new Turn") + .await + } + + async fn reconcile_restored_session( + &self, + session_id: &str, + restored: T, + ) -> BitFunResult { + let session_storage_path = self + .session_manager + .effective_session_storage_path(session_id) + .await + .ok_or_else(|| { + BitFunError::NotFound(format!("Session storage path not found: {session_id}")) + })?; + let _mutation_guard = self + .session_manager + .acquire_session_mutation(session_id) + .await?; + self.reconcile_session_revert_locked(&session_storage_path, session_id) + .await?; + Ok(restored) + } + + pub async fn restore_session_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult { + let session = self + .session_manager + .restore_session_from_storage_path(session_storage_path, session_id) + .await?; + self.reconcile_restored_session(session_id, session).await + } + + pub async fn restore_internal_session_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult { + let session = self + .session_manager + .restore_internal_session_from_storage_path(session_storage_path, session_id) + .await?; + self.reconcile_restored_session(session_id, session).await + } + + pub async fn restore_session_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult { + self.ensure_runtime_ownership( + &request.workspace_path, + request.remote_connection_id.as_deref(), + request.remote_ssh_host.as_deref(), + )?; + let session = self + .session_manager + .restore_session_for_workspace(request, session_id) + .await?; + self.reconcile_restored_session(session_id, session).await + } + + pub async fn restore_internal_session_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult { + self.ensure_runtime_ownership( + &request.workspace_path, + request.remote_connection_id.as_deref(), + request.remote_ssh_host.as_deref(), + )?; + let session = self + .session_manager + .restore_internal_session_for_workspace(request, session_id) + .await?; + self.reconcile_restored_session(session_id, session).await + } + + pub async fn restore_internal_session( + &self, + workspace_path: &Path, + session_id: &str, + ) -> BitFunResult { + self.ensure_runtime_ownership(workspace_path, None, None)?; + let session = self + .session_manager + .restore_internal_session(workspace_path, session_id) + .await?; + self.reconcile_restored_session(session_id, session).await + } + + /// Restore session and return the persisted turns read during restore. + pub async fn restore_session_with_turns( + &self, + workspace_path: &Path, + session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + self.ensure_runtime_ownership(workspace_path, None, None)?; + let restored = self + .session_manager + .restore_session_with_turns(workspace_path, session_id) + .await?; + self.reconcile_restored_session(session_id, restored).await + } + + pub async fn restore_session_with_turns_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + let restored = self + .session_manager + .restore_session_with_turns_from_storage_path(session_storage_path, session_id) + .await?; + self.reconcile_restored_session(session_id, restored).await + } + + pub async fn restore_internal_session_with_turns_from_storage_path( + &self, + session_storage_path: &Path, + session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + let restored = self + .session_manager + .restore_internal_session_with_turns_from_storage_path(session_storage_path, session_id) + .await?; + self.reconcile_restored_session(session_id, restored).await + } + + pub async fn restore_session_with_turns_for_workspace( + &self, + request: SessionStoragePathRequest, session_id: &str, ) -> BitFunResult<(Session, Vec)> { self.ensure_runtime_ownership( @@ -5510,9 +6009,28 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet request.remote_connection_id.as_deref(), request.remote_ssh_host.as_deref(), )?; - self.session_manager + let restored = self + .session_manager + .restore_session_with_turns_for_workspace(request, session_id) + .await?; + self.reconcile_restored_session(session_id, restored).await + } + + pub async fn restore_internal_session_with_turns_for_workspace( + &self, + request: SessionStoragePathRequest, + session_id: &str, + ) -> BitFunResult<(Session, Vec)> { + self.ensure_runtime_ownership( + &request.workspace_path, + request.remote_connection_id.as_deref(), + request.remote_ssh_host.as_deref(), + )?; + let restored = self + .session_manager .restore_internal_session_with_turns_for_workspace(request, session_id) - .await + .await?; + self.reconcile_restored_session(session_id, restored).await } pub async fn restore_internal_session_with_turns( @@ -5521,9 +6039,11 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet session_id: &str, ) -> BitFunResult<(Session, Vec)> { self.ensure_runtime_ownership(workspace_path, None, None)?; - self.session_manager + let restored = self + .session_manager .restore_internal_session_with_turns(workspace_path, session_id) - .await + .await?; + self.reconcile_restored_session(session_id, restored).await } /// Restore only the UI-visible persisted session view. @@ -7354,12 +7874,11 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet .to_string(), )); } - self.session_manager - .restore_internal_session_from_storage_path( - &binding.session_storage_dir(), - target_session_id, - ) - .await? + self.restore_internal_session_from_storage_path( + &binding.session_storage_dir(), + target_session_id, + ) + .await? } }; @@ -7503,8 +8022,7 @@ Update the persona files and delete BOOTSTRAP.md as soon as bootstrap is complet if needs_restore { let restore_path = self.restore_path_for_existing_session(session_id).await?; - self.session_manager - .restore_internal_session_from_storage_path(&restore_path, session_id) + self.restore_internal_session_from_storage_path(&restore_path, session_id) .await?; context_messages = self .session_manager @@ -9629,20 +10147,42 @@ impl bitfun_runtime_ports::AgentLocalCommandTurnPort for ConversationCoordinator ) -> bitfun_runtime_ports::PortResult<()> { self.ensure_session_runtime_ownership(&request.session_id, None) .map_err(runtime_port_error_preserving_message)?; + let mutation_guard = self + .session_manager + .acquire_session_mutation(&request.session_id) + .await + .map_err(runtime_port_error_preserving_message)?; + self.session_manager + .get_session(&request.session_id) + .ok_or_else(|| { + runtime_port_error_preserving_message(BitFunError::NotFound(format!( + "Session not found: {}", + request.session_id + ))) + })?; + self.commit_session_revert_before_persisted_turn_locked( + &request.session_id, + "Local command Turn", + ) + .await + .map_err(runtime_port_error_preserving_message)?; let metadata = if request.metadata.is_empty() { None } else { Some(serde_json::Value::Object(request.metadata)) }; - self.get_session_manager() - .append_completed_local_command_turn( + let result = self + .session_manager + .append_completed_local_command_turn_locked( &request.session_id, request.content, request.turn_id, request.timestamp_ms, metadata, ) - .await + .await; + drop(mutation_guard); + result .map(|_| ()) .map_err(runtime_port_error_preserving_message) } @@ -9857,15 +10397,14 @@ impl bitfun_runtime_ports::RemoteControlStatePort for ConversationCoordinator { } } -#[async_trait::async_trait] -impl bitfun_runtime_ports::SessionTranscriptReader for ConversationCoordinator { - async fn read_session_transcript( +impl ConversationCoordinator { + pub(crate) async fn read_session_transcript_locked( &self, request: bitfun_runtime_ports::SessionTranscriptRequest, ) -> bitfun_runtime_ports::PortResult { let messages = match self .session_manager - .load_persisted_transcript_turns(&request.session_id) + .load_persisted_transcript_turns_locked(&request.session_id) .await .map_err(runtime_port_error_preserving_message)? { @@ -9873,7 +10412,8 @@ impl bitfun_runtime_ports::SessionTranscriptReader for ConversationCoordinator { runtime_transcript_messages_from_turns(&turns, request.turn_id.as_deref()) } None => self - .get_messages(&request.session_id) + .session_manager + .get_context_messages(&request.session_id) .await .map_err(runtime_port_error_preserving_message)? .into_iter() @@ -9892,6 +10432,53 @@ impl bitfun_runtime_ports::SessionTranscriptReader for ConversationCoordinator { } } +#[async_trait::async_trait] +impl bitfun_runtime_ports::SessionTranscriptReader for ConversationCoordinator { + async fn read_session_transcript( + &self, + request: bitfun_runtime_ports::SessionTranscriptRequest, + ) -> bitfun_runtime_ports::PortResult { + let _mutation = self + .session_manager + .acquire_session_mutation(&request.session_id) + .await + .map_err(runtime_port_error_preserving_message)?; + if let Some(storage_path) = self + .session_manager + .effective_session_storage_path(&request.session_id) + .await + { + self.session_manager + .validate_session_storage_path_binding(&request.session_id, &storage_path) + .map_err(runtime_port_error_preserving_message)?; + if let Some(state) = self + .session_manager + .persistence_manager() + .load_session_revert_state(&storage_path, &request.session_id) + .await + .map_err(runtime_port_error_preserving_message)? + { + if state.phase != SessionRevertPhase::Staged { + if self + .session_manager + .get_session(&request.session_id) + .is_none() + { + return Err(bitfun_runtime_ports::PortError::new( + bitfun_runtime_ports::PortErrorKind::OutcomeUnknown, + "Session transcript is unavailable until the unfinished undo transition is restored", + )); + } + self.reconcile_session_revert_locked(&storage_path, &request.session_id) + .await + .map_err(runtime_port_error_preserving_message)?; + } + } + } + self.read_session_transcript_locked(request).await + } +} + async fn is_ai_session_title_generation_enabled() -> bool { match crate::service::config::get_global_config_service().await { Ok(service) => service @@ -10010,7 +10597,7 @@ mod tests { use crate::agentic::core::{ InternalReminderKind, Message, MessageContent, MessageRole, MessageSemanticKind, SessionConfig, SessionContinuationPolicy, SessionKind, SessionModelBindingPolicy, - SessionState, + SessionState, TurnStats, }; use crate::agentic::events::{AgenticEvent, EventQueue, EventQueueConfig, EventRouter}; use crate::agentic::execution::{ @@ -10053,7 +10640,8 @@ mod tests { use crate::service::config::{AgentModelDefaultsConfig, SubagentModelSelection}; use crate::service::remote_ssh::workspace_state::init_remote_workspace_manager; use crate::service::session::{ - DialogTurnData, DialogTurnKind, SessionMetadata, SessionStatus, TurnStatus, UserMessageData, + DialogTurnData, DialogTurnKind, SessionMetadata, SessionRelationship, SessionStatus, + TurnStatus, UserMessageData, }; use crate::service::workspace::WorkspaceKind; use bitfun_agent_runtime::permission::AUTO_APPROVE_ASK_CONTEXT_KEY; @@ -10061,11 +10649,12 @@ mod tests { SessionExecutionTarget, SessionExecutionTargetKind, WorktreeLifecycle, }; use bitfun_runtime_ports::{ - AgentSessionArchiveRequest, AgentSessionCreateRequest, AgentSessionManagementPort, - AgentSessionRenameRequest, AgentSubmissionPort, AgentSubmissionRequest, - AgentSubmissionSource, AgentThreadGoalGetRequest, AgentThreadGoalManagementPort, - DelegationPolicy, PermissionEffect, PermissionRule, PermissionRuntimeCeiling, - SessionStoragePathRequest, SubagentContextMode, ThreadGoal, ThreadGoalStatus, + AgentLocalCommandTurnPort, AgentLocalCommandTurnRecordRequest, AgentSessionArchiveRequest, + AgentSessionCreateRequest, AgentSessionManagementPort, AgentSessionRenameRequest, + AgentSubmissionPort, AgentSubmissionRequest, AgentSubmissionSource, + AgentThreadGoalGetRequest, AgentThreadGoalManagementPort, DelegationPolicy, + PermissionEffect, PermissionRule, PermissionRuntimeCeiling, SessionStoragePathRequest, + SubagentContextMode, ThreadGoal, ThreadGoalStatus, }; use std::collections::HashMap; use std::path::PathBuf; @@ -10868,6 +11457,463 @@ mod tests { test_coordinator_with_max_active_sessions(100) } + async fn create_two_turn_session( + session_manager: &SessionManager, + workspace: &std::path::Path, + session_id: &str, + ) -> PathBuf { + session_manager + .create_session_with_id( + Some(session_id.to_string()), + "Reverted".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.to_string_lossy().into_owned()), + ..Default::default() + }, + ) + .await + .expect("create session"); + for (turn_id, prompt) in [("turn-0", "first"), ("turn-1", "second")] { + session_manager + .start_dialog_turn( + session_id, + "agentic".to_string(), + prompt.to_string(), + Some(turn_id.to_string()), + None, + None, + ) + .await + .expect("start persisted turn"); + session_manager + .complete_dialog_turn( + session_id, + turn_id, + format!("reply to {prompt}"), + &[], + TurnStats::default(), + ) + .await + .expect("complete persisted turn"); + session_manager.reset_session_state_if_processing(session_id, turn_id); + } + let storage_path = session_manager + .effective_session_storage_path(session_id) + .await + .expect("session storage path"); + storage_path + } + + async fn create_staged_two_turn_session( + session_manager: &SessionManager, + workspace: &std::path::Path, + session_id: &str, + ) -> PathBuf { + let storage_path = create_two_turn_session(session_manager, workspace, session_id).await; + session_manager + .persistence_manager() + .save_session_revert_state( + &storage_path, + session_id, + &crate::agentic::session::revert::SessionRevertState { + schema_version: crate::agentic::session::revert::SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 1, + original_turn_end: 2, + phase: crate::agentic::session::revert::SessionRevertPhase::Staged, + workspace_checkpoint: Vec::new(), + }, + ) + .await + .expect("stage session revert"); + let mutation = session_manager + .acquire_session_mutation(session_id) + .await + .expect("session mutation"); + session_manager + .apply_staged_revert_context_locked(&storage_path, session_id, 1) + .await + .expect("apply staged context"); + drop(mutation); + storage_path + } + + #[tokio::test] + async fn staged_revert_is_committed_before_local_and_maintenance_turns() { + let (coordinator, session_manager) = test_persistent_coordinator(); + let workspace = tempfile::tempdir().expect("workspace"); + + let local_session_id = format!("local-revert-{}", uuid::Uuid::new_v4()); + let local_storage = create_staged_two_turn_session( + session_manager.as_ref(), + workspace.path(), + &local_session_id, + ) + .await; + let child_session_id = format!("{local_session_id}-child"); + let grandchild_session_id = format!("{local_session_id}-grandchild"); + let mut child = SessionMetadata::new( + child_session_id.clone(), + "Hidden child".to_string(), + "Explore".to_string(), + "model".to_string(), + ); + child.session_kind = SessionKind::Subagent; + child.relationship = Some(SessionRelationship { + kind: Some(SessionRelationshipKind::Subagent), + parent_session_id: Some(local_session_id.clone()), + parent_request_id: None, + parent_dialog_turn_id: Some("turn-1".to_string()), + parent_turn_index: Some(1), + parent_tool_call_id: Some("tool-child".to_string()), + subagent_type: Some("Explore".to_string()), + continuation_policy: None, + }); + child.workspace_path = Some(workspace.path().to_string_lossy().into_owned()); + session_manager + .persistence_manager() + .save_session_metadata(&local_storage, &child) + .await + .expect("hidden child metadata"); + let mut grandchild = SessionMetadata::new( + grandchild_session_id.clone(), + "Hidden grandchild".to_string(), + "Explore".to_string(), + "model".to_string(), + ); + grandchild.session_kind = SessionKind::Subagent; + grandchild.relationship = Some(SessionRelationship { + kind: Some(SessionRelationshipKind::Subagent), + parent_session_id: Some(child_session_id.clone()), + parent_request_id: None, + parent_dialog_turn_id: Some("child-turn".to_string()), + parent_turn_index: Some(0), + parent_tool_call_id: Some("tool-grandchild".to_string()), + subagent_type: Some("Explore".to_string()), + continuation_policy: None, + }); + grandchild.workspace_path = Some(workspace.path().to_string_lossy().into_owned()); + session_manager + .persistence_manager() + .save_session_metadata(&local_storage, &grandchild) + .await + .expect("hidden grandchild metadata"); + AgentLocalCommandTurnPort::record_completed_local_command_turn( + &coordinator, + AgentLocalCommandTurnRecordRequest { + session_id: local_session_id.clone(), + content: "/usage".to_string(), + turn_id: Some("local-turn".to_string()), + timestamp_ms: None, + metadata: serde_json::Map::new(), + }, + ) + .await + .expect("record local command after staged undo"); + let local_turns = session_manager + .persistence_manager() + .load_session_turns(&local_storage, &local_session_id) + .await + .expect("load local turns"); + assert_eq!( + local_turns + .iter() + .map(|turn| turn.turn_id.as_str()) + .collect::>(), + vec!["turn-0", "local-turn"] + ); + assert!(session_manager + .persistence_manager() + .load_session_revert_state(&local_storage, &local_session_id) + .await + .expect("load local marker") + .is_none()); + for discarded_session_id in [&child_session_id, &grandchild_session_id] { + assert!(session_manager + .persistence_manager() + .load_session_metadata(&local_storage, discarded_session_id) + .await + .expect("discarded child metadata lookup") + .is_none()); + } + + let maintenance_session_id = format!("compact-revert-{}", uuid::Uuid::new_v4()); + let maintenance_storage = create_staged_two_turn_session( + session_manager.as_ref(), + workspace.path(), + &maintenance_session_id, + ) + .await; + let task = coordinator + .start_manual_compaction_task( + maintenance_session_id.clone(), + Some("maintenance-turn".to_string()), + ) + .await + .expect("start maintenance after staged undo"); + let maintenance_turns = session_manager + .persistence_manager() + .load_session_turns(&maintenance_storage, &maintenance_session_id) + .await + .expect("load maintenance turns"); + assert_eq!( + maintenance_turns + .iter() + .map(|turn| turn.turn_id.as_str()) + .collect::>(), + vec!["turn-0", "maintenance-turn"] + ); + assert!(session_manager + .persistence_manager() + .load_session_revert_state(&maintenance_storage, &maintenance_session_id) + .await + .expect("load maintenance marker") + .is_none()); + coordinator + .cancel_dialog_turn(&maintenance_session_id, &task.turn_id) + .await + .expect("cancel maintenance task"); + let _ = tokio::time::timeout(Duration::from_secs(5), task.completion).await; + } + + #[tokio::test] + async fn mutating_restore_reconciles_a_marker_written_before_workspace_apply() { + let (coordinator, session_manager) = test_persistent_coordinator(); + let workspace = tempfile::tempdir().expect("workspace"); + let file_path = workspace.path().join("src/lib.rs"); + std::fs::create_dir_all(file_path.parent().expect("file parent")) + .expect("create file parent"); + tokio::fs::write(&file_path, "before\n") + .await + .expect("write original file"); + let session_id = format!("restore-revert-{}", uuid::Uuid::new_v4()); + let storage_path = + create_two_turn_session(session_manager.as_ref(), workspace.path(), &session_id).await; + let snapshot_manager = crate::service::snapshot::get_or_create_snapshot_manager( + workspace.path().to_path_buf(), + None, + ) + .await + .expect("snapshot manager"); + let operation_id = snapshot_manager + .record_file_change( + &session_id, + 1, + file_path.clone(), + crate::service::snapshot::types::OperationType::Modify, + "Edit".to_string(), + ) + .await + .expect("record file change"); + tokio::fs::write(&file_path, "after\n") + .await + .expect("write changed file"); + snapshot_manager + .get_snapshot_service() + .read() + .await + .complete_file_modification(&session_id, &operation_id, 1) + .await + .expect("complete file change"); + + let mut state = crate::agentic::session::revert::SessionRevertState { + schema_version: crate::agentic::session::revert::SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 1, + original_turn_end: 2, + phase: crate::agentic::session::revert::SessionRevertPhase::Applying, + workspace_checkpoint: Vec::new(), + }; + snapshot_manager + .prepare_workspace_revert(&session_id, &mut state) + .await + .expect("prepare staged checkpoint"); + session_manager + .persistence_manager() + .save_session_revert_state(&storage_path, &session_id, &state) + .await + .expect("persist marker before workspace apply"); + + coordinator + .restore_session_from_storage_path(&storage_path, &session_id) + .await + .expect("restore should reconcile staged workspace"); + + assert_eq!( + tokio::fs::read_to_string(&file_path) + .await + .expect("read reconciled file"), + "before\n" + ); + assert_eq!( + session_manager + .get_session(&session_id) + .expect("restored session") + .dialog_turn_ids, + vec!["turn-0"] + ); + let staged = session_manager + .persistence_manager() + .load_session_revert_state(&storage_path, &session_id) + .await + .expect("load staged marker") + .expect("staged marker should remain"); + assert_eq!( + staged.phase, + crate::agentic::session::revert::SessionRevertPhase::Staged + ); + + tokio::fs::write(&file_path, "external edit\n") + .await + .expect("write external edit after successful undo"); + coordinator + .commit_session_revert_before_submission(&session_id) + .await + .expect("commit stable staged boundary"); + assert_eq!( + tokio::fs::read_to_string(&file_path) + .await + .expect("read external edit after commit"), + "external edit\n" + ); + assert!(session_manager + .persistence_manager() + .load_session_revert_state(&storage_path, &session_id) + .await + .expect("load committed marker") + .is_none()); + } + + #[tokio::test] + async fn coordinator_delete_reconciles_an_unfinished_revert_before_cleanup() { + let (coordinator, session_manager) = test_persistent_coordinator(); + let workspace = tempfile::tempdir().expect("workspace"); + let session_id = format!("delete-revert-{}", uuid::Uuid::new_v4()); + let storage_path = + create_two_turn_session(session_manager.as_ref(), workspace.path(), &session_id).await; + let state = crate::agentic::session::revert::SessionRevertState { + schema_version: crate::agentic::session::revert::SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 1, + original_turn_end: 2, + phase: crate::agentic::session::revert::SessionRevertPhase::Applying, + workspace_checkpoint: Vec::new(), + }; + session_manager + .persistence_manager() + .save_session_revert_state(&storage_path, &session_id, &state) + .await + .expect("pending marker"); + + coordinator + .delete_session(workspace.path(), &session_id) + .await + .expect("coordinator should reconcile before deleting"); + + assert!(session_manager.get_session(&session_id).is_none()); + assert!(session_manager + .persistence_manager() + .load_session_revert_state(&storage_path, &session_id) + .await + .expect("deleted marker load") + .is_none()); + } + + #[tokio::test] + async fn transcript_read_waits_for_session_history_mutation_before_loading_turns() { + let (coordinator, session_manager) = test_persistent_coordinator(); + let coordinator = Arc::new(coordinator); + let workspace = tempfile::tempdir().expect("workspace"); + let session_id = format!("transcript-mutation-{}", uuid::Uuid::new_v4()); + let storage_path = + create_two_turn_session(session_manager.as_ref(), workspace.path(), &session_id).await; + + let mutation = session_manager + .acquire_session_mutation(&session_id) + .await + .expect("simulated revert mutation"); + let reader = coordinator.clone(); + let read_session_id = session_id.clone(); + let transcript_task = tokio::spawn(async move { + bitfun_runtime_ports::SessionTranscriptReader::read_session_transcript( + reader.as_ref(), + bitfun_runtime_ports::SessionTranscriptRequest { + session_id: read_session_id, + turn_id: None, + }, + ) + .await + }); + tokio::task::yield_now().await; + assert!( + !transcript_task.is_finished(), + "transcript reads must share the session history mutation boundary" + ); + + session_manager + .persistence_manager() + .delete_turns_from(&storage_path, &session_id, 1) + .await + .expect("commit simulated suffix deletion"); + drop(mutation); + + let transcript = transcript_task + .await + .expect("transcript task") + .expect("transcript after mutation"); + assert!(transcript + .messages + .iter() + .all(|message| message.turn_id.as_deref() != Some("turn-1"))); + assert!(transcript + .messages + .iter() + .any(|message| message.turn_id.as_deref() == Some("turn-0"))); + } + + #[tokio::test] + async fn transient_transcript_locked_fallback_does_not_reenter_session_mutation() { + let (coordinator, session_manager) = test_coordinator(); + let workspace = tempfile::tempdir().expect("transient workspace"); + let session = session_manager + .create_session( + "Transient transcript".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().into_owned()), + ..Default::default() + }, + ) + .await + .expect("transient session"); + session_manager + .add_message( + &session.session_id, + Message::user("visible transient context".to_string()), + ) + .await + .expect("transient context message"); + + let transcript = tokio::time::timeout( + Duration::from_secs(1), + bitfun_runtime_ports::SessionTranscriptReader::read_session_transcript( + &coordinator, + bitfun_runtime_ports::SessionTranscriptRequest { + session_id: session.session_id, + turn_id: None, + }, + ), + ) + .await + .expect("transient transcript must not deadlock") + .expect("transient transcript"); + + assert_eq!(transcript.messages.len(), 1); + assert!(matches!( + &transcript.messages[0].content, + bitfun_runtime_ports::TranscriptContent::Text(text) + if text == "visible transient context" + )); + } + #[tokio::test] async fn create_session_checks_runtime_ownership_before_persisting() { let ownership_root = tempfile::tempdir().expect("ownership root"); diff --git a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs index 7d0d4f4f1..6cebf6d00 100644 --- a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs +++ b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs @@ -343,6 +343,13 @@ pub struct DialogScheduler { /// performs maintenance that must not overlap turn dispatch. pub(crate) struct SessionMaintenancePermit { _operation_guard: KeyedAsyncLockGuard, + retired_turn_ids: Vec, +} + +impl SessionMaintenancePermit { + pub(crate) fn retired_turn_ids(&self) -> &[String] { + &self.retired_turn_ids + } } fn take_active_turn_for_outcome( @@ -948,7 +955,7 @@ impl DialogScheduler { ) .await .map_err(|error| error.to_string())?; - self.session_manager + self.coordinator .restore_session_from_storage_path(&restore_path, session_id) .await .map_err(|error| error.to_string())? @@ -1046,6 +1053,13 @@ impl DialogScheduler { )); } + // OpenCode-compatible semantics: accepting a new prompt while history + // is staged permanently discards the hidden suffix before the Turn starts. + self.coordinator + .commit_session_revert_before_submission(&session_id) + .await + .map_err(SchedulerSubmitError::Core)?; + match action { DialogSubmitQueueAction::StartImmediately => { let tid = self.start_turn(&session_id, &queued_turn).await?; @@ -1059,7 +1073,7 @@ impl DialogScheduler { } DialogSubmitQueueAction::ClearQueueAndStartImmediately => { - self.clear_queue(&session_id).await; + let _ = self.clear_queue(&session_id).await; let tid = self.start_turn(&session_id, &queued_turn).await?; queued_turn.accept_settlement(); self.record_last_submitted_agent_type(&session_id, &queued_turn.agent_type) @@ -1264,9 +1278,11 @@ impl DialogScheduler { let operation_guard = self.lock_session_operation(session_id).await; self.session_manager .validate_session_storage_path_binding(session_id, requested_storage_path)?; - if self.queue_depth(session_id) > 0 { - self.clear_queue(session_id).await; - } + let mut retired_turn_ids = if self.queue_depth(session_id) > 0 { + self.clear_queue(session_id).await + } else { + Vec::new() + }; abort_thread_goal_continuation_for_session(session_id); let deadline = Instant::now() + wait_timeout; let cancelled_before_parent = self @@ -1283,7 +1299,8 @@ impl DialogScheduler { self.maintenance_background_sessions .insert(session_id.to_string(), subagent_session_ids.clone()); } - self.coordinator + let cancelled_turn_id = self + .coordinator .cancel_active_turn_for_session( session_id, deadline.saturating_duration_since(Instant::now()), @@ -1313,9 +1330,15 @@ impl DialogScheduler { ) .await?; self.maintenance_background_sessions.remove(session_id); - self.retire_active_turn_for_maintenance(session_id); + let scheduler_turn_id = self.retire_active_turn_for_maintenance(session_id); + for retired_turn_id in [cancelled_turn_id, scheduler_turn_id].into_iter().flatten() { + if !retired_turn_ids.contains(&retired_turn_id) { + retired_turn_ids.push(retired_turn_id); + } + } Ok(SessionMaintenancePermit { _operation_guard: operation_guard, + retired_turn_ids, }) } @@ -1329,9 +1352,9 @@ impl DialogScheduler { .await } - fn retire_active_turn_for_maintenance(&self, session_id: &str) { + fn retire_active_turn_for_maintenance(&self, session_id: &str) -> Option { let Some(active_turn) = self.active_turns.remove(session_id) else { - return; + return None; }; let turn_id = active_turn.turn_id().to_string(); self.retired_maintenance_outcomes.mark(session_id, &turn_id); @@ -1343,6 +1366,7 @@ impl DialogScheduler { "Retired active turn before destructive session maintenance: session_id={}, turn_id={}", session_id, turn_id ); + Some(turn_id) } // ── Private helpers ────────────────────────────────────────────────────── @@ -1368,13 +1392,15 @@ impl DialogScheduler { Ok(()) } - async fn clear_queue(&self, session_id: &str) { + async fn clear_queue(&self, session_id: &str) -> Vec { let cleared_turns = self.queues.clear(session_id); let count = cleared_turns.len(); + let mut retired_turn_ids = Vec::new(); for queued_turn in cleared_turns { match queued_turn.execution { QueuedTurnExecution::Standard => { if let Some(turn_id) = queued_turn.turn_id { + retired_turn_ids.push(turn_id.clone()); self.coordinator .emit_event(AgenticEvent::DialogTurnCancelled { session_id: session_id.to_string(), @@ -1409,6 +1435,7 @@ impl DialogScheduler { count, session_id ); } + retired_turn_ids } fn dequeue_next(&self, session_id: &str) -> Option { @@ -1826,7 +1853,7 @@ impl DialogScheduler { "Turn {}, clearing queue: session_id={}", lifecycle_plan.status, session_id ); - self.clear_queue(&session_id).await; + let _ = self.clear_queue(&session_id).await; } (active_turn, active_internal_turn, lifecycle_plan) }; @@ -2389,6 +2416,7 @@ mod tests { use crate::agentic::persistence::PersistenceManager; use crate::agentic::session::{ compression::{CompressionConfig, ContextCompressor}, + revert::{SessionRevertPhase, SessionRevertState, SESSION_REVERT_SCHEMA_VERSION}, PromptCachePolicy, SessionContextStore, SessionManagerConfig, }; use crate::agentic::tools::registry::ToolRegistry; @@ -2480,6 +2508,67 @@ mod tests { )); } + #[tokio::test] + async fn submission_preflight_commits_a_persisted_revert_marker() { + let (scheduler, session_manager, _, root) = test_scheduler(); + let session_id = "reverted-session"; + let workspace = root.path().join("workspace"); + std::fs::create_dir_all(&workspace).expect("workspace"); + session_manager + .create_session_with_id( + Some(session_id.to_string()), + "Reverted".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.to_string_lossy().into_owned()), + ..Default::default() + }, + ) + .await + .expect("create session"); + let storage_path = session_manager + .effective_session_storage_path(session_id) + .await + .expect("storage path"); + session_manager + .persistence_manager() + .save_session_revert_state( + &storage_path, + session_id, + &SessionRevertState { + schema_version: SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 0, + original_turn_end: 1, + phase: SessionRevertPhase::Staged, + workspace_checkpoint: Vec::new(), + }, + ) + .await + .expect("persist staged revert"); + + scheduler + .coordinator + .commit_session_revert_before_submission(session_id) + .await + .expect("commit staged revert"); + + assert!(session_manager + .persistence_manager() + .load_session_revert_state(&storage_path, session_id) + .await + .expect("load revert marker") + .is_none()); + let source = include_str!("scheduler.rs"); + let submission = source + .split_once("async fn submit_queued_turn_locked(") + .expect("submission method") + .1 + .split_once("async fn record_last_submitted_agent_type(") + .expect("submission method boundary") + .0; + assert!(submission.contains("commit_session_revert_before_submission(&session_id)")); + } + #[tokio::test] async fn background_bash_result_injects_into_its_running_parent_turn() { let (scheduler, session_manager, _, root) = test_scheduler(); @@ -3118,6 +3207,58 @@ mod tests { .matches_turn(session_id, "turn-active")); } + #[tokio::test] + async fn maintenance_retires_scheduler_state_even_when_core_cancel_returns_a_turn_id() { + let (scheduler, session_manager, _, root) = test_scheduler(); + let session_id = "session-maintenance-retire"; + let turn_id = "turn-active"; + let workspace = root.path().join("workspace-maintenance-retire"); + std::fs::create_dir_all(&workspace).expect("workspace"); + session_manager + .create_session_with_id( + Some(session_id.to_string()), + "Maintenance retire".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.to_string_lossy().to_string()), + ..Default::default() + }, + ) + .await + .expect("create session"); + session_manager + .update_session_state( + session_id, + SessionState::Processing { + current_turn_id: turn_id.to_string(), + phase: ProcessingPhase::ToolCalling, + }, + ) + .await + .expect("mark processing"); + scheduler + .active_turns + .insert(session_id, desktop_active_turn(turn_id)); + let storage_path = session_manager + .storage_path_binding_for_test(session_id) + .expect("storage binding"); + + let maintenance = scheduler + .begin_session_maintenance(session_id, &storage_path, Duration::from_secs(1)) + .await + .expect("maintenance"); + + assert_eq!(maintenance.retired_turn_ids(), &[turn_id.to_string()]); + assert!(!scheduler.active_turns.matches_turn(session_id, turn_id)); + assert!(take_active_turn_for_outcome( + &scheduler.active_turns, + &scheduler.retired_maintenance_outcomes, + session_id, + turn_id, + ) + .is_none()); + } + #[test] fn retired_maintenance_outcome_cannot_mutate_a_recreated_session_generation() { let active_turns = ActiveDialogTurnStore::default(); diff --git a/src/crates/assembly/core/src/agentic/insights/collector.rs b/src/crates/assembly/core/src/agentic/insights/collector.rs index f466ca8de..95c2ade1f 100644 --- a/src/crates/assembly/core/src/agentic/insights/collector.rs +++ b/src/crates/assembly/core/src/agentic/insights/collector.rs @@ -1,3 +1,4 @@ +use crate::agentic::coordination::get_global_coordinator; use crate::agentic::core::{Message, MessageContent, MessageRole, ToolCall, ToolResult}; use crate::agentic::insights::session_paths::collect_effective_session_storage_targets; use crate::agentic::insights::types::*; @@ -11,7 +12,7 @@ use crate::service::session_usage::{ build_session_usage_report_from_turns, SessionUsageReportRequest, }; use crate::service::snapshot::get_snapshot_manager_for_workspace; -use crate::util::errors::BitFunResult; +use crate::util::errors::{BitFunError, BitFunResult}; use bitfun_agent_tools::ResolvedToolInvocation; use chrono::{DateTime, Local, Utc}; use log::{debug, warn}; @@ -46,6 +47,9 @@ impl InsightsCollector { pub async fn collect(days: u32) -> BitFunResult<(BaseStats, Vec)> { let path_manager = get_path_manager_arc(); let pm = PersistenceManager::new(path_manager)?; + let coordinator = get_global_coordinator().ok_or_else(|| { + BitFunError::service("Core coordinator is unavailable for Insights history reads") + })?; let now = SystemTime::now(); let now_ms = system_time_to_unix_ms(now); let cutoff_ms = now_ms.saturating_sub(days as u64 * 86_400_000); @@ -85,19 +89,33 @@ impl InsightsCollector { continue; } - let (session, parent_turns) = match pm - .load_session_with_turns(ws_path, &summary.session_id) + let mut session = match pm.load_session(ws_path, &summary.session_id).await { + Ok(value) => value, + Err(e) => { + warn!( + "Skipping session {}: load failed: {}", + summary.session_id, e + ); + continue; + } + }; + let parent_turns = match coordinator + .load_visible_persisted_session_turns(ws_path, &summary.session_id) .await { Ok(value) => value, Err(e) => { warn!( - "Skipping session {}: load failed: {}", + "Skipping session {}: visible history load failed: {}", summary.session_id, e ); continue; } }; + session.dialog_turn_ids = parent_turns + .iter() + .map(|turn| turn.turn_id.clone()) + .collect(); let parent_turn_ids = parent_turns .iter() @@ -136,7 +154,10 @@ impl InsightsCollector { ); let mut selected_turns = selected_parent_turns.clone(); for hidden_session_id in hidden_session_ids { - match pm.load_session_turns(ws_path, &hidden_session_id).await { + match coordinator + .load_visible_persisted_session_turns(ws_path, &hidden_session_id) + .await + { Ok(turns) => { selected_turns .extend(filter_turns_for_window(&turns, cutoff_ms, now_ms)); diff --git a/src/crates/assembly/core/src/agentic/memories/service.rs b/src/crates/assembly/core/src/agentic/memories/service.rs index 9ee1bd8b8..825a7235e 100644 --- a/src/crates/assembly/core/src/agentic/memories/service.rs +++ b/src/crates/assembly/core/src/agentic/memories/service.rs @@ -456,10 +456,17 @@ async fn process_single_session( source.turn_count, format_unix_secs(source.last_finished_unix_secs) ); - let turns = match persistence - .load_session_turns(session_storage_path, &source.session_id) - .await - { + let turns_result = match crate::agentic::coordination::get_global_coordinator() { + Some(coordinator) => { + coordinator + .load_visible_persisted_session_turns(session_storage_path, &source.session_id) + .await + } + None => Err(BitFunError::service( + "Core coordinator is unavailable for AI Memory history reads", + )), + }; + let turns = match turns_result { Ok(turns) => turns, Err(error) => { record_failure( diff --git a/src/crates/assembly/core/src/agentic/persistence/manager.rs b/src/crates/assembly/core/src/agentic/persistence/manager.rs index 55de2e24b..a8277c5fe 100644 --- a/src/crates/assembly/core/src/agentic/persistence/manager.rs +++ b/src/crates/assembly/core/src/agentic/persistence/manager.rs @@ -9,6 +9,7 @@ use crate::agentic::core::{ }; use crate::agentic::memories::db::{MemoryDatabase, MEMORY_PHASE2_GLOBAL_JOB_KEY}; use crate::agentic::memories::external_context::dialog_turn_uses_external_context; +use crate::agentic::session::revert::{SessionRevertState, SESSION_REVERT_SCHEMA_VERSION}; use crate::agentic::session::transcript_render::{ render_transcript, rendered_turn_char_count, transcript_fingerprint, }; @@ -471,6 +472,12 @@ impl PersistenceManager { .join("token-anchors.json") } + fn session_revert_path(&self, workspace_path: &Path, session_id: &str) -> PathBuf { + self.session_layout(workspace_path) + .session_dir(session_id) + .join("session-revert.json") + } + fn turns_dir(&self, workspace_path: &Path, session_id: &str) -> PathBuf { self.session_layout(workspace_path).turns_dir(session_id) } @@ -1441,6 +1448,59 @@ impl PersistenceManager { } } + pub(crate) async fn load_session_revert_state( + &self, + workspace_path: &Path, + session_id: &str, + ) -> BitFunResult> { + Self::validate_session_id(session_id)?; + let state = self + .read_json_optional::( + &self.session_revert_path(workspace_path, session_id), + ) + .await?; + if let Some(state) = state.as_ref() { + if state.schema_version != SESSION_REVERT_SCHEMA_VERSION { + return Err(BitFunError::Deserialization(format!( + "Unsupported Session revert schema version: session_id={}, version={}", + session_id, state.schema_version + ))); + } + } + Ok(state) + } + + pub(crate) async fn save_session_revert_state( + &self, + workspace_path: &Path, + session_id: &str, + state: &SessionRevertState, + ) -> BitFunResult<()> { + Self::validate_session_id(session_id)?; + let _session_write = self.lock_session_write_operation(workspace_path, session_id)?; + self.ensure_runtime_for_write(workspace_path).await?; + self.ensure_session_dir(workspace_path, session_id).await?; + self.write_json_atomic(&self.session_revert_path(workspace_path, session_id), state) + .await + } + + pub(crate) async fn delete_session_revert_state( + &self, + workspace_path: &Path, + session_id: &str, + ) -> BitFunResult<()> { + Self::validate_session_id(session_id)?; + let _session_write = self.lock_session_write_operation(workspace_path, session_id)?; + match fs::remove_file(self.session_revert_path(workspace_path, session_id)).await { + Ok(()) => Ok(()), + Err(error) if error.kind() == ErrorKind::NotFound => Ok(()), + Err(error) => Err(BitFunError::io(format!( + "Failed to delete staged Session revert for {}: {}", + session_id, error + ))), + } + } + pub async fn load_token_anchors( &self, workspace_path: &Path, @@ -1544,6 +1604,16 @@ impl PersistenceManager { &self, workspace_path: &Path, session_id: &str, + ) -> BitFunResult)>> { + self.load_latest_turn_context_snapshot_before(workspace_path, session_id, usize::MAX) + .await + } + + pub(crate) async fn load_latest_turn_context_snapshot_before( + &self, + workspace_path: &Path, + session_id: &str, + exclusive_turn_index: usize, ) -> BitFunResult)>> { Self::validate_session_id(session_id)?; let started_at = Instant::now(); @@ -1576,7 +1646,9 @@ impl PersistenceManager { }; if let Ok(index) = index_str.parse::() { snapshot_file_count += 1; - latest = Some(latest.map(|value| value.max(index)).unwrap_or(index)); + if index < exclusive_turn_index { + latest = Some(latest.map(|value| value.max(index)).unwrap_or(index)); + } } } let scan_duration = scan_started_at.elapsed(); @@ -2355,6 +2427,25 @@ impl PersistenceManager { && self .turn_path(workspace_path, &turn.session_id, turn.turn_index) .exists(); + if let Some(revert) = self + .read_json_optional::( + &self.session_revert_path(workspace_path, &turn.session_id), + ) + .await? + { + if revert.schema_version != SESSION_REVERT_SCHEMA_VERSION { + return Err(BitFunError::Deserialization(format!( + "Unsupported Session revert schema version: session_id={}, version={}", + turn.session_id, revert.schema_version + ))); + } + if turn.turn_index >= revert.boundary_turn { + return Err(BitFunError::Validation(format!( + "Cannot persist a Turn over the staged Session suffix: session_id={}, turn_index={}, boundary_turn={}", + turn.session_id, turn.turn_index, revert.boundary_turn + ))); + } + } let file = StoredDialogTurnFile { schema_version: SESSION_STORAGE_SCHEMA_VERSION, @@ -2561,6 +2652,32 @@ impl PersistenceManager { Ok(turns) } + /// Load the product-visible Session history while retaining the current + /// process's persisted writer lease across the marker and Turn reads. + /// + /// Runtime owners that reconcile, redo, or permanently discard a staged + /// suffix must use [`Self::load_session_turns`] instead. Passive product + /// consumers must enter through Core's per-Session mutation owner before + /// using this projection; the persistence lease supplies cross-process, + /// not in-process, ordering. + pub async fn load_visible_session_turns( + &self, + workspace_path: &Path, + session_id: &str, + ) -> BitFunResult> { + Self::validate_session_id(session_id)?; + let _session_write = self.lock_session_write_operation(workspace_path, session_id)?; + let boundary_turn = self + .load_session_revert_state(workspace_path, session_id) + .await? + .map(|state| state.boundary_turn); + let mut turns = self.load_session_turns(workspace_path, session_id).await?; + if let Some(boundary_turn) = boundary_turn { + turns.retain(|turn| turn.turn_index < boundary_turn); + } + Ok(turns) + } + pub async fn load_session_tail_turns( &self, workspace_path: &Path, @@ -2709,7 +2826,9 @@ impl PersistenceManager { count: usize, ) -> BitFunResult> { Self::validate_session_id(session_id)?; - let turns = self.load_session_turns(workspace_path, session_id).await?; + let turns = self + .load_visible_session_turns(workspace_path, session_id) + .await?; let start = turns.len().saturating_sub(count); Ok(turns[start..].to_vec()) } @@ -3024,7 +3143,14 @@ impl PersistenceManager { }), }; - let all_turns = self.load_session_turns(workspace_path, session_id).await?; + let revert_boundary = self + .load_session_revert_state(workspace_path, session_id) + .await? + .map(|state| state.boundary_turn); + let mut all_turns = self.load_session_turns(workspace_path, session_id).await?; + if let Some(boundary_turn) = revert_boundary { + all_turns.retain(|turn| turn.turn_index < boundary_turn); + } let selected_indices = parsed_turn_selectors .as_ref() .map(|selectors| Self::transcript_select_turn_indices(all_turns.len(), selectors)) @@ -3132,7 +3258,7 @@ impl PersistenceManager { turns: None, }; let all_turns = self - .load_session_turns(reference_workspace_path, reference_session_id) + .load_visible_session_turns(reference_workspace_path, reference_session_id) .await?; // Pick complete turns backwards from the newest one. The first turn @@ -3300,6 +3426,9 @@ mod tests { }; use crate::agentic::core::{Message, Session, SessionConfig, SessionKind, ToolResult}; use crate::agentic::memories::db::{MemoryDatabase, MemoryRow, MEMORY_PHASE2_GLOBAL_JOB_KEY}; + use crate::agentic::session::revert::{ + SessionRevertPhase, SessionRevertState, SESSION_REVERT_SCHEMA_VERSION, + }; use crate::agentic::session::{TokenAnchor, TokenAnchorInput}; use crate::agentic::skill_agent_snapshot::{ AgentSnapshotEntry, SkillSnapshotEntry, TurnSkillAgentSnapshot, @@ -3345,6 +3474,121 @@ mod tests { } } + #[tokio::test] + async fn staged_session_revert_state_round_trips_and_clears_independently() { + let workspace = TestWorkspace::new(); + let manager = + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"); + let state = SessionRevertState { + schema_version: SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 2, + original_turn_end: 5, + phase: SessionRevertPhase::Staged, + workspace_checkpoint: Vec::new(), + }; + + manager + .save_session_revert_state(workspace.path(), "session-1", &state) + .await + .expect("staged revert should persist"); + assert_eq!( + manager + .load_session_revert_state(workspace.path(), "session-1") + .await + .expect("staged revert should load"), + Some(state) + ); + + manager + .delete_session_revert_state(workspace.path(), "session-1") + .await + .expect("staged revert should clear"); + assert!(manager + .load_session_revert_state(workspace.path(), "session-1") + .await + .expect("cleared staged revert should stay absent") + .is_none()); + } + + #[tokio::test] + async fn staged_revert_rejects_overwriting_a_hidden_turn_index() { + let workspace = TestWorkspace::new(); + let manager = + PersistenceManager::new(workspace.path_manager()).expect("persistence manager"); + let session_id = "session-hidden-suffix"; + manager + .save_session_metadata( + workspace.path(), + &SessionMetadata::new( + session_id.to_string(), + "Hidden suffix".to_string(), + "agentic".to_string(), + "model-a".to_string(), + ), + ) + .await + .expect("session metadata should persist"); + + for index in 0..=1 { + let turn = DialogTurnData::new( + format!("turn-{index}"), + index, + session_id.to_string(), + UserMessageData { + id: format!("user-{index}"), + content: format!("prompt {index}"), + timestamp: index as u64, + metadata: None, + }, + ); + manager + .save_dialog_turn(workspace.path(), &turn) + .await + .expect("fixture turn should persist"); + } + manager + .save_session_revert_state( + workspace.path(), + session_id, + &SessionRevertState { + schema_version: SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 1, + original_turn_end: 2, + phase: SessionRevertPhase::Staged, + workspace_checkpoint: Vec::new(), + }, + ) + .await + .expect("staged revert should persist"); + + let replacement = DialogTurnData::new( + "local-command".to_string(), + 1, + session_id.to_string(), + UserMessageData { + id: "local-command-user".to_string(), + content: "usage report".to_string(), + timestamp: 3, + metadata: None, + }, + ); + let error = manager + .save_dialog_turn(workspace.path(), &replacement) + .await + .expect_err("a staged hidden turn must not be overwritten"); + assert!( + error.to_string().contains("staged Session suffix"), + "{error}" + ); + + let preserved = manager + .load_dialog_turn(workspace.path(), session_id, 1) + .await + .expect("hidden turn should remain readable") + .expect("hidden turn should remain present"); + assert_eq!(preserved.turn_id, "turn-1"); + } + #[test] fn unfinished_session_directory_is_removed_when_creation_is_cancelled() { let workspace = TestWorkspace::new(); @@ -3604,6 +3848,36 @@ mod tests { .save_dialog_turn(workspace.path(), &turn) .await .expect("turn should save"); + let mut hidden_turn = DialogTurnData::new( + "turn-hidden".to_string(), + 1, + session_id.clone(), + UserMessageData { + id: "user-hidden".to_string(), + content: "hidden transcript payload".to_string(), + timestamp: 1, + metadata: None, + }, + ); + hidden_turn.mark_completed(); + manager + .save_dialog_turn(workspace.path(), &hidden_turn) + .await + .expect("hidden turn should save"); + manager + .save_session_revert_state( + workspace.path(), + &session_id, + &SessionRevertState { + schema_version: SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 1, + original_turn_end: 2, + phase: SessionRevertPhase::Staged, + workspace_checkpoint: Vec::new(), + }, + ) + .await + .expect("staged marker should save"); let export = manager .export_session_transcript( @@ -3621,6 +3895,24 @@ mod tests { .expect("transcript file should be readable"); assert!(transcript.contains("## Turn 0")); assert!(transcript.contains("hello transcript")); + assert!(!transcript.contains("hidden transcript payload")); + + let selected = manager + .export_session_transcript( + workspace.path(), + &session_id, + &SessionTranscriptExportOptions { + turns: Some(vec!["-1".to_string()]), + ..Default::default() + }, + ) + .await + .expect("visible-relative transcript selection should succeed"); + assert_eq!(selected.turn_count, 1); + let selected_transcript = std::fs::read_to_string(&selected.transcript_path) + .expect("selected transcript should be readable"); + assert!(selected_transcript.contains("hello transcript")); + assert!(!selected_transcript.contains("hidden transcript payload")); } #[tokio::test] diff --git a/src/crates/assembly/core/src/agentic/session/mod.rs b/src/crates/assembly/core/src/agentic/session/mod.rs index 030f88091..5d8219553 100644 --- a/src/crates/assembly/core/src/agentic/session/mod.rs +++ b/src/crates/assembly/core/src/agentic/session/mod.rs @@ -7,6 +7,7 @@ pub mod context_store; pub mod evidence_ledger; pub mod file_read_state; pub mod prompt_cache; +pub(crate) mod revert; pub mod session_manager; pub mod session_store_port; pub mod token_anchor; diff --git a/src/crates/assembly/core/src/agentic/session/revert.rs b/src/crates/assembly/core/src/agentic/session/revert.rs new file mode 100644 index 000000000..b19226652 --- /dev/null +++ b/src/crates/assembly/core/src/agentic/session/revert.rs @@ -0,0 +1,261 @@ +use crate::service::session::{DialogTurnData, DialogTurnKind}; +use serde::{Deserialize, Serialize}; +use std::path::PathBuf; + +pub(crate) const SESSION_REVERT_SCHEMA_VERSION: u32 = 1; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub(crate) enum SessionRevertPhase { + Applying, + Staged, + Clearing, + Committing, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct SessionWorkspaceCheckpoint { + pub(crate) path: PathBuf, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub(crate) snapshot_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct SessionRevertState { + pub(crate) schema_version: u32, + pub(crate) boundary_turn: usize, + pub(crate) original_turn_end: usize, + pub(crate) phase: SessionRevertPhase, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub(crate) workspace_checkpoint: Vec, +} + +impl SessionRevertState { + fn initial(boundary_turn: usize, original_turn_end: usize) -> Self { + Self { + schema_version: SESSION_REVERT_SCHEMA_VERSION, + boundary_turn, + original_turn_end, + phase: SessionRevertPhase::Applying, + workspace_checkpoint: Vec::new(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub(crate) enum SessionRevertTransition { + Stage { + state: SessionRevertState, + replacement_prompt: Option, + hidden_turn_count: usize, + }, + Clear { + previous_state: SessionRevertState, + }, +} + +#[cfg(test)] +impl SessionRevertTransition { + pub(crate) fn boundary_turn(&self) -> Option { + match self { + Self::Stage { state, .. } => Some(state.boundary_turn), + Self::Clear { .. } => None, + } + } + + pub(crate) fn replacement_prompt(&self) -> Option<&str> { + match self { + Self::Stage { + replacement_prompt, .. + } => replacement_prompt.as_deref(), + Self::Clear { .. } => None, + } + } + + pub(crate) fn hidden_turn_count(&self) -> usize { + match self { + Self::Stage { + hidden_turn_count, .. + } => *hidden_turn_count, + Self::Clear { .. } => 0, + } + } + + pub(crate) fn staged_state(&self) -> Option { + match self { + Self::Stage { state, .. } => Some(state.clone()), + Self::Clear { .. } => None, + } + } +} + +pub(crate) fn resolve_undo( + turns: &[DialogTurnData], + current: Option<&SessionRevertState>, +) -> Option { + let current_boundary = current + .map(|state| state.boundary_turn) + .unwrap_or_else(|| original_turn_end(turns)); + let target = turns + .iter() + .filter(|turn| turn.kind == DialogTurnKind::UserDialog) + .filter(|turn| turn.turn_index < current_boundary) + .max_by_key(|turn| turn.turn_index)?; + + let mut state = current.cloned().unwrap_or_else(|| { + SessionRevertState::initial(target.turn_index, original_turn_end(turns)) + }); + state.boundary_turn = target.turn_index; + state.phase = SessionRevertPhase::Applying; + let hidden_turn_count = turns + .iter() + .filter(|turn| turn.turn_index >= state.boundary_turn) + .count(); + + Some(SessionRevertTransition::Stage { + state, + replacement_prompt: Some( + target + .user_message + .metadata + .as_ref() + .and_then(|metadata| metadata.get("original_text")) + .and_then(serde_json::Value::as_str) + .unwrap_or(&target.user_message.content) + .to_string(), + ), + hidden_turn_count, + }) +} + +pub(crate) fn resolve_redo( + turns: &[DialogTurnData], + current: Option<&SessionRevertState>, +) -> Option { + let current = current?; + let next = turns + .iter() + .filter(|turn| turn.kind == DialogTurnKind::UserDialog) + .filter(|turn| { + turn.turn_index > current.boundary_turn && turn.turn_index < current.original_turn_end + }) + .min_by_key(|turn| turn.turn_index); + + if let Some(next) = next { + let mut state = current.clone(); + state.boundary_turn = next.turn_index; + state.phase = SessionRevertPhase::Applying; + let hidden_turn_count = turns + .iter() + .filter(|turn| turn.turn_index >= state.boundary_turn) + .count(); + return Some(SessionRevertTransition::Stage { + state, + replacement_prompt: None, + hidden_turn_count, + }); + } + + Some(SessionRevertTransition::Clear { + previous_state: current.clone(), + }) +} + +fn original_turn_end(turns: &[DialogTurnData]) -> usize { + turns + .iter() + .map(|turn| turn.turn_index.saturating_add(1)) + .max() + .unwrap_or(0) +} + +#[cfg(test)] +mod tests { + use super::{resolve_redo, resolve_undo, SessionRevertTransition}; + use crate::service::session::{DialogTurnData, DialogTurnKind, TurnStatus, UserMessageData}; + + fn user_turn(index: usize, prompt: &str) -> DialogTurnData { + DialogTurnData { + turn_id: format!("turn-{index}"), + turn_index: index, + session_id: "session-1".to_string(), + timestamp: index as u64, + kind: DialogTurnKind::UserDialog, + agent_type: Some("agentic".to_string()), + user_message: UserMessageData { + id: format!("message-{index}"), + content: prompt.to_string(), + timestamp: index as u64, + metadata: None, + }, + model_rounds: Vec::new(), + start_time: index as u64, + end_time: Some(index as u64), + duration_ms: Some(0), + token_usage: None, + finish_reason: None, + has_final_response: Some(true), + error: None, + error_detail: None, + status: TurnStatus::Completed, + } + } + + fn maintenance_turn(index: usize) -> DialogTurnData { + let mut turn = user_turn(index, "internal compact"); + turn.kind = DialogTurnKind::ManualCompaction; + turn.agent_type = None; + turn + } + + #[test] + fn undo_walks_visible_user_prompts_and_restores_each_prompt() { + let turns = vec![user_turn(0, "first"), maintenance_turn(1), { + let mut turn = user_turn(2, "wrapped second prompt"); + turn.user_message.metadata = Some(serde_json::json!({ + "original_text": "second" + })); + turn + }]; + + let first = resolve_undo(&turns, None).expect("latest prompt should be undoable"); + assert_eq!(first.boundary_turn(), Some(2)); + assert_eq!(first.replacement_prompt(), Some("second")); + assert_eq!(first.hidden_turn_count(), 1); + + let state = first.staged_state().expect("undo should stage a boundary"); + let second = resolve_undo(&turns, Some(&state)).expect("earlier prompt should be undoable"); + assert_eq!(second.boundary_turn(), Some(0)); + assert_eq!(second.replacement_prompt(), Some("first")); + assert_eq!(second.hidden_turn_count(), 3); + assert!(resolve_undo(&turns, second.staged_state().as_ref()).is_none()); + } + + #[test] + fn redo_advances_to_the_next_user_prompt_then_clears_the_stage() { + let turns = vec![ + user_turn(0, "first"), + maintenance_turn(1), + user_turn(2, "second"), + ]; + let state = resolve_undo(&turns, None) + .and_then(|transition| transition.staged_state()) + .expect("undo should stage latest prompt"); + let earliest = resolve_undo(&turns, Some(&state)) + .and_then(|transition| transition.staged_state()) + .expect("second undo should stage first prompt"); + + let advance = resolve_redo(&turns, Some(&earliest)).expect("redo should advance"); + assert_eq!(advance.boundary_turn(), Some(2)); + assert!(advance.replacement_prompt().is_none()); + assert_eq!(advance.hidden_turn_count(), 1); + + let clear = resolve_redo(&turns, advance.staged_state().as_ref()) + .expect("final redo should clear the stage"); + assert!(matches!(clear, SessionRevertTransition::Clear { .. })); + assert_eq!(clear.hidden_turn_count(), 0); + assert!(resolve_redo(&turns, None).is_none()); + } +} 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 87cd1505c..add3dab66 100644 --- a/src/crates/assembly/core/src/agentic/session/session_manager.rs +++ b/src/crates/assembly/core/src/agentic/session/session_manager.rs @@ -12,6 +12,7 @@ use crate::agentic::image_analysis::ImageContextData; use crate::agentic::keyed_lock::{KeyedAsyncLock, KeyedAsyncLockGuard}; use crate::agentic::memories::db::{MemoryDatabase, MEMORY_PHASE2_GLOBAL_JOB_KEY}; use crate::agentic::persistence::{MaterializedSessionReferenceTranscript, PersistenceManager}; +use crate::agentic::session::revert::SessionRevertPhase; use crate::agentic::session::session_store_port::CoreSessionStorePort; use crate::agentic::session::{ prompt_cache_persist_action, reconcile_prompt_cache_restore, CachedSystemPrompt, @@ -37,7 +38,9 @@ use crate::service::session::{ SessionRelationship, SessionStatus, TextItemData, ThinkingItemData, ToolCallData, ToolItemData, ToolResultData, TranscriptLineRange, TurnStatus, UserMessageData, }; -use crate::service::snapshot::ensure_snapshot_manager_for_workspace; +use crate::service::snapshot::{ + ensure_snapshot_manager_for_workspace, get_or_create_snapshot_manager, +}; use crate::service::workspace::{get_global_workspace_service, WorkspaceInfo, WorkspaceKind}; use crate::util::errors::{BitFunError, BitFunResult}; use crate::util::sanitize_plain_model_output; @@ -864,7 +867,10 @@ impl SessionManager { .unwrap_or_else(|| workspace_path.to_path_buf()) } - async fn resolve_storage_path_for_workspace_path(&self, workspace_path: &Path) -> PathBuf { + pub(crate) async fn resolve_storage_path_for_workspace_path( + &self, + workspace_path: &Path, + ) -> PathBuf { let storage_path_started_at = Instant::now(); let session_storage_path = self .effective_storage_path_for_workspace_path(workspace_path) @@ -927,7 +933,7 @@ impl SessionManager { /// Resolve the effective storage path for a session by ID. /// For remote workspaces, maps the remote path to a local session storage path. - async fn effective_session_storage_path(&self, session_id: &str) -> Option { + pub(crate) async fn effective_session_storage_path(&self, session_id: &str) -> Option { let config = self.sessions.get(session_id)?.config.clone(); self.effective_storage_path_for_config(&config).await } @@ -955,9 +961,10 @@ impl SessionManager { "Session storage path not found: {parent_session_id}" )) })?; + let _related_history_read = self.acquire_session_mutation(related_session_id).await?; Ok(self .persistence_manager - .load_session_turns(&storage_path, related_session_id) + .load_visible_session_turns(&storage_path, related_session_id) .await? .into_iter() .find(|turn| turn.turn_id == dialog_turn_id)) @@ -1116,6 +1123,8 @@ impl SessionManager { ))); } + let _reference_history_read = self.acquire_session_mutation(&reference.session_id).await?; + let transcript = self .persistence_manager .materialize_session_reference_transcript( @@ -1452,7 +1461,7 @@ impl SessionManager { ) -> BitFunResult> { let turns = self .persistence_manager - .load_session_turns(workspace_path, session_id) + .load_visible_session_turns(workspace_path, session_id) .await?; Ok(Self::build_messages_from_turns(&turns)) } @@ -3632,9 +3641,33 @@ impl SessionManager { } let _mutation_guard = self.acquire_session_mutation(session_id).await?; + let mut has_persisted_turns = false; + if self.should_persist_session_id(session_id) { + if let Some(storage_path) = self.effective_session_storage_path(session_id).await { + if self + .persistence_manager + .load_session_revert_state(&storage_path, session_id) + .await? + .is_some() + { + return Err(SessionExecutionBindingError::Busy( + "Worktree isolation cannot change while Session undo or redo is pending" + .to_string(), + )); + } + has_persisted_turns = !self + .persistence_manager + .load_session_turns(&storage_path, session_id) + .await? + .is_empty(); + } + } if let Some(mut session) = self.sessions.get_mut(session_id) { - if !session.dialog_turn_ids.is_empty() || !matches!(session.state, SessionState::Idle) { + if has_persisted_turns + || !session.dialog_turn_ids.is_empty() + || !matches!(session.state, SessionState::Idle) + { return Err(SessionExecutionBindingError::Busy( "Worktree isolation can only be changed before the session's first message" .to_string(), @@ -3738,6 +3771,14 @@ impl SessionManager { ) -> BitFunResult<()> { bitfun_core_types::validate_session_id(session_id).map_err(BitFunError::Validation)?; let _mutation_guard = self.lock_session_mutation(session_id).await; + self.delete_session_locked(workspace_path, session_id).await + } + + pub(crate) async fn delete_session_locked( + &self, + workspace_path: &Path, + session_id: &str, + ) -> BitFunResult<()> { let session_storage_path = self .resolve_storage_path_for_workspace_path(workspace_path) .await; @@ -4137,6 +4178,18 @@ impl SessionManager { // before mutating loaded runtime state so a storage failure leaves the // active session usable and retryable. if self.config.enable_persistence { + let revert_state = self + .persistence_manager + .load_session_revert_state(session_storage_path, session_id) + .await?; + if revert_state + .as_ref() + .is_some_and(|state| state.phase != SessionRevertPhase::Staged) + { + return Err(BitFunError::OutcomeUnknown(format!( + "Session deletion cannot discard an unfinished revert transition: session_id={session_id}" + ))); + } let persistence_stage_started_at = Instant::now(); debug!( "Session deletion stage starting: session_id={}, stage=persistence_delete", @@ -4150,6 +4203,30 @@ impl SessionManager { session_id, elapsed_ms_u64(persistence_stage_started_at) ); + if let Some(revert_state) = revert_state { + match get_or_create_snapshot_manager( + cleanup_workspace_path.to_path_buf(), + None, + ) + .await + { + Ok(snapshot_manager) => { + if let Err(error) = snapshot_manager + .delete_workspace_revert_checkpoint(&revert_state) + .await + { + warn!( + "Failed to delete Session revert checkpoint after Session deletion: session_id={}, error={}", + session_id, error + ); + } + } + Err(error) => warn!( + "Failed to initialize snapshot cleanup for deleted Session revert checkpoint: session_id={}, error={}", + session_id, error + ), + } + } } self.cleanup_session_owned_resources( @@ -4551,9 +4628,17 @@ impl SessionManager { session_id, visibility_metadata_duration_ms ); + let staged_revert = self + .persistence_manager + .load_session_revert_state(session_storage_path, session_id) + .await?; + let effective_tail_turn_count = staged_revert + .as_ref() + .map(|_| None) + .unwrap_or(tail_turn_count); let session_started_at = Instant::now(); - let (mut session, persisted_turns, total_turn_count, turn_load) = - if let Some(tail_turn_count) = tail_turn_count { + let (mut session, mut persisted_turns, mut total_turn_count, turn_load) = + if let Some(tail_turn_count) = effective_tail_turn_count { self.persistence_manager .load_session_with_tail_turns_timed( session_storage_path, @@ -4569,6 +4654,10 @@ impl SessionManager { let total_turn_count = turns.len(); (session, turns, total_turn_count, timing) }; + if let Some(revert) = staged_revert.as_ref() { + persisted_turns.retain(|turn| turn.turn_index < revert.boundary_turn); + total_turn_count = persisted_turns.len(); + } let load_session_with_turns_duration_ms = elapsed_ms_u64(session_started_at); debug!( "Session view restore phase completed: session_id={}, phase=load_session_with_turns, turn_count={}, total_turn_count={}, tail_turn_count={:?}, duration_ms={}", @@ -4804,10 +4893,17 @@ impl SessionManager { // 1. Load session and turns from storage in one pass let session_started_at = Instant::now(); - let (mut session, persisted_turns) = self + let (mut session, mut persisted_turns) = self .persistence_manager .load_session_with_turns(session_storage_path, session_id) .await?; + let staged_revert = self + .persistence_manager + .load_session_revert_state(session_storage_path, session_id) + .await?; + if let Some(revert) = staged_revert.as_ref() { + persisted_turns.retain(|turn| turn.turn_index < revert.boundary_turn); + } debug!( "Session restore phase completed: session_id={}, phase=load_session_with_turns, turn_count={}, duration_ms={}", session_id, @@ -4916,11 +5012,23 @@ impl SessionManager { ); let mut latest_turn_index: Option = None; let context_snapshot_started_at = Instant::now(); - let mut messages = match self - .persistence_manager - .load_latest_turn_context_snapshot(session_storage_path, session_id) - .await? - { + let latest_context_snapshot = match staged_revert.as_ref() { + Some(revert) => { + self.persistence_manager + .load_latest_turn_context_snapshot_before( + session_storage_path, + session_id, + revert.boundary_turn, + ) + .await? + } + None => { + self.persistence_manager + .load_latest_turn_context_snapshot(session_storage_path, session_id) + .await? + } + }; + let mut messages = match latest_context_snapshot { Some((turn_index, msgs)) => { latest_turn_index = Some(turn_index); self.sanitize_listing_diff_context_snapshot_if_needed( @@ -4944,17 +5052,19 @@ impl SessionManager { ); if let Some(snapshot_turn_index) = latest_turn_index { - let delta_start = snapshot_turn_index.saturating_add(1); - if delta_start < persisted_turns.len() { + let delta_turns = persisted_turns + .iter() + .filter(|turn| turn.turn_index > snapshot_turn_index) + .cloned() + .collect::>(); + if !delta_turns.is_empty() { warn!( "Context snapshot is behind persisted turns, rebuilding delta: session_id={}, snapshot_turn_index={}, persisted_turn_count={}", session_id, snapshot_turn_index, persisted_turns.len() ); - messages.extend(Self::build_messages_from_turns( - &persisted_turns[delta_start..], - )); + messages.extend(Self::build_messages_from_turns(&delta_turns)); } }; @@ -4965,10 +5075,14 @@ impl SessionManager { ); } - let recoverable_turn_count = latest_turn_index - .map(|turn_index| turn_index + 1) - .unwrap_or(0) - .max(persisted_turns.len()); + let recoverable_turn_count = if staged_revert.is_some() { + persisted_turns.len() + } else { + latest_turn_index + .map(|turn_index| turn_index + 1) + .unwrap_or(0) + .max(persisted_turns.len()) + }; if session.dialog_turn_ids.len() < persisted_turns.len() { warn!( @@ -5110,6 +5224,150 @@ impl SessionManager { .await } + /// Move the loaded Session to a persisted staged-revert boundary without + /// deleting any turn, context snapshot, or compression artifact. The + /// durable `session-revert.json` remains the authoritative visibility fact. + pub(crate) async fn apply_staged_revert_context_locked( + &self, + session_storage_path: &Path, + session_id: &str, + boundary_turn: usize, + ) -> BitFunResult<()> { + let turns = self + .persistence_manager + .load_session_turns(session_storage_path, session_id) + .await?; + let visible_turns = turns + .iter() + .filter(|turn| turn.turn_index < boundary_turn) + .cloned() + .collect::>(); + let listing_baseline_rebuild_turn_index = self + .persistence_manager + .load_session_metadata(session_storage_path, session_id) + .await? + .as_ref() + .and_then(|metadata| { + Self::listing_baseline_rebuild_turn_index_from_metadata(Some(metadata)) + }); + + let messages = if boundary_turn == 0 { + Vec::new() + } else if let Some((snapshot_turn_index, snapshot_messages)) = self + .persistence_manager + .load_latest_turn_context_snapshot_before( + session_storage_path, + session_id, + boundary_turn, + ) + .await? + { + let mut messages = self + .sanitize_listing_diff_context_snapshot_if_needed( + session_storage_path, + session_id, + snapshot_turn_index, + snapshot_messages, + listing_baseline_rebuild_turn_index, + "staged_revert_context_snapshot", + ) + .await; + let delta = visible_turns + .iter() + .filter(|turn| turn.turn_index > snapshot_turn_index) + .cloned() + .collect::>(); + messages.extend(Self::build_messages_from_turns(&delta)); + messages + } else { + Self::build_messages_from_turns(&visible_turns) + }; + + self.context_store.replace_context(session_id, messages); + self.file_read_state_store.clear_session(session_id); + let fallback_agent_type = self + .sessions + .get(session_id) + .map(|session| session.agent_type.clone()); + let last_user_dialog_agent_type = Self::derive_last_user_dialog_agent_type_from_turns( + &visible_turns, + fallback_agent_type.as_deref(), + ); + if let Some(mut session) = self.sessions.get_mut(session_id) { + session.dialog_turn_ids = visible_turns + .iter() + .map(|turn| turn.turn_id.clone()) + .collect(); + session.last_user_dialog_agent_type = last_user_dialog_agent_type; + session.state = SessionState::Idle; + session.updated_at = SystemTime::now(); + session.last_activity_at = SystemTime::now(); + } + Ok(()) + } + + /// Permanently discard the hidden suffix after a staged revert. Callers + /// persist the `Committing` phase before entering this idempotent cleanup. + pub(crate) async fn commit_staged_revert_context_locked( + &self, + session_storage_path: &Path, + session_id: &str, + boundary_turn: usize, + ) -> BitFunResult<()> { + self.apply_staged_revert_context_locked(session_storage_path, session_id, boundary_turn) + .await?; + + let session_snapshot = self.sessions.get(session_id).and_then(|session| { + (self.should_persist_session(&session) && self.config.enable_persistence) + .then(|| session.clone()) + }); + if let Some(session) = session_snapshot { + self.persistence_manager + .save_session(session_storage_path, &session) + .await?; + } + // A durable revert marker means persisted Session artifacts exist even + // when automatic Session persistence is disabled for the current + // runtime (for example, an adapter restoring an explicitly selected + // history). Committing that marker must therefore always prune the + // persisted suffix; `enable_persistence` only controls automatic + // Session writes, not explicit history mutations. + self.persistence_manager + .delete_dialog_turns_from(session_storage_path, session_id, boundary_turn) + .await?; + self.persistence_manager + .delete_turn_context_snapshots_from(session_storage_path, session_id, boundary_turn) + .await?; + self.persistence_manager + .delete_compression_transcripts_from(session_storage_path, session_id, boundary_turn) + .await?; + self.truncate_listing_baseline_rebuild_turn_index_after_rollback( + session_storage_path, + session_id, + boundary_turn, + ) + .await?; + self.turn_skill_agent_snapshot_store + .remove_from(session_id, boundary_turn); + let surviving_turn_ids = self + .sessions + .get(session_id) + .map(|session| { + session + .dialog_turn_ids + .iter() + .cloned() + .collect::>() + }) + .unwrap_or_default(); + self.rollback_edit_constraint_state_to_turns(session_id, &surviving_turn_ids) + .await; + let messages = self.context_store.get_context_messages(session_id); + self.prune_token_anchors_to_messages(session_id, &messages) + .await; + Ok(()) + } + pub(crate) async fn rollback_context_to_turn_start_locked( &self, session_storage_path: &Path, @@ -5266,6 +5524,16 @@ impl SessionManager { if !self.config.enable_persistence { return Ok(()); } + if self + .persistence_manager + .load_session_revert_state(session_storage_path, session_id) + .await? + .is_some() + { + return Err(BitFunError::OutcomeUnknown(format!( + "Legacy context rollback cannot overlap a staged Session undo: session_id={session_id}" + ))); + } self.persistence_manager .load_session_metadata(session_storage_path, session_id) .await?; @@ -5561,6 +5829,26 @@ impl SessionManager { // ============ Dialog Turn Management ============ + async fn ensure_persisted_turn_append_allowed(&self, session_id: &str) -> BitFunResult<()> { + if !self.should_persist_session_id(session_id) { + return Ok(()); + } + let Some(storage_path) = self.effective_session_storage_path(session_id).await else { + return Ok(()); + }; + if let Some(revert) = self + .persistence_manager + .load_session_revert_state(&storage_path, session_id) + .await? + { + return Err(BitFunError::Validation(format!( + "Cannot append a persisted Turn while a Session revert is {:?}: session_id={}, boundary_turn={}", + revert.phase, session_id, revert.boundary_turn + ))); + } + Ok(()) + } + #[allow(clippy::too_many_arguments)] async fn start_persisted_turn( &self, @@ -5574,6 +5862,34 @@ impl SessionManager { user_message_metadata: Option, ) -> BitFunResult { let _mutation_guard = self.acquire_session_mutation(session_id).await?; + self.start_persisted_turn_locked( + session_id, + kind, + agent_type, + user_input, + turn_id, + context_messages, + processing_phase, + user_message_metadata, + ) + .await + } + + /// Start a persisted Turn while the caller owns the Session mutation lock. + #[allow(clippy::too_many_arguments)] + async fn start_persisted_turn_locked( + &self, + session_id: &str, + kind: DialogTurnKind, + agent_type: Option, + user_input: String, + turn_id: Option, + context_messages: Vec, + processing_phase: ProcessingPhase, + user_message_metadata: Option, + ) -> BitFunResult { + self.ensure_persisted_turn_append_allowed(session_id) + .await?; let session = self .get_session(session_id) .ok_or_else(|| BitFunError::NotFound(format!("Session not found: {}", session_id)))?; @@ -5822,6 +6138,29 @@ impl SessionManager { Ok(turn_id) } + pub(crate) async fn start_maintenance_turn_locked( + &self, + session_id: &str, + display_message: String, + turn_id: Option, + user_message_metadata: Option, + ) -> BitFunResult { + let turn_id = self + .start_persisted_turn_locked( + session_id, + DialogTurnKind::ManualCompaction, + None, + display_message, + turn_id, + Vec::new(), + ProcessingPhase::Compacting, + user_message_metadata, + ) + .await?; + debug!("Starting maintenance turn: turn_id={}", turn_id); + Ok(turn_id) + } + /// Append a completed local command turn that should be persisted in user-facing /// history without entering model-visible runtime context. pub async fn append_completed_local_command_turn( @@ -5833,6 +6172,26 @@ impl SessionManager { user_message_metadata: Option, ) -> BitFunResult { let _mutation_guard = self.lock_session_mutation(session_id).await; + self.append_completed_local_command_turn_locked( + session_id, + content, + turn_id, + timestamp_ms, + user_message_metadata, + ) + .await + } + + pub(crate) async fn append_completed_local_command_turn_locked( + &self, + session_id: &str, + content: String, + turn_id: Option, + timestamp_ms: Option, + user_message_metadata: Option, + ) -> BitFunResult { + self.ensure_persisted_turn_append_allowed(session_id) + .await?; let session = self .get_session(session_id) .ok_or_else(|| BitFunError::NotFound(format!("Session not found: {}", session_id)))?; @@ -6524,6 +6883,7 @@ impl SessionManager { pub async fn get_messages(&self, session_id: &str) -> BitFunResult> { if self.config.enable_persistence { if let Some(workspace_path) = self.effective_session_storage_path(session_id).await { + let _history_read = self.acquire_session_mutation(session_id).await?; let messages = self .rebuild_messages_from_turns(&workspace_path, session_id) .await?; @@ -6541,7 +6901,7 @@ impl SessionManager { /// This is intentionally separate from `get_messages`: runtime context /// reconstruction excludes model-invisible maintenance turns, while a /// transcript must be able to project those turns for the UI. - pub(crate) async fn load_persisted_transcript_turns( + pub(crate) async fn load_persisted_transcript_turns_locked( &self, session_id: &str, ) -> BitFunResult>> { @@ -6551,10 +6911,18 @@ impl SessionManager { let Some(workspace_path) = self.effective_session_storage_path(session_id).await else { return Ok(None); }; - self.persistence_manager + let mut turns = self + .persistence_manager .load_session_turns(&workspace_path, session_id) - .await - .map(Some) + .await?; + if let Some(revert) = self + .persistence_manager + .load_session_revert_state(&workspace_path, session_id) + .await? + { + turns.retain(|turn| turn.turn_index < revert.boundary_turn); + } + Ok(Some(turns)) } /// Get a paginated best-effort message view for the session. @@ -7041,6 +7409,7 @@ mod tests { }; use crate::agentic::persistence::PersistenceManager; use crate::agentic::session::{ + revert::{SessionRevertPhase, SessionRevertState, SESSION_REVERT_SCHEMA_VERSION}, PromptCachePolicy, PromptCacheScope, SessionContextStore, SystemPromptCacheIdentity, UserContextCacheIdentity, }; @@ -7365,6 +7734,104 @@ mod tests { assert_eq!(restored.config.workspace_id.as_deref(), Some("workspace-2")); } + #[tokio::test] + async fn execution_binding_rejects_a_boundary_zero_revert_after_explicit_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( + "Boundary zero binding".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + project_workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..SessionConfig::default() + }, + ) + .await + .expect("session should be created"); + let storage_path = manager + .effective_session_storage_path(&session.session_id) + .await + .expect("storage path"); + persistence_manager + .save_dialog_turn( + &storage_path, + &DialogTurnData::new( + "turn-0".to_string(), + 0, + session.session_id.clone(), + UserMessageData { + id: "user-0".to_string(), + content: "first prompt".to_string(), + timestamp: 1, + metadata: None, + }, + ), + ) + .await + .expect("persist first turn"); + persistence_manager + .save_session_revert_state( + &storage_path, + &session.session_id, + &SessionRevertState { + schema_version: SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 0, + original_turn_end: 1, + phase: SessionRevertPhase::Staged, + workspace_checkpoint: Vec::new(), + }, + ) + .await + .expect("persist boundary zero marker"); + assert!(manager + .unload_session_from_memory(&session.session_id) + .await + .expect("unload session")); + let restored = manager + .restore_session_from_storage_path(&storage_path, &session.session_id) + .await + .expect("restore from the known Session storage path"); + assert!( + restored.dialog_turn_ids.is_empty(), + "boundary-zero restore must project an empty visible history" + ); + + let target_path = workspace.path().join("managed-worktree"); + let error = manager + .update_session_execution_binding( + &session.session_id, + SessionExecutionBindingUpdate { + workspace_path: target_path.to_string_lossy().to_string(), + project_workspace_path: workspace.path().to_string_lossy().to_string(), + workspace_id: Some("workspace-2".to_string()), + execution_target: SessionExecutionTarget::local( + target_path.to_string_lossy().to_string(), + ), + }, + ) + .await + .expect_err("a staged revert must pin its snapshot workspace"); + + assert!( + matches!(error, SessionExecutionBindingError::Busy(_)), + "expected staged revert to reject rebinding as busy, got {error:?}" + ); + assert_eq!( + manager + .get_session(&session.session_id) + .expect("session remains restored after rejection") + .config + .workspace_path + .as_deref(), + Some(workspace.path().to_string_lossy().as_ref()) + ); + } + #[tokio::test] async fn unloading_a_session_releases_capacity_without_deleting_persistence() { let workspace = TestWorkspace::new(); @@ -10075,8 +10542,12 @@ mod tests { .await .expect("maintenance failure should persist"); + let _mutation = manager + .acquire_session_mutation(&session.session_id) + .await + .expect("transcript mutation"); let turns = manager - .load_persisted_transcript_turns(&session.session_id) + .load_persisted_transcript_turns_locked(&session.session_id) .await .expect("turns should load") .expect("persistence should be enabled"); @@ -10455,6 +10926,138 @@ mod tests { ); } + #[tokio::test] + async fn staged_revert_filters_runtime_restore_and_transcript_without_deleting_turns() { + use crate::agentic::session::revert::{ + SessionRevertPhase, SessionRevertState, SESSION_REVERT_SCHEMA_VERSION, + }; + + 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( + "Staged revert".to_string(), + "agentic".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().to_string()), + ..Default::default() + }, + ) + .await + .expect("session should create"); + + for index in 0..3 { + let turn = DialogTurnData::new( + format!("turn-{index}"), + index, + session.session_id.clone(), + UserMessageData { + id: format!("turn-{index}-user"), + content: format!("prompt {index}"), + timestamp: index as u64, + metadata: None, + }, + ); + persistence_manager + .save_dialog_turn(workspace.path(), &turn) + .await + .expect("turn should save"); + let messages = (0..=index) + .map(|message_index| { + crate::agentic::core::Message::user(format!("prompt {message_index}")) + }) + .collect::>(); + persistence_manager + .save_turn_context_snapshot(workspace.path(), &session.session_id, index, &messages) + .await + .expect("context snapshot should save"); + } + manager + .sessions + .get_mut(&session.session_id) + .expect("session should be active") + .dialog_turn_ids = vec![ + "turn-0".to_string(), + "turn-1".to_string(), + "turn-2".to_string(), + ]; + let state = SessionRevertState { + schema_version: SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 2, + original_turn_end: 3, + phase: SessionRevertPhase::Staged, + workspace_checkpoint: Vec::new(), + }; + persistence_manager + .save_session_revert_state(workspace.path(), &session.session_id, &state) + .await + .expect("staged revert should persist"); + + let legacy_error = manager + .validate_rollback_context_to_turn_start_locked( + workspace.path(), + &session.session_id, + 1, + ) + .await + .expect_err("legacy rollback must not truncate a staged Session suffix"); + assert!(matches!(legacy_error, BitFunError::OutcomeUnknown(_))); + + manager + .apply_staged_revert_context_locked( + workspace.path(), + &session.session_id, + state.boundary_turn, + ) + .await + .expect("staged context should apply"); + assert_eq!( + manager + .get_session(&session.session_id) + .expect("session should remain active") + .dialog_turn_ids, + vec!["turn-0".to_string(), "turn-1".to_string()] + ); + assert_eq!( + persistence_manager + .load_session_turns(workspace.path(), &session.session_id) + .await + .expect("hidden turns should remain persisted") + .len(), + 3 + ); + + manager.evict_loaded_session_for_test(&session.session_id); + let restored = manager + .restore_session(workspace.path(), &session.session_id) + .await + .expect("staged session should restore"); + assert_eq!(restored.dialog_turn_ids.len(), 2); + assert_eq!( + manager + .context_store + .get_context_messages(&session.session_id) + .len(), + 2 + ); + let _mutation = manager + .acquire_session_mutation(&session.session_id) + .await + .expect("transcript mutation"); + assert_eq!( + manager + .load_persisted_transcript_turns_locked(&session.session_id) + .await + .expect("transcript should load") + .expect("transcript should be persisted") + .len(), + 2 + ); + } + #[tokio::test] async fn rollback_context_failure_preserves_turn_history() { let workspace = TestWorkspace::new(); @@ -11276,6 +11879,52 @@ mod tests { .contains_key(&session.session_id)); } + #[tokio::test] + async fn direct_delete_rejects_an_unfinished_revert_transition() { + 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( + "Pending revert delete".to_string(), + "agent".to_string(), + SessionConfig { + workspace_path: Some(workspace.path().to_string_lossy().into_owned()), + ..Default::default() + }, + ) + .await + .expect("session should create"); + let state = SessionRevertState { + schema_version: SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 0, + original_turn_end: 1, + phase: SessionRevertPhase::Applying, + workspace_checkpoint: Vec::new(), + }; + persistence_manager + .save_session_revert_state(workspace.path(), &session.session_id, &state) + .await + .expect("pending marker should persist"); + + let error = manager + .delete_session(workspace.path(), &session.session_id) + .await + .expect_err("direct deletion must not discard recovery state"); + + assert!(matches!(error, BitFunError::OutcomeUnknown(_))); + assert!(manager.get_session(&session.session_id).is_some()); + assert_eq!( + persistence_manager + .load_session_revert_state(workspace.path(), &session.session_id) + .await + .expect("marker load"), + Some(state) + ); + } + #[test] fn build_messages_from_turns_skips_model_invisible_turns() { use crate::service::session::{DialogTurnData, DialogTurnKind, UserMessageData}; diff --git a/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs b/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs index 9b1e795cb..a6a7de379 100644 --- a/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs +++ b/src/crates/assembly/core/src/agentic/tools/implementations/session_history_tool.rs @@ -1,15 +1,12 @@ -use crate::agentic::persistence::PersistenceManager; use crate::agentic::tools::framework::{ Tool, ToolExposure, ToolRenderOptions, ToolResult, ToolUseContext, ValidationResult, }; -use crate::infrastructure::PathManager; use crate::service::session::SessionTranscriptExportOptions; use crate::service_agent_runtime::CoreServiceAgentRuntime; use crate::util::errors::{BitFunError, BitFunResult}; use async_trait::async_trait; use serde::Deserialize; use serde_json::{json, Value}; -use std::sync::Arc; /// SessionHistory tool - export a grep-friendly transcript file for a session. pub struct SessionHistoryTool; @@ -237,9 +234,12 @@ Examples: )) })?; let display_workspace = display_workspace.to_string_lossy().into_owned(); - let manager = PersistenceManager::new(Arc::new(PathManager::new()?))?; - let transcript = manager - .export_session_transcript( + let coordinator = + crate::agentic::coordination::get_global_coordinator().ok_or_else(|| { + BitFunError::service("Core coordinator is unavailable for SessionHistory export") + })?; + let transcript = coordinator + .export_visible_persisted_session_transcript( &session_storage_dir, &session_id, &SessionTranscriptExportOptions { diff --git a/src/crates/assembly/core/src/product_runtime.rs b/src/crates/assembly/core/src/product_runtime.rs index aa2d766e9..f57eacb91 100644 --- a/src/crates/assembly/core/src/product_runtime.rs +++ b/src/crates/assembly/core/src/product_runtime.rs @@ -40,8 +40,12 @@ use crate::agentic::persistence::session_branch::SessionBranchRequest; use crate::agentic::persistence::{PersistenceManager, SessionMetadataPage}; use crate::agentic::session::{CoreSessionStorePort, PromptCacheScope}; use crate::agentic::tools::implementations::skills::SkillRegistry; -use crate::service::session::{DialogTurnData, SessionMetadata}; -use crate::service::session_usage::{generate_session_usage_report, SessionUsageReport}; +use crate::service::session::{ + DialogTurnData, SessionMetadata, SessionTranscriptExport, SessionTranscriptExportOptions, +}; +use crate::service::session_usage::{ + generate_session_usage_report_from_storage_path, SessionUsageReport, +}; use crate::service::snapshot::{ get_snapshot_manager_for_workspace, initialize_snapshot_manager_for_workspace, open_snapshot_manager_for_view, SnapshotError, SnapshotManager, @@ -198,6 +202,21 @@ pub struct CoreSessionMutationPermit { storage_path: PathBuf, } +/// Holds a stable persisted-session visibility boundary while a product reads +/// facts that are also changed by Session undo/redo. +pub struct CoreSessionReadPermit { + _guard: KeyedAsyncLockGuard, + visible_turn_end: Option, +} + +impl CoreSessionReadPermit { + /// Exclusive turn end for product facts. `None` means the complete session + /// is visible because no undo is staged. + pub fn visible_turn_end(&self) -> Option { + self.visible_turn_end + } +} + /// Holds Core's scheduler boundary while a product compatibility operation /// mutates session state that must not overlap turn dispatch. pub struct CoreSessionMaintenancePermit { @@ -208,6 +227,44 @@ fn validate_persisted_session_id(session_id: &str) -> BitFunResult<()> { bitfun_core_types::validate_session_id(session_id).map_err(BitFunError::Validation) } +async fn begin_consistent_persisted_session_read( + coordinator: &ConversationCoordinator, + persistence: &PersistenceManager, + storage_path: &Path, + session_id: &str, +) -> BitFunResult { + validate_persisted_session_id(session_id)?; + let session_manager = coordinator.get_session_manager(); + let guard = session_manager.acquire_session_mutation(session_id).await?; + session_manager.validate_session_storage_path_binding(session_id, storage_path)?; + + if let Some(state) = persistence + .load_session_revert_state(storage_path, session_id) + .await? + { + if state.phase != crate::agentic::session::revert::SessionRevertPhase::Staged { + if session_manager.get_session(session_id).is_none() { + return Err(BitFunError::OutcomeUnknown( + "Session facts are unavailable until the unfinished undo transition is restored" + .to_string(), + )); + } + coordinator + .reconcile_session_revert_locked(storage_path, session_id) + .await?; + } + } + + let visible_turn_end = persistence + .load_session_revert_state(storage_path, session_id) + .await? + .map(|state| state.boundary_turn); + Ok(CoreSessionReadPermit { + _guard: guard, + visible_turn_end, + }) +} + fn latest_persisted_turn_id(turns: &[DialogTurnData]) -> BitFunResult { turns .last() @@ -220,10 +277,37 @@ fn latest_persisted_turn_id(turns: &[DialogTurnData]) -> BitFunResult { async fn generate_core_session_usage_report( persistence: &PersistenceManager, token_usage_service: &TokenUsageService, + session_storage_path: &Path, request: AgentSessionUsageRequest, ) -> BitFunResult { validate_persisted_session_id(&request.session_id)?; - generate_session_usage_report(persistence, Some(token_usage_service), request).await + generate_session_usage_report_from_storage_path( + persistence, + Some(token_usage_service), + session_storage_path, + request, + ) + .await +} + +async fn load_visible_persisted_session_turns( + persistence: &PersistenceManager, + storage_path: &Path, + session_id: &str, + visible_turn_end: Option, + limit: Option, +) -> BitFunResult> { + let mut turns = persistence + .load_session_turns(storage_path, session_id) + .await?; + if let Some(end) = visible_turn_end { + turns.retain(|turn| turn.turn_index < end); + } + if let Some(limit) = limit { + let start = turns.len().saturating_sub(limit); + turns = turns.split_off(start); + } + Ok(turns) } fn snapshot_port_error(error: SnapshotError) -> PortError { @@ -318,7 +402,7 @@ impl LocalWorkspaceSnapshotPort for CoreLocalWorkspaceSnapshot { validate_persisted_session_id(&request.session_id).map_err(runtime_port_error)?; let manager = local_snapshot_manager_for_view(&request.workspace_path).await?; manager - .get_session_files(&request.session_id) + .get_session_files_before(&request.session_id, request.max_turn_exclusive) .await .map_err(snapshot_port_error) } @@ -330,7 +414,7 @@ impl LocalWorkspaceSnapshotPort for CoreLocalWorkspaceSnapshot { validate_persisted_session_id(&request.session_id).map_err(runtime_port_error)?; let manager = local_snapshot_manager_for_view(&request.workspace_path).await?; let stats = manager - .get_session_stats_fact(&request.session_id) + .get_session_stats_fact_before(&request.session_id, request.max_turn_exclusive) .await .map_err(snapshot_port_error)?; Ok(LocalWorkspaceSnapshotStats { @@ -865,6 +949,45 @@ impl CoreAgentRuntimeCompatibility { }) } + /// Admit an external history import only when no undo/redo transaction is + /// active. The returned permit must remain alive for the complete batch. + /// Importers intentionally fail closed instead of guessing how a remote + /// history should merge with a locally staged branch. + pub async fn begin_external_persisted_history_write( + &self, + storage_path: &Path, + session_id: &str, + ) -> BitFunResult { + let permit = self + .begin_persisted_session_mutation(storage_path, session_id) + .await?; + if self + .persistence + .load_session_revert_state(storage_path, session_id) + .await? + .is_some() + { + return Err(BitFunError::SessionInUse { + session_id: session_id.to_string(), + }); + } + Ok(permit) + } + + pub async fn begin_persisted_session_read( + &self, + storage_path: &Path, + session_id: &str, + ) -> BitFunResult { + begin_consistent_persisted_session_read( + self.coordinator.as_ref(), + self.persistence.as_ref(), + storage_path, + session_id, + ) + .await + } + pub async fn begin_session_maintenance( &self, storage_path: &Path, @@ -933,22 +1056,89 @@ impl CoreAgentRuntimeCompatibility { .await } + /// Commits any staged Session undo before a snapshot mutation starts a + /// different destructive history branch. The caller must retain the + /// mutation permit through the complete snapshot operation. + pub async fn commit_session_revert_before_snapshot_mutation( + &self, + permit: &CoreSessionMutationPermit, + ) -> BitFunResult<()> { + self.coordinator + .commit_session_revert_locked(&permit.storage_path, &permit.session_id) + .await + } + + /// Snapshot recording belongs to an already admitted turn and must never + /// append operation history behind a staged Session undo boundary. + pub async fn ensure_snapshot_record_allowed( + &self, + permit: &CoreSessionMutationPermit, + ) -> BitFunResult<()> { + if self + .persistence + .load_session_revert_state(&permit.storage_path, &permit.session_id) + .await? + .is_some() + { + return Err(BitFunError::OutcomeUnknown(format!( + "Snapshot recording is not allowed while a Session undo is staged: session_id={}", + permit.session_id + ))); + } + Ok(()) + } + pub async fn load_persisted_session_turns( &self, - workspace_path: &Path, + storage_path: &Path, session_id: &str, limit: Option, ) -> BitFunResult> { - validate_persisted_session_id(session_id)?; - if let Some(limit) = limit { - self.persistence - .load_recent_turns(workspace_path, session_id, limit) - .await - } else { - self.persistence - .load_session_turns(workspace_path, session_id) - .await - } + let read = self + .begin_persisted_session_read(storage_path, session_id) + .await?; + load_visible_persisted_session_turns( + self.persistence.as_ref(), + storage_path, + session_id, + read.visible_turn_end(), + limit, + ) + .await + } + + pub async fn load_persisted_session_turns_for_mutation( + &self, + permit: &CoreSessionMutationPermit, + limit: Option, + ) -> BitFunResult> { + let visible_turn_end = self + .persistence + .load_session_revert_state(&permit.storage_path, &permit.session_id) + .await? + .map(|state| state.boundary_turn); + load_visible_persisted_session_turns( + self.persistence.as_ref(), + &permit.storage_path, + &permit.session_id, + visible_turn_end, + limit, + ) + .await + } + + pub async fn export_persisted_session_transcript( + &self, + storage_path: &Path, + session_id: &str, + options: &SessionTranscriptExportOptions, + ) -> BitFunResult { + let _read = self + .begin_persisted_session_read(storage_path, session_id) + .await?; + self.persistence + .export_session_transcript(storage_path, session_id, options) + .await } pub async fn touch_persisted_session( @@ -964,12 +1154,35 @@ impl CoreAgentRuntimeCompatibility { pub async fn save_persisted_dialog_turn( &self, - workspace_path: &Path, + permit: &CoreSessionMutationPermit, turn: &DialogTurnData, ) -> BitFunResult<()> { validate_persisted_session_id(&turn.session_id)?; + if turn.session_id != permit.session_id { + return Err(BitFunError::Validation(format!( + "Turn session does not match the active mutation: turn_session_id={}, mutation_session_id={}", + turn.session_id, permit.session_id + ))); + } + let session = self + .coordinator + .get_session_manager() + .get_session(&permit.session_id) + .ok_or_else(|| { + BitFunError::OutcomeUnknown(format!( + "Session must be loaded before saving a projected Turn: session_id={}", + permit.session_id + )) + })?; + let expected_turn_id = session.dialog_turn_ids.get(turn.turn_index); + if expected_turn_id != Some(&turn.turn_id) { + return Err(BitFunError::Validation(format!( + "Turn does not belong to the active Session history branch: session_id={}, turn_index={}, turn_id={}", + permit.session_id, turn.turn_index, turn.turn_id + ))); + } self.persistence - .save_dialog_turn(workspace_path, turn) + .save_dialog_turn(&permit.storage_path, turn) .await } @@ -1010,7 +1223,7 @@ impl CoreSessionOperationsPort { } } - async fn resolve_fork_storage_path( + async fn resolve_session_storage_path( &self, workspace_path: String, remote_connection_id: Option, @@ -1030,15 +1243,83 @@ impl CoreSessionOperationsPort { &self, storage_path: &Path, source_session_id: String, - source_turn_id: String, + source_turn_id: Option, boundary: SessionBranchBoundary, ) -> PortResult { - if source_turn_id.trim().is_empty() { + if source_turn_id + .as_deref() + .is_some_and(|id| id.trim().is_empty()) + { return Err(PortError::new( PortErrorKind::InvalidRequest, "source turn id is required", )); } + if !self + .coordinator + .get_session_manager() + .is_session_loaded_from_storage_path(storage_path, &source_session_id) + .map_err(runtime_port_error)? + { + self.coordinator + .restore_session_from_storage_path(storage_path, &source_session_id) + .await + .map_err(runtime_port_error)?; + } + let session_manager = self.coordinator.get_session_manager(); + let _mutation_guard = session_manager + .acquire_session_mutation(&source_session_id) + .await + .map_err(runtime_port_error)?; + session_manager + .validate_session_storage_path_binding(&source_session_id, storage_path) + .map_err(runtime_port_error)?; + self.coordinator + .reconcile_session_revert_locked(storage_path, &source_session_id) + .await + .map_err(runtime_port_error)?; + let revert_boundary = self + .persistence + .load_session_revert_state(storage_path, &source_session_id) + .await + .map_err(runtime_port_error)? + .map(|state| state.boundary_turn); + let source_turns = self + .persistence + .load_session_turns(storage_path, &source_session_id) + .await + .map_err(runtime_port_error)?; + let source_turn_id = match source_turn_id { + Some(source_turn_id) => { + let source_turn = source_turns + .iter() + .find(|turn| turn.turn_id == source_turn_id) + .ok_or_else(|| { + PortError::new( + PortErrorKind::NotFound, + format!("Source turn not found in persisted session: {source_turn_id}"), + ) + })?; + if revert_boundary + .is_some_and(|boundary_turn| source_turn.turn_index >= boundary_turn) + { + return Err(PortError::new( + PortErrorKind::InvalidRequest, + "Cannot fork from a turn hidden by the current Session undo boundary", + )); + } + source_turn_id + } + None => latest_persisted_turn_id( + &source_turns + .into_iter() + .filter(|turn| { + revert_boundary.is_none_or(|boundary_turn| turn.turn_index < boundary_turn) + }) + .collect::>(), + ) + .map_err(runtime_port_error)?, + }; let source_session_id_for_coordination = source_session_id.clone(); let result = self .persistence @@ -1135,18 +1416,12 @@ impl AgentSessionForkPort for CoreSessionOperationsPort { ) .map_err(runtime_port_error)?; let storage_path = self - .resolve_fork_storage_path(workspace_path, remote_connection_id, remote_ssh_host) + .resolve_session_storage_path(workspace_path, remote_connection_id, remote_ssh_host) .await?; - let (_, turns, _) = self - .coordinator - .restore_session_view_from_storage_path_timed(&storage_path, &source_session_id) - .await - .map_err(runtime_port_error)?; - let source_turn_id = latest_persisted_turn_id(&turns).map_err(runtime_port_error)?; self.fork_at_persisted_turn( &storage_path, source_session_id, - source_turn_id, + None, SessionBranchBoundary::ThroughTurn, ) .await @@ -1168,7 +1443,7 @@ impl AgentSessionForkPort for CoreSessionOperationsPort { ) .map_err(runtime_port_error)?; let storage_path = self - .resolve_fork_storage_path( + .resolve_session_storage_path( request.workspace_path, request.remote_connection_id, request.remote_ssh_host, @@ -1177,7 +1452,7 @@ impl AgentSessionForkPort for CoreSessionOperationsPort { self.fork_at_persisted_turn( &storage_path, request.source_session_id, - request.source_turn_id, + Some(request.source_turn_id), SessionBranchBoundary::ThroughTurn, ) .await @@ -1199,7 +1474,7 @@ impl AgentSessionForkPort for CoreSessionOperationsPort { ) .map_err(runtime_port_error)?; let storage_path = self - .resolve_fork_storage_path( + .resolve_session_storage_path( request.workspace_path, request.remote_connection_id, request.remote_ssh_host, @@ -1208,7 +1483,7 @@ impl AgentSessionForkPort for CoreSessionOperationsPort { self.fork_at_persisted_turn( &storage_path, request.source_session_id, - request.source_turn_id, + Some(request.source_turn_id), SessionBranchBoundary::BeforeTurn, ) .await @@ -1221,9 +1496,31 @@ impl AgentSessionUsagePort for CoreSessionOperationsPort { &self, request: AgentSessionUsageRequest, ) -> PortResult { + let workspace_path = request.workspace_path.clone().ok_or_else(|| { + PortError::new( + PortErrorKind::InvalidRequest, + "Workspace path is required for usage reports", + ) + })?; + let storage_path = self + .resolve_session_storage_path( + workspace_path, + request.remote_connection_id.clone(), + request.remote_ssh_host.clone(), + ) + .await?; + let _read = begin_consistent_persisted_session_read( + self.coordinator.as_ref(), + self.persistence.as_ref(), + &storage_path, + &request.session_id, + ) + .await + .map_err(runtime_port_error)?; generate_core_session_usage_report( self.persistence.as_ref(), self.token_usage_service.as_ref(), + &storage_path, request, ) .await @@ -1260,6 +1557,7 @@ mod tests { use std::sync::Arc; use std::time::Duration; + use crate::service::session::SessionTranscriptExportOptions; use bitfun_agent_runtime::sdk::{AgentEventSource, AgentRuntime}; use bitfun_harness::HarnessRegistry; use bitfun_runtime_ports::{ @@ -1300,7 +1598,8 @@ mod tests { }; use crate::util::errors::BitFunError; use bitfun_agent_runtime::sdk::{ - AgentSessionForkPort, AgentSessionForkRequest, AgentSessionUsageRequest, PortErrorKind, + AgentSessionForkAtTurnRequest, AgentSessionForkPort, AgentSessionForkRequest, + AgentSessionUsagePort, AgentSessionUsageRequest, PortErrorKind, }; use bitfun_events::AgenticEvent; use tokio::sync::RwLock as TokioRwLock; @@ -1614,6 +1913,7 @@ mod tests { let request = LocalWorkspaceSnapshotSessionRequest { workspace_path: workspace.path().to_path_buf(), session_id: "session-empty".to_string(), + max_turn_exclusive: None, }; assert!(port .get_session_files(request.clone()) @@ -1646,6 +1946,7 @@ mod tests { let request = LocalWorkspaceSnapshotSessionRequest { workspace_path: workspace.path().to_path_buf(), session_id: "session-view-only".to_string(), + max_turn_exclusive: None, }; assert!(get_snapshot_manager_for_workspace(workspace.path()).is_none()); @@ -1673,6 +1974,7 @@ mod tests { .get_session_files(LocalWorkspaceSnapshotSessionRequest { workspace_path: workspace.path().to_path_buf(), session_id: "../other-session".to_string(), + max_turn_exclusive: None, }) .await .expect_err("path-like session ids must be rejected"); @@ -1770,7 +2072,7 @@ mod tests { } #[tokio::test] - async fn latest_turn_fork_restores_from_the_resolved_storage_path() { + async fn session_fork_reconciles_pending_revert_and_rejects_hidden_explicit_turns() { let workspace = TestWorkspace::new(); let workspace_root = workspace.path().join("project"); std::fs::create_dir_all(&workspace_root).expect("workspace root"); @@ -1810,7 +2112,7 @@ mod tests { ExecutionEngineConfig::default(), )); let coordinator = Arc::new(ConversationCoordinator::new( - session_manager, + session_manager.clone(), execution_engine, tool_pipeline, event_queue, @@ -1831,23 +2133,24 @@ mod tests { .await .expect("token usage service"), ); - let port = CoreSessionOperationsPort::new(coordinator, token_usage_service); + let scheduler = DialogScheduler::new(coordinator.clone(), session_manager.clone()); + let compatibility = CoreAgentRuntimeCompatibility::build(coordinator.clone(), scheduler); + let port = CoreSessionOperationsPort::new(coordinator.clone(), token_usage_service); let session_id = "session-latest-fork"; + let mut metadata = SessionMetadata::new( + session_id.to_string(), + "Latest fork".to_string(), + "agentic".to_string(), + "model-a".to_string(), + ); + metadata.workspace_path = Some(workspace_root.to_string_lossy().into_owned()); persistence - .save_session_metadata( - &storage_path, - &SessionMetadata::new( - session_id.to_string(), - "Latest fork".to_string(), - "agentic".to_string(), - "model-a".to_string(), - ), - ) + .save_session_metadata(&storage_path, &metadata) .await .expect("session metadata"); - let mut turn = DialogTurnData::new( - "turn-latest".to_string(), + let mut visible_turn = DialogTurnData::new( + "turn-visible".to_string(), 0, session_id.to_string(), UserMessageData { @@ -1857,11 +2160,47 @@ mod tests { metadata: None, }, ); - turn.mark_completed(); + visible_turn.mark_completed(); + persistence + .save_dialog_turn(&storage_path, &visible_turn) + .await + .expect("visible turn"); + let mut hidden_turn = DialogTurnData::new( + "turn-hidden".to_string(), + 1, + session_id.to_string(), + UserMessageData { + id: "user-hidden".to_string(), + content: "hidden by undo".to_string(), + timestamp: 2, + metadata: None, + }, + ); + hidden_turn.mark_completed(); persistence - .save_dialog_turn(&storage_path, &turn) + .save_dialog_turn(&storage_path, &hidden_turn) + .await + .expect("hidden turn"); + persistence + .save_session_revert_state( + &storage_path, + session_id, + &crate::agentic::session::revert::SessionRevertState { + schema_version: crate::agentic::session::revert::SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 1, + original_turn_end: 2, + phase: crate::agentic::session::revert::SessionRevertPhase::Applying, + workspace_checkpoint: Vec::new(), + }, + ) + .await + .expect("pending revert marker"); + + let cold_read_error = coordinator + .load_visible_persisted_session_turns(&storage_path, session_id) .await - .expect("persisted turn"); + .expect_err("cold readers must fail closed on an unfinished undo transition"); + assert!(matches!(cold_read_error, BitFunError::OutcomeUnknown(_))); let result = port .fork_session(AgentSessionForkRequest { @@ -1875,6 +2214,196 @@ mod tests { assert_ne!(result.session_id, session_id); assert_eq!(result.agent_type, "agentic"); + assert_eq!( + persistence + .load_session_turns(&storage_path, &result.session_id) + .await + .expect("fork turns") + .iter() + .map(|turn| turn.turn_id.as_str()) + .collect::>(), + vec!["turn-visible"] + ); + assert_eq!( + persistence + .load_session_revert_state(&storage_path, session_id) + .await + .expect("source marker") + .expect("reconciled marker remains staged") + .phase, + crate::agentic::session::revert::SessionRevertPhase::Staged + ); + let visible_tail = compatibility + .load_persisted_session_turns(&storage_path, session_id, Some(1)) + .await + .expect("staged undo must filter before applying the recent-turn limit"); + assert_eq!( + visible_tail + .iter() + .map(|turn| turn.turn_id.as_str()) + .collect::>(), + vec!["turn-visible"] + ); + assert_eq!( + persistence + .load_visible_session_turns(&storage_path, session_id) + .await + .expect("passive consumers share the persisted visible history") + .iter() + .map(|turn| turn.turn_id.as_str()) + .collect::>(), + vec!["turn-visible"] + ); + let passive_read_mutation = session_manager + .acquire_session_mutation(session_id) + .await + .expect("simulated concurrent Session mutation"); + let passive_reader = coordinator.clone(); + let passive_storage = storage_path.clone(); + let passive_read = tokio::spawn(async move { + passive_reader + .load_visible_persisted_session_turns(&passive_storage, session_id) + .await + }); + tokio::task::yield_now().await; + assert!( + !passive_read.is_finished(), + "passive persisted readers must wait for Core's Session mutation owner" + ); + drop(passive_read_mutation); + assert_eq!( + passive_read + .await + .expect("passive read task") + .expect("passive visible history") + .iter() + .map(|turn| turn.turn_id.as_str()) + .collect::>(), + vec!["turn-visible"] + ); + let export_mutation = session_manager + .acquire_session_mutation(session_id) + .await + .expect("simulated concurrent export mutation"); + let export_reader = coordinator.clone(); + let export_storage = storage_path.clone(); + let transcript_export = tokio::spawn(async move { + export_reader + .export_visible_persisted_session_transcript( + &export_storage, + session_id, + &SessionTranscriptExportOptions { + tools: false, + tool_inputs: false, + thinking: false, + turns: None, + }, + ) + .await + }); + tokio::task::yield_now().await; + assert!( + !transcript_export.is_finished(), + "SessionHistory export must wait for Core's Session mutation owner" + ); + drop(export_mutation); + let transcript_export = transcript_export + .await + .expect("transcript export task") + .expect("visible transcript export"); + let transcript_text = tokio::fs::read_to_string(&transcript_export.transcript_path) + .await + .expect("visible transcript artifact"); + assert!(transcript_text.contains("fork here")); + assert!(!transcript_text.contains("hidden by undo")); + + let hidden_error = port + .fork_session_at_turn(AgentSessionForkAtTurnRequest { + workspace_path: workspace_root.to_string_lossy().into_owned(), + source_session_id: session_id.to_string(), + source_turn_id: "turn-hidden".to_string(), + remote_connection_id: None, + remote_ssh_host: None, + }) + .await + .expect_err("explicit fork cannot target a hidden turn"); + assert_eq!(hidden_error.kind, PortErrorKind::InvalidRequest); + + compatibility + .ensure_session_loaded_from_storage_path(&storage_path, session_id, false) + .await + .expect("load source before branch validation"); + let commit_mutation = compatibility + .begin_persisted_session_mutation(&storage_path, session_id) + .await + .expect("commit mutation"); + compatibility + .commit_session_revert_before_snapshot_mutation(&commit_mutation) + .await + .expect("commit staged suffix"); + let stale_save = compatibility + .save_persisted_dialog_turn(&commit_mutation, &hidden_turn) + .await + .expect_err("a delayed save cannot revive a committed suffix Turn"); + assert!(matches!(stale_save, BitFunError::Validation(_))); + drop(commit_mutation); + assert_eq!( + persistence + .load_session_turns(&storage_path, session_id) + .await + .expect("turns after stale save rejection") + .iter() + .map(|turn| turn.turn_id.as_str()) + .collect::>(), + vec!["turn-visible"] + ); + + persistence + .delete_session_revert_state(&storage_path, session_id) + .await + .expect("clear marker before controlled usage interleaving"); + let mutation_guard = session_manager + .acquire_session_mutation(session_id) + .await + .expect("simulated undo mutation"); + let usage_port = port.clone(); + let usage_workspace_path = workspace_root.to_string_lossy().into_owned(); + let usage_task = tokio::spawn(async move { + usage_port + .generate_session_usage(AgentSessionUsageRequest { + session_id: session_id.to_string(), + workspace_path: Some(usage_workspace_path), + remote_connection_id: None, + remote_ssh_host: None, + include_hidden_subagents: false, + }) + .await + }); + tokio::task::yield_now().await; + assert!( + !usage_task.is_finished(), + "usage must wait for the same mutation boundary as undo" + ); + persistence + .save_session_revert_state( + &storage_path, + session_id, + &crate::agentic::session::revert::SessionRevertState { + schema_version: crate::agentic::session::revert::SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 1, + original_turn_end: 2, + phase: crate::agentic::session::revert::SessionRevertPhase::Staged, + workspace_checkpoint: Vec::new(), + }, + ) + .await + .expect("stage marker while usage is blocked"); + drop(mutation_guard); + let usage = usage_task + .await + .expect("usage task") + .expect("usage after staged marker"); + assert_eq!(usage.scope.turn_count, 1); } #[tokio::test] @@ -1889,6 +2418,7 @@ mod tests { let invalid = generate_core_session_usage_report( &persistence, &token_usage_service, + workspace.path(), AgentSessionUsageRequest { session_id: "../other-session".to_string(), workspace_path: Some(workspace.path().to_string_lossy().into_owned()), @@ -1951,6 +2481,7 @@ mod tests { let report = generate_core_session_usage_report( &persistence, &token_usage_service, + workspace.path(), AgentSessionUsageRequest { session_id: session_id.to_string(), workspace_path: Some(workspace.path().to_string_lossy().into_owned()), diff --git a/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs b/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs index 34fe3db3d..54ab238a3 100644 --- a/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs +++ b/src/crates/assembly/core/src/service/remote_connect/bot/command_router.rs @@ -1796,18 +1796,14 @@ async fn load_last_dialog_pair_from_turns( workspace: Option<&BotWorkspaceRef>, session_id: &str, ) -> Option<(String, String)> { - use crate::agentic::persistence::PersistenceManager; - use crate::infrastructure::PathManager; - const MAX_USER_LEN: usize = 200; const MAX_AI_LEN: usize = 400; let workspace = workspace?; let storage_path = resolve_bot_session_storage_path_for_ref(workspace).await?; - let pm = std::sync::Arc::new(PathManager::new().ok()?); - let store = PersistenceManager::new(pm).ok()?; - let turns = store - .load_session_turns(&storage_path, session_id) + let coordinator = crate::agentic::coordination::get_global_coordinator()?; + let turns = coordinator + .load_visible_persisted_session_turns(&storage_path, session_id) .await .ok()?; let turn = turns.last()?; diff --git a/src/crates/assembly/core/src/service/session_usage/mod.rs b/src/crates/assembly/core/src/service/session_usage/mod.rs index 7d0415b4c..fbf5a7911 100644 --- a/src/crates/assembly/core/src/service/session_usage/mod.rs +++ b/src/crates/assembly/core/src/service/session_usage/mod.rs @@ -16,5 +16,6 @@ pub use bitfun_services_core::session_usage::{ }; pub use service::{ build_session_usage_report_from_sources, build_session_usage_report_from_turns, - generate_session_usage_report, SessionUsageReportRequest, + generate_session_usage_report, generate_session_usage_report_from_storage_path, + SessionUsageReportRequest, }; diff --git a/src/crates/assembly/core/src/service/session_usage/service.rs b/src/crates/assembly/core/src/service/session_usage/service.rs index 0668883f5..cf9753173 100644 --- a/src/crates/assembly/core/src/service/session_usage/service.rs +++ b/src/crates/assembly/core/src/service/session_usage/service.rs @@ -29,19 +29,39 @@ pub async fn generate_session_usage_report( .workspace_path .clone() .ok_or_else(|| BitFunError::validation("Workspace path is required for usage reports"))?; - let turns = persistence_manager - .load_session_turns(Path::new(&workspace_path), &request.session_id) - .await?; - let (session_ids, subagent_scope_complete) = token_usage_session_scope( + generate_session_usage_report_from_storage_path( persistence_manager, + token_usage_service, Path::new(&workspace_path), - &request, - &turns, + request, ) - .await; + .await +} + +/// Build a usage report after a product adapter has resolved the Session +/// storage directory. The request workspace remains the source of snapshot +/// facts and user-facing identity; persisted Session and subagent facts use +/// `session_storage_path`. +pub async fn generate_session_usage_report_from_storage_path( + persistence_manager: &PersistenceManager, + token_usage_service: Option<&TokenUsageService>, + session_storage_path: &Path, + request: SessionUsageReportRequest, +) -> BitFunResult { + let revert_boundary = persistence_manager + .load_session_revert_state(session_storage_path, &request.session_id) + .await? + .map(|state| state.boundary_turn); + let mut turns = persistence_manager + .load_session_turns(session_storage_path, &request.session_id) + .await?; + retain_usage_turns_before_boundary(&mut turns, revert_boundary); + let (session_ids, subagent_scope_complete) = + token_usage_session_scope(persistence_manager, session_storage_path, &request, &turns) + .await; let (token_turn_scope, turn_scope_complete) = token_usage_turn_scope( persistence_manager, - Path::new(&workspace_path), + session_storage_path, &request, &turns, &session_ids, @@ -58,7 +78,7 @@ pub async fn generate_session_usage_report( Vec::new() }; - let snapshot_facts = load_snapshot_facts(&request).await; + let snapshot_facts = load_snapshot_facts(&request, revert_boundary).await; Ok(build_session_usage_report_from_sources_with_scope( request, @@ -170,7 +190,7 @@ fn extend_token_usage_turn_scope( } async fn token_usage_turn_scope( - persistence_manager: &PersistenceManager, + _persistence_manager: &PersistenceManager, workspace_path: &Path, request: &SessionUsageReportRequest, parent_turns: &[DialogTurnData], @@ -187,10 +207,18 @@ async fn token_usage_turn_scope( if session_id == &request.session_id { continue; } - match persistence_manager - .load_session_turns(workspace_path, session_id) - .await - { + let turns = match crate::agentic::coordination::get_global_coordinator() { + Some(coordinator) => { + coordinator + .load_visible_persisted_session_turns(workspace_path, session_id) + .await + } + None => { + complete = false; + continue; + } + }; + match turns { Ok(turns) => extend_token_usage_turn_scope(&mut scope, session_id, &turns), Err(_) => complete = false, } @@ -299,7 +327,10 @@ fn build_session_usage_report_from_sources_with_scope( report } -async fn load_snapshot_facts(request: &SessionUsageReportRequest) -> UsageSnapshotFacts { +async fn load_snapshot_facts( + request: &SessionUsageReportRequest, + revert_boundary: Option, +) -> UsageSnapshotFacts { let Some(workspace_path) = request.workspace_path.as_deref() else { return UsageSnapshotFacts::default(); }; @@ -314,6 +345,7 @@ async fn load_snapshot_facts(request: &SessionUsageReportRequest) -> UsageSnapsh operations: session .operations .into_iter() + .filter(|operation| usage_index_is_visible(operation.turn_index, revert_boundary)) .map(snapshot_operation_from_file_operation) .collect(), }, @@ -321,6 +353,17 @@ async fn load_snapshot_facts(request: &SessionUsageReportRequest) -> UsageSnapsh } } +fn retain_usage_turns_before_boundary( + turns: &mut Vec, + revert_boundary: Option, +) { + turns.retain(|turn| usage_index_is_visible(turn.turn_index, revert_boundary)); +} + +fn usage_index_is_visible(turn_index: usize, revert_boundary: Option) -> bool { + revert_boundary.is_none_or(|boundary| turn_index < boundary) +} + fn is_reportable_usage_turn(turn: &DialogTurnData) -> bool { turn.kind != DialogTurnKind::LocalCommand } @@ -1878,6 +1921,54 @@ mod tests { assert_eq!(report.files.files[0].operation_count, 1); } + #[test] + fn staged_revert_boundary_excludes_hidden_turn_token_tool_and_snapshot_usage() { + let request = test_request(None); + let mut turns = vec![ + test_turn("turn-visible", 0, DialogTurnKind::UserDialog), + test_turn("turn-hidden", 1, DialogTurnKind::UserDialog), + ]; + retain_usage_turns_before_boundary(&mut turns, Some(1)); + + let mut visible_record = test_token_record("model-a", 100, 20, 0); + visible_record.turn_id = "turn-visible".to_string(); + let mut hidden_record = test_token_record("model-a", 900, 100, 0); + hidden_record.turn_id = "turn-hidden".to_string(); + let snapshot_facts = test_snapshot_facts( + [ + test_snapshot_operation( + "op-visible", + 0, + "D:/workspace/bitfun/src/visible.rs", + 2, + 0, + ), + test_snapshot_operation("op-hidden", 1, "D:/workspace/bitfun/src/hidden.rs", 30, 4), + ] + .into_iter() + .filter(|operation| usage_index_is_visible(operation.turn_index, Some(1))) + .collect(), + ); + + let report = build_session_usage_report_from_sources( + request, + &turns, + &[visible_record, hidden_record], + &snapshot_facts, + 1_778_347_200_000, + ); + + assert_eq!(report.scope.turn_count, 1); + assert_eq!(report.scope.to_turn_id.as_deref(), Some("turn-visible")); + assert_eq!(report.tokens.total_tokens, Some(120)); + assert_eq!( + report.tools.iter().map(|tool| tool.call_count).sum::(), + 1 + ); + assert_eq!(report.files.files.len(), 1); + assert_eq!(report.files.files[0].path_label, "src/visible.rs"); + } + #[test] fn report_classifies_deferred_calls_by_effective_identity_and_nested_args() { let request = test_request(None); diff --git a/src/crates/assembly/core/src/service/snapshot/manager.rs b/src/crates/assembly/core/src/service/snapshot/manager.rs index 8dc144d96..6507a89aa 100644 --- a/src/crates/assembly/core/src/service/snapshot/manager.rs +++ b/src/crates/assembly/core/src/service/snapshot/manager.rs @@ -90,6 +90,53 @@ impl SnapshotManager { .await } + pub(crate) async fn prepare_workspace_revert( + &self, + session_id: &str, + state: &mut crate::agentic::session::revert::SessionRevertState, + ) -> SnapshotResult<()> { + self.snapshot_service + .read() + .await + .prepare_workspace_revert(session_id, state) + .await + } + + pub(crate) async fn apply_workspace_revert( + &self, + session_id: &str, + state: &crate::agentic::session::revert::SessionRevertState, + ) -> SnapshotResult> { + self.snapshot_service + .read() + .await + .apply_workspace_revert(session_id, state) + .await + } + + pub(crate) async fn commit_workspace_revert( + &self, + session_id: &str, + state: &crate::agentic::session::revert::SessionRevertState, + ) -> SnapshotResult<()> { + self.snapshot_service + .read() + .await + .commit_workspace_revert(session_id, state) + .await + } + + pub(crate) async fn delete_workspace_revert_checkpoint( + &self, + state: &crate::agentic::session::revert::SessionRevertState, + ) -> SnapshotResult<()> { + self.snapshot_service + .read() + .await + .delete_workspace_revert_checkpoint(state) + .await + } + /// Accepts all changes in a session. pub async fn accept_session(&self, session_id: &str) -> SnapshotResult<()> { let snapshot_service = self.snapshot_service.read().await; @@ -116,14 +163,34 @@ impl SnapshotManager { /// Returns the list of files affected by a session. pub async fn get_session_files(&self, session_id: &str) -> SnapshotResult> { + self.get_session_files_before(session_id, None).await + } + + pub async fn get_session_files_before( + &self, + session_id: &str, + max_turn_exclusive: Option, + ) -> SnapshotResult> { let snapshot_service = self.snapshot_service.read().await; - snapshot_service.get_session_files(session_id).await + snapshot_service + .get_session_files_before(session_id, max_turn_exclusive) + .await } /// Returns the list of turns for a session. pub async fn get_session_turns(&self, session_id: &str) -> SnapshotResult> { + self.get_session_turns_before(session_id, None).await + } + + pub async fn get_session_turns_before( + &self, + session_id: &str, + max_turn_exclusive: Option, + ) -> SnapshotResult> { let snapshot_service = self.snapshot_service.read().await; - snapshot_service.get_session_turns(session_id).await + snapshot_service + .get_session_turns_before(session_id, max_turn_exclusive) + .await } /// Returns the list of files modified in a turn. @@ -131,10 +198,20 @@ impl SnapshotManager { &self, session_id: &str, turn_index: usize, + ) -> SnapshotResult> { + self.get_turn_files_before(session_id, turn_index, None) + .await + } + + pub async fn get_turn_files_before( + &self, + session_id: &str, + turn_index: usize, + max_turn_exclusive: Option, ) -> SnapshotResult> { let snapshot_service = self.snapshot_service.read().await; snapshot_service - .get_turn_files(session_id, turn_index) + .get_turn_files_before(session_id, turn_index, max_turn_exclusive) .await } @@ -144,11 +221,27 @@ impl SnapshotManager { session_id: &str, file_path: &str, anchor_operation_id: Option<&str>, + ) -> SnapshotResult { + self.get_file_diff_before(session_id, file_path, anchor_operation_id, None) + .await + } + + pub async fn get_file_diff_before( + &self, + session_id: &str, + file_path: &str, + anchor_operation_id: Option<&str>, + max_turn_exclusive: Option, ) -> SnapshotResult { let snapshot_service = self.snapshot_service.read().await; let file_path = std::path::Path::new(file_path); let (original, modified, anchor_line) = snapshot_service - .get_file_diff_with_anchor(session_id, file_path, anchor_operation_id) + .get_file_diff_with_anchor_before( + session_id, + file_path, + anchor_operation_id, + max_turn_exclusive, + ) .await?; Ok(serde_json::json!({ @@ -163,11 +256,21 @@ impl SnapshotManager { &self, session_id: &str, file_path: &str, + ) -> SnapshotResult { + self.get_session_file_diff_stats_before(session_id, file_path, None) + .await + } + + pub async fn get_session_file_diff_stats_before( + &self, + session_id: &str, + file_path: &str, + max_turn_exclusive: Option, ) -> SnapshotResult { let snapshot_service = self.snapshot_service.read().await; let file_path = std::path::Path::new(file_path); snapshot_service - .get_session_file_diff_stats(session_id, file_path) + .get_session_file_diff_stats_before(session_id, file_path, max_turn_exclusive) .await } @@ -175,10 +278,20 @@ impl SnapshotManager { &self, session_id: &str, operation_id: &str, + ) -> SnapshotResult { + self.get_operation_summary_before(session_id, operation_id, None) + .await + } + + pub async fn get_operation_summary_before( + &self, + session_id: &str, + operation_id: &str, + max_turn_exclusive: Option, ) -> SnapshotResult { let snapshot_service = self.snapshot_service.read().await; let op = snapshot_service - .get_operation_summary(session_id, operation_id) + .get_operation_summary_before(session_id, operation_id, max_turn_exclusive) .await?; Ok(serde_json::json!({ "operation_id": op.operation_id, @@ -196,26 +309,49 @@ impl SnapshotManager { pub async fn get_session( &self, session_id: &str, + ) -> SnapshotResult { + self.get_session_before(session_id, None).await + } + + pub async fn get_session_before( + &self, + session_id: &str, + max_turn_exclusive: Option, ) -> SnapshotResult { let snapshot_service = self.snapshot_service.read().await; - snapshot_service.get_session(session_id).await + snapshot_service + .get_session_before(session_id, max_turn_exclusive) + .await } /// Returns session statistics. pub async fn get_session_stats(&self, session_id: &str) -> SnapshotResult { - let stats = self.get_session_stats_fact(session_id).await?; + self.get_session_stats_before(session_id, None).await + } + + pub async fn get_session_stats_before( + &self, + session_id: &str, + max_turn_exclusive: Option, + ) -> SnapshotResult { + let stats = self + .get_session_stats_fact_before(session_id, max_turn_exclusive) + .await?; serde_json::to_value(stats).map_err(|e| { SnapshotError::ConfigError(format!("Failed to serialize statistics: {}", e)) }) } - pub(crate) async fn get_session_stats_fact( + pub(crate) async fn get_session_stats_fact_before( &self, session_id: &str, + max_turn_exclusive: Option, ) -> SnapshotResult { let snapshot_service = self.snapshot_service.read().await; - snapshot_service.get_session_stats(session_id).await + snapshot_service + .get_session_stats_before(session_id, max_turn_exclusive) + .await } /// Returns system statistics. diff --git a/src/crates/assembly/core/src/service/snapshot/service.rs b/src/crates/assembly/core/src/service/snapshot/service.rs index 9fbd6fb5c..7b7ef454a 100644 --- a/src/crates/assembly/core/src/service/snapshot/service.rs +++ b/src/crates/assembly/core/src/service/snapshot/service.rs @@ -200,11 +200,27 @@ impl SnapshotService { session_id: &str, file_path: &Path, anchor_operation_id: Option<&str>, + ) -> SnapshotResult<(String, String, Option)> { + self.get_file_diff_with_anchor_before(session_id, file_path, anchor_operation_id, None) + .await + } + + pub async fn get_file_diff_with_anchor_before( + &self, + session_id: &str, + file_path: &Path, + anchor_operation_id: Option<&str>, + max_turn_exclusive: Option, ) -> SnapshotResult<(String, String, Option)> { self.ensure_initialized().await?; let snapshot_core = self.snapshot_core.read().await; snapshot_core - .get_file_diff_with_anchor(file_path, session_id, anchor_operation_id) + .get_file_diff_with_anchor_before( + file_path, + session_id, + anchor_operation_id, + max_turn_exclusive, + ) .await } @@ -212,11 +228,21 @@ impl SnapshotService { &self, session_id: &str, file_path: &Path, + ) -> SnapshotResult { + self.get_session_file_diff_stats_before(session_id, file_path, None) + .await + } + + pub async fn get_session_file_diff_stats_before( + &self, + session_id: &str, + file_path: &Path, + max_turn_exclusive: Option, ) -> SnapshotResult { self.ensure_initialized().await?; let snapshot_core = self.snapshot_core.read().await; snapshot_core - .get_session_file_diff_stats(session_id, file_path) + .get_session_file_diff_stats_before(session_id, file_path, max_turn_exclusive) .await } @@ -224,10 +250,20 @@ impl SnapshotService { &self, session_id: &str, operation_id: &str, + ) -> SnapshotResult { + self.get_operation_summary_before(session_id, operation_id, None) + .await + } + + pub async fn get_operation_summary_before( + &self, + session_id: &str, + operation_id: &str, + max_turn_exclusive: Option, ) -> SnapshotResult { self.ensure_initialized().await?; let snapshot_core = self.snapshot_core.read().await; - snapshot_core.get_operation(session_id, operation_id) + snapshot_core.get_operation_before(session_id, operation_id, max_turn_exclusive) } /// Complete a file modification (after snapshot + diff summary). @@ -321,6 +357,57 @@ impl SnapshotService { snapshot_core.rollback_to_turn(session_id, turn_index).await } + pub(crate) async fn prepare_workspace_revert( + &self, + session_id: &str, + state: &mut crate::agentic::session::revert::SessionRevertState, + ) -> SnapshotResult<()> { + self.ensure_writable().await?; + self.snapshot_core + .write() + .await + .prepare_workspace_revert(session_id, state) + .await + } + + pub(crate) async fn apply_workspace_revert( + &self, + session_id: &str, + state: &crate::agentic::session::revert::SessionRevertState, + ) -> SnapshotResult> { + self.ensure_writable().await?; + self.snapshot_core + .read() + .await + .apply_workspace_revert(session_id, state) + .await + } + + pub(crate) async fn commit_workspace_revert( + &self, + session_id: &str, + state: &crate::agentic::session::revert::SessionRevertState, + ) -> SnapshotResult<()> { + self.ensure_writable().await?; + self.snapshot_core + .write() + .await + .commit_workspace_revert(session_id, state) + .await + } + + pub(crate) async fn delete_workspace_revert_checkpoint( + &self, + state: &crate::agentic::session::revert::SessionRevertState, + ) -> SnapshotResult<()> { + self.ensure_writable().await?; + self.snapshot_core + .write() + .await + .delete_workspace_revert_checkpoint(state) + .await + } + pub async fn accept_session(&self, session_id: &str) -> SnapshotResult<()> { self.ensure_writable().await?; info!("Accepting session changes: session_id={}", session_id); @@ -379,25 +466,51 @@ impl SnapshotService { } pub async fn get_session_files(&self, session_id: &str) -> SnapshotResult> { + self.get_session_files_before(session_id, None).await + } + + pub async fn get_session_files_before( + &self, + session_id: &str, + max_turn_exclusive: Option, + ) -> SnapshotResult> { self.ensure_initialized().await?; let snapshot_core = self.snapshot_core.read().await; - Ok(snapshot_core.get_session_files(session_id)) + Ok(snapshot_core.get_session_files_before(session_id, max_turn_exclusive)) } pub async fn get_session_turns(&self, session_id: &str) -> SnapshotResult> { + self.get_session_turns_before(session_id, None).await + } + + pub async fn get_session_turns_before( + &self, + session_id: &str, + max_turn_exclusive: Option, + ) -> SnapshotResult> { self.ensure_initialized().await?; let snapshot_core = self.snapshot_core.read().await; - Ok(snapshot_core.get_session_turns(session_id)) + Ok(snapshot_core.get_session_turns_before(session_id, max_turn_exclusive)) } pub async fn get_turn_files( &self, session_id: &str, turn_index: usize, + ) -> SnapshotResult> { + self.get_turn_files_before(session_id, turn_index, None) + .await + } + + pub async fn get_turn_files_before( + &self, + session_id: &str, + turn_index: usize, + max_turn_exclusive: Option, ) -> SnapshotResult> { self.ensure_initialized().await?; let snapshot_core = self.snapshot_core.read().await; - Ok(snapshot_core.get_turn_files(session_id, turn_index)) + Ok(snapshot_core.get_turn_files_before(session_id, turn_index, max_turn_exclusive)) } pub async fn get_file_diff( @@ -411,9 +524,17 @@ impl SnapshotService { } pub async fn get_session_stats(&self, session_id: &str) -> SnapshotResult { + self.get_session_stats_before(session_id, None).await + } + + pub async fn get_session_stats_before( + &self, + session_id: &str, + max_turn_exclusive: Option, + ) -> SnapshotResult { self.ensure_initialized().await?; let snapshot_core = self.snapshot_core.read().await; - Ok(snapshot_core.get_session_stats(session_id)) + Ok(snapshot_core.get_session_stats_before(session_id, max_turn_exclusive)) } pub async fn get_system_stats(&self) -> SnapshotResult { @@ -450,9 +571,18 @@ impl SnapshotService { } pub async fn get_session(&self, session_id: &str) -> SnapshotResult { + self.get_session_before(session_id, None).await + } + + pub async fn get_session_before( + &self, + session_id: &str, + max_turn_exclusive: Option, + ) -> SnapshotResult { self.ensure_initialized().await?; let snapshot_core = self.snapshot_core.read().await; - let operations = snapshot_core.get_session_operations(session_id); + let operations = + snapshot_core.get_session_operations_before(session_id, max_turn_exclusive); Ok(SessionInfo { session_id: session_id.to_string(), operations, diff --git a/src/crates/assembly/core/src/service/snapshot/snapshot_core.rs b/src/crates/assembly/core/src/service/snapshot/snapshot_core.rs index bb3e6f608..a247b7970 100644 --- a/src/crates/assembly/core/src/service/snapshot/snapshot_core.rs +++ b/src/crates/assembly/core/src/service/snapshot/snapshot_core.rs @@ -1,3 +1,4 @@ +use crate::agentic::session::revert::{SessionRevertState, SessionWorkspaceCheckpoint}; use crate::service::snapshot::snapshot_system::FileSnapshotSystem; use crate::service::snapshot::types::{ DiffSummary, FileOperation, OperationType, SessionFileDiffStats, SnapshotError, SnapshotResult, @@ -150,8 +151,21 @@ impl SnapshotCore { tool_input: serde_json::Value, operation_id_override: Option, ) -> SnapshotResult { + let operation_id = operation_id_override + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| Uuid::new_v4().to_string()); + if self.operation_index.contains_key(&operation_id) { + return Err(SnapshotError::ConfigError(format!( + "operation_id already exists: {}", + operation_id + ))); + } let before_snapshot_id = if file_path.exists() { - Some(self.snapshot_system.create_snapshot(&file_path).await?) + Some( + self.snapshot_system + .create_owned_snapshot(&file_path) + .await?, + ) } else { None }; @@ -207,15 +221,6 @@ impl SnapshotCore { .or_insert_with(|| SessionHistory::new(session_id.to_string())); let turn = session.ensure_turn_mut(turn_index); let seq_in_turn = turn.operations.len(); - let operation_id = operation_id_override - .filter(|s| !s.is_empty()) - .unwrap_or_else(|| Uuid::new_v4().to_string()); - if self.operation_index.contains_key(&operation_id) { - return Err(SnapshotError::ConfigError(format!( - "operation_id already exists: {}", - operation_id - ))); - } turn.operations.push(FileOperation { operation_id: operation_id.clone(), @@ -251,6 +256,15 @@ impl SnapshotCore { &self, session_id: &str, operation_id: &str, + ) -> SnapshotResult { + self.get_operation_before(session_id, operation_id, None) + } + + pub fn get_operation_before( + &self, + session_id: &str, + operation_id: &str, + max_turn_exclusive: Option, ) -> SnapshotResult { let Some((sid, turn_index, seq)) = self.operation_index.get(operation_id).cloned() else { return Err(SnapshotError::OperationNotFound(operation_id.to_string())); @@ -261,6 +275,9 @@ impl SnapshotCore { operation_id, session_id, sid ))); } + if max_turn_exclusive.is_some_and(|end| turn_index >= end) { + return Err(SnapshotError::OperationNotFound(operation_id.to_string())); + } let session = self .sessions .get(session_id) @@ -311,7 +328,11 @@ impl SnapshotCore { op.tool_context.execution_time_ms = execution_time_ms; let after_snapshot_id = if op.file_path.exists() { - Some(self.snapshot_system.create_snapshot(&op.file_path).await?) + Some( + self.snapshot_system + .create_owned_snapshot(&op.file_path) + .await?, + ) } else { None }; @@ -348,13 +369,38 @@ impl SnapshotCore { } pub fn get_session_turns(&self, session_id: &str) -> Vec { + self.get_session_turns_before(session_id, None) + } + + pub fn get_session_turns_before( + &self, + session_id: &str, + max_turn_exclusive: Option, + ) -> Vec { let Some(session) = self.sessions.get(session_id) else { return Vec::new(); }; - session.turns.keys().cloned().collect() + session + .turns + .keys() + .copied() + .filter(|turn| max_turn_exclusive.is_none_or(|end| *turn < end)) + .collect() } pub fn get_turn_files(&self, session_id: &str, turn_index: usize) -> Vec { + self.get_turn_files_before(session_id, turn_index, None) + } + + pub fn get_turn_files_before( + &self, + session_id: &str, + turn_index: usize, + max_turn_exclusive: Option, + ) -> Vec { + if max_turn_exclusive.is_some_and(|end| turn_index >= end) { + return Vec::new(); + } let Some(session) = self.sessions.get(session_id) else { return Vec::new(); }; @@ -370,22 +416,43 @@ impl SnapshotCore { } pub fn get_session_files(&self, session_id: &str) -> Vec { + self.get_session_files_before(session_id, None) + } + + pub fn get_session_files_before( + &self, + session_id: &str, + max_turn_exclusive: Option, + ) -> Vec { let Some(session) = self.sessions.get(session_id) else { return Vec::new(); }; unique_paths( session .all_operations_iter() + .filter(|op| operation_is_before(op, max_turn_exclusive)) .filter(|op| operation_is_completed_for_session_file(op)) .map(|op| op.file_path.clone()), ) } pub fn get_session_operations(&self, session_id: &str) -> Vec { + self.get_session_operations_before(session_id, None) + } + + pub fn get_session_operations_before( + &self, + session_id: &str, + max_turn_exclusive: Option, + ) -> Vec { let Some(session) = self.sessions.get(session_id) else { return Vec::new(); }; - session.all_operations_iter().cloned().collect() + session + .all_operations_iter() + .filter(|op| operation_is_before(op, max_turn_exclusive)) + .cloned() + .collect() } pub fn get_all_modified_files(&self) -> Vec { @@ -402,14 +469,24 @@ impl SnapshotCore { } pub fn get_session_stats(&self, session_id: &str) -> SessionStats { + self.get_session_stats_before(session_id, None) + } + + pub fn get_session_stats_before( + &self, + session_id: &str, + max_turn_exclusive: Option, + ) -> SessionStats { let ops: Vec = self - .get_session_operations(session_id) + .get_session_operations_before(session_id, max_turn_exclusive) .into_iter() .filter(operation_is_completed_for_session_file) .collect(); let total_changes = ops.len(); let total_files = unique_paths(ops.iter().map(|op| op.file_path.clone())).len(); - let total_turns = self.get_session_turns(session_id).len(); + let total_turns = self + .get_session_turns_before(session_id, max_turn_exclusive) + .len(); SessionStats { session_id: session_id.to_string(), total_files, @@ -488,12 +565,21 @@ impl SnapshotCore { &self, file_path: &Path, session_id: &str, + ) -> SnapshotResult<(String, String)> { + self.get_file_diff_before(file_path, session_id, None).await + } + + pub async fn get_file_diff_before( + &self, + file_path: &Path, + session_id: &str, + max_turn_exclusive: Option, ) -> SnapshotResult<(String, String)> { let Some(session) = self.sessions.get(session_id) else { return Err(SnapshotError::SessionNotFound(session_id.to_string())); }; - let Some(boundary) = session_file_boundary(session, file_path) else { + let Some(boundary) = session_file_boundary(session, file_path, max_turn_exclusive) else { debug!( "No completed session file operation found for diff: file_path={:?} session_id={}", file_path, session_id @@ -531,13 +617,26 @@ impl SnapshotCore { session_id: &str, anchor_operation_id: Option<&str>, ) -> SnapshotResult<(String, String, Option)> { - let (before, after) = self.get_file_diff(file_path, session_id).await?; + self.get_file_diff_with_anchor_before(file_path, session_id, anchor_operation_id, None) + .await + } + + pub async fn get_file_diff_with_anchor_before( + &self, + file_path: &Path, + session_id: &str, + anchor_operation_id: Option<&str>, + max_turn_exclusive: Option, + ) -> SnapshotResult<(String, String, Option)> { + let (before, after) = self + .get_file_diff_before(file_path, session_id, max_turn_exclusive) + .await?; let Some(operation_id) = anchor_operation_id.filter(|s| !s.is_empty()) else { return Ok((before, after, None)); }; - let op = self.get_operation(session_id, operation_id)?; + let op = self.get_operation_before(session_id, operation_id, max_turn_exclusive)?; if op.file_path != file_path { return Ok((before, after, None)); } @@ -575,12 +674,22 @@ impl SnapshotCore { &self, session_id: &str, file_path: &Path, + ) -> SnapshotResult { + self.get_session_file_diff_stats_before(session_id, file_path, None) + .await + } + + pub async fn get_session_file_diff_stats_before( + &self, + session_id: &str, + file_path: &Path, + max_turn_exclusive: Option, ) -> SnapshotResult { let Some(session) = self.sessions.get(session_id) else { return Err(SnapshotError::SessionNotFound(session_id.to_string())); }; - let Some(boundary) = session_file_boundary(session, file_path) else { + let Some(boundary) = session_file_boundary(session, file_path, max_turn_exclusive) else { return Ok(SessionFileDiffStats { file_path: file_path.to_string_lossy().to_string(), lines_added: 0, @@ -604,7 +713,8 @@ impl SnapshotCore { || before_bytes > SESSION_FILE_DIFF_STATS_MAX_SOURCE_BYTES; if too_large { - let agg = aggregate_operations_diff_summary_for_file(session, file_path); + let agg = + aggregate_operations_diff_summary_for_file(session, file_path, max_turn_exclusive); let change_kind = change_kind_from_session_boundary(&boundary); debug!( "get_session_file_diff_stats: approximate session_id={} file_path={:?} after_bytes={} before_bytes={} lines_added={} lines_removed={}", @@ -624,7 +734,9 @@ impl SnapshotCore { }); } - let (before, after) = self.get_file_diff(file_path, session_id).await?; + let (before, after) = self + .get_file_diff_before(file_path, session_id, max_turn_exclusive) + .await?; let summary = compute_diff_summary(&before, &after); let change_kind = change_kind_from_session_boundary(&boundary); debug!( @@ -735,36 +847,199 @@ impl SnapshotCore { Ok(restored) } - pub async fn cleanup_session(&mut self, session_id: &str) -> SnapshotResult<()> { - let snapshot_ids_to_delete: Vec = - if let Some(session) = self.sessions.get(session_id) { - session - .all_operations_iter() - .flat_map(|op| { - let mut ids = Vec::new(); - if let Some(ref id) = op.before_snapshot_id { - if !id.starts_with("empty_snapshot_") { - ids.push(id.clone()); - } - } - if let Some(ref id) = op.after_snapshot_id { - if !id.starts_with("empty_snapshot_") { - ids.push(id.clone()); - } - } - ids - }) - .collect() + /// Extends the pre-undo workspace checkpoint for every path affected at or + /// after the requested boundary. Repeated undo can safely discover earlier + /// paths because paths not already checkpointed are still at their original + /// pre-undo state before that earlier boundary is applied. + pub(crate) async fn prepare_workspace_revert( + &mut self, + session_id: &str, + state: &mut SessionRevertState, + ) -> SnapshotResult<()> { + let Some(session) = self.sessions.get(session_id) else { + return Ok(()); + }; + let existing = state + .workspace_checkpoint + .iter() + .map(|checkpoint| checkpoint.path.clone()) + .collect::>(); + let mut affected = session + .all_operations_iter() + .filter(|operation| operation.turn_index >= state.boundary_turn) + .flat_map(|operation| { + [ + operation + .path_before + .clone() + .unwrap_or_else(|| operation.file_path.clone()), + operation + .path_after + .clone() + .unwrap_or_else(|| operation.file_path.clone()), + ] + }) + .filter(|path| !existing.contains(path)) + .collect::>(); + affected.sort(); + affected.dedup(); + + for path in affected { + let snapshot_id = if path.exists() { + Some(self.snapshot_system.create_owned_snapshot(&path).await?) } else { - Vec::new() + None }; + state + .workspace_checkpoint + .push(SessionWorkspaceCheckpoint { path, snapshot_id }); + } + state + .workspace_checkpoint + .sort_by(|left, right| left.path.cmp(&right.path)); + Ok(()) + } + + /// Applies a staged Session boundary without deleting operation history. + /// Restoring the original checkpoint first makes retries and redo idempotent. + pub(crate) async fn apply_workspace_revert( + &self, + session_id: &str, + state: &SessionRevertState, + ) -> SnapshotResult> { + let mut restored = self.restore_workspace_revert(state).await?; + let Some(session) = self.sessions.get(session_id) else { + return Ok(restored); + }; + let mut operations = session + .all_operations_iter() + .filter(|operation| operation.turn_index >= state.boundary_turn) + .cloned() + .collect::>(); + operations.sort_by_key(|operation| (operation.turn_index, operation.seq_in_turn)); + operations.reverse(); + restored.extend( + self.apply_rollback_ops_with_policy(&operations, true) + .await?, + ); + Ok(unique_paths(restored.into_iter())) + } + + pub(crate) async fn restore_workspace_revert( + &self, + state: &SessionRevertState, + ) -> SnapshotResult> { + let mut restored = Vec::new(); + for checkpoint in &state.workspace_checkpoint { + match checkpoint.snapshot_id.as_deref() { + Some(snapshot_id) => { + self.snapshot_system + .restore_file(snapshot_id, &checkpoint.path) + .await?; + restored.push(checkpoint.path.clone()); + } + None if checkpoint.path.exists() => { + tokio::fs::remove_file(&checkpoint.path).await?; + restored.push(checkpoint.path.clone()); + } + None => {} + } + } + Ok(unique_paths(restored.into_iter())) + } + + pub(crate) async fn commit_workspace_revert( + &mut self, + session_id: &str, + state: &SessionRevertState, + ) -> SnapshotResult<()> { + let discarded_snapshot_ids = self + .sessions + .get(session_id) + .into_iter() + .flat_map(|session| { + session + .turns + .range(state.boundary_turn..) + .flat_map(|(_, turn)| turn.operations.iter()) + }) + .flat_map(Self::operation_snapshot_ids) + .collect::>(); + let retained_snapshot_ids = self + .sessions + .iter() + .flat_map(|(candidate_session_id, session)| { + session.all_operations_iter().filter(move |operation| { + candidate_session_id != session_id || operation.turn_index < state.boundary_turn + }) + }) + .flat_map(Self::operation_snapshot_ids) + .collect::>(); + // Release suffix-owned handles before pruning their durable history. + // A retry can still rediscover every handle if deletion or persistence + // fails part-way through this committing phase. + for snapshot_id in discarded_snapshot_ids.difference(&retained_snapshot_ids) { + match self.snapshot_system.delete_snapshot(snapshot_id).await { + Ok(()) | Err(SnapshotError::SnapshotNotFound(_)) => {} + Err(error) => return Err(error), + } + } + if let Some(session) = self.sessions.get_mut(session_id) { + session + .turns + .retain(|turn_index, _| *turn_index < state.boundary_turn); + session.last_updated = SystemTime::now(); + self.persist_session(session_id).await?; + self.rebuild_operation_index(); + } + self.delete_workspace_revert_checkpoint(state).await + } + + pub(crate) async fn delete_workspace_revert_checkpoint( + &mut self, + state: &SessionRevertState, + ) -> SnapshotResult<()> { + let operation_snapshot_ids = self.referenced_operation_snapshot_ids(); + for checkpoint in &state.workspace_checkpoint { + if let Some(snapshot_id) = checkpoint.snapshot_id.as_deref() { + // Markers written by older builds could borrow a deduplicated + // operation snapshot ID. Preserve that handle until its + // operation history owner releases it. + if operation_snapshot_ids.contains(snapshot_id) { + continue; + } + match self.snapshot_system.delete_snapshot(snapshot_id).await { + Ok(()) | Err(SnapshotError::SnapshotNotFound(_)) => {} + Err(error) => return Err(error), + } + } + } + Ok(()) + } + + pub async fn cleanup_session(&mut self, session_id: &str) -> SnapshotResult<()> { + let owned_snapshot_ids = self + .sessions + .get(session_id) + .into_iter() + .flat_map(SessionHistory::all_operations_iter) + .flat_map(Self::operation_snapshot_ids) + .collect::>(); + + self.sessions.remove(session_id); + let retained_snapshot_ids = self.referenced_operation_snapshot_ids(); + let snapshot_ids_to_delete = owned_snapshot_ids + .difference(&retained_snapshot_ids) + .cloned() + .collect::>(); for snapshot_id in &snapshot_ids_to_delete { - if let Err(e) = self.snapshot_system.delete_snapshot(snapshot_id).await { - warn!( + match self.snapshot_system.delete_snapshot(snapshot_id).await { + Ok(()) | Err(SnapshotError::SnapshotNotFound(_)) => {} + Err(error) => warn!( "Failed to delete snapshot: snapshot_id={} error={}", - snapshot_id, e - ); + snapshot_id, error + ), } } @@ -776,8 +1051,6 @@ impl SnapshotCore { ); } - self.sessions.remove(session_id); - self.delete_session_file(session_id).await?; self.rebuild_operation_index(); @@ -785,6 +1058,25 @@ impl SnapshotCore { Ok(()) } + fn operation_snapshot_ids(operation: &FileOperation) -> impl Iterator + '_ { + [ + operation.before_snapshot_id.as_ref(), + operation.after_snapshot_id.as_ref(), + ] + .into_iter() + .flatten() + .filter(|snapshot_id| !snapshot_id.starts_with("empty_snapshot_")) + .cloned() + } + + fn referenced_operation_snapshot_ids(&self) -> HashSet { + self.sessions + .values() + .flat_map(SessionHistory::all_operations_iter) + .flat_map(Self::operation_snapshot_ids) + .collect() + } + pub async fn cleanup_file_session( &mut self, session_id: &str, @@ -834,6 +1126,14 @@ impl SnapshotCore { } async fn apply_rollback_ops(&self, ops: &[FileOperation]) -> SnapshotResult> { + self.apply_rollback_ops_with_policy(ops, false).await + } + + async fn apply_rollback_ops_with_policy( + &self, + ops: &[FileOperation], + fail_on_delete_error: bool, + ) -> SnapshotResult> { let mut restored_files: Vec = Vec::new(); for op in ops { @@ -850,6 +1150,9 @@ impl SnapshotCore { if before_path != after_path && after_path.exists() { if let Err(e) = tokio::fs::remove_file(&after_path).await { + if fail_on_delete_error { + return Err(SnapshotError::Io(e)); + } warn!( "Failed to delete after_path: path={} error={}", after_path.display(), @@ -862,6 +1165,9 @@ impl SnapshotCore { None => { if after_path.exists() { if let Err(e) = tokio::fs::remove_file(&after_path).await { + if fail_on_delete_error { + return Err(SnapshotError::Io(e)); + } warn!( "Failed to delete file: path={} error={}", after_path.display(), @@ -874,7 +1180,11 @@ impl SnapshotCore { } Some(snapshot_id) if snapshot_id.starts_with("empty_snapshot_") => { if after_path.exists() { - let _ = tokio::fs::remove_file(&after_path).await; + if let Err(error) = tokio::fs::remove_file(&after_path).await { + if fail_on_delete_error { + return Err(SnapshotError::Io(error)); + } + } restored_files.push(after_path.clone()); } } @@ -1014,12 +1324,18 @@ fn operation_is_completed_for_session_file(op: &FileOperation) -> bool { || op.diff_summary.lines_modified > 0 } +fn operation_is_before(op: &FileOperation, max_turn_exclusive: Option) -> bool { + max_turn_exclusive.is_none_or(|end| op.turn_index < end) +} + fn completed_session_operations_for_file<'a>( session: &'a SessionHistory, file_path: &Path, + max_turn_exclusive: Option, ) -> Vec<&'a FileOperation> { let mut operations: Vec<&FileOperation> = session .all_operations_iter() + .filter(|op| operation_is_before(op, max_turn_exclusive)) .filter(|op| SnapshotCore::operation_matches_file_path(op, file_path)) .filter(|op| operation_is_completed_for_session_file(op)) .collect(); @@ -1031,8 +1347,9 @@ fn completed_session_operations_for_file<'a>( fn session_file_boundary( session: &SessionHistory, file_path: &Path, + max_turn_exclusive: Option, ) -> Option { - let operations = completed_session_operations_for_file(session, file_path); + let operations = completed_session_operations_for_file(session, file_path, max_turn_exclusive); let first = operations.first()?; let last = operations.last()?; @@ -1048,10 +1365,12 @@ fn session_file_boundary( fn aggregate_operations_diff_summary_for_file( session: &SessionHistory, file_path: &Path, + max_turn_exclusive: Option, ) -> DiffSummary { let mut out = DiffSummary::default(); for op in session.all_operations_iter() { - if SnapshotCore::operation_matches_file_path(op, file_path) + if operation_is_before(op, max_turn_exclusive) + && SnapshotCore::operation_matches_file_path(op, file_path) && operation_is_completed_for_session_file(op) { out.lines_added += op.diff_summary.lines_added; @@ -1178,6 +1497,9 @@ fn find_anchor_in_current( #[cfg(test)] mod tests { use super::*; + use crate::agentic::session::revert::{ + SessionRevertPhase, SessionRevertState, SESSION_REVERT_SCHEMA_VERSION, + }; use crate::service::snapshot::snapshot_system::FileSnapshotSystem; use crate::service::workspace_runtime::{WorkspaceRuntimeContext, WorkspaceRuntimeTarget}; use serde_json::json; @@ -1274,6 +1596,121 @@ mod tests { assert_eq!(after, "base\nsession\n"); } + #[tokio::test] + async fn bounded_session_view_hides_suffix_operations_from_every_projection() { + let mut runtime = make_test_runtime("bounded_session_view").await; + let file_path = runtime.workspace.join("src/lib.rs"); + let hidden_only_path = runtime.workspace.join("src/hidden.rs"); + fs::create_dir_all(file_path.parent().unwrap()).unwrap(); + tokio::fs::write(&file_path, "base\n").await.unwrap(); + + let visible_operation = runtime + .core + .start_file_operation( + "session-1", + 0, + file_path.clone(), + OperationType::Modify, + "Edit".to_string(), + json!({}), + None, + ) + .await + .unwrap(); + tokio::fs::write(&file_path, "visible\n").await.unwrap(); + runtime + .core + .complete_file_operation("session-1", &visible_operation, 1) + .await + .unwrap(); + + let hidden_operation = runtime + .core + .start_file_operation( + "session-1", + 1, + file_path.clone(), + OperationType::Modify, + "Edit".to_string(), + json!({}), + None, + ) + .await + .unwrap(); + tokio::fs::write(&file_path, "hidden\n").await.unwrap(); + runtime + .core + .complete_file_operation("session-1", &hidden_operation, 1) + .await + .unwrap(); + + let hidden_only_operation = runtime + .core + .start_file_operation( + "session-1", + 1, + hidden_only_path.clone(), + OperationType::Create, + "Write".to_string(), + json!({}), + None, + ) + .await + .unwrap(); + tokio::fs::write(&hidden_only_path, "hidden only\n") + .await + .unwrap(); + runtime + .core + .complete_file_operation("session-1", &hidden_only_operation, 1) + .await + .unwrap(); + + assert_eq!( + runtime.core.get_session_turns_before("session-1", Some(1)), + vec![0] + ); + assert!(runtime + .core + .get_turn_files_before("session-1", 1, Some(1)) + .is_empty()); + assert_eq!( + runtime.core.get_session_files_before("session-1", Some(1)), + vec![file_path.clone()] + ); + assert_eq!( + runtime + .core + .get_session_operations_before("session-1", Some(1)) + .len(), + 1 + ); + let stats = runtime.core.get_session_stats_before("session-1", Some(1)); + assert_eq!( + (stats.total_turns, stats.total_files, stats.total_changes), + (1, 1, 1) + ); + + let (before, after) = runtime + .core + .get_file_diff_before(&file_path, "session-1", Some(1)) + .await + .unwrap(); + assert_eq!((before.as_str(), after.as_str()), ("base\n", "visible\n")); + assert!(matches!( + runtime + .core + .get_file_diff_with_anchor_before( + &file_path, + "session-1", + Some(&hidden_operation), + Some(1), + ) + .await, + Err(SnapshotError::OperationNotFound(_)) + )); + } + #[tokio::test] async fn session_files_ignore_unfinished_operations() { let mut runtime = make_test_runtime("unfinished_ops").await; @@ -1308,4 +1745,176 @@ mod tests { assert_eq!(stats.lines_added, 0); assert_eq!(stats.lines_removed, 0); } + + #[tokio::test] + async fn legacy_checkpoint_cleanup_preserves_an_operation_owned_snapshot() { + let mut runtime = make_test_runtime("legacy_checkpoint_owner").await; + let file_path = runtime.workspace.join("src/lib.rs"); + fs::create_dir_all(file_path.parent().unwrap()).unwrap(); + tokio::fs::write(&file_path, "base\n").await.unwrap(); + let operation_id = runtime + .core + .start_file_operation( + "session-1", + 0, + file_path.clone(), + OperationType::Modify, + "Edit".to_string(), + json!({ "file_path": "src/lib.rs" }), + None, + ) + .await + .unwrap(); + let snapshot_id = runtime + .core + .get_operation("session-1", &operation_id) + .unwrap() + .before_snapshot_id + .expect("before snapshot"); + let legacy_state = SessionRevertState { + schema_version: SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 0, + original_turn_end: 1, + phase: SessionRevertPhase::Staged, + workspace_checkpoint: vec![SessionWorkspaceCheckpoint { + path: file_path, + snapshot_id: Some(snapshot_id.clone()), + }], + }; + + runtime + .core + .delete_workspace_revert_checkpoint(&legacy_state) + .await + .expect("legacy checkpoint cleanup"); + + assert_eq!( + runtime + .core + .snapshot_system + .get_snapshot_content(&snapshot_id) + .await + .expect("operation snapshot must survive"), + "base\n" + ); + } + + #[tokio::test] + async fn staged_workspace_revert_moves_between_boundaries_and_restores_original_state() { + let mut runtime = make_test_runtime("staged_revert").await; + let file_path = runtime.workspace.join("src/lib.rs"); + fs::create_dir_all(file_path.parent().unwrap()).unwrap(); + tokio::fs::write(&file_path, "base\n").await.unwrap(); + + // A -> B -> A exercises content deduplication: the Session checkpoint + // must not borrow and later delete the first operation's A snapshot. + for (turn_index, next) in [(0, "first\n"), (2, "base\n")] { + let operation_id = runtime + .core + .start_file_operation( + "session-1", + turn_index, + file_path.clone(), + OperationType::Modify, + "Edit".to_string(), + json!({ "file_path": "src/lib.rs" }), + None, + ) + .await + .unwrap(); + tokio::fs::write(&file_path, next).await.unwrap(); + runtime + .core + .complete_file_operation("session-1", &operation_id, 1) + .await + .unwrap(); + } + + let mut state = SessionRevertState { + schema_version: SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 2, + original_turn_end: 3, + phase: SessionRevertPhase::Staged, + workspace_checkpoint: Vec::new(), + }; + runtime + .core + .prepare_workspace_revert("session-1", &mut state) + .await + .expect("latest workspace state should be checkpointed"); + runtime + .core + .apply_workspace_revert("session-1", &state) + .await + .expect("latest turn should be reverted"); + assert_eq!( + tokio::fs::read_to_string(&file_path).await.unwrap(), + "first\n" + ); + + state.boundary_turn = 0; + runtime + .core + .prepare_workspace_revert("session-1", &mut state) + .await + .expect("earlier affected files should join the checkpoint"); + runtime + .core + .apply_workspace_revert("session-1", &state) + .await + .expect("all turns should be reverted"); + assert_eq!( + tokio::fs::read_to_string(&file_path).await.unwrap(), + "base\n" + ); + + state.boundary_turn = 2; + runtime + .core + .apply_workspace_revert("session-1", &state) + .await + .expect("redo should advance to the later boundary"); + assert_eq!( + tokio::fs::read_to_string(&file_path).await.unwrap(), + "first\n" + ); + + runtime + .core + .restore_workspace_revert(&state) + .await + .expect("final redo should restore the pre-undo workspace"); + assert_eq!( + tokio::fs::read_to_string(&file_path).await.unwrap(), + "base\n" + ); + + runtime + .core + .delete_workspace_revert_checkpoint(&state) + .await + .expect("full redo should release only its owned checkpoint handle"); + + let mut repeated = SessionRevertState { + schema_version: SESSION_REVERT_SCHEMA_VERSION, + boundary_turn: 0, + original_turn_end: 3, + phase: SessionRevertPhase::Staged, + workspace_checkpoint: Vec::new(), + }; + runtime + .core + .prepare_workspace_revert("session-1", &mut repeated) + .await + .expect("a later undo should create a fresh checkpoint"); + runtime + .core + .apply_workspace_revert("session-1", &repeated) + .await + .expect("operation snapshots must survive a previous checkpoint release"); + assert_eq!( + tokio::fs::read_to_string(&file_path).await.unwrap(), + "base\n" + ); + } } diff --git a/src/crates/assembly/core/src/service/snapshot/snapshot_system.rs b/src/crates/assembly/core/src/service/snapshot/snapshot_system.rs index 88aee46c7..aaa79c3e0 100644 --- a/src/crates/assembly/core/src/service/snapshot/snapshot_system.rs +++ b/src/crates/assembly/core/src/service/snapshot/snapshot_system.rs @@ -330,6 +330,24 @@ impl FileSnapshotSystem { /// Creates a file snapshot. pub async fn create_snapshot(&mut self, file_path: &Path) -> SnapshotResult { + self.create_snapshot_with_ownership(file_path, false).await + } + + /// Creates a snapshot with an independent metadata handle while retaining + /// content-addressed blob deduplication. The caller owns this handle and may + /// delete it without invalidating another operation's snapshot reference. + pub(crate) async fn create_owned_snapshot( + &mut self, + file_path: &Path, + ) -> SnapshotResult { + self.create_snapshot_with_ownership(file_path, true).await + } + + async fn create_snapshot_with_ownership( + &mut self, + file_path: &Path, + independent_handle: bool, + ) -> SnapshotResult { debug!("Creating snapshot: file_path={}", file_path.display()); if !file_path.exists() { @@ -356,7 +374,10 @@ impl FileSnapshotSystem { let content_hash = self.calculate_content_hash(&content); - if self.dedup_enabled && self.hash_to_path.contains_key(&content_hash) { + if !independent_handle + && self.dedup_enabled + && self.hash_to_path.contains_key(&content_hash) + { if let Some(snapshot_id) = self.find_snapshot_by_hash(&content_hash) { debug!( "Found duplicate content, reusing existing snapshot: content_hash={}", @@ -917,4 +938,42 @@ mod tests { fs::remove_dir_all(&context.runtime_root).expect("cleanup runtime root"); } + + #[tokio::test] + async fn owned_snapshot_handle_does_not_delete_an_existing_deduplicated_snapshot() { + let context = test_runtime_context(); + create_runtime_dirs(&context); + let file_path = context.runtime_root.join("workspace").join("shared.txt"); + fs::create_dir_all(file_path.parent().expect("file has parent")).expect("create parent"); + fs::write(&file_path, "same content").expect("write fixture"); + + let mut snapshot_system = FileSnapshotSystem::new(context.clone()); + snapshot_system + .initialize() + .await + .expect("initialize snapshots"); + let existing = snapshot_system + .create_snapshot(&file_path) + .await + .expect("create existing snapshot"); + let owned = snapshot_system + .create_owned_snapshot(&file_path) + .await + .expect("create independent checkpoint handle"); + + assert_ne!(existing, owned); + snapshot_system + .delete_snapshot(&owned) + .await + .expect("delete owned handle"); + assert_eq!( + snapshot_system + .get_snapshot_content(&existing) + .await + .expect("existing snapshot must survive"), + "same content" + ); + + fs::remove_dir_all(&context.runtime_root).expect("cleanup runtime root"); + } } diff --git a/src/crates/assembly/core/src/service_agent_runtime.rs b/src/crates/assembly/core/src/service_agent_runtime.rs index 234536f40..861414893 100644 --- a/src/crates/assembly/core/src/service_agent_runtime.rs +++ b/src/crates/assembly/core/src/service_agent_runtime.rs @@ -8,17 +8,19 @@ use bitfun_agent_runtime::sdk::{ AgentEventSource, AgentInteractionResponsePort, AgentRuntime, AgentRuntimeBuilder, AgentSessionCompactionPort, AgentSessionForkPort, AgentSessionModePort, AgentSessionModelPort, - AgentSessionModelUpdateRequest, AgentSessionRestorePort, AgentSessionUsagePort, - AgentTurnSettlementPort, RuntimeError, + AgentSessionModelUpdateRequest, AgentSessionRestorePort, AgentSessionRevertPort, + AgentSessionUsagePort, AgentTurnSettlementPort, RuntimeError, }; +use bitfun_events::AgenticEvent; use bitfun_runtime_ports::{ AgentDialogTurnPort, AgentDialogTurnRequest, AgentInputAttachment, AgentLifecycleDeliveryPort, AgentLocalCommandTurnPort, AgentSessionClosePort, AgentSessionCreateRequest, - AgentSessionManagementPort, AgentSubmissionPort, AgentSubmissionSource, - AgentThreadGoalManagementPort, AgentTurnCancellationPort, AgentTurnCancellationRequest, - PermissionPolicyPreset, RemoteControlStatePort, RemoteControlStateRequest, - RemoteControlStateSnapshot, RemoteSessionWorkspaceIdentity, RuntimeServiceCapability, - RuntimeServicePort, SessionStoragePathRequest, SessionStorePort, ToolPermissionConfig, + AgentSessionManagementPort, AgentSessionRevertRequest, AgentSessionRevertResult, + AgentSubmissionPort, AgentSubmissionSource, AgentThreadGoalManagementPort, + AgentTurnCancellationPort, AgentTurnCancellationRequest, PermissionPolicyPreset, + RemoteControlStatePort, RemoteControlStateRequest, RemoteControlStateSnapshot, + RemoteSessionWorkspaceIdentity, RuntimeServiceCapability, RuntimeServicePort, + SessionStoragePathRequest, SessionStorePort, ToolPermissionConfig, }; use bitfun_services_integrations::remote_connect::{ agent_input_attachment_from_remote_image_context, build_remote_chat_messages, @@ -450,6 +452,114 @@ impl ScheduledSessionManagementPort { scheduler, } } + + async fn apply_session_revert( + &self, + request: AgentSessionRevertRequest, + undo: bool, + ) -> bitfun_runtime_ports::PortResult { + bitfun_core_types::validate_session_id(&request.session_id).map_err(|message| { + bitfun_runtime_ports::PortError::new( + bitfun_runtime_ports::PortErrorKind::InvalidRequest, + message, + ) + })?; + if request.remote_connection_id.is_some() || request.remote_ssh_host.is_some() { + return Err(bitfun_runtime_ports::PortError::new( + bitfun_runtime_ports::PortErrorKind::NotAvailable, + "Session undo and redo are unavailable for remote workspaces", + )); + } + self.coordinator + .local_revert_workspace(&request.session_id) + .map_err(|error| { + if matches!(&error, crate::util::errors::BitFunError::Validation(message) if message == "Session undo and redo are unavailable for remote workspaces") + { + bitfun_runtime_ports::PortError::new( + bitfun_runtime_ports::PortErrorKind::NotAvailable, + error.to_string(), + ) + } else { + map_session_close_error(error) + } + })?; + let storage_path = CoreSessionStorePort::default() + .resolve_session_storage_path(SessionStoragePathRequest { + workspace_path: std::path::PathBuf::from(&request.workspace_path), + remote_connection_id: None, + remote_ssh_host: None, + }) + .await + .map(|resolution| resolution.effective_storage_path)?; + let session_manager = self.coordinator.get_session_manager(); + session_manager + .validate_session_storage_path_binding(&request.session_id, &storage_path) + .map_err(map_session_close_error)?; + let maintenance = self + .scheduler + .begin_session_maintenance(&request.session_id, &storage_path, Duration::from_secs(30)) + .await + .map_err(map_session_close_error)?; + let _mutation = session_manager + .acquire_session_mutation(&request.session_id) + .await + .map_err(map_session_close_error)?; + session_manager + .validate_session_storage_path_binding(&request.session_id, &storage_path) + .map_err(map_session_close_error)?; + let (composer, changed, hidden_turn_count) = self + .coordinator + .apply_session_revert_locked(&storage_path, &request.session_id, undo) + .await + .map_err(map_session_close_error)?; + if changed { + self.coordinator + .emit_event(AgenticEvent::SessionHistoryChanged { + session_id: request.session_id.clone(), + }) + .await; + } + let transcript = self + .coordinator + .read_session_transcript_locked(bitfun_runtime_ports::SessionTranscriptRequest { + session_id: request.session_id.clone(), + turn_id: None, + }) + .await + .map_err(|error| { + bitfun_runtime_ports::PortError::new( + bitfun_runtime_ports::PortErrorKind::OutcomeUnknown, + format!( + "Session revert completed but the authoritative transcript could not be read: {error}" + ), + ) + })?; + Ok(AgentSessionRevertResult { + session_id: request.session_id, + transcript, + composer, + retired_turn_ids: maintenance.retired_turn_ids().to_vec(), + changed, + hidden_turn_count, + }) + } +} + +#[async_trait::async_trait] +impl AgentSessionRevertPort for ScheduledSessionManagementPort { + async fn undo_session( + &self, + request: AgentSessionRevertRequest, + ) -> bitfun_runtime_ports::PortResult { + self.apply_session_revert(request, true).await + } + + async fn redo_session( + &self, + request: AgentSessionRevertRequest, + ) -> bitfun_runtime_ports::PortResult { + self.apply_session_revert(request, false).await + } } /// ACP accepts one prompt at a time per session. Keep that protocol-specific @@ -683,6 +793,13 @@ fn scheduled_session_close_port( Arc::new(ScheduledSessionManagementPort::new(coordinator, scheduler)) } +fn scheduled_session_revert_port( + coordinator: Arc, + scheduler: Arc, +) -> Arc { + Arc::new(ScheduledSessionManagementPort::new(coordinator, scheduler)) +} + pub(crate) struct CoreServiceAgentRuntime; impl CoreServiceAgentRuntime { @@ -780,21 +897,15 @@ impl CoreServiceAgentRuntime { pub(crate) async fn load_remote_chat_messages( session_storage_dir: &std::path::Path, session_id: &str, - ) -> (Vec, bool) { - let Ok(pm) = crate::infrastructure::PathManager::new() else { - return (vec![], false); - }; - let pm = std::sync::Arc::new(pm); - let Ok(store) = crate::agentic::persistence::PersistenceManager::new(pm) else { - return (vec![], false); - }; - let Ok(turns) = store - .load_session_turns(session_storage_dir, session_id) + ) -> Result<(Vec, bool), String> { + let coordinator = get_global_coordinator().ok_or_else(|| { + "Core coordinator is unavailable for remote history reads".to_string() + })?; + let turns = coordinator + .load_visible_persisted_session_turns(session_storage_dir, session_id) .await - else { - return (vec![], false); - }; - (remote_chat_messages_from_turns(&turns), false) + .map_err(|error| error.to_string())?; + Ok((remote_chat_messages_from_turns(&turns), false)) } pub(crate) async fn load_remote_model_catalog( @@ -973,6 +1084,7 @@ impl CoreServiceAgentRuntime { let session_management = scheduled_session_management_port(coordinator.clone(), scheduler.clone()); let session_close = scheduled_session_close_port(coordinator.clone(), scheduler.clone()); + let session_revert = scheduled_session_revert_port(coordinator.clone(), scheduler.clone()); let session_mode: Arc = coordinator.clone(); let session_model: Arc = coordinator.clone(); let session_restore: Arc = coordinator.clone(); @@ -999,6 +1111,7 @@ impl CoreServiceAgentRuntime { interaction_response, )? .with_session_close_port(session_close) + .with_session_revert_port(session_revert) .with_dialog_turn_port(dialog_turn) .with_lifecycle_delivery_port(lifecycle_delivery) .build() @@ -1012,6 +1125,7 @@ impl CoreServiceAgentRuntime { let submission: Arc = coordinator.clone(); let session_management = scheduled_session_management_port(coordinator.clone(), scheduler.clone()); + let session_revert = scheduled_session_revert_port(coordinator.clone(), scheduler.clone()); let session_mode: Arc = coordinator.clone(); let session_model: Arc = coordinator.clone(); let session_restore: Arc = coordinator.clone(); @@ -1036,6 +1150,7 @@ impl CoreServiceAgentRuntime { cancellation, interaction_response, )? + .with_session_revert_port(session_revert) .with_lifecycle_delivery_port(lifecycle_delivery) .build() .map_err(|error| error.to_string()) @@ -1052,6 +1167,7 @@ impl CoreServiceAgentRuntime { let submission: Arc = coordinator.clone(); let session_management = scheduled_session_management_port(coordinator.clone(), scheduler.clone()); + let session_revert = scheduled_session_revert_port(coordinator.clone(), scheduler.clone()); let session_model: Arc = coordinator.clone(); let session_compaction: Arc = coordinator.clone(); let interaction_response: Arc = coordinator; @@ -1061,6 +1177,7 @@ impl CoreServiceAgentRuntime { AgentRuntimeBuilder::new() .with_submission_port(submission) .with_session_management_port(session_management) + .with_session_revert_port(session_revert) .with_session_model_port(session_model) .with_session_compaction_port(session_compaction) .with_dialog_turn_port(dialog_turn) @@ -1082,6 +1199,7 @@ impl CoreServiceAgentRuntime { let submission: Arc = coordinator.clone(); let session_management = scheduled_session_management_port(coordinator.clone(), scheduler.clone()); + let session_revert = scheduled_session_revert_port(coordinator.clone(), scheduler.clone()); let session_mode: Arc = coordinator.clone(); let session_model: Arc = coordinator.clone(); let session_restore: Arc = coordinator.clone(); @@ -1107,6 +1225,7 @@ impl CoreServiceAgentRuntime { cancellation, interaction_response, )? + .with_session_revert_port(session_revert) .with_dialog_turn_port(dialog_turn) .with_lifecycle_delivery_port(lifecycle_delivery) .build() @@ -1198,6 +1317,7 @@ impl CoreServiceAgentRuntime { let session_management = scheduled_session_management_port(coordinator.clone(), scheduler.clone()); let session_close = scheduled_session_close_port(coordinator.clone(), scheduler.clone()); + let session_revert = scheduled_session_revert_port(coordinator.clone(), scheduler.clone()); let session_mode: Arc = coordinator.clone(); let session_model: Arc = coordinator.clone(); let session_restore: Arc = coordinator.clone(); @@ -1224,6 +1344,7 @@ impl CoreServiceAgentRuntime { interaction_response, )? .with_session_close_port(session_close) + .with_session_revert_port(session_revert) .with_dialog_turn_port(dialog_turn) .with_lifecycle_delivery_port(lifecycle_delivery); let builder = match event_source { @@ -1765,7 +1886,7 @@ impl RemoteSessionRuntimeHost for CoreRemoteSessionRuntimeHost { &self, session_storage_dir: &std::path::Path, session_id: &str, - ) -> (Vec, bool) { + ) -> Result<(Vec, bool), String> { CoreServiceAgentRuntime::load_remote_chat_messages(session_storage_dir, session_id).await } @@ -1848,7 +1969,7 @@ impl RemotePollRuntimeHost for CoreRemotePollRuntimeHost<'_> { &self, session_storage_dir: &std::path::Path, session_id: &str, - ) -> (Vec, bool) { + ) -> Result<(Vec, bool), String> { CoreServiceAgentRuntime::load_remote_chat_messages(session_storage_dir, session_id).await } } diff --git a/src/crates/contracts/events/src/agentic.rs b/src/crates/contracts/events/src/agentic.rs index 3d6ba336e..3d2cb1918 100644 --- a/src/crates/contracts/events/src/agentic.rs +++ b/src/crates/contracts/events/src/agentic.rs @@ -102,6 +102,14 @@ pub enum AgenticEvent { new_state: String, }, + /// The authoritative visible history for a session changed outside the + /// append-only turn lifecycle (for example, after restoring a checkpoint). + /// Consumers should invalidate cached transcript projections for this + /// session and reload them from the owning runtime. + SessionHistoryChanged { + session_id: String, + }, + SessionDeleted { session_id: String, }, @@ -599,6 +607,7 @@ impl AgenticEvent { match self { Self::SessionCreated { session_id, .. } | Self::SessionStateChanged { session_id, .. } + | Self::SessionHistoryChanged { session_id } | Self::SessionDeleted { session_id } | Self::SessionTitleGenerated { session_id, .. } | Self::ImageAnalysisStarted { session_id, .. } @@ -626,6 +635,29 @@ impl AgenticEvent { } } + /// Get the dialog Turn identity carried by a Turn-scoped event. + pub fn turn_id(&self) -> Option<&str> { + match self { + Self::DialogTurnStarted { turn_id, .. } + | Self::DialogTurnCompleted { turn_id, .. } + | Self::DialogTurnCancelled { turn_id, .. } + | Self::DialogTurnFailed { turn_id, .. } + | Self::TokenUsageUpdated { turn_id, .. } + | Self::ContextCompressionStarted { turn_id, .. } + | Self::ContextCompressionCompleted { turn_id, .. } + | Self::ContextCompressionFailed { turn_id, .. } + | Self::ModelRoundStarted { turn_id, .. } + | Self::ModelRoundAttemptSuperseded { turn_id, .. } + | Self::ModelRoundCompleted { turn_id, .. } + | Self::TextChunk { turn_id, .. } + | Self::ThinkingChunk { turn_id, .. } + | Self::ToolEvent { turn_id, .. } + | Self::DeepReviewQueueStateChanged { turn_id, .. } + | Self::UserSteeringInjected { turn_id, .. } => Some(turn_id), + _ => None, + } + } + /// Get the default priority pub fn default_priority(&self) -> AgenticEventPriority { match self { @@ -634,6 +666,7 @@ impl AgenticEvent { | Self::DialogTurnCancelled { .. } => AgenticEventPriority::Critical, Self::SessionStateChanged { .. } + | Self::SessionHistoryChanged { .. } | Self::SessionTitleGenerated { .. } | Self::SessionModelAutoMigrated { .. } | Self::SubagentSessionLinked { .. } diff --git a/src/crates/contracts/events/src/frontend_projection.rs b/src/crates/contracts/events/src/frontend_projection.rs index a9fb3ca0c..561c05707 100644 --- a/src/crates/contracts/events/src/frontend_projection.rs +++ b/src/crates/contracts/events/src/frontend_projection.rs @@ -358,6 +358,10 @@ pub fn project_agentic_frontend_event(event: AgenticEvent) -> Option Some(AgenticFrontendEvent::new( + "agentic://session-history-changed", + json!({ "sessionId": session_id }), + )), AgenticEvent::SessionModelAutoMigrated { session_id, previous_model_id, diff --git a/src/crates/contracts/runtime-ports/src/lib.rs b/src/crates/contracts/runtime-ports/src/lib.rs index 0e259686e..6772e0e81 100644 --- a/src/crates/contracts/runtime-ports/src/lib.rs +++ b/src/crates/contracts/runtime-ports/src/lib.rs @@ -2072,6 +2072,57 @@ pub trait AgentSessionCompactionPort: Send + Sync { ) -> PortResult; } +/// Local Session history mutation requested by an interactive product surface. +/// +/// This is deliberately narrower than a generic checkpoint or workspace rewind +/// API: one Core-owned operation keeps the persisted transcript, model context, +/// and tracked workspace files on the same staged boundary. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentSessionRevertRequest { + pub workspace_path: String, + pub session_id: String, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_connection_id: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub remote_ssh_host: Option, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum AgentSessionComposerUpdate { + Preserve, + Replace { text: String }, + Clear, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct AgentSessionRevertResult { + pub session_id: String, + pub transcript: SessionTranscript, + pub composer: AgentSessionComposerUpdate, + /// Active and queued Turns retired by the maintenance boundary. Consumers + /// use this as an event-stream fence after replacing their local projection. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub retired_turn_ids: Vec, + pub changed: bool, + pub hidden_turn_count: usize, +} + +#[async_trait::async_trait] +pub trait AgentSessionRevertPort: Send + Sync { + async fn undo_session( + &self, + request: AgentSessionRevertRequest, + ) -> PortResult; + + async fn redo_session( + &self, + request: AgentSessionRevertRequest, + ) -> PortResult; +} + #[async_trait::async_trait] pub trait AgentSessionForkPort: Send + Sync { async fn fork_session( @@ -2399,6 +2450,35 @@ mod tests { use super::*; use std::sync::Mutex; + #[test] + fn session_revert_contract_preserves_authoritative_transcript_and_composer_intent() { + let result = AgentSessionRevertResult { + session_id: "session-1".to_string(), + transcript: SessionTranscript { + session_id: "session-1".to_string(), + messages: Vec::new(), + }, + composer: AgentSessionComposerUpdate::Replace { + text: "restore this prompt".to_string(), + }, + retired_turn_ids: vec!["turn-2".to_string(), "turn-queued".to_string()], + changed: true, + hidden_turn_count: 2, + }; + + let value = serde_json::to_value(&result).expect("session revert result should serialize"); + assert_eq!(value["sessionId"], "session-1"); + assert_eq!(value["composer"]["kind"], "replace"); + assert_eq!(value["composer"]["text"], "restore this prompt"); + assert_eq!(value["hiddenTurnCount"], 2); + assert_eq!(value["retiredTurnIds"][1], "turn-queued"); + assert_eq!( + serde_json::from_value::(value) + .expect("session revert result should deserialize"), + result + ); + } + #[test] fn context_reload_contract_is_closed_and_target_specific() { let cases = [ diff --git a/src/crates/contracts/runtime-ports/src/local_workspace_snapshot.rs b/src/crates/contracts/runtime-ports/src/local_workspace_snapshot.rs index a32e1b936..a014941f0 100644 --- a/src/crates/contracts/runtime-ports/src/local_workspace_snapshot.rs +++ b/src/crates/contracts/runtime-ports/src/local_workspace_snapshot.rs @@ -12,6 +12,9 @@ use crate::PortResult; pub struct LocalWorkspaceSnapshotSessionRequest { pub workspace_path: PathBuf, pub session_id: String, + /// Exclusive visible turn end supplied by the session lifecycle owner. + /// `None` exposes the complete snapshot history. + pub max_turn_exclusive: Option, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -65,6 +68,7 @@ mod tests { let session = LocalWorkspaceSnapshotSessionRequest { workspace_path: PathBuf::from("workspace"), session_id: "session-1".to_string(), + max_turn_exclusive: Some(5), }; let turn = LocalWorkspaceSnapshotTurnRequest { workspace_path: session.workspace_path.clone(), diff --git a/src/crates/execution/agent-runtime/src/runtime.rs b/src/crates/execution/agent-runtime/src/runtime.rs index 7d94741a5..fbcc9ea91 100644 --- a/src/crates/execution/agent-runtime/src/runtime.rs +++ b/src/crates/execution/agent-runtime/src/runtime.rs @@ -19,7 +19,8 @@ use bitfun_runtime_ports::{ AgentSessionForkBeforeTurnRequest, AgentSessionForkPort, AgentSessionForkRequest, AgentSessionForkResult, AgentSessionListRequest, AgentSessionManagementPort, AgentSessionModePort, AgentSessionModeUpdateRequest, AgentSessionModelPort, - AgentSessionModelUpdateRequest, AgentSessionRenameRequest, AgentSessionSummary, + AgentSessionModelUpdateRequest, AgentSessionRenameRequest, AgentSessionRevertPort, + AgentSessionRevertRequest, AgentSessionRevertResult, AgentSessionSummary, AgentSessionUsagePort, AgentSessionUsageRequest, AgentSessionWorkspaceBinding, AgentSessionWorkspaceRequest, AgentSubmissionPort, AgentSubmissionRequest, AgentSubmissionResult, AgentSubmissionSource, AgentThreadGoalCreateRequest, @@ -190,6 +191,7 @@ pub struct AgentRuntime { session_mode: Option>, session_model: Option>, session_compaction: Option>, + session_revert: Option>, session_fork: Option>, session_usage: Option>, turn_settlement: Option>, @@ -258,6 +260,13 @@ impl std::fmt::Debug for AgentRuntime { .as_ref() .map(|_| ""), ) + .field( + "session_revert", + &self + .session_revert + .as_ref() + .map(|_| ""), + ) .field( "session_usage", &self @@ -382,6 +391,7 @@ pub struct AgentRuntimeBuilder { session_mode: Option>, session_model: Option>, session_compaction: Option>, + session_revert: Option>, session_fork: Option>, session_usage: Option>, turn_settlement: Option>, @@ -445,6 +455,11 @@ impl AgentRuntimeBuilder { self } + pub fn with_session_revert_port(mut self, port: Arc) -> Self { + self.session_revert = Some(port); + self + } + pub fn with_session_fork_port(mut self, port: Arc) -> Self { self.session_fork = Some(port); self @@ -571,6 +586,7 @@ impl AgentRuntimeBuilder { session_mode, session_model, session_compaction, + session_revert, session_fork, session_usage, turn_settlement, @@ -604,6 +620,7 @@ impl AgentRuntimeBuilder { session_mode, session_model, session_compaction, + session_revert, session_fork, session_usage, turn_settlement, @@ -1094,6 +1111,32 @@ impl AgentRuntime { .map_err(RuntimeError::from) } + pub async fn undo_session( + &self, + request: AgentSessionRevertRequest, + ) -> Result { + let port = self.session_revert.as_ref().ok_or_else(|| { + RuntimeError::Port(PortError::new( + PortErrorKind::NotAvailable, + "agent session revert port is not registered", + )) + })?; + port.undo_session(request).await.map_err(RuntimeError::from) + } + + pub async fn redo_session( + &self, + request: AgentSessionRevertRequest, + ) -> Result { + let port = self.session_revert.as_ref().ok_or_else(|| { + RuntimeError::Port(PortError::new( + PortErrorKind::NotAvailable, + "agent session revert port is not registered", + )) + })?; + port.redo_session(request).await.map_err(RuntimeError::from) + } + pub async fn fork_session( &self, request: AgentSessionForkRequest, @@ -1422,6 +1465,7 @@ mod tests { AgentSessionCompactionPort, AgentSessionCompactionRequest, AgentSessionCompactionResult, AgentSessionCreateResult, AgentSessionDeleteRequest, AgentSessionListRequest, AgentSessionManagementPort, AgentSessionModePort, AgentSessionModeUpdateRequest, + AgentSessionRevertPort, AgentSessionRevertRequest, AgentSessionRevertResult, AgentSessionSummary, AgentSessionWorkspaceRequest, AgentSubmissionResult, AgentThreadGoalDeliveryKind, AgentThreadGoalDeliveryRequest, AgentThreadGoalManagementPort, AgentTurnCancellationResult, ClockPort, DialogQueuePriority, DialogSubmissionPolicy, @@ -1449,6 +1493,7 @@ mod tests { local_command_turns: Mutex>, restored_sessions: Mutex>, mode_updates: Mutex>, + undo_requests: Mutex>, transcript_requests: Mutex>, workspace_binding_requests: Mutex>, thread_goal_gets: Mutex>, @@ -1637,6 +1682,34 @@ mod tests { } } + #[async_trait::async_trait] + impl AgentSessionRevertPort for FakeAgentRuntimePorts { + async fn undo_session( + &self, + request: AgentSessionRevertRequest, + ) -> PortResult { + self.undo_requests.lock().unwrap().push(request.clone()); + Ok(AgentSessionRevertResult { + session_id: request.session_id.clone(), + transcript: SessionTranscript { + session_id: request.session_id, + messages: Vec::new(), + }, + composer: bitfun_runtime_ports::AgentSessionComposerUpdate::Preserve, + retired_turn_ids: Vec::new(), + changed: true, + hidden_turn_count: 1, + }) + } + + async fn redo_session( + &self, + request: AgentSessionRevertRequest, + ) -> PortResult { + self.undo_session(request).await + } + } + #[async_trait::async_trait] impl AgentLocalCommandTurnPort for FakeAgentRuntimePorts { async fn record_completed_local_command_turn( @@ -1933,6 +2006,30 @@ mod tests { ); } + #[tokio::test] + async fn session_undo_forwards_through_the_optional_runtime_port() { + let ports = Arc::new(FakeAgentRuntimePorts::default()); + let runtime = AgentRuntimeBuilder::new() + .with_submission_port(ports.clone()) + .with_session_revert_port(ports.clone()) + .build() + .expect("runtime"); + let request = AgentSessionRevertRequest { + workspace_path: "/workspace/project".to_string(), + session_id: "session-1".to_string(), + remote_connection_id: None, + remote_ssh_host: None, + }; + + let result = runtime + .undo_session(request.clone()) + .await + .expect("undo session"); + + assert_eq!(result.hidden_turn_count, 1); + assert_eq!(ports.undo_requests.lock().unwrap().as_slice(), &[request]); + } + #[tokio::test] async fn builder_keeps_plugin_runtime_disabled_by_default() { let ports = Arc::new(FakeAgentRuntimePorts::default()); diff --git a/src/crates/execution/agent-runtime/src/sdk.rs b/src/crates/execution/agent-runtime/src/sdk.rs index a4a4940b6..755939362 100644 --- a/src/crates/execution/agent-runtime/src/sdk.rs +++ b/src/crates/execution/agent-runtime/src/sdk.rs @@ -62,25 +62,26 @@ pub use bitfun_runtime_ports::{ AgentInputAttachment, AgentLifecycleDeliveryPort, AgentLocalCommandTurnPort, AgentLocalCommandTurnRecordRequest, AgentSessionArchiveRequest, AgentSessionArchiveStateRequest, AgentSessionClosePort, AgentSessionCompactionPort, - AgentSessionCompactionRequest, AgentSessionCompactionResult, AgentSessionCreateRequest, - AgentSessionCreateResult, AgentSessionDeleteRequest, AgentSessionForkAtTurnRequest, - AgentSessionForkBeforeTurnRequest, AgentSessionForkPort, AgentSessionForkRequest, - AgentSessionForkResult, AgentSessionListRequest, AgentSessionManagementPort, - AgentSessionModePort, AgentSessionModeUpdateRequest, AgentSessionModelPort, - AgentSessionModelUpdateRequest, AgentSessionRenameRequest, AgentSessionSummary, - AgentSessionUsagePort, AgentSessionUsageRequest, AgentSessionWorkspaceBinding, - AgentSessionWorkspaceRequest, AgentSubmissionPort, AgentSubmissionRequest, - AgentSubmissionResult, AgentSubmissionSource, AgentThreadGoalCreateRequest, - AgentThreadGoalDeliveryRequest, AgentThreadGoalGetRequest, AgentThreadGoalManagementPort, - AgentThreadGoalUpdateStatusRequest, AgentTransientSessionDiscardRequest, - AgentTurnCancellationPort, AgentTurnCancellationRequest, AgentTurnCancellationResult, - AgentTurnSettlementPort, AgentTurnSettlementRequest, ClockPort, DialogSubmissionPolicy, - DialogSubmitOutcome, FileSystemPort, GitPort, McpCatalogPort, NetworkPort, - PermissionAuditRecord, PermissionDelegationContext, PermissionGrant, PermissionGrantKey, - PermissionReply, PermissionReplySource, PermissionRequest, PermissionRequestEvent, - PermissionRequestSource, PermissionRequestSourceKind, PortError, PortErrorKind, PortResult, - RemoteAssistantWorkspaceFacts, RemoteCapabilityPort, RemoteConnectionPort, - RemoteProjectionPort, RemoteRecentWorkspaceFacts, RemoteWorkspaceFacts, + AgentSessionCompactionRequest, AgentSessionCompactionResult, AgentSessionComposerUpdate, + AgentSessionCreateRequest, AgentSessionCreateResult, AgentSessionDeleteRequest, + AgentSessionForkAtTurnRequest, AgentSessionForkBeforeTurnRequest, AgentSessionForkPort, + AgentSessionForkRequest, AgentSessionForkResult, AgentSessionListRequest, + AgentSessionManagementPort, AgentSessionModePort, AgentSessionModeUpdateRequest, + AgentSessionModelPort, AgentSessionModelUpdateRequest, AgentSessionRenameRequest, + AgentSessionRevertPort, AgentSessionRevertRequest, AgentSessionRevertResult, + AgentSessionSummary, AgentSessionUsagePort, AgentSessionUsageRequest, + AgentSessionWorkspaceBinding, AgentSessionWorkspaceRequest, AgentSubmissionPort, + AgentSubmissionRequest, AgentSubmissionResult, AgentSubmissionSource, + AgentThreadGoalCreateRequest, AgentThreadGoalDeliveryRequest, AgentThreadGoalGetRequest, + AgentThreadGoalManagementPort, AgentThreadGoalUpdateStatusRequest, + AgentTransientSessionDiscardRequest, AgentTurnCancellationPort, AgentTurnCancellationRequest, + AgentTurnCancellationResult, AgentTurnSettlementPort, AgentTurnSettlementRequest, ClockPort, + DialogSubmissionPolicy, DialogSubmitOutcome, FileSystemPort, GitPort, McpCatalogPort, + NetworkPort, PermissionAuditRecord, PermissionDelegationContext, PermissionGrant, + PermissionGrantKey, PermissionReply, PermissionReplySource, PermissionRequest, + PermissionRequestEvent, PermissionRequestSource, PermissionRequestSourceKind, PortError, + PortErrorKind, PortResult, RemoteAssistantWorkspaceFacts, RemoteCapabilityPort, + RemoteConnectionPort, RemoteProjectionPort, RemoteRecentWorkspaceFacts, RemoteWorkspaceFacts, RemoteWorkspaceFileRuntimeHost, RemoteWorkspaceKind, RemoteWorkspacePort, RemoteWorkspaceRuntimeHost, RemoteWorkspaceUpdate, RuntimeEventEnvelope, RuntimeEventSink, RuntimeEventType, RuntimeServiceCapability, RuntimeServicePort, SessionStorageKind, @@ -152,6 +153,11 @@ impl AgentRuntimeBuilder { self } + pub fn with_session_revert_port(mut self, port: Arc) -> Self { + self.inner = self.inner.with_session_revert_port(port); + self + } + pub fn with_session_fork_port(mut self, port: Arc) -> Self { self.inner = self.inner.with_session_fork_port(port); self @@ -478,6 +484,20 @@ impl AgentRuntime { self.inner.start_session_compaction(request).await } + pub async fn undo_session( + &self, + request: AgentSessionRevertRequest, + ) -> Result { + self.inner.undo_session(request).await + } + + pub async fn redo_session( + &self, + request: AgentSessionRevertRequest, + ) -> Result { + self.inner.redo_session(request).await + } + pub async fn fork_session( &self, request: AgentSessionForkRequest, diff --git a/src/crates/services/services-integrations/src/remote_connect.rs b/src/crates/services/services-integrations/src/remote_connect.rs index a3915b635..50aea1142 100644 --- a/src/crates/services/services-integrations/src/remote_connect.rs +++ b/src/crates/services/services-integrations/src/remote_connect.rs @@ -1206,7 +1206,7 @@ pub trait RemoteSessionRuntimeHost: Send + Sync { &self, session_storage_dir: &Path, session_id: &str, - ) -> (Vec, bool); + ) -> Result<(Vec, bool), String>; async fn delete_session( &self, session_storage_dir: &Path, @@ -1379,9 +1379,13 @@ where ), }; }; - let (chat_messages, has_more) = host + let (chat_messages, has_more) = match host .load_remote_chat_messages(&session_storage_dir, session_id) - .await; + .await + { + Ok(messages) => messages, + Err(message) => return RemoteResponse::Error { message }, + }; remote_messages_response(session_id.clone(), chat_messages, has_more) } RemoteCommand::DeleteSession { session_id } => { @@ -1421,7 +1425,7 @@ pub trait RemotePollRuntimeHost: Send + Sync { &self, session_storage_dir: &Path, session_id: &str, - ) -> (Vec, bool); + ) -> Result<(Vec, bool), String>; } pub async fn handle_remote_poll_command(host: &H, command: &RemoteCommand) -> RemoteResponse @@ -1446,12 +1450,17 @@ where let current_model_catalog = host.load_model_catalog(session_id).await; let model_catalog_delta = remote_model_catalog_poll_delta(current_model_catalog, *known_model_catalog_version); + let persistence_dirty = tracker.is_persistence_dirty(); - if *since_version == current_version && *since_version > 0 && !model_catalog_delta.changed { + if *since_version == current_version + && *since_version > 0 + && !model_catalog_delta.changed + && !persistence_dirty + { return remote_no_change_poll_response(current_version); } - let needs_persistence = *since_version == 0 || tracker.is_persistence_dirty(); + let needs_persistence = *since_version == 0 || persistence_dirty; if !needs_persistence { return remote_snapshot_poll_response( &tracker, @@ -1468,9 +1477,13 @@ where ), }; }; - let (all_chat_messages, _) = host + let (all_chat_messages, _) = match host .load_remote_chat_messages(&session_storage_dir, session_id) - .await; + .await + { + Ok(messages) => messages, + Err(message) => return RemoteResponse::Error { message }, + }; let total_msg_count = all_chat_messages.len(); let new_messages = all_chat_messages .into_iter() @@ -2777,6 +2790,20 @@ impl RemoteSessionStateTracker { self.state.write().unwrap().persistence_dirty = false; } + fn invalidate_history_projection(&self) { + let mut state = self.state.write().unwrap(); + state.turn_id = None; + state.turn_status.clear(); + state.accumulated_text.clear(); + state.accumulated_thinking.clear(); + state.active_tools.clear(); + state.active_items.clear(); + state.session_state = "idle".to_string(); + state.persistence_dirty = true; + drop(state); + self.bump_version(); + } + fn find_mergeable_item( items: &[ChatMessageItem], target_type: &str, @@ -3219,6 +3246,9 @@ impl RemoteSessionStateTracker { drop(state); self.bump_version(); } + AE::SessionHistoryChanged { .. } if is_direct => { + self.invalidate_history_projection(); + } AE::SessionTitleGenerated { title, .. } if is_direct => { let mut state = self.state.write().unwrap(); state.title = title.clone(); @@ -3382,7 +3412,7 @@ pub fn remote_persisted_poll_response( let (send_messages, send_total) = if turn_finished && !has_assistant_msg { (None, None) } else { - if !new_messages.is_empty() { + if !new_messages.is_empty() || active_turn.is_none() { tracker.mark_persistence_clean(); } (Some(new_messages), Some(total_msg_count)) @@ -3413,6 +3443,7 @@ fn non_empty_title(title: String) -> Option { #[cfg(test)] mod tests { use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Mutex; struct FakeWorkspaceHost; @@ -3521,6 +3552,7 @@ mod tests { created_requests: Mutex>, list_identities: Mutex>, removed_trackers: Mutex>, + history_error: Option, } #[async_trait::async_trait] @@ -3606,8 +3638,11 @@ mod tests { &self, _session_storage_dir: &Path, _session_id: &str, - ) -> (Vec, bool) { - ( + ) -> Result<(Vec, bool), String> { + if let Some(error) = self.history_error.as_ref() { + return Err(error.clone()); + } + Ok(( vec![ChatMessage { id: "message-1".to_string(), role: "user".to_string(), @@ -3620,7 +3655,7 @@ mod tests { items: None, }], false, - ) + )) } async fn delete_session( @@ -3735,8 +3770,36 @@ mod tests { ); } + #[tokio::test] + async fn remote_session_handler_propagates_history_outcome_unknown() { + let host = FakeSessionHost { + history_error: Some("Session history restore is incomplete".to_string()), + ..Default::default() + }; + + let response = handle_remote_session_command( + &host, + &RemoteCommand::GetSessionMessages { + session_id: "session-a".to_string(), + limit: None, + before_message_id: None, + }, + ) + .await; + + assert_eq!( + response, + RemoteResponse::Error { + message: "Session history restore is incomplete".to_string(), + } + ); + } + struct FakePollHost { tracker: Arc, + storage_dir: Option, + messages: Vec, + history_read_count: Arc, } #[async_trait::async_trait] @@ -3750,15 +3813,16 @@ mod tests { } async fn resolve_session_storage_dir(&self, _session_id: &str) -> Option { - None + self.storage_dir.clone() } async fn load_remote_chat_messages( &self, _session_storage_dir: &Path, _session_id: &str, - ) -> (Vec, bool) { - (Vec::new(), false) + ) -> Result<(Vec, bool), String> { + self.history_read_count.fetch_add(1, Ordering::Relaxed); + Ok((self.messages.clone(), false)) } } @@ -3766,6 +3830,9 @@ mod tests { async fn remote_poll_handler_preserves_missing_workspace_error() { let host = FakePollHost { tracker: Arc::new(RemoteSessionStateTracker::new("session-a".to_string())), + storage_dir: None, + messages: Vec::new(), + history_read_count: Arc::new(AtomicUsize::new(0)), }; let response = handle_remote_poll_command( @@ -3788,6 +3855,72 @@ mod tests { ); } + #[tokio::test] + async fn history_change_invalidates_clean_poll_cache_and_reports_a_shorter_projection() { + let tracker = Arc::new(RemoteSessionStateTracker::new("session-a".to_string())); + tracker.handle_agentic_event(&AgenticEvent::DialogTurnStarted { + session_id: "session-a".to_string(), + turn_id: "turn-hidden".to_string(), + turn_index: 1, + user_input: "hidden".to_string(), + original_user_input: None, + user_message_metadata: None, + }); + tracker.mark_persistence_clean(); + let previous_version = tracker.version(); + let history_read_count = Arc::new(AtomicUsize::new(0)); + let host = FakePollHost { + tracker: tracker.clone(), + storage_dir: Some(PathBuf::from("/workspace/project/.bitfun/sessions")), + messages: vec![ChatMessage { + id: "message-visible".to_string(), + role: "user".to_string(), + content: "visible".to_string(), + timestamp: "1".to_string(), + metadata: None, + tools: None, + thinking: None, + items: None, + images: None, + }], + history_read_count: history_read_count.clone(), + }; + + tracker.handle_agentic_event(&AgenticEvent::SessionHistoryChanged { + session_id: "session-a".to_string(), + }); + + assert_eq!(tracker.version(), previous_version + 1); + assert!(tracker.is_persistence_dirty()); + assert!(tracker.snapshot_active_turn().is_none()); + let response = handle_remote_poll_command( + &host, + &RemoteCommand::PollSession { + session_id: "session-a".to_string(), + since_version: previous_version, + known_msg_count: 2, + known_model_catalog_version: None, + }, + ) + .await; + + assert_eq!(history_read_count.load(Ordering::Relaxed), 1); + assert_eq!( + response, + RemoteResponse::SessionPoll { + version: previous_version + 1, + changed: true, + session_state: Some("idle".to_string()), + title: None, + new_messages: Some(Vec::new()), + total_msg_count: Some(1), + active_turn: None, + model_catalog: Box::new(None), + } + ); + assert!(!tracker.is_persistence_dirty()); + } + #[derive(Default)] struct FakeInteractionHost; diff --git a/src/crates/services/services-integrations/src/remote_connect/account.rs b/src/crates/services/services-integrations/src/remote_connect/account.rs index c7063206f..5bb19ef94 100644 --- a/src/crates/services/services-integrations/src/remote_connect/account.rs +++ b/src/crates/services/services-integrations/src/remote_connect/account.rs @@ -8,6 +8,7 @@ use anyhow::{anyhow, Result}; use argon2::{Algorithm, Argon2, Params, Version}; use base64::{engine::general_purpose::STANDARD as BASE64, Engine}; +use bitfun_services_core::session::SessionMetadata; #[cfg(test)] use rand::RngCore; use serde::{Deserialize, Serialize}; @@ -76,6 +77,73 @@ pub struct AccountSession { pub master_key: [u8; MASTER_KEY_LEN], } +const RELAY_TURNS_IMPORT_STATE_KEY: &str = "relayTurnsImportState"; +const RELAY_TURNS_IMPORT_PENDING: &str = "pending"; +const RELAY_TURNS_IMPORT_COMPLETE: &str = "complete"; + +/// Return the durable import state used by account-backed Session history. +/// A missing marker denotes a locally-created Session. +pub fn relay_session_history_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) +} + +pub fn relay_session_history_import_is_complete(metadata: &SessionMetadata) -> bool { + relay_session_history_import_state(metadata) == Some(RELAY_TURNS_IMPORT_COMPLETE) +} + +/// Fail closed while an imported Session has not completed its local turn +/// batch. Uploading a metadata-only or partial prefix would overwrite the +/// account's authoritative full-history bundle. +pub fn ensure_relay_session_history_exportable( + metadata: &SessionMetadata, +) -> std::result::Result<(), String> { + match relay_session_history_import_state(metadata) { + None | Some(RELAY_TURNS_IMPORT_COMPLETE) => Ok(()), + Some(state) => Err(format!( + "session {} history import is incomplete ({state})", + metadata.session_id + )), + } +} + +/// Align exported metadata with the visible history projection. This matters +/// while a staged Session restore hides a physical suffix from every consumer. +pub fn relay_session_export_metadata( + metadata: &SessionMetadata, + visible_turn_count: usize, +) -> SessionMetadata { + let mut exported = metadata.clone(); + exported.turn_count = visible_turn_count; + exported +} + +fn set_relay_session_history_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)); +} + +pub fn mark_relay_session_history_import_pending(metadata: &mut SessionMetadata) { + set_relay_session_history_import_state(metadata, RELAY_TURNS_IMPORT_PENDING); +} + +pub fn mark_relay_session_history_import_complete(metadata: &mut SessionMetadata) { + set_relay_session_history_import_state(metadata, RELAY_TURNS_IMPORT_COMPLETE); +} + /// A delegated token for a paired client (mobile-web / IM bot). /// The desktop requests this from the relay and transmits it along /// with the master_key to the paired client via the E2E room channel. @@ -1086,4 +1154,32 @@ mod tests { "ws://127.0.0.1:3000/relay/ws" ); } + + #[test] + fn account_history_export_waits_for_the_durable_import_marker() { + let mut metadata = SessionMetadata::new( + "session".to_string(), + "Session".to_string(), + "agentic".to_string(), + "auto".to_string(), + ); + metadata.turn_count = 3; + + assert!(ensure_relay_session_history_exportable(&metadata).is_ok()); + mark_relay_session_history_import_pending(&mut metadata); + assert!(!relay_session_history_import_is_complete(&metadata)); + assert!(ensure_relay_session_history_exportable(&metadata).is_err()); + + mark_relay_session_history_import_complete(&mut metadata); + assert!(relay_session_history_import_is_complete(&metadata)); + assert!(ensure_relay_session_history_exportable(&metadata).is_ok()); + let exported = relay_session_export_metadata(&metadata, 2); + assert_eq!(metadata.turn_count, 3); + assert_eq!(exported.turn_count, 2); + + metadata.custom_metadata = Some(serde_json::json!({ + "relayTurnsImportState": "unknown" + })); + assert!(ensure_relay_session_history_exportable(&metadata).is_err()); + } }