fix: await async callbacks in CommandHandle.wait()

.changeset/await-async-callbacks.md

@@ -0,0 +1,5 @@
+---
+"e2b": patch
+---
+
+fix: await async callbacks in CommandHandle.wait()

packages/js-sdk/src/sandbox/commands/commandHandle.ts

@@ -226,11 +226,11 @@ export class CommandHandle
     try {
       for await (const [stdout, stderr, pty] of this.iterateEvents()) {
         if (stdout !== null) {
-          this.onStdout?.(stdout)
+          await this.onStdout?.(stdout)
         } else if (stderr !== null) {
-          this.onStderr?.(stderr)
+          await this.onStderr?.(stderr)
         } else if (pty) {
-          this.onPty?.(pty)
+          await this.onPty?.(pty)
         }
       }
     } catch (e) {

packages/js-sdk/tests/sandbox/commands/commandHandle.test.ts

@@ -0,0 +1,97 @@
+import { describe, expect, it, vi } from 'vitest'
+
+import { CommandHandle } from '../../../src/sandbox/commands/commandHandle'
+
+type EventKind = 'stdout' | 'stderr' | 'pty'
+
+function createEvents(kind: EventKind): AsyncIterable<any> {
+  async function* events() {
+    if (kind === 'pty') {
+      yield {
+        event: {
+          event: {
+            case: 'data',
+            value: {
+              output: {
+                case: 'pty',
+                value: new Uint8Array([1, 2, 3]),
+              },
+            },
+          },
+        },
+      }
+    } else {
+      yield {
+        event: {
+          event: {
+            case: 'data',
+            value: {
+              output: {
+                case: kind,
+                value: new TextEncoder().encode(kind),
+              },
+            },
+          },
+        },
+      }
+    }
+
+    yield {
+      event: {
+        event: {
+          case: 'end',
+          value: {
+            exitCode: 0,
+            error: undefined,
+          },
+        },
+      },
+    }
+  }
+
+  return events()
+}
+
+describe('CommandHandle', () => {
+  it.each<EventKind>(['stdout', 'stderr', 'pty'])(
+    'wait awaits async %s callbacks',
+    async (kind) => {
+      let callbackStarted = false
+      let releaseCallback: (() => void) | undefined
+
+      const callbackBlocked = new Promise<void>((resolve) => {
+        releaseCallback = resolve
+      })
+
+      const callback = async () => {
+        callbackStarted = true
+        await callbackBlocked
+      }
+
+      const handle = new CommandHandle(
+        1,
+        () => {},
+        async () => true,
+        createEvents(kind),
+        kind === 'stdout' ? callback : undefined,
+        kind === 'stderr' ? callback : undefined,
+        kind === 'pty' ? callback : undefined
+      )
+
+      let waitResolved = false
+      const waitPromise = handle.wait().then(() => {
+        waitResolved = true
+      })
+
+      await vi.waitFor(() => {
+        expect(callbackStarted).toBe(true)
+      })
+      expect(waitResolved).toBe(false)
+
+      releaseCallback?.()
+      await waitPromise
+
+      expect(waitResolved).toBe(true)
+    }
+  )
+})