fix(peer-device): reconcile live session snapshots (#1708)

Author: bobleer

docs/architecture/peer-device-mode.md

@@ -75,6 +75,14 @@ FS) and must not be mixed with Peer Device Mode.
   controllers; controller re-emits the same event names locally. This includes
   SSH-backed remote PTY Ready / Data / Exit events created on B, not only B's
   local terminal service events.
+- Because DeviceEvent delivery has no ACK/replay contract, the controller also
+  reconciles its active chat session from the Peer Host every 3s, immediately
+  after session/visibility changes, and after detecting a dropped data event.
+  Realtime events remain the primary path; snapshot reconciliation repairs a
+  controller that attached after turn/round lifecycle events or crossed a
+  transient relay gap. The host overlays its authoritative in-memory session
+  state onto the persisted view so an executing turn is not misclassified as
+  interrupted history.
 - CLI Peer Host forwards only turns submitted through Peer Host and linked
   child turns. A background-result follow-up inherits ownership only when its
   Core-internal metadata identifies the exact tracked parent and source child
@@ -96,8 +104,7 @@ FS) and must not be mixed with Peer Device Mode.
   delivery lease serializes detach or offline removal with the local Relay
   enqueue attempt. An explicit disconnect still restores the local controller
   UI, but reports a warning when host cancellation was not confirmed. This
-  boundary does not change the Relay envelope or add ACK, replay, or reconnect
-  recovery.
+  boundary does not change the Relay envelope or add ACK or replay.
 - Relay `POST /api/devices/:id/rpc` waits up to **120s** for the peer response;
   reverse proxies in front of the relay must use a matching (or higher) read
   timeout or they will return 504 first.

src/apps/cli/src/peer_host/commands/session.rs

@@ -69,6 +69,15 @@ fn session_to_json(session: Session, turn_count: usize) -> Value {
     })
 }
 
