feat(remote-connect): support targeted session rollback and edit-resend from remote control

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<SessionMaintenancePermit> {
+        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<SessionMaintenancePermit> {
         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();

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<usize>,
+    ) -> Result<RemoteSessionRollbackOutcome, String> {
+        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]

src/crates/contracts/runtime-ports/src/agent_api.rs

@@ -1742,6 +1742,10 @@ pub struct AgentSessionRollbackToTurnRequest {
     pub workspace_hostname: Option<String>,
     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<usize>,
     #[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()),