diff --git a/rivetkit-typescript/packages/rivetkit/src/registry/index.ts b/rivetkit-typescript/packages/rivetkit/src/registry/index.ts
index 088f49af8e..9e9fe4dfed 100644
--- a/rivetkit-typescript/packages/rivetkit/src/registry/index.ts
+++ b/rivetkit-typescript/packages/rivetkit/src/registry/index.ts
@@ -507,7 +507,7 @@ export class Registry {
const application = opts.application
? createApplicationFetch(opts.application, runtime)
: undefined;
- await runtime.serveListener(
+ const listenerPromise = runtime.serveListener(
registry,
{
port,
@@ -517,6 +517,8 @@ export class Registry {
},
serveConfig,
);
+ this.#applicationListenerPromise = listenerPromise;
+ await listenerPromise;
}
/**
diff --git a/rivetkit-typescript/packages/rivetkit/tests/registry-shutdown.test.ts b/rivetkit-typescript/packages/rivetkit/tests/registry-shutdown.test.ts
index e398904e14..8be944fdef 100644
--- a/rivetkit-typescript/packages/rivetkit/tests/registry-shutdown.test.ts
+++ b/rivetkit-typescript/packages/rivetkit/tests/registry-shutdown.test.ts
@@ -29,6 +29,8 @@ function makeGate(): Gate {
interface FakeState {
/** Number of times the injected registry builder was invoked. */
builderCalls: number;
+ /** Number of times serveListener was invoked. */
+ listenerCalls: number;
/** Registry handles passed to `shutdownRegistry`, in call order. */
shutdownRegistries: RegistryHandle[];
/** Value returned by `registryActorStopThresholdMs`. */
@@ -37,6 +39,8 @@ interface FakeState {
hangShutdown: boolean;
/** When set, `shutdownRegistry` blocks on this gate before resolving. */
gate: Gate | null;
+ /** When set, `serveListener` blocks on this gate before resolving. */
+ listenerGate: Gate | null;
}
interface Fake {
@@ -52,15 +56,23 @@ interface Fake {
function createFake(): Fake {
const state: FakeState = {
builderCalls: 0,
+ listenerCalls: 0,
shutdownRegistries: [],
stopThresholdMs: undefined,
hangShutdown: false,
gate: null,
+ listenerGate: null,
};
const runtime = {
kind: "napi",
serveRegistry: async () => {},
+ serveListener: async () => {
+ state.listenerCalls += 1;
+ if (state.listenerGate) {
+ await state.listenerGate.promise;
+ }
+ },
shutdownRegistry: async (registry: RegistryHandle) => {
state.shutdownRegistries.push(registry);
if (state.hangShutdown) {
@@ -261,4 +273,29 @@ describe("Registry.shutdown", () => {
await drained;
expect(settled).toBe(true);
});
+
+ test("waits for in-flight serverless application listener before resolving", async () => {
+ const { deps, state } = createFake();
+ const gate = makeGate();
+ state.listenerGate = gate;
+
+ const registry = makeRegistry(deps, {
+ shutdown: { gracePeriodMs: 60_000 },
+ });
+ const listenPromise = registry.listen();
+
+ let settled = false;
+ const drained = registry.shutdown().then(() => {
+ settled = true;
+ });
+
+ await vi.advanceTimersByTimeAsync(0);
+ expect(state.listenerCalls).toBe(1);
+ expect(settled).toBe(false);
+
+ gate.release();
+ await vi.advanceTimersByTimeAsync(0);
+ await Promise.all([listenPromise, drained]);
+ expect(settled).toBe(true);
+ });
});