diff --git a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs index f8da52705b..08372b8d81 100644 --- a/src/crates/assembly/core/src/agentic/coordination/scheduler.rs +++ b/src/crates/assembly/core/src/agentic/coordination/scheduler.rs @@ -1712,12 +1712,39 @@ impl DialogScheduler { session_id: &str, requested_storage_path: &std::path::Path, wait_timeout: Duration, + ) -> OpenBitFunResult { + self.begin_session_maintenance_with_policy( + session_id, + requested_storage_path, + wait_timeout, + false, + ) + .await + } + + pub(crate) async fn begin_session_maintenance_with_policy( + &self, + session_id: &str, + requested_storage_path: &std::path::Path, + wait_timeout: Duration, + require_idle: bool, ) -> OpenBitFunResult { openbitfun_core_types::validate_session_id(session_id) .map_err(OpenBitFunError::Validation)?; let operation_guard = self.lock_session_operation(session_id).await; self.session_manager .validate_session_storage_path_binding(session_id, requested_storage_path)?; + // Check only after admission is locked, before cancelling or retiring anything. + // A controller's stream can lag another controller's accepted submission. + if require_idle + && (self.is_session_busy_or_queued(session_id) + || self.queue_depth(session_id) > 0 + || self.round_injection_buffer.pending_count(session_id) > 0) + { + return Err(OpenBitFunError::Validation( + "Session rollback requires an idle session with an empty queue".to_string(), + )); + } let mut retired_turn_ids = if self.queue_depth(session_id) > 0 { self.clear_queue(session_id).await } else { @@ -4544,6 +4571,68 @@ mod tests { .matches_turn("session-a", "shared-turn")); } + #[tokio::test] + async fn idle_only_maintenance_preserves_work_accepted_before_lock_acquisition() { + let (scheduler, session_manager, _, root) = test_scheduler(); + let session_id = "remote-rollback-busy"; + let storage = root.path().join("sessions"); + session_manager + .ensure_session_storage_path(session_id, &storage) + .unwrap(); + // The phone may have observed idle before another controller was admitted. + let guard = scheduler.lock_session_operation(session_id).await; + let request = scheduler.begin_session_maintenance_with_policy( + session_id, + &storage, + Duration::ZERO, + true, + ); + tokio::pin!(request); + assert!( + tokio::time::timeout(Duration::from_millis(10), &mut request) + .await + .is_err() + ); + scheduler + .queues + .enqueue( + session_id, + standard_queued_turn("queued"), + DialogQueuePriority::Normal, + ) + .unwrap(); + scheduler + .active_turns + .insert(session_id, desktop_active_turn("active")); + drop(guard); + assert!(request + .await + .err() + .unwrap() + .to_string() + .contains("idle session")); + assert!(scheduler.active_turns.matches_turn(session_id, "active")); + assert_eq!(scheduler.queue_depth(session_id), 1); + + scheduler.active_turns.remove(session_id); + assert!( + scheduler + .begin_session_maintenance_with_policy(session_id, &storage, Duration::ZERO, true) + .await + .is_err(), + "idle but queued must also be rejected" + ); + assert_eq!(scheduler.queue_depth(session_id), 1); + scheduler.clear_queue(session_id).await; + assert!( + scheduler + .begin_session_maintenance_with_policy(session_id, &storage, Duration::ZERO, true) + .await + .is_ok(), + "empty idle session can be maintained" + ); + } + #[tokio::test] async fn wrong_workspace_deletion_leaves_active_and_queued_turns_untouched() { let (scheduler, session_manager, _, root) = test_scheduler(); diff --git a/src/crates/assembly/core/src/service_agent_runtime.rs b/src/crates/assembly/core/src/service_agent_runtime.rs index 6eb2db285a..a4b5a95748 100644 --- a/src/crates/assembly/core/src/service_agent_runtime.rs +++ b/src/crates/assembly/core/src/service_agent_runtime.rs @@ -22,11 +22,11 @@ use openbitfun_agent_runtime::sdk::{ use openbitfun_events::AgenticEvent; #[cfg(feature = "remote-connect")] use openbitfun_runtime_ports::{ - AgentDialogSteerRequest, AgentInputAttachment, AgentSubmissionSource, - AgentTurnCancellationRequest, DialogSteerOutcome, PermissionPolicyPreset, - RemoteControlStatePort, RemoteControlStateRequest, RemoteControlStateSnapshot, - RemoteSessionWorkspaceIdentity, RuntimeServiceCapability, RuntimeServicePort, - ToolPermissionConfig, + AgentDialogSteerRequest, AgentInputAttachment, AgentSessionComposerUpdate, + AgentSubmissionSource, AgentTurnCancellationRequest, DialogSteerOutcome, + PermissionPolicyPreset, RemoteControlStatePort, RemoteControlStateRequest, + RemoteControlStateSnapshot, RemoteSessionWorkspaceIdentity, RuntimeServiceCapability, + RuntimeServicePort, ToolPermissionConfig, }; use openbitfun_runtime_ports::{ AgentDialogTurnPort, AgentDialogTurnRequest, AgentLifecycleDeliveryPort, @@ -54,10 +54,11 @@ use openbitfun_services_integrations::remote_connect::{ RemoteInitialSyncRuntimeHost, RemoteInteractionRuntimeHost, RemoteModelCapabilityFact, RemoteModelCatalog, RemoteModelCatalogFacts, RemoteModelFacts, RemotePermissionMode, RemotePollRuntimeHost, RemoteRecentWorkspaceFacts, RemoteSessionMetadata, - RemoteSessionModelSelection, RemoteSessionRuntimeHost, RemoteSessionStateTracker, - RemoteSessionTrackerHost, RemoteTerminalPrewarmRequest, RemoteWorkspaceFacts, - RemoteWorkspaceFileRuntimeHost, RemoteWorkspaceKind as RemoteConnectWorkspaceKind, - RemoteWorkspaceRuntimeHost, RemoteWorkspaceUpdate, + RemoteSessionModelSelection, RemoteSessionRollbackOutcome, RemoteSessionRuntimeHost, + RemoteSessionStateTracker, RemoteSessionTrackerHost, RemoteTerminalPrewarmRequest, + RemoteWorkspaceFacts, RemoteWorkspaceFileRuntimeHost, + RemoteWorkspaceKind as RemoteConnectWorkspaceKind, RemoteWorkspaceRuntimeHost, + RemoteWorkspaceUpdate, }; use std::sync::Arc; use std::time::Duration; @@ -866,6 +867,7 @@ fn remote_chat_history_turn_from_core_turn( RemoteChatHistoryTurn { turn_id: turn.turn_id.clone(), + turn_index: turn.turn_index, user_message_id: turn.user_message.id.clone(), user_display_content: user_projection.content, user_timestamp_ms: turn.user_message.timestamp, @@ -1265,7 +1267,12 @@ impl AgentSessionRevertPort for ScheduledSessionManagementPort { .map_err(map_session_close_error)?; let maintenance = self .scheduler - .begin_session_maintenance(&request.session_id, &storage_path, Duration::from_secs(30)) + .begin_session_maintenance_with_policy( + &request.session_id, + &storage_path, + Duration::from_secs(30), + request.require_idle, + ) .await .map_err(map_session_close_error)?; let _mutation = session_manager @@ -3411,6 +3418,61 @@ impl RemoteSessionRuntimeHost for CoreRemoteSessionRuntimeHost { CoreServiceAgentRuntime::load_remote_chat_messages(session_storage_dir, session_id).await } + /// Reuse the desktop targeted-rollback transaction so a remote client + /// retires turns and restores files for real instead of hiding messages in + /// its own transcript. + async fn rollback_session_to_turn( + &self, + session_id: &str, + target_turn_id: &str, + expected_storage_turn_index: Option, + ) -> Result { + let binding = CoreServiceAgentRuntime::resolve_session_workspace_binding(session_id) + .await + .ok_or_else(|| { + format!("Session workspace binding not available for session: {session_id}") + })?; + if binding.is_remote() { + return Err("Session rollback is unavailable for remote workspaces".to_string()); + } + ensure_remote_binding_runtime_ownership(self.coordinator.as_ref(), &binding).await?; + + let outcome = self + .runtime + .rollback_session_to_turn(AgentSessionRollbackToTurnRequest { + workspace_path: binding.logical_workspace_path_string(), + workspace_id: binding.workspace_id.clone(), + workspace_hostname: Some(binding.session_identity.hostname.clone()), + session_id: session_id.to_string(), + target_turn_id: target_turn_id.to_string(), + require_idle: true, + expected_storage_turn_index, + expected_catalog_revision: None, + remote_connection_id: None, + remote_ssh_host: None, + }) + .await + .map_err(|error| error.into_message())?; + + match outcome { + AgentSessionRollbackToTurnOutcome::Completed { result } => { + Ok(RemoteSessionRollbackOutcome { + retired_turn_ids: result.retired_turn_ids, + restored_files: result.restored_files, + composer_text: match result.composer { + AgentSessionComposerUpdate::Replace { text } => Some(text), + AgentSessionComposerUpdate::Preserve + | AgentSessionComposerUpdate::Clear => None, + }, + changed: result.changed, + }) + } + AgentSessionRollbackToTurnOutcome::RecoveryRequired { reason, .. } => Err(format!( + "Session rollback requires recovery before it can continue: {reason}" + )), + } + } + async fn delete_session( &self, session_storage_dir: &std::path::Path, @@ -3991,6 +4053,14 @@ mod tests { .and_then(|source| source.split("fn remove_tracker").next()) .expect("remote session delete"); assert!(delete.contains("ensure_remote_binding_runtime_ownership")); + + let rollback = remote_session_host + .split("async fn rollback_session_to_turn") + .nth(1) + .and_then(|source| source.split("async fn delete_session").next()) + .expect("remote session rollback"); + assert!(rollback.contains("ensure_remote_binding_runtime_ownership")); + assert!(rollback.contains("binding.is_remote()")); } #[test] diff --git a/src/crates/contracts/runtime-ports/src/agent_api.rs b/src/crates/contracts/runtime-ports/src/agent_api.rs index 45b50bc523..1c701b1712 100644 --- a/src/crates/contracts/runtime-ports/src/agent_api.rs +++ b/src/crates/contracts/runtime-ports/src/agent_api.rs @@ -1742,6 +1742,10 @@ pub struct AgentSessionRollbackToTurnRequest { pub workspace_hostname: Option, pub session_id: String, pub target_turn_id: String, + /// Reject active or queued work under the host scheduling lock before any mutation. + /// Older callers retain the existing cancel-and-drain maintenance policy. + #[serde(default)] + pub require_idle: bool, #[serde(default, skip_serializing_if = "Option::is_none")] pub expected_storage_turn_index: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -2221,6 +2225,7 @@ mod tests { workspace_hostname: Some("localhost".to_string()), session_id: "session-1".to_string(), target_turn_id: "turn-7".to_string(), + require_idle: true, expected_storage_turn_index: Some(7), expected_catalog_revision: Some("catalog-3".to_string()), remote_connection_id: None, @@ -2246,6 +2251,7 @@ mod tests { "targetTurnId": "turn-7" })) .expect("deserialize pre-workspace-identity rollback request"); + assert!(!legacy_request.require_idle); assert_eq!(legacy_request.workspace_id, None); assert_eq!(legacy_request.workspace_hostname, None); // IDs are opaque. SSH metadata is an old projection, not type authority. @@ -3408,6 +3414,7 @@ mod tests { parent_tool_call_id: Some("tool_1".to_string()), subagent_type: Some("explore".to_string()), agent_id: Some("parser-review".to_string()), + workspace_id: None, workspace_path: Some("/workspace/project".to_string()), remote_connection_id: Some("conn-1".to_string()), remote_ssh_host: Some("host-1".to_string()), diff --git a/src/crates/services/services-integrations/src/remote_connect.rs b/src/crates/services/services-integrations/src/remote_connect.rs index af9e166285..de9249ad17 100644 --- a/src/crates/services/services-integrations/src/remote_connect.rs +++ b/src/crates/services/services-integrations/src/remote_connect.rs @@ -557,6 +557,7 @@ pub const REMOTE_FILE_MAX_CHUNK_BYTES: u64 = 3 * 1024 * 1024; pub const REMOTE_CAPABILITY_HARNESS_PROFILES_V1: &str = "harness_profiles_v1"; pub const REMOTE_CAPABILITY_DIALOG_STEER_V1: &str = "dialog_steer_v1"; pub const REMOTE_CAPABILITY_USER_QUESTION_INTERACTION_V1: &str = "user_question_interaction_v1"; +pub const REMOTE_CAPABILITY_SESSION_ROLLBACK_V1: &str = "session_rollback_v1"; pub const REMOTE_CAPABILITY_PLAN_BUILD_V1: &str = "plan_build_v1"; pub use host_stream::REMOTE_CAPABILITY_HOST_STREAM_V1; @@ -566,6 +567,7 @@ fn remote_host_capabilities() -> Vec { REMOTE_CAPABILITY_HARNESS_PROFILES_V1.to_string(), REMOTE_CAPABILITY_DIALOG_STEER_V1.to_string(), REMOTE_CAPABILITY_PLAN_BUILD_V1.to_string(), + REMOTE_CAPABILITY_SESSION_ROLLBACK_V1.to_string(), REMOTE_CAPABILITY_USER_QUESTION_INTERACTION_V1.to_string(), REMOTE_CAPABILITY_HOST_STREAM_V1.to_string(), ] @@ -1368,6 +1370,32 @@ pub fn remote_session_deleted_response(session_id: impl Into) -> RemoteR } } +/// What the host observed while rolling a session back to a turn. +/// +/// This mirrors the fields the desktop rollback result exposes that a remote +/// transcript actually needs, so this crate keeps describing the wire contract +/// without depending on the runtime port types. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct RemoteSessionRollbackOutcome { + pub retired_turn_ids: Vec, + pub restored_files: Vec, + pub composer_text: Option, + pub changed: bool, +} + +pub fn remote_session_rolled_back_response( + session_id: impl Into, + outcome: RemoteSessionRollbackOutcome, +) -> RemoteResponse { + RemoteResponse::SessionRolledBack { + session_id: session_id.into(), + retired_turn_ids: outcome.retired_turn_ids, + restored_files: outcome.restored_files, + composer_text: outcome.composer_text, + changed: outcome.changed, + } +} + #[async_trait::async_trait] pub trait RemoteSessionRuntimeHost: Send + Sync { async fn workspace_by_id(&self, _workspace_id: &str) -> Result { @@ -1413,6 +1441,16 @@ pub trait RemoteSessionRuntimeHost: Send + Sync { session_storage_dir: &Path, session_id: &str, ) -> Result<(Vec, bool), String>; + /// Roll the session back to `target_turn_id`, retiring that turn and the + /// ones after it and restoring the files those turns wrote. Hosts are + /// expected to reuse the same targeted-rollback path the desktop uses rather + /// than hiding turns, and to reject sessions whose workspace they do not own. + async fn rollback_session_to_turn( + &self, + session_id: &str, + target_turn_id: &str, + expected_storage_turn_index: Option, + ) -> Result; async fn delete_session( &self, session_storage_dir: &Path, @@ -1690,6 +1728,29 @@ where Err(message) => RemoteResponse::Error { message }, } } + RemoteCommand::RollbackSessionToTurn { + session_id, + target_turn_id, + expected_storage_turn_index, + } => { + if target_turn_id.trim().is_empty() { + return RemoteResponse::Error { + message: "Rollback requires a target turn id".into(), + }; + } + + match host + .rollback_session_to_turn( + session_id, + target_turn_id.trim(), + *expected_storage_turn_index, + ) + .await + { + Ok(outcome) => remote_session_rolled_back_response(session_id.clone(), outcome), + Err(message) => RemoteResponse::Error { message }, + } + } _ => RemoteResponse::Error { message: "Unknown session command".into(), }, @@ -2202,6 +2263,11 @@ pub struct ChatMessage { pub metadata: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub turn_id: Option, + /// Storage turn index for `turn_id`, carried so a remote rollback can send + /// the same optimistic-concurrency guard the desktop sends and cannot + /// retire a different turn after the transcript moved underneath it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub turn_index: Option, #[serde(default, skip_serializing_if = "Option::is_none")] pub status: Option, #[serde(default, skip_serializing_if = "Option::is_none")] @@ -2235,6 +2301,7 @@ pub struct ChatMessageItem { #[derive(Debug, Clone, PartialEq)] pub struct RemoteChatHistoryTurn { pub turn_id: String, + pub turn_index: usize, pub user_message_id: String, pub user_display_content: String, pub user_timestamp_ms: u64, @@ -2300,6 +2367,7 @@ pub fn build_remote_chat_messages(turns: Vec) -> Vec) -> Vec, + }, CancelTool { tool_id: String, reason: Option, @@ -2913,6 +2997,19 @@ pub enum RemoteResponse { SessionDeleted { session_id: String, }, + SessionRolledBack { + session_id: String, + /// Turns the host retired, including the target turn itself. Reported so + /// a client can tell the user what was withdrawn; the transcript itself + /// is repaired by the next poll's authoritative `message_snapshot`. + retired_turn_ids: Vec, + restored_files: Vec, + /// Prompt text the desktop would have put back into its composer, so a + /// remote client can offer the same "edit and resend" continuation. + #[serde(skip_serializing_if = "Option::is_none")] + composer_text: Option, + changed: bool, + }, InitialSync { #[serde(default, skip_serializing_if = "Option::is_none")] workspace_id: Option, @@ -3097,7 +3194,8 @@ where | RemoteCommand::SetSessionModel { .. } | RemoteCommand::UpdateSessionTitle { .. } | RemoteCommand::GetSessionMessages { .. } - | RemoteCommand::DeleteSession { .. } => host.handle_session_command(command).await, + | RemoteCommand::DeleteSession { .. } + | RemoteCommand::RollbackSessionToTurn { .. } => host.handle_session_command(command).await, RemoteCommand::PollSession { .. } => host.handle_poll_command(command).await, @@ -4610,6 +4708,7 @@ mod tests { model_updates: Mutex>)>>, removed_trackers: Mutex>, history_error: Option, + rollback_requests: Mutex)>>, } fn fake_workspace( @@ -4797,7 +4896,8 @@ mod tests { content: "hello".to_string(), timestamp: "1".to_string(), metadata: None, - turn_id: None, + turn_id: Some("turn-1".to_string()), + turn_index: Some(0), status: None, error: None, images: None, @@ -4809,6 +4909,25 @@ mod tests { )) } + async fn rollback_session_to_turn( + &self, + session_id: &str, + target_turn_id: &str, + expected_storage_turn_index: Option, + ) -> Result { + self.rollback_requests.lock().unwrap().push(( + session_id.to_string(), + target_turn_id.to_string(), + expected_storage_turn_index, + )); + Ok(RemoteSessionRollbackOutcome { + retired_turn_ids: vec!["turn-2".to_string()], + restored_files: vec!["src/main.rs".to_string()], + composer_text: Some("previous prompt".to_string()), + changed: true, + }) + } + async fn delete_session( &self, _session_storage_dir: &Path, @@ -5092,6 +5211,59 @@ mod tests { ); } + #[tokio::test] + async fn remote_session_handler_forwards_rollback_target_and_guard() { + let host = FakeSessionHost::default(); + + let response = handle_remote_session_command( + &host, + &RemoteCommand::RollbackSessionToTurn { + session_id: "session-a".to_string(), + target_turn_id: " turn-1 ".to_string(), + expected_storage_turn_index: Some(3), + }, + ) + .await; + + assert_eq!( + response, + RemoteResponse::SessionRolledBack { + session_id: "session-a".to_string(), + retired_turn_ids: vec!["turn-2".to_string()], + restored_files: vec!["src/main.rs".to_string()], + composer_text: Some("previous prompt".to_string()), + changed: true, + } + ); + assert_eq!( + host.rollback_requests.lock().unwrap().as_slice(), + [("session-a".to_string(), "turn-1".to_string(), Some(3))] + ); + } + + #[tokio::test] + async fn remote_session_handler_rejects_rollback_without_target_turn() { + let host = FakeSessionHost::default(); + + let response = handle_remote_session_command( + &host, + &RemoteCommand::RollbackSessionToTurn { + session_id: "session-a".to_string(), + target_turn_id: " ".to_string(), + expected_storage_turn_index: None, + }, + ) + .await; + + assert_eq!( + response, + RemoteResponse::Error { + message: "Rollback requires a target turn id".to_string(), + } + ); + assert!(host.rollback_requests.lock().unwrap().is_empty()); + } + #[tokio::test] async fn remote_session_handler_propagates_history_outcome_unknown() { let host = FakeSessionHost { @@ -5132,6 +5304,7 @@ mod tests { timestamp: "1".to_string(), metadata: None, turn_id: Some(turn_id.to_string()), + turn_index: None, status: status.map(str::to_string), error: None, tools: None, @@ -5203,6 +5376,7 @@ mod tests { timestamp: "1".to_string(), metadata: None, turn_id: None, + turn_index: None, status: None, error: None, tools: None, @@ -5268,6 +5442,7 @@ mod tests { timestamp: "1".to_string(), metadata: None, turn_id: None, + turn_index: None, status: None, error: None, tools: None, @@ -5317,6 +5492,7 @@ mod tests { timestamp: "1".to_string(), metadata: None, turn_id: None, + turn_index: None, status: None, error: None, tools: None, diff --git a/src/crates/services/services-integrations/tests/remote_connect_contracts.rs b/src/crates/services/services-integrations/tests/remote_connect_contracts.rs index fb2403d108..585188135f 100644 --- a/src/crates/services/services-integrations/tests/remote_connect_contracts.rs +++ b/src/crates/services/services-integrations/tests/remote_connect_contracts.rs @@ -49,7 +49,8 @@ use openbitfun_services_integrations::remote_connect::{ RemoteWorkspaceFileContent, RemoteWorkspaceFileInfo, RemoteWorkspaceFileRuntimeHost, RemoteWorkspaceKind, RemoteWorkspaceUpdate, TrackerEvent, REMOTE_CAPABILITY_DIALOG_STEER_V1, REMOTE_CAPABILITY_HARNESS_PROFILES_V1, REMOTE_CAPABILITY_HOST_STREAM_V1, - REMOTE_CAPABILITY_PLAN_BUILD_V1, REMOTE_FILE_MAX_CHUNK_BYTES, REMOTE_FILE_MAX_READ_BYTES, + REMOTE_CAPABILITY_PLAN_BUILD_V1, REMOTE_CAPABILITY_SESSION_ROLLBACK_V1, + REMOTE_FILE_MAX_CHUNK_BYTES, REMOTE_FILE_MAX_READ_BYTES, }; use std::path::PathBuf; use std::sync::{Arc, Mutex}; @@ -542,6 +543,8 @@ fn remote_chat_history_assembly_preserves_message_shape_and_item_order() { assert_eq!(messages[0].role, "user"); assert_eq!(messages[0].content, "original question"); assert_eq!(messages[0].timestamp, "1"); + assert_eq!(messages[0].turn_id.as_deref(), Some("turn-1")); + assert_eq!(messages[0].turn_index, Some(4)); assert_eq!( messages[0].images.as_ref().unwrap()[0], ChatImageAttachment { @@ -552,6 +555,7 @@ fn remote_chat_history_assembly_preserves_message_shape_and_item_order() { assert_eq!(messages[1].id, "turn-1_assistant"); assert_eq!(messages[1].role, "assistant"); + assert_eq!(messages[1].turn_index, None); assert_eq!(messages[1].content, "visible text"); assert_eq!(messages[1].timestamp, "1"); assert_eq!(messages[1].thinking.as_deref(), Some("visible thought")); @@ -664,6 +668,7 @@ fn remote_connect_cancel_and_restore_policy_preserve_runtime_decisions() { fn remote_history_contract_turn(is_in_progress: bool) -> RemoteChatHistoryTurn { RemoteChatHistoryTurn { turn_id: "turn-1".to_string(), + turn_index: 4, user_message_id: "user-1".to_string(), user_display_content: "original question".to_string(), user_timestamp_ms: 1_000, @@ -1365,6 +1370,25 @@ async fn remote_connect_command_owner_preserves_cancel_and_group_routing() { ); } +#[tokio::test] +async fn remote_connect_command_owner_routes_rollback_to_the_session_group() { + let host = RecordingCommandHost::default(); + + let response = handle_remote_command( + &host, + &RemoteCommand::RollbackSessionToTurn { + session_id: "session-1".to_string(), + target_turn_id: "turn-7".to_string(), + expected_storage_turn_index: Some(6), + }, + RemoteConnectSubmissionSource::Relay, + ) + .await; + + assert!(matches!(response, RemoteResponse::SessionCreated { .. })); + assert_eq!(host.events(), vec!["session"]); +} + #[tokio::test] async fn remote_connect_dialog_runtime_owns_restore_prewarm_and_submit_order() { let host = RecordingDialogHost::new(false, Some("D:/workspace/project")); @@ -2200,6 +2224,7 @@ fn remote_connect_workspace_response_helpers_own_wire_shape() { REMOTE_CAPABILITY_HARNESS_PROFILES_V1, REMOTE_CAPABILITY_DIALOG_STEER_V1, REMOTE_CAPABILITY_PLAN_BUILD_V1, + REMOTE_CAPABILITY_SESSION_ROLLBACK_V1, "user_question_interaction_v1", REMOTE_CAPABILITY_HOST_STREAM_V1 ]) @@ -2445,6 +2470,7 @@ fn remote_connect_session_response_helpers_own_pagination_and_timestamps() { REMOTE_CAPABILITY_HARNESS_PROFILES_V1, REMOTE_CAPABILITY_DIALOG_STEER_V1, REMOTE_CAPABILITY_PLAN_BUILD_V1, + REMOTE_CAPABILITY_SESSION_ROLLBACK_V1, "user_question_interaction_v1", REMOTE_CAPABILITY_HOST_STREAM_V1 ]) @@ -2556,6 +2582,7 @@ fn remote_connect_message_dtos_keep_current_wire_shape() { timestamp: "1".to_string(), metadata: None, turn_id: Some("turn-1".to_string()), + turn_index: None, status: Some("done".to_string()), error: None, tools: Some(vec![RemoteToolStatus { @@ -2803,6 +2830,82 @@ fn remote_connect_response_wire_shape_lives_in_owner_contract() { assert_eq!(title_updated["title"], "Renamed session"); } +#[test] +fn remote_connect_rollback_wire_shape_lives_in_owner_contract() { + let with_guard: RemoteCommand = serde_json::from_value(serde_json::json!({ + "cmd": "rollback_session_to_turn", + "session_id": "session-1", + "target_turn_id": "turn-7", + "expected_storage_turn_index": 6 + })) + .expect("deserialize rollback command with guard"); + let without_guard: RemoteCommand = serde_json::from_value(serde_json::json!({ + "cmd": "rollback_session_to_turn", + "session_id": "session-1", + "target_turn_id": "turn-7" + })) + .expect("deserialize rollback command without guard"); + + assert!(matches!( + with_guard, + RemoteCommand::RollbackSessionToTurn { + expected_storage_turn_index: Some(6), + ref target_turn_id, + .. + } if target_turn_id == "turn-7" + )); + assert!(matches!( + without_guard, + RemoteCommand::RollbackSessionToTurn { + expected_storage_turn_index: None, + .. + } + )); + + let rolled_back = serde_json::to_value(RemoteResponse::SessionRolledBack { + session_id: "session-1".to_string(), + retired_turn_ids: vec!["turn-8".to_string(), "turn-9".to_string()], + restored_files: vec!["src/main.rs".to_string()], + composer_text: Some("original question".to_string()), + changed: true, + }) + .expect("serialize rollback response"); + + assert_eq!(rolled_back["resp"], "session_rolled_back"); + assert_eq!(rolled_back["retired_turn_ids"][1], "turn-9"); + assert_eq!(rolled_back["restored_files"][0], "src/main.rs"); + assert_eq!(rolled_back["composer_text"], "original question"); + assert_eq!(rolled_back["changed"], true); + + let without_composer = serde_json::to_value(RemoteResponse::SessionRolledBack { + session_id: "session-1".to_string(), + retired_turn_ids: Vec::new(), + restored_files: Vec::new(), + composer_text: None, + changed: false, + }) + .expect("serialize rollback response without composer text"); + assert!(without_composer.get("composer_text").is_none()); +} + +#[test] +fn remote_connect_user_message_carries_rollback_turn_identity_on_the_wire() { + let messages = build_remote_chat_messages(vec![remote_history_contract_turn(false)]); + let user = serde_json::to_value(&messages[0]).expect("serialize user message"); + let assistant = serde_json::to_value(&messages[1]).expect("serialize assistant message"); + + assert_eq!(user["turn_id"], "turn-1"); + assert_eq!(user["turn_index"], 4); + assert!(assistant.get("turn_index").is_none()); + + // Existing hosts and saved messages omit the additive storage index. + let mut legacy = user; + legacy.as_object_mut().unwrap().remove("turn_index"); + let decoded: ChatMessage = serde_json::from_value(legacy.clone()).unwrap(); + assert_eq!(decoded.turn_index, None); + assert_eq!(serde_json::to_value(decoded).unwrap(), legacy); +} + fn sample_remote_model_catalog(version: u64) -> RemoteModelCatalog { RemoteModelCatalog { version, @@ -3453,6 +3556,7 @@ fn remote_connect_poll_helpers_preserve_delta_and_completion_policy() { timestamp: "2".to_string(), metadata: None, turn_id: Some("turn-1".to_string()), + turn_index: None, status: Some("done".to_string()), error: None, tools: None, diff --git a/src/mobile-web/src/components/ChatMessageActions.tsx b/src/mobile-web/src/components/ChatMessageActions.tsx index d2149471d8..faa8b9cd59 100644 --- a/src/mobile-web/src/components/ChatMessageActions.tsx +++ b/src/mobile-web/src/components/ChatMessageActions.tsx @@ -1,4 +1,4 @@ -import { Copy as LucideCopy, RotateCw as LucideRotateCw, Trash2 as LucideTrash2 } from 'lucide-react'; +import { Copy as LucideCopy, Edit3 as LucideEdit3, RotateCcw as LucideRotateCcw, RotateCw as LucideRotateCw, Trash2 as LucideTrash2 } from 'lucide-react'; import React from 'react'; import { MobileActionSheet, type MobileActionSheetItem } from '@openbitfun/ui/mobile'; import { useI18n } from '../i18n'; @@ -7,21 +7,56 @@ import type { ChatMessage } from '../services/RemoteSessionManager'; interface ChatMessageActionsProps { deleting: boolean; message: ChatMessage | null; + streaming?: boolean; + rollbackSupported?: boolean; onClose: () => void; onCopy: () => void; onDelete: () => void; onResend: () => void; + onEdit?: () => void; + onRollback?: () => void; } const CopyIcon = () =>