feat(devbox): add devbox.onEvict eviction notifications (SSE) (#820)

Co-authored-by: Reflex <noreply@anthropic.com>
Co-authored-by: Reflex <reflex@runloop.ai>

Author: 3 people

src/sdk/devbox.ts

@@ -22,6 +22,7 @@ import { LongPollRequestOptions, PollingOptions } from '../lib/polling';
 import { Snapshot } from './snapshot';
 import { Execution } from './execution';
 import { ExecutionResult } from './execution-result';
+import { EvictionCallback, getEvictionMonitor } from './eviction';
 import { uuidv7 } from 'uuidv7';
 
 // Re-export Execution and ExecutionResult for Devbox namespace
@@ -1029,6 +1030,35 @@ export class Devbox {
   async keepAlive(options?: Core.RequestOptions): Promise<DevboxKeepAliveResponse> {
     return this.client.devboxes.keepAlive(this._id, options);
   }
+
+  /**
+   * Register a callback fired once if this devbox gets a pending infrastructure eviction.
+   *
+   * The first `onEvict` across any devbox on this client opens a single account-wide
+   * notification stream; it closes automatically once every registered devbox has been
+   * notified. The callback runs at most once and is invoked as
+   * `callback(devbox, evictionDeadlineMs)` — this devbox and the Unix millisecond
+   * deadline by which it will be suspended. Use it to run cleanup before suspend.
+   *
+   * @param {EvictionCallback} callback - Invoked with this devbox and its eviction deadline (ms).
+   *
+   * @example
+   * ```typescript
+   * devbox.onEvict((devbox, evictionDeadlineMs) => {
+   *   console.log(`${devbox.id} will be suspended at ${evictionDeadlineMs}`);
+   * });
+   * ```
+   */
+  onEvict(callback: EvictionCallback): void {
+    getEvictionMonitor(this.client).register(this, callback);
+  }
+
+  /**
+   * Withdraw this devbox's eviction interest registered via {@link onEvict}.
+   */
+  cancelOnEvict(): void {
+    getEvictionMonitor(this.client).unregister(this._id);
+  }
 }
 
 /**

src/sdk/eviction.ts

@@ -0,0 +1,126 @@
+import { Runloop } from '../index';
+import { Stream } from '../streaming';
+import type { DevboxEvictionEventView } from '../resources/devboxes/devboxes';
+import type { Devbox } from './devbox';
+
+/**
+ * Invoked once when the devbox it was registered for has a pending eviction.
+ *
+ * @param devbox - The devbox with a pending eviction.
+ * @param evictionDeadlineMs - Unix timestamp (ms) by which the devbox will be suspended.
+ */
+export type EvictionCallback = (devbox: Devbox, evictionDeadlineMs: number) => void;
+
+/**
+ * Fans account-wide eviction notifications out to per-devbox callbacks.
+ *
+ * One monitor is shared by every {@link Devbox} built from the same generated
+ * client (see {@link getEvictionMonitor}), so all registered devboxes are served
+ * by a single SSE connection. The stream opens the moment the first devbox
+ * registers interest and closes as soon as the last interested devbox has been
+ * notified.
+ *
+ * Delivery contract:
+ * - The server replays every currently-pending eviction on connect, so a devbox
+ *   that registers after its eviction was scheduled is still notified.
+ * - Notifications for devboxes not in the interest set are discarded.
+ * - A devbox is removed from the interest set *before* its callback runs, so the
+ *   callback fires at most once even if the server repeats the notification.
+ */
+export class EvictionMonitor {
+  private entries = new Map<string, { devbox: Devbox; callback: EvictionCallback }>();
+  private stream: Stream<DevboxEvictionEventView> | null = null;
+  private running = false;
+
+  constructor(private client: Runloop) {}
+
+  /** Add `devbox` to the interest set, opening the stream if idle. */
+  register(devbox: Devbox, callback: EvictionCallback): void {
+    this.entries.set(devbox.id, { devbox, callback });
+    if (!this.running) {
+      this.running = true;
+      void this.run();
+    }
+  }
+
+  /** Drop `devboxId`; close the stream if it was the last interested devbox. */
+  unregister(devboxId: string): void {
+    this.entries.delete(devboxId);
+    if (this.entries.size === 0) {
+      this.close();
+    }
+  }
+
+  /** Clear all interest and tear down the stream. */
+  close(): void {
+    this.entries.clear();
+    this.stream?.controller.abort();
+    this.stream = null;
+    this.running = false;
+  }
+
+  private async run(): Promise<void> {
+    // The server force-closes the stream on purpose (leader change / slow consumer) and a
+    // long-lived HTTP/2 stream can drop; the client is expected to reconnect and re-read the
+    // snapshot, which re-delivers anything missed. So reconnect (with backoff) until no devbox is
+    // still interested.
+    const INITIAL_BACKOFF_MS = 500;
+    const MAX_BACKOFF_MS = 30_000;
+    let backoff = INITIAL_BACKOFF_MS;
+    try {
+      while (this.entries.size > 0) {
+        try {
+          // Force the SSE Accept header: the endpoint only streams for text/event-stream; the
+          // generated client's default (application/json) gets an empty text/plain response, so the
+          // feed would silently deliver nothing.
+          const stream = await this.client.devboxes.watchEvictions({
+            headers: { Accept: 'text/event-stream' },
+          });
+          this.stream = stream;
+          for await (const event of stream) {
+            this.dispatch(event);
+            if (this.entries.size === 0) return;
+          }
+          // Clean end: reset backoff and reconnect if still interested.
+          backoff = INITIAL_BACKOFF_MS;
+        } catch (error) {
+          // Aborting the stream on close surfaces as an AbortError; that is an intentional teardown.
+          if (error instanceof Error && error.name === 'AbortError') return;
+          if (this.entries.size === 0) return;
+          // Otherwise a routine disconnect — fall through to backoff + reconnect.
+        }
+        if (this.entries.size === 0) return;
+        await new Promise((resolve) => setTimeout(resolve, backoff));
+        backoff = Math.min(backoff * 2, MAX_BACKOFF_MS);
+      }
+    } finally {
+      this.stream = null;
+      this.running = false;
+    }
+  }
+
+  private dispatch(event: DevboxEvictionEventView): void {
+    const entry = this.entries.get(event.devbox_id);
+    if (!entry) return;
+    // Remove before signaling so a duplicate notification for the same devbox is
+    // discarded and the callback fires at most once.
+    this.entries.delete(event.devbox_id);
+    try {
+      entry.callback(entry.devbox, event.eviction_deadline_ms);
+    } catch (error) {
+      console.error(`Error in eviction callback for devbox ${event.devbox_id}:`, error);
+    }
+  }
+}
+
+const monitors = new WeakMap<Runloop, EvictionMonitor>();
+
+/** Return the shared {@link EvictionMonitor} for `client`, creating it once. */
+export function getEvictionMonitor(client: Runloop): EvictionMonitor {
+  let monitor = monitors.get(client);
+  if (!monitor) {
+    monitor = new EvictionMonitor(client);
+    monitors.set(client, monitor);
+  }
+  return monitor;
+}

src/sdk/index.ts

@@ -1,4 +1,5 @@
 export { Devbox, DevboxCmdOps, DevboxFileOps, DevboxNetOps, type ExecuteStreamingCallbacks } from './devbox';
+export { EvictionMonitor, getEvictionMonitor, type EvictionCallback } from './eviction';
 export { Blueprint } from './blueprint';
 export { Snapshot } from './snapshot';
 export { StorageObject } from './storage-object';

tests/smoketests/object-oriented/eviction.test.ts

@@ -0,0 +1,105 @@
+import { Devbox } from '@runloop/api-client/sdk';
+import type { DevboxEvictionEventView } from '@runloop/api-client/resources/devboxes/devboxes';
+
+// Hermetic tests: they drive the shared EvictionMonitor through an in-memory fake
+// SSE stream instead of a live devbox eviction, so `devbox.onEvict` / `cancelOnEvict`
+// (and the monitor's dispatch + teardown) are exercised deterministically.
+
+// The generated watch_evictions stream: an async-iterable of eviction events plus an
+// AbortController the monitor aborts on close().
+type FakeStream = {
+  controller: AbortController;
+  [Symbol.asyncIterator](): AsyncIterator<DevboxEvictionEventView>;
+};
+
+function streamOf(events: DevboxEvictionEventView[]): FakeStream {
+  return {
+    controller: new AbortController(),
+    async *[Symbol.asyncIterator]() {
+      for (const event of events) {
+        yield event;
+      }
+    },
+  };
+}
+
+// A stream that emits nothing until the monitor aborts it (i.e. on cancelOnEvict).
+function blockingStream(): FakeStream {
+  const controller = new AbortController();
+  return {
+    controller,
+    async *[Symbol.asyncIterator]() {
+      await new Promise<void>((resolve) => {
+        if (controller.signal.aborted) {
+          resolve();
+        } else {
+          controller.signal.addEventListener('abort', () => resolve(), { once: true });
+        }
+      });
+    },
+  };
+}
+
+// Type the fake as whatever Devbox.fromId expects, without importing the generated client.
+type Client = Parameters<typeof Devbox.fromId>[0];
+function fakeClient(watchEvictions: () => Promise<FakeStream>): Client {
+  return { devboxes: { watchEvictions } } as unknown as Client;
+}
+
+const tick = (ms = 20) => new Promise((resolve) => setTimeout(resolve, ms));
+
+(process.env['RUN_SMOKETESTS'] ? describe : describe.skip)('object-oriented eviction notifications', () => {
+  test('onEvict fires once for a matching devbox and ignores others', async () => {
+    const events: DevboxEvictionEventView[] = [
+      { devbox_id: 'dbx_other', eviction_deadline_ms: 1 },
+      { devbox_id: 'dbx_match', eviction_deadline_ms: 1_720_000_000_000 },
+    ];
+    const devbox = Devbox.fromId(
+      fakeClient(async () => streamOf(events)),
+      'dbx_match',
+    );
+
+    let receivedDevbox: Devbox | undefined;
+    const calls: Array<{ id: string; deadline: number }> = [];
+    const fired = new Promise<void>((resolve) => {
+      devbox.onEvict((evicted, evictionDeadlineMs) => {
+        receivedDevbox = evicted;
+        calls.push({ id: evicted.id, deadline: evictionDeadlineMs });
+        resolve();
+      });
+    });
+
+    let guard: ReturnType<typeof setTimeout>;
+    const timeout = new Promise<never>((_resolve, reject) => {
+      guard = setTimeout(() => reject(new Error('onEvict callback never fired')), 2000);
+    });
+    try {
+      await Promise.race([fired, timeout]);
+    } finally {
+      clearTimeout(guard!);
+    }
+    await tick();
+
+    // Fired once, with the devbox object and its deadline; the unmatched id was discarded.
+    expect(receivedDevbox).toBe(devbox);
+    expect(calls).toEqual([{ id: 'dbx_match', deadline: 1_720_000_000_000 }]);
+  });
+
+  test('cancelOnEvict stops watching and aborts the stream', async () => {
+    const stream = blockingStream();
+    const devbox = Devbox.fromId(
+      fakeClient(async () => stream),
+      'dbx_cancel',
+    );
+
+    const callback = jest.fn();
+    devbox.onEvict(callback);
+    await tick(); // let the monitor open the stream
+
+    devbox.cancelOnEvict();
+
+    expect(stream.controller.signal.aborted).toBe(true);
+    await tick();
+    expect(callback).not.toHaveBeenCalled();
+  });
+});