feat(server): abort ai stream jobs that outlive the shutdown drain budget (#22517)

## Why #22514 makes SIGTERM drain workers, but `worker.close()` waits for active jobs with **no upper bound** (BullMQ semantics). An AI stream job can run for 10 minutes; a deploy would either hang the rollout or hit the pod's termination grace deadline and get SIGKILLed anyway — back to the frozen-stream + stalled-rerun failure this series eliminates. ## What On shutdown, `aiStreamQueue` gets a bounded drain: active stream jobs have `AI_STREAM_SHUTDOWN_DRAIN_MS` (60s) to finish naturally; stragglers are then aborted and terminate exactly like a stream failure — `lastStreamError` persisted with the new typed `AiExceptionCode.STREAM_INTERRUPTED`, the pinned `stream-error` → `queue-updated` terminal sequence published, claim released. The client shows the interrupted state with Retry (#22434) within seconds instead of a stream frozen mid-sentence. This is deliberately **not** the user-cancel path, which resolves cleanly and persists no error. Mechanism — evaluated BullMQ 5.78's native cancellation vs a parallel in-process registry, and picked native: - The driver's processor now declares the 3-arg signature, which makes BullMQ create a per-job `AbortController` (`processor.length >= 3` is the trigger), and the signal is handed to job handlers as an optional `MessageQueueJobContext`. - `worker.cancelAllJobs()` is purely cooperative: it aborts the signal and nothing else, so the job's own persist/publish/cleanup still runs to completion and `worker.close()` still waits for it — no force-fail race, no second signaling channel to maintain, and the timer lives inside the same `closeWorker()` call so there is no dependence on Nest module-destroy ordering. - The stream job maps the shutdown signal onto its **existing** AbortController (the one already wired through the AI SDK for user cancel), with an `AiException(STREAM_INTERRUPTED)` reason to tell the two apart. One abort path end to end, no new infrastructure. Error-type choice: the job throws a plain `AiException`, not BullMQ's `UnrecoverableError`. Stream jobs are enqueued with `attempts: 1` (no `retryLimit`), so there is no BullMQ retry to suppress — retryability for this queue lives at the app layer (`lastStreamError` + client Retry), and an `UnrecoverableError` would only obscure the typed exception. `STREAM_INTERRUPTED` also replaces the string constant introduced on the base branch (#22482's reap now uses the same enum member) — one code, two producers (reap for dead workers, abort for live shutdowns), identical client behavior. Notes: - 60s is a static constant mirroring `AI_STREAM_LOCK_DURATION_MS` rather than an env var — it has to move in lockstep with the worker's `terminationGracePeriodSeconds` (120s, twenty-infra PR) anyway, and we ship multiple releases a day. Happy to lift it into a config variable if you want runtime tunability. - The `onModuleDestroy` scaffolding (drain logs, workers-close-before-queues) deliberately matches #22514; whichever lands second rebases clean. Stacked on #22482 (needs the heartbeat/reap base). Merge order: #22482 → this. Depends on #22514 for SIGTERM to reach the driver at all. ## Validation - `stream-agent-chat.job.spec.ts`: shutdown-abort persists `STREAM_INTERRUPTED`, publishes `stream-error` before `queue-updated`, releases the claim, skips the queued-message flush; user-cancel semantics unchanged with a wired-but-idle shutdown signal (60/60 ai-chat tests green). - Local end-to-end: real AI stream mid-flight, SIGTERM the worker → drain window → abort → interrupted state persisted, process exits on its own. (Transcript in the PR conversation.)  <a href="https://cubic.dev/pr/twentyhq/twenty/pull/22517?utm_source=github" target="_blank" rel="noopener noreferrer" data-no-image-dialog="true"><picture><source media="(prefers-color-scheme: dark)" srcset="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"><source media="(prefers-color-scheme: light)" srcset="https://www.cubic.dev/buttons/review-in-cubic-light.svg"><img alt="Review in cubic" src="https://www.cubic.dev/buttons/review-in-cubic-dark.svg"></picture></a>
2026-07-05 14:15:24 +02:00
parent 904957ea1e
commit 3a21089e0a
17 changed files with 261 additions and 35 deletions
@@ -30,6 +30,7 @@ import { getJobKey } from 'src/engine/core-modules/message-queue/utils/get-job-k
 import { type MetricsService } from 'src/engine/core-modules/metrics/metrics.service';
 import { MetricsKeys } from 'src/engine/core-modules/metrics/types/metrics-keys.type';
 import { applyWorkspaceSentryContextFromJobData } from 'src/engine/core-modules/sentry/utils/apply-workspace-sentry-context-from-job-data.util';
+import { type TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';

 export type BullMQDriverOptions = QueueOptions;

@@ -47,10 +48,14 @@ export class BullMQDriver
    MessageQueue,
    Worker
  >;
+  private workerOptionsMap: Partial<
+    Record<MessageQueue, MessageQueueWorkerOptions>
+  > = {};

  constructor(
    private options: BullMQDriverOptions,
    private metricsService: MetricsService,
+    private twentyConfigService: TwentyConfigService,
  ) {}

  onModuleInit() {
@@ -89,7 +94,7 @@ export class BullMQDriver
  }

  async onModuleDestroy() {
-    const workers = Object.entries(this.workerMap);
+    const workers = Object.entries(this.workerMap) as [MessageQueue, Worker][];
    const queues = Object.values(this.queueMap);

    if (workers.length > 0) {
@@ -101,7 +106,11 @@ export class BullMQDriver
    let workerCloseError: unknown;

    try {
-      await Promise.all(workers.map(([, worker]) => worker.close()));
+      await Promise.all(
+        workers.map(([queueName, worker]) =>
+          this.closeWorker(queueName, worker),
+        ),
+      );
    } catch (error) {
      workerCloseError = error;
    }
@@ -125,6 +134,34 @@ export class BullMQDriver
    this.logger.log('Message queue shutdown complete');
  }

+  private async closeWorker(
+    queueName: MessageQueue,
+    worker: Worker,
+  ): Promise<void> {
+    if (!this.workerOptionsMap[queueName]?.boundedShutdownDrain) {
+      await worker.close();
+
+      return;
+    }
+
+    const shutdownTimeoutMs = this.twentyConfigService.get(
+      'AI_STREAM_SHUTDOWN_DRAIN_MS',
+    );
+
+    const abortTimer = setTimeout(() => {
+      this.logger.warn(
+        `Queue ${queueName} still has active jobs after draining for ${shutdownTimeoutMs}ms, aborting them`,
+      );
+      worker.cancelAllJobs('worker shutdown');
+    }, shutdownTimeoutMs);
+
+    try {
+      await worker.close();
+    } finally {
+      clearTimeout(abortTimer);
+    }
+  }
+
  work<T>(
    queueName: MessageQueue,
    handler: (job: MessageQueueJob<T>) => Promise<void>,
@@ -138,15 +175,20 @@ export class BullMQDriver
      ...(isDefined(options?.lockDuration)
        ? { lockDuration: options.lockDuration }
        : {}),
+      ...(isDefined(options?.maxStalledCount)
+        ? { maxStalledCount: options.maxStalledCount }
+        : {}),
      metrics: {
        maxDataPoints: MetricsTime.ONE_WEEK,
        collectInterval: 60000,
      },
    };

+    this.workerOptionsMap[queueName] = options;
+
    this.workerMap[queueName] = new Worker(
      queueName,
-      async (job) =>
+      async (job, _token, abortSignal) =>
        Sentry.withIsolationScope(async () => {
          applyWorkspaceSentryContextFromJobData(job.data);

@@ -169,7 +211,12 @@ export class BullMQDriver
          this.logger.log(
            `Processing job ${job.id} with name ${job.name} on queue ${queueName}${workspaceSuffix}`,
          );
-          await handler({ data: job.data, id: job.id ?? '', name: job.name });
+          await handler({
+            data: job.data,
+            id: job.id ?? '',
+            name: job.name,
+            abortSignal,
+          });
          const timeEnd = performance.now();
          const executionTime = timeEnd - timeStart;

@@ -3,6 +3,11 @@ export interface MessageQueueJob<T = any> {
  id: string;
  name: string;
  data: T;
+  abortSignal?: AbortSignal;
+}
+
+export interface MessageQueueJobContext {
+  abortSignal?: AbortSignal;
 }

 export interface MessageQueueCronJobData<
@@ -1,5 +1,6 @@
 import { type BullMQDriverOptions } from 'src/engine/core-modules/message-queue/drivers/bullmq.driver';
 import { type MetricsService } from 'src/engine/core-modules/metrics/metrics.service';
+import { type TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';

 export enum MessageQueueDriverType {
  BullMQ = 'bull-mq',
@@ -10,6 +11,7 @@ export interface BullMQDriverFactoryOptions {
  type: MessageQueueDriverType.BullMQ;
  options: BullMQDriverOptions;
  metricsService: MetricsService;
+  twentyConfigService: TwentyConfigService;
 }

 export interface SyncDriverFactoryOptions {
@@ -1,4 +1,6 @@
 export interface MessageQueueWorkerOptions {
  concurrency?: number;
  lockDuration?: number;
+  maxStalledCount?: number;
+  boundedShutdownDrain?: boolean;
 }
@@ -96,7 +96,11 @@ export class MessageQueueCoreModule extends ConfigurableModuleClass {
  static async createDriver(config: typeof OPTIONS_TYPE) {
    switch (config.type) {
      case MessageQueueDriverType.BullMQ: {
-        return new BullMQDriver(config.options, config.metricsService);
+        return new BullMQDriver(
+          config.options,
+          config.metricsService,
+          config.twentyConfigService,
+        );
      }
      case MessageQueueDriverType.Sync: {
        return new SyncDriver();
@@ -8,6 +8,8 @@ export const QUEUE_WORKER_OPTIONS: Partial<
  [MessageQueue.aiStreamQueue]: {
    concurrency: 20,
    lockDuration: AI_STREAM_LOCK_DURATION_MS,
+    maxStalledCount: 0,
+    boundedShutdownDrain: true,
  },
  [MessageQueue.logicFunctionQueue]: { concurrency: 10 },
 };
@@ -208,7 +208,9 @@ export class MessageQueueExplorer implements OnModuleInit {
    for (const processMethodName of processMethodNames) {
      try {
        // @ts-expect-error legacy noImplicitAny
-        await instance[processMethodName].call(instance, job.data);
+        await instance[processMethodName].call(instance, job.data, {
+          abortSignal: job.abortSignal,
+        });
      } catch (err) {
        if (shouldCaptureException(err)) {
          this.exceptionHandlerService.captureExceptions([err]);
@@ -15,7 +15,7 @@ import { type TwentyConfigService } from 'src/engine/core-modules/twenty-config/
 * @param metricsService
 */
 export const messageQueueModuleFactory = async (
-  _twentyConfigService: TwentyConfigService,
+  twentyConfigService: TwentyConfigService,
  redisClientService: RedisClientService,
  metricsService: MetricsService,
 ): Promise<MessageQueueModuleOptions> => {
@@ -29,6 +29,7 @@ export const messageQueueModuleFactory = async (
          connection: redisClientService.getQueueClient(),
        },
        metricsService,
+        twentyConfigService,
      } satisfies BullMQDriverFactoryOptions;
    }
    default:
@@ -1280,6 +1280,17 @@ export class ConfigVariables {
  @IsOptional()
  SERVER_KEEP_ALIVE_TIMEOUT_MS = 65000;

+  @ConfigVariablesMetadata({
+    group: ConfigVariablesGroup.SERVER_CONFIG,
+    description:
+      'How long (ms) a worker shutdown waits for active AI chat stream jobs to finish before aborting them into a retryable interrupted state. Must be lower than the pod terminationGracePeriodSeconds so the abort and clean exit fit before SIGKILL (default: 300000)',
+    type: ConfigVariableType.NUMBER,
+    isEnvOnly: true,
+  })
+  @CastToPositiveNumber()
+  @IsOptional()
+  AI_STREAM_SHUTDOWN_DRAIN_MS = 300_000;
+
  @ConfigVariablesMetadata({
    group: ConfigVariablesGroup.SERVER_CONFIG,
    description: 'Base URL for the server',