Apply remote ACK recovery snapshots on desktop clients

Co-authored-by: Orca <help@stably.ai>

src/main/runtime/rpc/methods/terminal.ts

@@ -1031,6 +1031,9 @@ export const TERMINAL_METHODS: RpcAnyMethod[] = [
           const displayMode = runtime.getMobileDisplayMode(stream.ptyId)
           // Why: dropped ACK-pending output means live frames are no longer a
           // complete replay. Send a fresh model snapshot before resuming output.
+          // Why: truncated marks an unusable snapshot, and clients discard
+          // those. The recovery snapshot must be applied to cover dropped
+          // output, so it is only truncated when serialization failed.
           sendSnapshotFrames((opcode, payload) => sendFrame(stream.streamId, opcode, payload), {
             kind: 'scrollback',
             cols: serialized?.cols ?? size?.cols ?? 80,
@@ -1039,10 +1042,23 @@ export const TERMINAL_METHODS: RpcAnyMethod[] = [
             reason: 'ack-pending-overflow',
             seq: serialized?.seq,
             source: serialized?.source,
-            truncated: true,
+            truncated: !serialized,
             truncatedByByteBudget: serialized?.truncatedByByteBudget,
             data: serialized?.data ?? ''
           })
+          if (serialized && typeof serialized.seq === 'number') {
+            // Why: retained chunks queued before the snapshot serialized are
+            // already contained in it; replaying them would duplicate output.
+            const snapshotSeq = serialized.seq
+            const retained = stream.ackPendingOutput.filter(
+              (chunk) => !(typeof chunk.seq === 'number' && chunk.seq <= snapshotSeq)
+            )
+            stream.ackPendingOutput = retained
+            stream.ackPendingOutputBytes = retained.reduce(
+              (total, chunk) => total + chunk.bytes.byteLength,
+              0
+            )
+          }
           stream.ackPendingOutputOverflowed = false
         } catch (error) {
           sendStreamError(

src/main/runtime/rpc/terminal-multiplex.test.ts

@@ -867,6 +867,12 @@ describe('terminal multiplex RPC', () => {
       (frame) => frame?.opcode === TerminalStreamOpcode.Output
     )
     expect(recoveryStartIndex).toBeGreaterThanOrEqual(0)
+    // Why: clients discard truncated snapshots; a usable recovery snapshot
+    // must not be marked truncated or the dropped output gap is permanent.
+    expect(
+      decodeTerminalStreamJson<{ truncated?: boolean }>(drainFrames[recoveryStartIndex]!.payload)
+        ?.truncated
+    ).toBe(false)
     expect(firstOutputAfterAckIndex).toBeGreaterThan(recoveryStartIndex)
     expect(
       drainFrames
@@ -884,6 +890,154 @@ describe('terminal multiplex RPC', () => {
     await dispatchPromise
   })
 
+  it('trims recovery-covered ACK pending output instead of replaying it', async () => {
+    const messages: string[] = []
+    const binaryFrames: Uint8Array<ArrayBufferLike>[] = []
+    const handlers = new Map<
+      number,
+      (frame: NonNullable<ReturnType<typeof decodeTerminalStreamFrame>>) => void
+    >()
+    const cleanups = new Map<string, () => void>()
+    const dataListenerRef: {
+      current?: (data: string, meta?: { seq?: number; rawLength?: number }) => void
+    } = {}
+    const floodedChars = 3 * 1024 * 1024
+    const runtime = stubRuntime({
+      resolveLeafForHandle: vi.fn().mockReturnValue({ ptyId: 'pty-1' }),
+      readTerminal: vi.fn().mockResolvedValue({ tail: [], truncated: false }),
+      serializeTerminalBuffer: vi
+        .fn()
+        .mockResolvedValueOnce({ data: 'initial snapshot', cols: 120, rows: 40 })
+        // Why: the recovery snapshot seq covers the entire flood, so every
+        // retained pending chunk is already contained in the snapshot.
+        .mockResolvedValue({ data: 'recovered snapshot', cols: 120, rows: 40, seq: floodedChars }),
+      getTerminalSize: vi.fn().mockReturnValue({ cols: 120, rows: 40 }),
+      getMobileDisplayMode: vi.fn().mockReturnValue('auto'),
+      getLayout: vi.fn().mockReturnValue({ seq: 1 }),
+      subscribeToTerminalData: vi.fn(
+        (
+          _: string,
+          listener: (data: string, meta?: { seq?: number; rawLength?: number }) => void
+        ) => {
+          dataListenerRef.current = listener
+          return vi.fn()
+        }
+      ),
+      subscribeToTerminalResize: vi.fn().mockReturnValue(vi.fn()),
+      subscribeToFitOverrideChanges: vi.fn().mockReturnValue(vi.fn()),
+      subscribeToDriverChanges: vi.fn().mockReturnValue(vi.fn()),
+      getTerminalFitOverride: vi.fn().mockReturnValue(null),
+      getDriver: vi.fn().mockReturnValue({ kind: 'idle' }),
+      registerSubscriptionCleanup: vi.fn((id: string, cleanup: () => void) => {
+        cleanups.set(id, cleanup)
+      }),
+      cleanupSubscription: vi.fn((id: string) => {
+        const cleanup = cleanups.get(id)
+        cleanups.delete(id)
+        cleanup?.()
+      }),
+      waitForTerminal: vi.fn(() => new Promise<RuntimeTerminalWait>(() => {})),
+      sendTerminal: vi.fn().mockResolvedValue({ accepted: true }),
+      updateDesktopViewport: vi.fn().mockResolvedValue(true)
+    })
+    const dispatcher = new RpcDispatcher({ runtime, methods: TERMINAL_METHODS })
+
+    const dispatchPromise = dispatcher.dispatchStreaming(
+      makeRequest('terminal.multiplex', {}),
+      (msg) => messages.push(msg),
+      {
+        connectionId: 'conn-ack-trim',
+        sendBinary: (bytes) => binaryFrames.push(bytes),
+        registerBinaryStreamHandler: (streamId, handler) => {
+          handlers.set(streamId, handler)
+          return () => handlers.delete(streamId)
+        }
+      }
+    )
+
+    await vi.waitFor(() =>
+      expect(messages.some((msg) => JSON.parse(msg).result?.type === 'ready')).toBe(true)
+    )
+    handlers.get(0)?.(
+      decodeTerminalStreamFrame(
+        encodeTerminalStreamFrame({
+          opcode: TerminalStreamOpcode.Subscribe,
+          streamId: 0,
+          seq: 1,
+          payload: encodeTerminalStreamJson({
+            streamId: 31,
+            terminal: 'terminal-1',
+            client: { id: 'desktop-1', type: 'desktop' },
+            viewport: { cols: 120, rows: 40 },
+            capabilities: { ackOutput: 1 }
+          })
+        })
+      )!
+    )
+    await vi.waitFor(() =>
+      expect(messages.some((msg) => JSON.parse(msg).result?.type === 'subscribed')).toBe(true)
+    )
+    binaryFrames.splice(0)
+
+    const output = 'x'.repeat(floodedChars)
+    dataListenerRef.current?.(output, { seq: floodedChars, rawLength: floodedChars })
+    const initialBytes = binaryFrames
+      .map((frame) => decodeTerminalStreamFrame(frame))
+      .filter((frame) => frame?.opcode === TerminalStreamOpcode.Output)
+      .reduce((total, frame) => total + (frame?.payload.byteLength ?? 0), 0)
+    expect(initialBytes).toBeLessThanOrEqual(512 * 1024)
+
+    binaryFrames.splice(0)
+    handlers.get(31)?.(
+      decodeTerminalStreamFrame(
+        encodeTerminalStreamFrame({
+          opcode: TerminalStreamOpcode.Ack,
+          streamId: 31,
+          seq: 2,
+          payload: encodeTerminalStreamJson({ bytes: initialBytes })
+        })
+      )!
+    )
+    await vi.waitFor(() =>
+      expect(
+        binaryFrames
+          .map((frame) => decodeTerminalStreamFrame(frame))
+          .some((frame) => frame?.opcode === TerminalStreamOpcode.SnapshotEnd)
+      ).toBe(true)
+    )
+
+    const framesAfterRecovery = binaryFrames.map((frame) => decodeTerminalStreamFrame(frame))
+    expect(
+      framesAfterRecovery
+        .filter((frame) => frame?.opcode === TerminalStreamOpcode.SnapshotChunk)
+        .map((frame) => (frame ? decodeTerminalStreamText(frame.payload) : ''))
+        .join('')
+    ).toBe('recovered snapshot')
+    // Why: every retained chunk is covered by the recovery snapshot seq;
+    // replaying any of them would duplicate snapshot content.
+    expect(
+      framesAfterRecovery.filter((frame) => frame?.opcode === TerminalStreamOpcode.Output)
+    ).toEqual([])
+
+    binaryFrames.splice(0)
+    const fresh = 'fresh-after-recovery\r\n'
+    dataListenerRef.current?.(fresh, {
+      seq: floodedChars + fresh.length,
+      rawLength: fresh.length
+    })
+    await vi.waitFor(() => {
+      const freshOutput = binaryFrames
+        .map((frame) => decodeTerminalStreamFrame(frame))
+        .filter((frame) => frame?.opcode === TerminalStreamOpcode.Output)
+        .map((frame) => (frame ? decodeTerminalStreamText(frame.payload) : ''))
+        .join('')
+      expect(freshOutput).toBe(fresh)
+    })
+
+    runtime.cleanupSubscription('terminal-multiplex:conn-ack-trim')
+    await dispatchPromise
+  })
+
   it('marks multiplex fallback snapshots truncated when the uncursored read is limited', async () => {
     const messages: string[] = []
     const binaryFrames: Uint8Array<ArrayBufferLike>[] = []

src/renderer/src/runtime/remote-runtime-terminal-multiplexer.ts

@@ -76,7 +76,7 @@ type RemoteRuntimeMultiplexedTerminalState = {
   snapshotChunks: Uint8Array<ArrayBufferLike>[]
   snapshotBytes: number
   snapshotOverflowed: boolean
-  snapshotTarget: 'initial' | 'request'
+  snapshotTarget: 'initial' | 'request' | 'recovery'
   snapshotInfo: RemoteRuntimeSnapshotInfo | null
   initialSnapshotReceived: boolean
   pendingSnapshotRequest: RemoteRuntimeSnapshotRequest | null
@@ -426,7 +426,9 @@ class RemoteRuntimeTerminalMultiplexer {
         typeof requestId === 'number' ||
         (stream.initialSnapshotReceived && stream.pendingSnapshotRequest)
           ? 'request'
-          : 'initial'
+          : stream.initialSnapshotReceived
+            ? 'recovery'
+            : 'initial'
       return
     }
     if (frame.opcode === TerminalStreamOpcode.SnapshotChunk) {
@@ -469,6 +471,10 @@ class RemoteRuntimeTerminalMultiplexer {
           clearPendingSnapshotRequest(stream)
         } else if (target === 'initial') {
           stream.callbacks.onSnapshot(data ?? '')
+        } else if (target === 'recovery' && data) {
+          // Why: a server-pushed recovery snapshot replaces terminal state
+          // mid-session; clear the screen and scrollback before applying it.
+          stream.callbacks.onSnapshot(`\x1b[2J\x1b[3J\x1b[H${data}`)
         }
       } else if (matchesPendingRequest) {
         pendingRequest.resolve(null)

src/renderer/src/runtime/runtime-terminal-stream.test.ts

@@ -4,6 +4,7 @@ import {
   decodeTerminalStreamFrame,
   decodeTerminalStreamJson,
   encodeTerminalStreamFrame,
+  encodeTerminalStreamJson,
   encodeTerminalStreamText
 } from '../../../shared/terminal-stream-protocol'
 import {
@@ -332,4 +333,74 @@ describe('remote runtime terminal multiplex ACK gate', () => {
     heldTerminal.close()
     liveTerminal.close()
   })
+
+  it('applies mid-session recovery snapshots without re-subscribing', async () => {
+    const { getRemoteRuntimeTerminalMultiplexer } =
+      await import('./remote-runtime-terminal-multiplexer')
+    resetRemoteRuntimeTerminalMultiplexersForTests()
+
+    const multiplexer = getRemoteRuntimeTerminalMultiplexer('env-recovery')
+    const onSnapshot = vi.fn()
+    const onSubscribed = vi.fn()
+    const stream = await multiplexer.subscribeTerminal({
+      terminal: 'terminal-recovery',
+      client: { id: 'desktop-recovery', type: 'desktop' },
+      callbacks: {
+        onData: vi.fn(),
+        onSnapshot,
+        onSubscribed
+      }
+    })
+    await vi.waitFor(() => expect(sendBinary).toHaveBeenCalled())
+    const streamId = stream.streamId
+
+    const injectSnapshot = (info: Record<string, unknown>, text: string): void => {
+      callbacks?.onBinary?.(
+        encodeTerminalStreamFrame({
+          opcode: TerminalStreamOpcode.SnapshotStart,
+          streamId,
+          seq: 1,
+          payload: encodeTerminalStreamJson(info)
+        })
+      )
+      callbacks?.onBinary?.(
+        encodeTerminalStreamFrame({
+          opcode: TerminalStreamOpcode.SnapshotChunk,
+          streamId,
+          seq: 2,
+          payload: encodeTerminalStreamText(text)
+        })
+      )
+      callbacks?.onBinary?.(
+        encodeTerminalStreamFrame({
+          opcode: TerminalStreamOpcode.SnapshotEnd,
+          streamId,
+          seq: 3,
+          payload: new Uint8Array(0)
+        })
+      )
+    }
+
+    injectSnapshot({ kind: 'scrollback', cols: 120, rows: 40, truncated: false }, 'initial state')
+    expect(onSnapshot).toHaveBeenCalledWith('initial state')
+    expect(onSubscribed).toHaveBeenCalledTimes(1)
+
+    injectSnapshot(
+      {
+        kind: 'scrollback',
+        cols: 120,
+        rows: 40,
+        reason: 'ack-pending-overflow',
+        truncated: false
+      },
+      'recovered state'
+    )
+    // Why: an unsolicited recovery snapshot replaces terminal state, so it
+    // clears screen and scrollback first and must not replay the subscribe
+    // lifecycle.
+    expect(onSnapshot).toHaveBeenCalledWith(`\x1b[2J\x1b[3J\x1b[H${'recovered state'}`)
+    expect(onSubscribed).toHaveBeenCalledTimes(1)
+
+    stream.close()
+  })
 })