@parall/agent-core 1.18.1 → 1.19.0

package/dist/gateway-base.d.ts

@@ -18,6 +18,7 @@ export type ParallGatewayOptions = {
     dispatchAdapter: DispatchAdapter;
     log?: GatewayLogger;
     coldStartWindowMs?: number;
+    shutdownDeadlineMs?: number;
     stepIdFilePathForSession?: (sessionKey: string) => string | undefined;
     onConfigUpdate?: (data: AgentConfigUpdateData) => Promise<void> | void;
     onSessionReady?: (state: {
@@ -27,6 +28,7 @@ export type ParallGatewayOptions = {
     }) => Promise<void> | void;
     onBeforeDisconnect?: () => Promise<void> | void;
 };
+export declare function parseShutdownDeadlineMs(raw: string | undefined): number | undefined;
 export declare class ParallAgentGateway {
     private readonly opts;
     private readonly chatInfoMap;
@@ -41,8 +43,12 @@ export declare class ParallAgentGateway {
     private hadSuccessfulHello;
     private lastHeartbeatAt;
     private draining;
+    private shuttingDown;
+    private inFlightDispatches;
+    private drainResolvers;
     private readonly DISPATCHED_MESSAGES_CAP;
     private readonly COLD_START_WINDOW_MS;
+    private readonly SHUTDOWN_DEADLINE_MS;
     constructor(opts: ParallGatewayOptions);
     run(abortSignal: AbortSignal): Promise<void>;
     private tryClaimMessage;
@@ -65,6 +71,7 @@ export declare class ParallAgentGateway {
     private handleTaskComment;
     private catchUpFromDispatch;
     private handleHello;
+    private waitForDrain;
     private shutdown;
 }
 //# sourceMappingURL=gateway-base.d.ts.map

package/dist/gateway-base.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"gateway-base.d.ts","sourceRoot":"","sources":["../src/gateway-base.ts"],"names":[],"mappings":"AAIA,OAAO,EAAuB,YAAY,EAAE,QAAQ,EAAE,MAAM,aAAa,CAAC;AAC1E,OAAO,KAAK,EACV,qBAAqB,EAUtB,MAAM,aAAa,CAAC;AAErB,OAAO,KAAK,EAEV,eAAe,EAGf,aAAa,EAEd,MAAM,uBAAuB,CAAC;AAoD/B,MAAM,MAAM,oBAAoB,GAAG;IACjC,SAAS,EAAE,MAAM,CAAC;IAClB,MAAM,EAAE,YAAY,CAAC;IACrB,EAAE,EAAE,QAAQ,CAAC;IACb,eAAe,CAAC,EAAE,MAAM,CAAC;IACzB,MAAM,EAAE;QACN,UAAU,EAAE,MAAM,CAAC;QACnB,OAAO,EAAE,MAAM,CAAC;QAChB,MAAM,EAAE,MAAM,CAAC;KAChB,CAAC;IACF,WAAW,EAAE,MAAM,CAAC;IACpB,WAAW,EAAE,MAAM,CAAC;IACpB,UAAU,EAAE,MAAM,CAAC;IACnB,UAAU,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;IACrC,eAAe,EAAE,eAAe,CAAC;IACjC,GAAG,CAAC,EAAE,aAAa,CAAC;IACpB,iBAAiB,CAAC,EAAE,MAAM,CAAC;IAC3B,wBAAwB,CAAC,EAAE,CAAC,UAAU,EAAE,MAAM,KAAK,MAAM,GAAG,SAAS,CAAC;IACtE,cAAc,CAAC,EAAE,CAAC,IAAI,EAAE,qBAAqB,KAAK,OAAO,CAAC,IAAI,CAAC,GAAG,IAAI,CAAC;IACvE,cAAc,CAAC,EAAE,CAAC,KAAK,EAAE;QAAE,eAAe,CAAC,EAAE,MAAM,CAAC;QAAC,EAAE,EAAE,QAAQ,CAAC;QAAC,UAAU,EAAE,MAAM,CAAA;KAAE,KAAK,OAAO,CAAC,IAAI,CAAC,GAAG,IAAI,CAAC;IACjH,kBAAkB,CAAC,EAAE,MAAM,OAAO,CAAC,IAAI,CAAC,GAAG,IAAI,CAAC;CACjD,CAAC;AAkCF,qBAAa,kBAAkB;IAuBjB,OAAO,CAAC,QAAQ,CAAC,IAAI;IAtBjC,OAAO,CAAC,QAAQ,CAAC,WAAW,CAA+B;IAC3D,OAAO,CAAC,QAAQ,CAAC,gBAAgB,CAAqC;IACtE,OAAO,CAAC,QAAQ,CAAC,eAAe,CAAqB;IACrD,OAAO,CAAC,QAAQ,CAAC,kBAAkB,CAAqB;IACxD,OAAO,CAAC,QAAQ,CAAC,UAAU,CAAsC;IACjE,OAAO,CAAC,QAAQ,CAAC,aAAa,CAK5B;IAEF,OAAO,CAAC,SAAS,CAAM;IACvB,OAAO,CAAC,eAAe,CAAqB;IAC5C,OAAO,CAAC,cAAc,CAA+C;IACrE,OAAO,CAAC,kBAAkB,CAAS;IACnC,OAAO,CAAC,eAAe,CAAc;IACrC,OAAO,CAAC,QAAQ,CAAS;IAEzB,OAAO,CAAC,QAAQ,CAAC,uBAAuB,CAAQ;IAChD,OAAO,CAAC,QAAQ,CAAC,oBAAoB,CAAS;gBAEjB,IAAI,EAAE,oBAAoB;IAIjD,GAAG,CAAC,WAAW,EAAE,WAAW,GAAG,OAAO,CAAC,IAAI,CAAC;IA2FlD,OAAO,CAAC,eAAe;IAavB,OAAO,CAAC,WAAW;IAcnB,OAAO,CAAC,UAAU;IAWlB,OAAO,CAAC,oBAAoB;YAmBd,eAAe;YAqBf,iBAAiB;IAsF/B,OAAO,CAAC,eAAe;IASvB,OAAO,CAAC,eAAe;YAQT,gCAAgC;YAMhC,WAAW;YAyDX,gBAAgB;YAmEhB,eAAe;YAmDf,kBAAkB;YAoElB,kBAAkB;YAkBlB,4BAA4B;YAoD5B,aAAa;YA8Bb,oBAAoB;YAmCpB,iBAAiB;YAqEjB,mBAAmB;YAsGnB,WAAW;YA0DX,QAAQ;CAsBvB"}
+{"version":3,"file":"gateway-base.d.ts","sourceRoot":"","sources":["../src/gateway-base.ts"],"names":[],"mappings":"AAIA,OAAO,EAAuB,YAAY,EAAE,QAAQ,EAAE,MAAM,aAAa,CAAC;AAC1E,OAAO,KAAK,EACV,qBAAqB,EAUtB,MAAM,aAAa,CAAC;AAErB,OAAO,KAAK,EAEV,eAAe,EAGf,aAAa,EAEd,MAAM,uBAAuB,CAAC;AAoD/B,MAAM,MAAM,oBAAoB,GAAG;IACjC,SAAS,EAAE,MAAM,CAAC;IAClB,MAAM,EAAE,YAAY,CAAC;IACrB,EAAE,EAAE,QAAQ,CAAC;IACb,eAAe,CAAC,EAAE,MAAM,CAAC;IACzB,MAAM,EAAE;QACN,UAAU,EAAE,MAAM,CAAC;QACnB,OAAO,EAAE,MAAM,CAAC;QAChB,MAAM,EAAE,MAAM,CAAC;KAChB,CAAC;IACF,WAAW,EAAE,MAAM,CAAC;IACpB,WAAW,EAAE,MAAM,CAAC;IACpB,UAAU,EAAE,MAAM,CAAC;IACnB,UAAU,CAAC,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC,CAAC;IACrC,eAAe,EAAE,eAAe,CAAC;IACjC,GAAG,CAAC,EAAE,aAAa,CAAC;IACpB,iBAAiB,CAAC,EAAE,MAAM,CAAC;IAI3B,kBAAkB,CAAC,EAAE,MAAM,CAAC;IAC5B,wBAAwB,CAAC,EAAE,CAAC,UAAU,EAAE,MAAM,KAAK,MAAM,GAAG,SAAS,CAAC;IACtE,cAAc,CAAC,EAAE,CAAC,IAAI,EAAE,qBAAqB,KAAK,OAAO,CAAC,IAAI,CAAC,GAAG,IAAI,CAAC;IACvE,cAAc,CAAC,EAAE,CAAC,KAAK,EAAE;QAAE,eAAe,CAAC,EAAE,MAAM,CAAC;QAAC,EAAE,EAAE,QAAQ,CAAC;QAAC,UAAU,EAAE,MAAM,CAAA;KAAE,KAAK,OAAO,CAAC,IAAI,CAAC,GAAG,IAAI,CAAC;IACjH,kBAAkB,CAAC,EAAE,MAAM,OAAO,CAAC,IAAI,CAAC,GAAG,IAAI,CAAC;CACjD,CAAC;AAUF,wBAAgB,uBAAuB,CAAC,GAAG,EAAE,MAAM,GAAG,SAAS,GAAG,MAAM,GAAG,SAAS,CAKnF;AAkCD,qBAAa,kBAAkB;IAmCjB,OAAO,CAAC,QAAQ,CAAC,IAAI;IAlCjC,OAAO,CAAC,QAAQ,CAAC,WAAW,CAA+B;IAC3D,OAAO,CAAC,QAAQ,CAAC,gBAAgB,CAAqC;IACtE,OAAO,CAAC,QAAQ,CAAC,eAAe,CAAqB;IACrD,OAAO,CAAC,QAAQ,CAAC,kBAAkB,CAAqB;IACxD,OAAO,CAAC,QAAQ,CAAC,UAAU,CAAsC;IACjE,OAAO,CAAC,QAAQ,CAAC,aAAa,CAK5B;IAEF,OAAO,CAAC,SAAS,CAAM;IACvB,OAAO,CAAC,eAAe,CAAqB;IAC5C,OAAO,CAAC,cAAc,CAA+C;IACrE,OAAO,CAAC,kBAAkB,CAAS;IACnC,OAAO,CAAC,eAAe,CAAc;IACrC,OAAO,CAAC,QAAQ,CAAS;IAMzB,OAAO,CAAC,YAAY,CAAS;IAC7B,OAAO,CAAC,kBAAkB,CAAK;IAC/B,OAAO,CAAC,cAAc,CAAyB;IAE/C,OAAO,CAAC,QAAQ,CAAC,uBAAuB,CAAQ;IAChD,OAAO,CAAC,QAAQ,CAAC,oBAAoB,CAAS;IAI9C,OAAO,CAAC,QAAQ,CAAC,oBAAoB,CAAS;gBAEjB,IAAI,EAAE,oBAAoB;IAKjD,GAAG,CAAC,WAAW,EAAE,WAAW,GAAG,OAAO,CAAC,IAAI,CAAC;IA2FlD,OAAO,CAAC,eAAe;IAavB,OAAO,CAAC,WAAW;IAcnB,OAAO,CAAC,UAAU;IAWlB,OAAO,CAAC,oBAAoB;YAmBd,eAAe;YAqBf,iBAAiB;IAsF/B,OAAO,CAAC,eAAe;IASvB,OAAO,CAAC,eAAe;YAQT,gCAAgC;YAUhC,WAAW;YAoFX,gBAAgB;YAkFhB,eAAe;YA6Ef,kBAAkB;YA2ElB,kBAAkB;YAkBlB,4BAA4B;YAoD5B,aAAa;YA+Bb,oBAAoB;YAoCpB,iBAAiB;YAsEjB,mBAAmB;YA2GnB,WAAW;IA6DzB,OAAO,CAAC,YAAY;YAiBN,QAAQ;CA4CvB"}

package/dist/gateway-base.js

@@ -6,6 +6,22 @@ import { MENTION_ALL_USER_ID } from "@parall/sdk";
 import { buildEventBody, buildForkResultPrefix } from "./event-format.js";
 import { routeTrigger } from "./routing.js";
 import { clearDispatchMessageId, clearDispatchNoReply, clearSessionMessageId, setDispatchMessageId, setDispatchNoReply, setSessionChatId, setSessionMessageId, } from "./session-state.js";
+// Parse PRLL_SHUTDOWN_DEADLINE_MS (or any string env value) into a positive
+// integer milliseconds value, or undefined if unset/invalid. Runtimes pass
+// the result into ParallGatewayOptions.shutdownDeadlineMs; leaving it
+// undefined falls back to the gateway default (60s). Operators can raise it
+// when long-running tools need more time to drain — the K8s
+// terminationGracePeriodSeconds (currently 90s) must remain larger than this
+// plus the post-drain cleanup window (~15s), or the kubelet will SIGKILL
+// before cleanup finishes.
+export function parseShutdownDeadlineMs(raw) {
+    if (!raw)
+        return undefined;
+    const n = Number(raw);
+    if (!Number.isFinite(n) || n <= 0)
+        return undefined;
+    return Math.floor(n);
+}
 function resolveStepTarget(event) {
     if (event.type === "task" || event.targetId.startsWith("tsk_")) {
         return { target_type: "task", target_id: event.targetId };
@@ -51,11 +67,23 @@ export class ParallAgentGateway {
     hadSuccessfulHello = false;
     lastHeartbeatAt = Date.now();
     draining = false;
+    // Graceful shutdown state. When SIGTERM / abort fires, `shuttingDown` flips
+    // to true so no new dispatches start, and `inFlightDispatches` counts runs
+    // still in progress. `shutdown()` awaits drain up to SHUTDOWN_DEADLINE_MS
+    // before tearing down the WS; see handleTermination caller.
+    shuttingDown = false;
+    inFlightDispatches = 0;
+    drainResolvers = [];
     DISPATCHED_MESSAGES_CAP = 5000;
     COLD_START_WINDOW_MS;
+    // SHUTDOWN_DEADLINE_MS is read by waitForDrain via the configured value
+    // below — kept as instance state so per-runtime configs can override it
+    // (see parseShutdownDeadlineMs and runtime entrypoints).
+    SHUTDOWN_DEADLINE_MS;
     constructor(opts) {
         this.opts = opts;
         this.COLD_START_WINDOW_MS = opts.coldStartWindowMs ?? 5 * 60_000;
+        this.SHUTDOWN_DEADLINE_MS = opts.shutdownDeadlineMs ?? 60_000;
     }
     async run(abortSignal) {
         const { ws, log } = this.opts;
@@ -329,13 +357,28 @@ export class ParallAgentGateway {
             await this.createInputStep(event);
         }
     }
+    // Returns true if the dispatch actually ran; false if skipped because we
+    // are shutting down. Callers MUST treat `false` as "not dispatched" and
+    // skip the server-side ack so the event stays in the dispatch queue for
+    // catch-up on the replacement pod — otherwise we silently drop work.
     async runDispatch(event, sessionKey, bodyForAgent, earlierEvents = []) {
+        // Graceful shutdown: once SIGTERM has fired we stop accepting new work.
+        // In-flight dispatches that started before the flag flipped keep running
+        // and are awaited by shutdown() up to SHUTDOWN_DEADLINE_MS.
+        if (this.shuttingDown) {
+            this.opts.log?.info(`parall[${this.opts.accountId}]: skipping dispatch for ${event.messageId} (shutting down) — leaving unacked for catch-up on replacement pod`);
+            return false;
+        }
         setSessionChatId(sessionKey, event.targetId);
         setSessionMessageId(sessionKey, event.messageId);
         setDispatchMessageId(sessionKey, event.messageId);
         setDispatchNoReply(sessionKey, event.noReply ?? false);
         const dispatchContext = this.buildDispatchContext(event, sessionKey);
         const stepIdFilePath = dispatchContext.stepIdFilePath;
+        // sync: no await between the shuttingDown check above and this increment
+        // — JS event loop is single-threaded, so shutdown() cannot interleave
+        // here and miss our in-flight count.
+        this.inFlightDispatches++;
         try {
             if (this.activeSessionId) {
                 try {
@@ -345,6 +388,13 @@ export class ParallAgentGateway {
                     this.opts.log?.warn(`parall[${this.opts.accountId}]: failed to set session active: ${String(err)}`);
                 }
             }
+            // Persist input steps for "earlier events" (batched events that arrived
+            // while a dispatch was in flight) inside the in-flight window so a
+            // shutdown short-circuit BEFORE this point cannot leave orphan input
+            // steps that the replacement pod would duplicate on replay.
+            if (earlierEvents.length > 0) {
+                await this.createInputStepsForEarlierEvents(earlierEvents);
+            }
             await this.createInputStep(event);
             for await (const runtimeEvent of this.opts.dispatchAdapter.dispatch({
                 event,
@@ -378,20 +428,43 @@ export class ParallAgentGateway {
                     this.opts.log?.warn(`parall[${this.opts.accountId}]: failed to set session idle: ${String(err)}`);
                 }
             }
+            this.inFlightDispatches--;
+            if (this.inFlightDispatches === 0 && this.drainResolvers.length > 0) {
+                const resolvers = this.drainResolvers.splice(0);
+                for (const resolve of resolvers)
+                    resolve();
+            }
         }
+        return true;
     }
     async runForkDrainLoop(fork) {
         try {
             while (fork.queue.length > 0) {
+                // Early shutdown guard — pure optimization to avoid wasted draining
+                // during the shutdown window. The structural correctness now lives in
+                // runDispatch: input AgentStep persistence happens behind its
+                // shuttingDown gate and inFlightDispatches counter, so a shutdown
+                // landing mid-batch cannot leave orphan steps that replay would
+                // duplicate.
+                if (this.shuttingDown) {
+                    for (const item of fork.queue.splice(0))
+                        item.resolve(false);
+                    break;
+                }
                 const items = fork.queue.splice(0);
                 const events = items.map((item) => item.event);
                 const last = events[events.length - 1];
                 const earlier = events.slice(0, -1);
                 try {
-                    if (earlier.length > 0) {
-                        await this.createInputStepsForEarlierEvents(earlier);
+                    const dispatched = await this.runDispatch(last, fork.fork.sessionKey, buildEventBody(last), earlier);
+                    if (!dispatched) {
+                        // Shutdown short-circuit — resolve un-acked so the server requeues
+                        // for the replacement pod and stop draining further items.
+                        for (const item of items) {
+                            item.resolve(false);
+                        }
+                        break;
                     }
-                    await this.runDispatch(last, fork.fork.sessionKey, buildEventBody(last), earlier);
                     fork.processedEvents.push(...events);
                     for (const item of items) {
                         item.resolve(true);
@@ -453,8 +526,17 @@ export class ParallAgentGateway {
         this.draining = true;
         try {
             while (this.dispatchState.mainBuffer.length > 0 || this.dispatchState.pendingForkResults.length > 0) {
+                // Early shutdown guard — same reasoning as runForkDrainLoop: stop
+                // before any input AgentSteps are persisted, leaving the buffer +
+                // pendingForkResults intact so dispatch catch-up on the replacement
+                // pod is the single source of truth for replay.
+                if (this.shuttingDown) {
+                    this.opts.log?.info(`parall[${this.opts.accountId}]: drainMainBuffer halted (shutting down) — ${this.dispatchState.mainBuffer.length} buffered, ${this.dispatchState.pendingForkResults.length} pending fork results left for catch-up`);
+                    break;
+                }
                 if (this.dispatchState.mainBuffer.length === 0 && this.dispatchState.pendingForkResults.length > 0) {
-                    const forkPrefix = buildForkResultPrefix(this.dispatchState.pendingForkResults.splice(0));
+                    const pending = this.dispatchState.pendingForkResults.splice(0);
+                    const forkPrefix = buildForkResultPrefix(pending);
                     const syntheticEvent = {
                         type: "message",
                         targetId: "_orchestrator",
@@ -465,7 +547,14 @@ export class ParallAgentGateway {
                         body: "[Orchestrator: fork session(s) completed — review results above]",
                     };
                     this.dispatchState.mainCurrentTargetId = undefined;
-                    await this.runDispatch(syntheticEvent, this.opts.runtimeKey, forkPrefix + buildEventBody(syntheticEvent));
+                    const dispatched = await this.runDispatch(syntheticEvent, this.opts.runtimeKey, forkPrefix + buildEventBody(syntheticEvent));
+                    if (!dispatched) {
+                        // Shutdown: put the fork results back at the head so the synthetic
+                        // event is regenerated on the replacement pod. Synthetic events are
+                        // not ack'd, so this is the only state we need to preserve.
+                        this.dispatchState.pendingForkResults.unshift(...pending);
+                        break;
+                    }
                     continue;
                 }
                 const targetId = this.dispatchState.mainBuffer[0].targetId;
@@ -475,12 +564,20 @@ export class ParallAgentGateway {
                 }
                 const event = events[events.length - 1];
                 const earlier = events.slice(0, -1);
-                const forkPrefix = buildForkResultPrefix(this.dispatchState.pendingForkResults.splice(0));
+                const pendingFork = this.dispatchState.pendingForkResults.splice(0);
+                const forkPrefix = buildForkResultPrefix(pendingFork);
                 this.dispatchState.mainCurrentTargetId = event.targetId;
-                if (earlier.length > 0) {
-                    await this.createInputStepsForEarlierEvents(earlier);
+                // Earlier-event input steps are persisted inside runDispatch (behind
+                // its shutdown gate) so a shutdown short-circuit cannot leave orphans.
+                const dispatched = await this.runDispatch(event, this.opts.runtimeKey, forkPrefix + buildEventBody(event), earlier);
+                if (!dispatched) {
+                    // Shutdown: skip the ack so the server redelivers these buffered
+                    // events to the replacement pod via dispatch catch-up. Put both the
+                    // buffered events and the fork results back so nothing is lost.
+                    this.dispatchState.mainBuffer.unshift(...events);
+                    this.dispatchState.pendingForkResults.unshift(...pendingFork);
+                    break;
                 }
-                await this.runDispatch(event, this.opts.runtimeKey, forkPrefix + buildEventBody(event), earlier);
                 for (const bufferedEvent of events) {
                     const sourceType = bufferedEvent.ackSourceType ?? (bufferedEvent.type === "task" ? "task_activity" : "message");
                     const sourceId = bufferedEvent.ackSourceId ?? bufferedEvent.messageId;
@@ -501,16 +598,23 @@ export class ParallAgentGateway {
         const disposition = routeTrigger(event, this.dispatchState);
         switch (disposition.action) {
             case "main": {
-                const forkPrefix = buildForkResultPrefix(this.dispatchState.pendingForkResults.splice(0));
+                const pendingFork = this.dispatchState.pendingForkResults.splice(0);
+                const forkPrefix = buildForkResultPrefix(pendingFork);
                 this.dispatchState.mainDispatching = true;
                 this.dispatchState.mainCurrentTargetId = event.targetId;
+                let dispatched = false;
                 try {
-                    await this.runDispatch(event, this.opts.runtimeKey, forkPrefix + buildEventBody(event));
+                    dispatched = await this.runDispatch(event, this.opts.runtimeKey, forkPrefix + buildEventBody(event));
+                    if (!dispatched) {
+                        // Shutdown short-circuit — restore the fork results so a future
+                        // pod can replay them, and return false so handleMessage skips ack.
+                        this.dispatchState.pendingForkResults.unshift(...pendingFork);
+                    }
                 }
                 finally {
                     await this.drainMainBuffer();
                 }
-                return true;
+                return dispatched;
             }
             case "buffer-main":
                 this.dispatchState.mainBuffer.push(event);
@@ -620,6 +724,8 @@ export class ParallAgentGateway {
         };
     }
     async handleMessage(data) {
+        if (this.shuttingDown)
+            return; // drain window — let server requeue via catch-up
         const chatId = data.chat_id;
         if (data.sender_id === this.opts.agentUserId)
             return;
@@ -655,6 +761,8 @@ export class ParallAgentGateway {
         }
     }
     async handleTaskAssignment(task, ackSourceId) {
+        if (this.shuttingDown)
+            return false; // drain window — let server requeue via catch-up
         const dedupeKey = `${task.id}:${task.updated_at}`;
         if (this.dispatchedTasks.has(dedupeKey)) {
             this.opts.log?.info(`parall[${this.opts.accountId}]: skipping already-dispatched task ${task.identifier ?? task.id}`);
@@ -689,6 +797,8 @@ export class ParallAgentGateway {
         return dispatched;
     }
     async handleTaskComment(commentId, taskId, actorId) {
+        if (this.shuttingDown)
+            return false; // drain window — let server requeue via catch-up
         const dedupeKey = `comment:${commentId}`;
         if (this.dispatchedTasks.has(dedupeKey))
             return false;
@@ -766,6 +876,10 @@ export class ParallAgentGateway {
                 cursor,
             });
             for (const item of page.data ?? []) {
+                // Shutdown short-circuit: stop fetching/building work that runDispatch
+                // will only reject. Items remain unacked for the replacement pod.
+                if (this.shuttingDown)
+                    break;
                 if (minAge > 0 && new Date(item.created_at).getTime() < minAge) {
                     this.opts.client.ackDispatchByID(this.opts.config.org_id, item.id).catch(() => { });
                     skippedOld++;
@@ -854,7 +968,9 @@ export class ParallAgentGateway {
                     this.opts.log?.warn(`parall[${this.opts.accountId}]: catch-up dispatch ${item.id} (${item.event_type}) failed: ${String(err)}`);
                 }
             }
-            cursor = page.has_more ? page.next_cursor : undefined;
+            // Stop paginating once shutdown begins — the inner loop already broke,
+            // and the next page would be wasted API work the replacement pod redoes.
+            cursor = !this.shuttingDown && page.has_more ? page.next_cursor : undefined;
         } while (cursor);
         if (processed > 0 || skippedOld > 0) {
             this.opts.log?.info(`parall[${this.opts.accountId}]: dispatch catch-up: processed ${processed}, skipped ${skippedOld} old item(s)`);
@@ -917,7 +1033,45 @@ export class ParallAgentGateway {
             log?.error(`parall[${this.opts.accountId}]: failed to fetch chats: ${String(err)}`);
         }
     }
+    // Resolves when in-flight dispatches hit 0 or the deadline elapses.
+    // Deadline arming the resolver rather than rejecting keeps shutdown() linear
+    // — callers don't need to try/catch.
+    waitForDrain(deadlineMs) {
+        if (this.inFlightDispatches === 0)
+            return Promise.resolve();
+        return new Promise((resolve) => {
+            const timer = setTimeout(() => {
+                // Remove this resolver so a late-completing dispatch doesn't call it.
+                const idx = this.drainResolvers.indexOf(onDrain);
+                if (idx >= 0)
+                    this.drainResolvers.splice(idx, 1);
+                resolve();
+            }, deadlineMs);
+            const onDrain = () => {
+                clearTimeout(timer);
+                resolve();
+            };
+            this.drainResolvers.push(onDrain);
+        });
+    }
     async shutdown() {
+        // Flip the flag first so any WS event that fires while we're draining
+        // short-circuits at runDispatch instead of starting new work.
+        this.shuttingDown = true;
+        // Drain: wait for in-flight dispatches to finish up to the configured
+        // deadline. If we time out, whatever is still running will be killed by
+        // the process exit — that's acceptable because K8s will retry the event
+        // via dispatch catch-up after the new pod starts.
+        if (this.inFlightDispatches > 0) {
+            this.opts.log?.info(`parall[${this.opts.accountId}]: draining ${this.inFlightDispatches} in-flight dispatch(es), deadline ${this.SHUTDOWN_DEADLINE_MS}ms`);
+            await this.waitForDrain(this.SHUTDOWN_DEADLINE_MS);
+            if (this.inFlightDispatches > 0) {
+                this.opts.log?.warn(`parall[${this.opts.accountId}]: drain deadline hit; ${this.inFlightDispatches} dispatch(es) still running — they will be killed by process exit`);
+            }
+            else {
+                this.opts.log?.info(`parall[${this.opts.accountId}]: drain complete`);
+            }
+        }
         if (this.heartbeatTimer)
             clearInterval(this.heartbeatTimer);
         for (const [, dispatch] of this.activeDispatches) {

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@parall/agent-core",
-  "version": "1.18.1",
+  "version": "1.19.0",
   "description": "Shared agent runtime orchestration helpers for Parall",
   "license": "MIT",
   "repository": {
@@ -22,7 +22,7 @@
     "src"
   ],
   "dependencies": {
-    "@parall/sdk": "1.18.0"
+    "@parall/sdk": "1.19.0"
   },
   "devDependencies": {
     "@types/node": "^22.0.0",

package/src/gateway-base.ts

@@ -92,12 +92,31 @@ export type ParallGatewayOptions = {
   dispatchAdapter: DispatchAdapter;
   log?: GatewayLogger;
   coldStartWindowMs?: number;
+  // Maximum time to wait for in-flight dispatches to finish after SIGTERM /
+  // abort before forcing WS disconnect. Pod termination grace period should
+  // be at least this + a few seconds for the remaining cleanup work.
+  shutdownDeadlineMs?: number;
   stepIdFilePathForSession?: (sessionKey: string) => string | undefined;
   onConfigUpdate?: (data: AgentConfigUpdateData) => Promise<void> | void;
   onSessionReady?: (state: { activeSessionId?: string; ws: ParallWs; runtimeKey: string }) => Promise<void> | void;
   onBeforeDisconnect?: () => Promise<void> | void;
 };
 
+// Parse PRLL_SHUTDOWN_DEADLINE_MS (or any string env value) into a positive
+// integer milliseconds value, or undefined if unset/invalid. Runtimes pass
+// the result into ParallGatewayOptions.shutdownDeadlineMs; leaving it
+// undefined falls back to the gateway default (60s). Operators can raise it
+// when long-running tools need more time to drain — the K8s
+// terminationGracePeriodSeconds (currently 90s) must remain larger than this
+// plus the post-drain cleanup window (~15s), or the kubelet will SIGKILL
+// before cleanup finishes.
+export function parseShutdownDeadlineMs(raw: string | undefined): number | undefined {
+  if (!raw) return undefined;
+  const n = Number(raw);
+  if (!Number.isFinite(n) || n <= 0) return undefined;
+  return Math.floor(n);
+}
+
 function resolveStepTarget(event: ParallEvent): { target_type: string; target_id?: string } {
   if (event.type === "task" || event.targetId.startsWith("tsk_")) {
     return { target_type: "task", target_id: event.targetId };
@@ -150,11 +169,24 @@ export class ParallAgentGateway {
   private lastHeartbeatAt = Date.now();
   private draining = false;
 
+  // Graceful shutdown state. When SIGTERM / abort fires, `shuttingDown` flips
+  // to true so no new dispatches start, and `inFlightDispatches` counts runs
+  // still in progress. `shutdown()` awaits drain up to SHUTDOWN_DEADLINE_MS
+  // before tearing down the WS; see handleTermination caller.
+  private shuttingDown = false;
+  private inFlightDispatches = 0;
+  private drainResolvers: Array<() => void> = [];
+
   private readonly DISPATCHED_MESSAGES_CAP = 5000;
   private readonly COLD_START_WINDOW_MS: number;
+  // SHUTDOWN_DEADLINE_MS is read by waitForDrain via the configured value
+  // below — kept as instance state so per-runtime configs can override it
+  // (see parseShutdownDeadlineMs and runtime entrypoints).
+  private readonly SHUTDOWN_DEADLINE_MS: number;
 
   constructor(private readonly opts: ParallGatewayOptions) {
     this.COLD_START_WINDOW_MS = opts.coldStartWindowMs ?? 5 * 60_000;
+    this.SHUTDOWN_DEADLINE_MS = opts.shutdownDeadlineMs ?? 60_000;
   }
 
   async run(abortSignal: AbortSignal): Promise<void> {
@@ -435,12 +467,26 @@ export class ParallAgentGateway {
     }
   }
 
+  // Returns true if the dispatch actually ran; false if skipped because we
+  // are shutting down. Callers MUST treat `false` as "not dispatched" and
+  // skip the server-side ack so the event stays in the dispatch queue for
+  // catch-up on the replacement pod — otherwise we silently drop work.
   private async runDispatch(
     event: ParallEvent,
     sessionKey: string,
     bodyForAgent: string,
     earlierEvents: ParallEvent[] = [],
-  ) {
+  ): Promise<boolean> {
+    // Graceful shutdown: once SIGTERM has fired we stop accepting new work.
+    // In-flight dispatches that started before the flag flipped keep running
+    // and are awaited by shutdown() up to SHUTDOWN_DEADLINE_MS.
+    if (this.shuttingDown) {
+      this.opts.log?.info(
+        `parall[${this.opts.accountId}]: skipping dispatch for ${event.messageId} (shutting down) — leaving unacked for catch-up on replacement pod`,
+      );
+      return false;
+    }
+
     setSessionChatId(sessionKey, event.targetId);
     setSessionMessageId(sessionKey, event.messageId);
     setDispatchMessageId(sessionKey, event.messageId);
@@ -449,6 +495,10 @@ export class ParallAgentGateway {
     const dispatchContext = this.buildDispatchContext(event, sessionKey);
     const stepIdFilePath = dispatchContext.stepIdFilePath;
 
+    // sync: no await between the shuttingDown check above and this increment
+    // — JS event loop is single-threaded, so shutdown() cannot interleave
+    // here and miss our in-flight count.
+    this.inFlightDispatches++;
     try {
       if (this.activeSessionId) {
         try {
@@ -458,6 +508,13 @@ export class ParallAgentGateway {
         }
       }
 
+      // Persist input steps for "earlier events" (batched events that arrived
+      // while a dispatch was in flight) inside the in-flight window so a
+      // shutdown short-circuit BEFORE this point cannot leave orphan input
+      // steps that the replacement pod would duplicate on replay.
+      if (earlierEvents.length > 0) {
+        await this.createInputStepsForEarlierEvents(earlierEvents);
+      }
       await this.createInputStep(event);
 
       for await (const runtimeEvent of this.opts.dispatchAdapter.dispatch({
@@ -489,21 +546,42 @@ export class ParallAgentGateway {
           this.opts.log?.warn(`parall[${this.opts.accountId}]: failed to set session idle: ${String(err)}`);
         }
       }
+      this.inFlightDispatches--;
+      if (this.inFlightDispatches === 0 && this.drainResolvers.length > 0) {
+        const resolvers = this.drainResolvers.splice(0);
+        for (const resolve of resolvers) resolve();
+      }
     }
+    return true;
   }
 
   private async runForkDrainLoop(fork: ActiveForkState) {
     try {
       while (fork.queue.length > 0) {
+        // Early shutdown guard — pure optimization to avoid wasted draining
+        // during the shutdown window. The structural correctness now lives in
+        // runDispatch: input AgentStep persistence happens behind its
+        // shuttingDown gate and inFlightDispatches counter, so a shutdown
+        // landing mid-batch cannot leave orphan steps that replay would
+        // duplicate.
+        if (this.shuttingDown) {
+          for (const item of fork.queue.splice(0)) item.resolve(false);
+          break;
+        }
         const items = fork.queue.splice(0);
         const events = items.map((item) => item.event);
         const last = events[events.length - 1];
         const earlier = events.slice(0, -1);
         try {
-          if (earlier.length > 0) {
-            await this.createInputStepsForEarlierEvents(earlier);
+          const dispatched = await this.runDispatch(last, fork.fork.sessionKey, buildEventBody(last), earlier);
+          if (!dispatched) {
+            // Shutdown short-circuit — resolve un-acked so the server requeues
+            // for the replacement pod and stop draining further items.
+            for (const item of items) {
+              item.resolve(false);
+            }
+            break;
           }
-          await this.runDispatch(last, fork.fork.sessionKey, buildEventBody(last), earlier);
           fork.processedEvents.push(...events);
           for (const item of items) {
             item.resolve(true);
@@ -564,8 +642,19 @@ export class ParallAgentGateway {
     this.draining = true;
     try {
       while (this.dispatchState.mainBuffer.length > 0 || this.dispatchState.pendingForkResults.length > 0) {
+        // Early shutdown guard — same reasoning as runForkDrainLoop: stop
+        // before any input AgentSteps are persisted, leaving the buffer +
+        // pendingForkResults intact so dispatch catch-up on the replacement
+        // pod is the single source of truth for replay.
+        if (this.shuttingDown) {
+          this.opts.log?.info(
+            `parall[${this.opts.accountId}]: drainMainBuffer halted (shutting down) — ${this.dispatchState.mainBuffer.length} buffered, ${this.dispatchState.pendingForkResults.length} pending fork results left for catch-up`,
+          );
+          break;
+        }
         if (this.dispatchState.mainBuffer.length === 0 && this.dispatchState.pendingForkResults.length > 0) {
-          const forkPrefix = buildForkResultPrefix(this.dispatchState.pendingForkResults.splice(0));
+          const pending = this.dispatchState.pendingForkResults.splice(0);
+          const forkPrefix = buildForkResultPrefix(pending);
           const syntheticEvent: ParallEvent = {
             type: "message",
             targetId: "_orchestrator",
@@ -576,7 +665,14 @@ export class ParallAgentGateway {
             body: "[Orchestrator: fork session(s) completed — review results above]",
           };
           this.dispatchState.mainCurrentTargetId = undefined;
-          await this.runDispatch(syntheticEvent, this.opts.runtimeKey, forkPrefix + buildEventBody(syntheticEvent));
+          const dispatched = await this.runDispatch(syntheticEvent, this.opts.runtimeKey, forkPrefix + buildEventBody(syntheticEvent));
+          if (!dispatched) {
+            // Shutdown: put the fork results back at the head so the synthetic
+            // event is regenerated on the replacement pod. Synthetic events are
+            // not ack'd, so this is the only state we need to preserve.
+            this.dispatchState.pendingForkResults.unshift(...pending);
+            break;
+          }
           continue;
         }
 
@@ -588,12 +684,20 @@ export class ParallAgentGateway {
 
         const event = events[events.length - 1];
         const earlier = events.slice(0, -1);
-        const forkPrefix = buildForkResultPrefix(this.dispatchState.pendingForkResults.splice(0));
+        const pendingFork = this.dispatchState.pendingForkResults.splice(0);
+        const forkPrefix = buildForkResultPrefix(pendingFork);
         this.dispatchState.mainCurrentTargetId = event.targetId;
-        if (earlier.length > 0) {
-          await this.createInputStepsForEarlierEvents(earlier);
+        // Earlier-event input steps are persisted inside runDispatch (behind
+        // its shutdown gate) so a shutdown short-circuit cannot leave orphans.
+        const dispatched = await this.runDispatch(event, this.opts.runtimeKey, forkPrefix + buildEventBody(event), earlier);
+        if (!dispatched) {
+          // Shutdown: skip the ack so the server redelivers these buffered
+          // events to the replacement pod via dispatch catch-up. Put both the
+          // buffered events and the fork results back so nothing is lost.
+          this.dispatchState.mainBuffer.unshift(...events);
+          this.dispatchState.pendingForkResults.unshift(...pendingFork);
+          break;
         }
-        await this.runDispatch(event, this.opts.runtimeKey, forkPrefix + buildEventBody(event), earlier);
         for (const bufferedEvent of events) {
           const sourceType = bufferedEvent.ackSourceType ?? (bufferedEvent.type === "task" ? "task_activity" : "message");
           const sourceId = bufferedEvent.ackSourceId ?? bufferedEvent.messageId;
@@ -615,15 +719,22 @@ export class ParallAgentGateway {
 
     switch (disposition.action) {
       case "main": {
-        const forkPrefix = buildForkResultPrefix(this.dispatchState.pendingForkResults.splice(0));
+        const pendingFork = this.dispatchState.pendingForkResults.splice(0);
+        const forkPrefix = buildForkResultPrefix(pendingFork);
         this.dispatchState.mainDispatching = true;
         this.dispatchState.mainCurrentTargetId = event.targetId;
+        let dispatched = false;
         try {
-          await this.runDispatch(event, this.opts.runtimeKey, forkPrefix + buildEventBody(event));
+          dispatched = await this.runDispatch(event, this.opts.runtimeKey, forkPrefix + buildEventBody(event));
+          if (!dispatched) {
+            // Shutdown short-circuit — restore the fork results so a future
+            // pod can replay them, and return false so handleMessage skips ack.
+            this.dispatchState.pendingForkResults.unshift(...pendingFork);
+          }
         } finally {
           await this.drainMainBuffer();
         }
-        return true;
+        return dispatched;
       }
 
       case "buffer-main":
@@ -749,6 +860,7 @@ export class ParallAgentGateway {
   }
 
   private async handleMessage(data: MessageNewData) {
+    if (this.shuttingDown) return; // drain window — let server requeue via catch-up
     const chatId = data.chat_id;
     if (data.sender_id === this.opts.agentUserId) return;
     if (data.message_type !== "text") return;
@@ -779,6 +891,7 @@ export class ParallAgentGateway {
   }
 
   private async handleTaskAssignment(task: Task, ackSourceId?: string): Promise<boolean> {
+    if (this.shuttingDown) return false; // drain window — let server requeue via catch-up
     const dedupeKey = `${task.id}:${task.updated_at}`;
     if (this.dispatchedTasks.has(dedupeKey)) {
       this.opts.log?.info(`parall[${this.opts.accountId}]: skipping already-dispatched task ${task.identifier ?? task.id}`);
@@ -814,6 +927,7 @@ export class ParallAgentGateway {
   }
 
   private async handleTaskComment(commentId: string, taskId: string, actorId: string | null): Promise<boolean> {
+    if (this.shuttingDown) return false; // drain window — let server requeue via catch-up
     const dedupeKey = `comment:${commentId}`;
     if (this.dispatchedTasks.has(dedupeKey)) return false;
     this.dispatchedTasks.add(dedupeKey);
@@ -895,6 +1009,9 @@ export class ParallAgentGateway {
       });
 
       for (const item of page.data ?? []) {
+        // Shutdown short-circuit: stop fetching/building work that runDispatch
+        // will only reject. Items remain unacked for the replacement pod.
+        if (this.shuttingDown) break;
         if (minAge > 0 && new Date(item.created_at).getTime() < minAge) {
           this.opts.client.ackDispatchByID(this.opts.config.org_id, item.id).catch(() => {});
           skippedOld++;
@@ -976,7 +1093,9 @@ export class ParallAgentGateway {
           this.opts.log?.warn(`parall[${this.opts.accountId}]: catch-up dispatch ${item.id} (${item.event_type}) failed: ${String(err)}`);
         }
       }
-      cursor = page.has_more ? page.next_cursor : undefined;
+      // Stop paginating once shutdown begins — the inner loop already broke,
+      // and the next page would be wasted API work the replacement pod redoes.
+      cursor = !this.shuttingDown && page.has_more ? page.next_cursor : undefined;
     } while (cursor);
 
     if (processed > 0 || skippedOld > 0) {
@@ -1042,7 +1161,49 @@ export class ParallAgentGateway {
     }
   }
 
+  // Resolves when in-flight dispatches hit 0 or the deadline elapses.
+  // Deadline arming the resolver rather than rejecting keeps shutdown() linear
+  // — callers don't need to try/catch.
+  private waitForDrain(deadlineMs: number): Promise<void> {
+    if (this.inFlightDispatches === 0) return Promise.resolve();
+    return new Promise<void>((resolve) => {
+      const timer = setTimeout(() => {
+        // Remove this resolver so a late-completing dispatch doesn't call it.
+        const idx = this.drainResolvers.indexOf(onDrain);
+        if (idx >= 0) this.drainResolvers.splice(idx, 1);
+        resolve();
+      }, deadlineMs);
+      const onDrain = () => {
+        clearTimeout(timer);
+        resolve();
+      };
+      this.drainResolvers.push(onDrain);
+    });
+  }
+
   private async shutdown() {
+    // Flip the flag first so any WS event that fires while we're draining
+    // short-circuits at runDispatch instead of starting new work.
+    this.shuttingDown = true;
+
+    // Drain: wait for in-flight dispatches to finish up to the configured
+    // deadline. If we time out, whatever is still running will be killed by
+    // the process exit — that's acceptable because K8s will retry the event
+    // via dispatch catch-up after the new pod starts.
+    if (this.inFlightDispatches > 0) {
+      this.opts.log?.info(
+        `parall[${this.opts.accountId}]: draining ${this.inFlightDispatches} in-flight dispatch(es), deadline ${this.SHUTDOWN_DEADLINE_MS}ms`,
+      );
+      await this.waitForDrain(this.SHUTDOWN_DEADLINE_MS);
+      if (this.inFlightDispatches > 0) {
+        this.opts.log?.warn(
+          `parall[${this.opts.accountId}]: drain deadline hit; ${this.inFlightDispatches} dispatch(es) still running — they will be killed by process exit`,
+        );
+      } else {
+        this.opts.log?.info(`parall[${this.opts.accountId}]: drain complete`);
+      }
+    }
+
     if (this.heartbeatTimer) clearInterval(this.heartbeatTimer);
     for (const [, dispatch] of this.activeDispatches) {
       clearInterval(dispatch.typingTimer);