fix(app): terminal disconnect and resync (#14004)

3 files changed, 91 insertions, 12 deletions

diff --git a/packages/app/src/components/terminal.tsx b/packages/app/src/components/terminal.tsx
index 14413dfda..9048c481c 100644
--- a/packages/app/src/components/terminal.tsx
+++ b/packages/app/src/components/terminal.tsx
@@ -346,7 +346,7 @@ export const Terminal = (props: TerminalProps) => {
       }
       ghostty = g
       term = t
-      output = terminalWriter((data) => t.write(data))
+      output = terminalWriter((data, done) => t.write(data, done))
 
       t.attachCustomKeyEventHandler((event) => {
         const key = event.key.toLowerCase()
@@ -520,9 +520,19 @@ export const Terminal = (props: TerminalProps) => {
     disposed = true
     if (fitFrame !== undefined) cancelAnimationFrame(fitFrame)
     if (sizeTimer !== undefined) clearTimeout(sizeTimer)
-    output?.flush()
-    persistTerminal({ term, addon: serializeAddon, cursor, pty: local.pty, onCleanup: props.onCleanup })
-    cleanup()
+    if (ws && ws.readyState !== WebSocket.CLOSED && ws.readyState !== WebSocket.CLOSING) ws.close()
+
+    const finalize = () => {
+      persistTerminal({ term, addon: serializeAddon, cursor, pty: local.pty, onCleanup: props.onCleanup })
+      cleanup()
+    }
+
+    if (!output) {
+      finalize()
+      return
+    }
+
+    output.flush(finalize)
   })
 
   return (
diff --git a/packages/app/src/utils/terminal-writer.test.ts b/packages/app/src/utils/terminal-writer.test.ts
index d48dd4f4e..c49702e39 100644
--- a/packages/app/src/utils/terminal-writer.test.ts
+++ b/packages/app/src/utils/terminal-writer.test.ts
@@ -6,7 +6,10 @@ describe("terminalWriter", () => {
     const calls: string[] = []
     const scheduled: VoidFunction[] = []
     const writer = terminalWriter(
-      (data) => calls.push(data),
+      (data, done) => {
+        calls.push(data)
+        done?.()
+      },
       (flush) => scheduled.push(flush),
     )
 
@@ -24,10 +27,38 @@ describe("terminalWriter", () => {
   test("flush is a no-op when empty", () => {
     const calls: string[] = []
     const writer = terminalWriter(
-      (data) => calls.push(data),
+      (data, done) => {
+        calls.push(data)
+        done?.()
+      },
       (flush) => flush(),
     )
     writer.flush()
     expect(calls).toEqual([])
   })
+
+  test("flush waits for pending write completion", () => {
+    const calls: string[] = []
+    let done: VoidFunction | undefined
+    const writer = terminalWriter(
+      (data, finish) => {
+        calls.push(data)
+        done = finish
+      },
+      (flush) => flush(),
+    )
+
+    writer.push("a")
+
+    let settled = false
+    writer.flush(() => {
+      settled = true
+    })
+
+    expect(calls).toEqual(["a"])
+    expect(settled).toBe(false)
+
+    done?.()
+    expect(settled).toBe(true)
+  })
 })
diff --git a/packages/app/src/utils/terminal-writer.ts b/packages/app/src/utils/terminal-writer.ts
index b6caff789..083f51de4 100644
--- a/packages/app/src/utils/terminal-writer.ts
+++ b/packages/app/src/utils/terminal-writer.ts
@@ -1,16 +1,42 @@
 export function terminalWriter(
-  write: (data: string) => void,
+  write: (data: string, done?: VoidFunction) => void,
   schedule: (flush: VoidFunction) => void = queueMicrotask,
 ) {
   let chunks: string[] | undefined
+  let waits: VoidFunction[] | undefined
   let scheduled = false
+  let writing = false
 
-  const flush = () => {
+  const settle = () => {
+    if (scheduled || writing || chunks?.length) return
+    const list = waits
+    if (!list?.length) return
+    waits = undefined
+    for (const fn of list) {
+      fn()
+    }
+  }
+
+  const run = () => {
+    if (writing) return
     scheduled = false
     const items = chunks
-    if (!items?.length) return
+    if (!items?.length) {
+      settle()
+      return
+    }
     chunks = undefined
-    write(items.join(""))
+    writing = true
+    write(items.join(""), () => {
+      writing = false
+      if (chunks?.length) {
+        if (scheduled) return
+        scheduled = true
+        schedule(run)
+        return
+      }
+      settle()
+    })
   }
 
   const push = (data: string) => {
@@ -18,9 +44,21 @@ export function terminalWriter(
     if (chunks) chunks.push(data)
     else chunks = [data]
 
-    if (scheduled) return
+    if (scheduled || writing) return
     scheduled = true
-    schedule(flush)
+    schedule(run)
+  }
+
+  const flush = (done?: VoidFunction) => {
+    if (!scheduled && !writing && !chunks?.length) {
+      done?.()
+      return
+    }
+    if (done) {
+      if (waits) waits.push(done)
+      else waits = [done]
+    }
+    run()
   }
 
   return { push, flush }