fix(typecheck): recreate missing CLI Transport interface (Gitlawb#1581)

* fix(typecheck): recreate missing CLI Transport interface

* fix(transports): implement async close in CLI transports

* fix(transports): harden async close cleanup

* fix: address async transport close review feedback

* test: isolate environment-sensitive suites

* fix(transports): drain uploader after close failure

* fix(transports): type guard CCR stream events

* test: remove unused auto-compact fixture helper

Author: chioarub

src/bridge/remoteBridgeCore.ts

@@ -490,7 +490,7 @@ export async function initEnvLessBridgeCore(
     flushGate.start()
     try {
       const seq = transport.getLastSequenceNum()
-      transport.close()
+      await transport.close()
       transport = await createV2ReplTransport({
         sessionUrl: buildCCRv2SdkUrl(fresh.api_base_url, sessionId),
         ingressToken: fresh.worker_jwt,
@@ -506,7 +506,7 @@ export async function initEnvLessBridgeCore(
         // Teardown fired during the async createV2ReplTransport window.
         // Don't wire/connect/schedule — we'd re-arm timers after cancelAll()
         // and fire onInboundMessage into a torn-down bridge.
-        transport.close()
+        await transport.close()
         return
       }
       wireTransportCallbacks()
@@ -717,7 +717,14 @@ export async function initEnvLessBridgeCore(
       }
     }
 
-    transport.close()
+    try {
+      await transport.close()
+    } catch (err) {
+      logForDebugging(
+        `[remote-bridge] Transport close threw during teardown: ${errorMessage(err)}`,
+        { level: 'error' },
+      )
+    }
 
     const archiveStatus: ArchiveTelemetryStatus =
       status === 'no_token'

src/bridge/replBridge.ts

@@ -626,6 +626,20 @@ export async function initBridgeCore(
     }
   }
 
+  async function closeTransportBestEffort(
+    transportToClose: ReplBridgeTransport,
+    reason: string,
+  ): Promise<void> {
+    try {
+      await transportToClose.close()
+    } catch (err) {
+      logForDebugging(
+        `[bridge:repl] Transport close threw ${reason}: ${errorMessage(err)}`,
+        { level: 'error' },
+      )
+    }
+  }
+
   async function doReconnect(): Promise<boolean> {
     environmentRecreations++
     // Invalidate any in-flight v2 handshake — the environment is being
@@ -652,7 +666,7 @@ export async function initBridgeCore(
       if (seq > lastTransportSequenceNum) {
         lastTransportSequenceNum = seq
       }
-      transport.close()
+      await closeTransportBestEffort(transport, 'during reconnect')
       transport = null
     }
     // Transport is gone — wake the poll loop out of its at-capacity
@@ -1055,7 +1069,7 @@ export async function initBridgeCore(
     // forwarding prompts → ~25-min dead window observed in daemon logs.
     // Kill the transport + work state so isAtCapacity()=false; the loop
     // fast-polls and picks up the server's re-dispatched work in seconds.
-    onHeartbeatFatal: (err: BridgeFatalError) => {
+    onHeartbeatFatal: async (err: BridgeFatalError) => {
       logForDebugging(
         `[bridge:repl] heartbeatWork fatal (status=${err.status}) — tearing down work item for fast re-dispatch`,
       )
@@ -1064,7 +1078,7 @@ export async function initBridgeCore(
         if (seq > lastTransportSequenceNum) {
           lastTransportSequenceNum = seq
         }
-        transport.close()
+        await closeTransportBestEffort(transport, 'after heartbeat fatal')
         transport = null
       }
       flushGate.drop()
@@ -1094,7 +1108,7 @@ export async function initBridgeCore(
       }
       return { environmentId, environmentSecret }
     },
-    onWorkReceived: (
+    onWorkReceived: async (
       workSessionId: string,
       ingressToken: string,
       workId: string,
@@ -1199,7 +1213,10 @@ export async function initBridgeCore(
         if (oldSeq > lastTransportSequenceNum) {
           lastTransportSequenceNum = oldSeq
         }
-        oldTransport.close()
+        await closeTransportBestEffort(
+          oldTransport,
+          'while replacing work transport',
+        )
       }
       // Reset flush state — the old flush (if any) is no longer relevant.
       // Preserve pending messages so they're drained after the new
@@ -1410,13 +1427,16 @@ export async function initBridgeCore(
           sessionId: workSessionId,
           initialSequenceNum: lastTransportSequenceNum,
         }).then(
-          t => {
+          async t => {
             // Teardown started while registerWorker was in flight. Teardown
             // saw transport === null and skipped close(); installing now
             // would leak CCRClient heartbeat timers and reset
             // teardownStarted via wireTransport's side effects.
             if (pollController.signal.aborted) {
-              t.close()
+              await closeTransportBestEffort(
+                t,
+                'while discarding aborted CCR v2 transport',
+              )
               return
             }
             // onWorkReceived may have fired again while registerWorker()
@@ -1429,7 +1449,10 @@ export async function initBridgeCore(
               logForDebugging(
                 `[bridge:repl] CCR v2: discarding stale handshake gen=${thisGen} current=${v2Generation}`,
               )
-              t.close()
+              await closeTransportBestEffort(
+                t,
+                'while discarding stale CCR v2 transport',
+              )
               return
             }
             wireTransport(t)
@@ -1668,8 +1691,10 @@ export async function initBridgeCore(
     // log their own success/failure internally.
     await Promise.all([stopWorkP, archiveSession(currentSessionId)])
 
-    teardownTransport?.close()
-    logForDebugging('[bridge:repl] Teardown: transport closed')
+    if (teardownTransport) {
+      await closeTransportBestEffort(teardownTransport, 'during teardown')
+      logForDebugging('[bridge:repl] Teardown: transport closed')
+    }
 
     await api.deregisterEnvironment(environmentId).catch((err: unknown) => {
       logForDebugging(
@@ -1892,7 +1917,7 @@ async function startWorkPollLoop({
     ingressToken: string,
     workId: string,
     useCodeSessions: boolean,
-  ) => void
+  ) => Promise<void>
   /** Called when the environment has been deleted. Returns new credentials or null. */
   onEnvironmentLost?: () => Promise<{
     environmentId: string
@@ -1935,7 +1960,7 @@ async function startWorkPollLoop({
    * ~10-minute dead window before recovery). When omitted, falls back to
    * the backoff sleep to avoid a tight poll+heartbeat loop.
    */
-  onHeartbeatFatal?: (err: BridgeFatalError) => void
+  onHeartbeatFatal?: (err: BridgeFatalError) => Promise<void>
 }): Promise<void> {
   const MAX_ENVIRONMENT_RECREATIONS = 3
 
@@ -2063,7 +2088,7 @@ async function startWorkPollLoop({
                   // for the server's re-dispatched work item. Without
                   // the hook, backoff to avoid tight poll+heartbeat loop.
                   if (onHeartbeatFatal) {
-                    onHeartbeatFatal(err)
+                    await onHeartbeatFatal(err)
                     logForDebugging(
                       `[bridge:repl:heartbeat] Fatal (status=${err.status}), work state cleared — fast-polling for re-dispatch`,
                     )
@@ -2197,7 +2222,7 @@ async function startWorkPollLoop({
           continue
         }
 
-        onWorkReceived(
+        await onWorkReceived(
           workSessionId,
           secret.session_ingress_token,
           work.id,

src/bridge/replBridgeTransport.ts

@@ -23,7 +23,7 @@ import { registerWorker } from './workSecret.js'
 export type ReplBridgeTransport = {
   write(message: StdoutMessage): Promise<void>
   writeBatch(messages: StdoutMessage[]): Promise<void>
-  close(): void
+  close(): Promise<void>
   isConnectedStatus(): boolean
   getStateLabel(): string
   setOnData(callback: (data: string) => void): void
@@ -210,20 +210,8 @@ export async function createV2ReplTransport(opts: {
       logForDebugging(
         '[bridge:repl] CCR v2: epoch superseded (409) — closing for poll-loop recovery',
       )
-      // Close resources in a try block so the throw always executes.
-      // If ccr.close() or sse.close() throw, we still need to unwind
-      // the caller (request()) — otherwise handleEpochMismatch's `never`
-      // return type is violated at runtime and control falls through.
-      try {
-        ccr.close()
-        sse.close()
-        onCloseCb?.(4090)
-      } catch (closeErr: unknown) {
-        logForDebugging(
-          `[bridge:repl] CCR v2: error during epoch-mismatch cleanup: ${errorMessage(closeErr)}`,
-          { level: 'error' },
-        )
-      }
+      closeResourcesBestEffort('epoch-mismatch cleanup')
+      onCloseCb?.(4090)
       // Don't return — the calling request() code continues after the 409
       // branch, so callers see the logged warning and a false return. We
       // throw to unwind; the uploaders catch it as a send failure.
@@ -267,6 +255,30 @@ export async function createV2ReplTransport(opts: {
   let ccrInitialized = false
   let closed = false
 
+  function closeResources(): Promise<void> {
+    closed = true
+    return Promise.all([ccr.close(), sse.close()]).then(() => {})
+  }
+
+  function closeResourcesBestEffort(reason: string): void {
+    void closeResources().catch((closeErr: unknown) => {
+      logForDebugging(
+        `[bridge:repl] CCR v2: error during ${reason}: ${errorMessage(closeErr)}`,
+        { level: 'error' },
+      )
+    })
+  }
+
+  function closeCcrBestEffort(reason: string): void {
+    closed = true
+    void ccr.close().catch((closeErr: unknown) => {
+      logForDebugging(
+        `[bridge:repl] CCR v2: error during ${reason}: ${errorMessage(closeErr)}`,
+        { level: 'error' },
+      )
+    })
+  }
+
   return {
     write(msg) {
       return ccr.writeEvent(msg)
@@ -281,11 +293,7 @@ export async function createV2ReplTransport(opts: {
         await ccr.writeEvent(m)
       }
     },
-    close() {
-      closed = true
-      ccr.close()
-      sse.close()
-    },
+    close: closeResources,
     isConnectedStatus() {
       // Write-readiness, not read-readiness — replBridge checks this
       // before calling writeBatch. SSE open state is orthogonal.
@@ -309,7 +317,7 @@ export async function createV2ReplTransport(opts: {
       // heartbeat timer before notifying replBridge. (sse.close() doesn't
       // invoke this, so the epoch-mismatch path above isn't double-firing.)
       sse.setOnClose(code => {
-        ccr.close()
+        closeCcrBestEffort('SSE close cleanup')
         cb(code ?? 4092)
       })
     },
@@ -360,8 +368,7 @@ export async function createV2ReplTransport(opts: {
           // so the poll loop can retry on the next work dispatch.
           // Without this callback, replBridge never learns the transport
           // failed to initialize and sits with transport === null forever.
-          ccr.close()
-          sse.close()
+          closeResourcesBestEffort('initialize failure cleanup')
           onCloseCb?.(4091) // 4091 = init failure, distinguishable from 4090 epoch mismatch
         },
       )

src/bridge/replBridgeTransport.types.test.ts

@@ -0,0 +1,99 @@
+import { expect, test } from 'bun:test'
+import { BridgeFatalError } from './bridgeApi.js'
+import { DEFAULT_POLL_CONFIG } from './pollConfigDefaults.js'
+import { _startWorkPollLoopForTesting } from './replBridge.js'
+import type { ReplBridgeTransport } from './replBridgeTransport.js'
+import type { BridgeApiClient } from './types.js'
+
+type AssertPromiseVoid<T extends Promise<void>> = T
+
+type _CloseReturnsPromise = AssertPromiseVoid<
+  ReturnType<ReplBridgeTransport['close']>
+>
+
+test('work poll loop waits for heartbeat-fatal cleanup before fast-polling again', async () => {
+  const abort = new AbortController()
+  let pollCount = 0
+  let resolveCleanup: (() => void) | undefined
+  const cleanupStarted = deferred<void>()
+  const secondPollStarted = deferred<void>()
+  const cleanupReleased = new Promise<void>(resolve => {
+    resolveCleanup = resolve
+  })
+
+  const api = {
+    pollForWork: async () => {
+      pollCount += 1
+      if (pollCount === 2) {
+        secondPollStarted.resolve()
+        abort.abort()
+      }
+      return null
+    },
+    heartbeatWork: async () => {
+      throw new BridgeFatalError('work item gone', 404)
+    },
+  } as unknown as BridgeApiClient
+
+  let atCapacity = true
+  const loop = _startWorkPollLoopForTesting({
+    api,
+    getCredentials: () => ({
+      environmentId: 'env-1',
+      environmentSecret: 'secret-1',
+    }),
+    signal: abort.signal,
+    isAtCapacity: () => atCapacity,
+    capacitySignal: createCapacitySignal,
+    getHeartbeatInfo: () => ({
+      environmentId: 'env-1',
+      workId: 'work-1',
+      sessionToken: 'token-1',
+    }),
+    getPollIntervalConfig: () => ({
+      ...DEFAULT_POLL_CONFIG,
+      poll_interval_ms_not_at_capacity: 1,
+      poll_interval_ms_at_capacity: 10_000,
+      non_exclusive_heartbeat_interval_ms: 1,
+      reclaim_older_than_ms: 0,
+    }),
+    onWorkReceived: async () => {},
+    onHeartbeatFatal: async () => {
+      atCapacity = false
+      cleanupStarted.resolve()
+      await cleanupReleased
+    },
+  })
+
+  await cleanupStarted.promise
+  await Promise.resolve()
+  await Promise.resolve()
+  expect(pollCount).toBe(1)
+
+  resolveCleanup?.()
+  await secondPollStarted.promise
+  await loop
+  expect(pollCount).toBe(2)
+})
+
+function createCapacitySignal(): {
+  signal: AbortSignal
+  cleanup: () => void
+} {
+  const controller = new AbortController()
+  return {
+    signal: controller.signal,
+    cleanup: () => {},
+  }
+}
+
+function deferred<T>(): {
+  promise: Promise<T>
+  resolve: (value: T | PromiseLike<T>) => void
+} {
+  let resolve!: (value: T | PromiseLike<T>) => void
+  const promise = new Promise<T>(res => {
+    resolve = res
+  })
+  return { promise, resolve }
+}

src/cli/remoteIO.ts

@@ -244,12 +244,12 @@ export class RemoteIO extends StructuredIO {
   /**
    * Clean up connections gracefully
    */
-  close(): void {
+  async close(): Promise<void> {
     if (this.keepAliveTimer) {
       clearInterval(this.keepAliveTimer)
       this.keepAliveTimer = null
     }
-    this.transport.close()
+    await this.transport.close()
     this.inputStream.end()
   }
 }