diff --git a/src/sdk/devbox.ts b/src/sdk/devbox.ts index e34a938d3..6d8627f58 100644 --- a/src/sdk/devbox.ts +++ b/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 { 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); + } } /** diff --git a/src/sdk/eviction.ts b/src/sdk/eviction.ts new file mode 100644 index 000000000..60eb93017 --- /dev/null +++ b/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(); + private stream: Stream | 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 { + // 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(); + +/** 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; +} diff --git a/src/sdk/index.ts b/src/sdk/index.ts index 088063279..4c8da4af8 100644 --- a/src/sdk/index.ts +++ b/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'; diff --git a/tests/smoketests/object-oriented/eviction.test.ts b/tests/smoketests/object-oriented/eviction.test.ts new file mode 100644 index 000000000..6cae58f2e --- /dev/null +++ b/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; +}; + +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((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[0]; +function fakeClient(watchEvictions: () => Promise): 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((resolve) => { + devbox.onEvict((evicted, evictionDeadlineMs) => { + receivedDevbox = evicted; + calls.push({ id: evicted.id, deadline: evictionDeadlineMs }); + resolve(); + }); + }); + + let guard: ReturnType; + const timeout = new Promise((_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(); + }); +});