+fn overlay_live_session_state(restored: &mut Session, live: Option<Session>) {
+    let Some(live) = live else {
+        return;
+    };
+    if live.session_id == restored.session_id {
+        restored.state = live.state;
+    }
+}
+
 fn restored_session_to_json(restored: AgentSessionRestoreResult) -> Value {
     let session = restored.session;
     json!({
@@ -165,7 +174,7 @@ pub(crate) async fn restore_session_view(
         .filter(|n| *n > 0)
         .map(|n| n.min(16));
 
-    let (session, turns, total_turn_count, timings) = state
+    let (mut session, turns, total_turn_count, timings) = state
         .compatibility
         .restore_session_view_for_workspace(
             storage_request,
@@ -175,6 +184,11 @@ pub(crate) async fn restore_session_view(
         )
         .await
         .map_err(|e| format!("Failed to restore session view: {e}"))?;
+    let live_session = state
+        .compatibility
+        .loaded_session_snapshot(&session_id)
+        .map_err(|e| format!("Failed to read live session state: {e}"))?;
+    overlay_live_session_state(&mut session, live_session);
 
     let loaded_turn_count = turns.len();
     let is_partial = loaded_turn_count < total_turn_count;
@@ -533,8 +547,13 @@ pub(crate) async fn save_session_turn(
 
 #[cfg(test)]
 mod tests {
-    use super::{restored_session_to_json, session_stats_validation_error};
+    use super::{
+        overlay_live_session_state, restored_session_to_json, session_stats_validation_error,
+    };
     use bitfun_agent_runtime::sdk::{AgentSessionRestoreResult, AgentSessionSummary, SessionState};
+    use bitfun_core::agentic::core::{
+        ProcessingPhase, Session as CoreSession, SessionConfig, SessionState as CoreSessionState,
+    };
 
     #[test]
     fn basic_restore_keeps_peer_host_session_shape() {
@@ -572,4 +591,29 @@ mod tests {
             "Failed to get session stats: Validation error: session_id cannot contain path separators"
         );
     }
+
+    #[test]
+    fn view_restore_uses_the_cli_hosts_live_processing_state() {
+        let mut restored = CoreSession::new_with_id(
+            "session_1".to_string(),
+            "Main".to_string(),
+            "agentic".to_string(),
+            SessionConfig::default(),
+        );
+        let mut live = restored.clone();
+        live.state = CoreSessionState::Processing {
+            current_turn_id: "turn_1".to_string(),
+            phase: ProcessingPhase::Streaming,
+        };
+
+        overlay_live_session_state(&mut restored, Some(live));
+
+        assert!(matches!(
+            restored.state,
+            CoreSessionState::Processing {
+                ref current_turn_id,
+                ..
+            } if current_turn_id == "turn_1"
+        ));
+    }
 }

src/apps/desktop/src/runtime/session_application.rs

@@ -76,6 +76,21 @@ pub(crate) struct DesktopSessionWithTurnsRestore {
     pub turns: Vec<DialogTurnData>,
 }
 
+fn overlay_live_session_state(restored: &mut Session, live: Option<Session>) {
+    let Some(live) = live else {
+        return;
+    };
+    if live.session_id != restored.session_id {
+        return;
+    }
+
+    // Disk state deliberately stores Processing as Idle so a process restart
+    // never revives work. A view served by the process that still owns the
+    // runtime must expose its live state, otherwise a Peer controller treats
+    // an executing turn as interrupted history and drops subsequent chunks.
+    restored.state = live.state;
+}
+
 #[derive(Clone)]
 struct ResolvedDesktopSessionScope {
     workspace_path: String,
@@ -510,7 +525,7 @@ impl DesktopSessionApplication {
         let resolve_storage_path_duration_ms =
             path_started_at.elapsed().as_millis().min(u64::MAX as u128) as u64;
         on_storage_path_resolved(resolve_storage_path_duration_ms);
-        let (session, turns, total_turn_count, mut timings) = self
+        let (mut session, turns, total_turn_count, mut timings) = self
             .compatibility
             .restore_session_view_from_storage_path(
                 &storage_path,
@@ -520,6 +535,11 @@ impl DesktopSessionApplication {
             )
             .await
             .map_err(|error| DesktopSessionApplicationError::Core(error.to_string()))?;
+        let live_session = self
+            .compatibility
+            .loaded_session_snapshot(session_id)
+            .map_err(|error| DesktopSessionApplicationError::Core(error.to_string()))?;
+        overlay_live_session_state(&mut session, live_session);
         timings.resolve_storage_path_duration_ms = resolve_storage_path_duration_ms;
         Ok(DesktopSessionViewRestore {
             session,
@@ -765,6 +785,48 @@ mod tests {
         );
     }
 
+    #[test]
+    fn live_processing_state_overlays_sanitized_view_state() {
+        let mut restored = Session::new_with_id(
+            "session-1".to_string(),
+            "Restored".to_string(),
+            "agentic".to_string(),
+            Default::default(),
+        );
+        let mut live = restored.clone();
+        live.state = bitfun_core::agentic::core::SessionState::Processing {
+            current_turn_id: "turn-1".to_string(),
+            phase: bitfun_core::agentic::core::ProcessingPhase::Streaming,
+        };
+
+        overlay_live_session_state(&mut restored, Some(live));
+
+        assert!(matches!(
+            restored.state,
+            bitfun_core::agentic::core::SessionState::Processing {
+                ref current_turn_id,
+                ..
+            } if current_turn_id == "turn-1"
+        ));
+    }
+
+    #[test]
+    fn missing_live_session_keeps_persisted_view_state() {
+        let mut restored = Session::new_with_id(
+            "session-1".to_string(),
+            "Restored".to_string(),
+            "agentic".to_string(),
+            Default::default(),
+        );
+
+        overlay_live_session_state(&mut restored, None);
+
+        assert!(matches!(
+            restored.state,
+            bitfun_core::agentic::core::SessionState::Idle
+        ));
+    }
+
     #[tokio::test]
     async fn local_session_storage_identity_survives_workspace_removal() {
         let root = std::env::temp_dir().join(format!(

src/crates/assembly/core/src/product_runtime.rs

@@ -541,6 +541,21 @@ impl CoreAgentRuntimeCompatibility {
             .is_some())
     }
 
+    /// Return the authoritative in-memory session snapshot when this process
+    /// currently owns the session runtime.
+    ///
+    /// Persisted session state intentionally sanitizes `Processing` to `Idle`
+    /// so a later process restart cannot resurrect work. Live read-only views
+    /// (for example Peer Device controllers) still need the current process
+    /// state to distinguish an executing session from interrupted history.
+    pub fn loaded_session_snapshot(&self, session_id: &str) -> BitFunResult<Option<Session>> {
+        validate_persisted_session_id(session_id)?;
+        Ok(self
+            .coordinator
+            .get_session_manager()
+            .get_session(session_id))
+    }
+
     pub async fn update_loaded_session_title(
         &self,
         session_id: &str,
@@ -1056,6 +1071,7 @@ mod tests {
         let _ = build;
         let _ = CoreAgentRuntimeCompatibility::list_persisted_sessions;
         let _ = CoreAgentRuntimeCompatibility::load_persisted_session_turns;
+        let _ = CoreAgentRuntimeCompatibility::loaded_session_snapshot;
         let _ = CoreAgentRuntimeCompatibility::unload_persisted_session;
     }
 

src/web-ui/src/flow_chat/services/FlowChatManager.ts

@@ -49,6 +49,7 @@ import {
   updateImageAnalysisItem as updateImageAnalysisItemModule,
   updateSessionMetadata,
 } from './flow-chat-manager';
+import { installPeerSessionRefresh } from './flow-chat-manager/PeerSessionRefreshModule';
 
 const log = createLogger('FlowChatManager');
 
@@ -61,6 +62,7 @@ export class FlowChatManager {
   private eventListenerCleanup: (() => void) | null = null;
   private initializationRequests = new Map<string, Promise<boolean>>();
   private latestInitializationRequestKey: string | null = null;
+  private peerSessionRefreshCleanup: (() => void) | null = null;
   private disposed = false;
 
   private constructor() {
@@ -88,6 +90,7 @@ export class FlowChatManager {
 
     this.agentService = AgentService.getInstance();
     installPendingQueueDrainListener(this.context);
+    this.peerSessionRefreshCleanup = installPeerSessionRefresh(this.context);
   }
 
   /** Public hook used by the queue panel "send now" fallback to drain head item. */
@@ -414,6 +417,8 @@ export class FlowChatManager {
     this.initializationRequests.clear();
     this.latestInitializationRequestKey = null;
     this.cleanupEventListeners();
+    this.peerSessionRefreshCleanup?.();
+    this.peerSessionRefreshCleanup = null;
     this.context.eventBatcher.destroy();
   }
 

src/web-ui/src/flow_chat/services/flow-chat-manager/EventHandlerModule.ts

@@ -80,6 +80,7 @@ import {
   clearRuntimeStatus,
   scheduleModelResponseStatus,
 } from './RuntimeStatusModule';
+import { requestPeerSessionRefresh } from './PeerSessionRefreshModule';
 
 const log = createLogger('EventHandlerModule');
 const TURN_COMPLETION_QUIET_WINDOW_MS = 500;
@@ -167,6 +168,7 @@ function logDroppedDataEvent(
   turnId: string | null,
   details: Record<string, unknown>
 ): void {
+  requestPeerSessionRefresh(sessionId);
   log.debug('Dropped agentic data event', {
     eventName,
     sessionId,
@@ -1623,6 +1625,7 @@ function handleTextChunk(context: FlowChatContext, event: any): void {
 
   const dialogTurn = session.dialogTurns.find((turn: DialogTurn) => turn.id === turnId);
   if (!dialogTurn) {
+    requestPeerSessionRefresh(sessionId);
     log.debug('Dialog turn not found', { turnId });
     return;
   }
@@ -1795,6 +1798,7 @@ function handleModelRoundStart(context: FlowChatContext, event: ModelRoundStarte
 
   const dialogTurn = session.dialogTurns.find((turn: DialogTurn) => turn.id === turnId);
   if (!dialogTurn) {
+    requestPeerSessionRefresh(sessionId);
     log.debug('Dialog turn not found (model round start)', { turnId });
     return;
   }