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',
@@ -1 +0,0 @@
-export const STREAM_INTERRUPTED_CODE = 'STREAM_INTERRUPTED';
@@ -99,9 +99,11 @@ describe('StreamAgentChatJob', () => {
    const publishedEvents: PublishedEvent[] = [];

    const threadRepository = {
-      findOne: jest
-        .fn()
-        .mockResolvedValue({ id: 'thread-id', deletedAt: null }),
+      findOne: jest.fn().mockResolvedValue({
+        id: 'thread-id',
+        deletedAt: null,
+        activeStreamId: 'stream-id',
+      }),
      update: jest.fn().mockImplementation((_workspaceId, _criteria, values) =>
        Promise.resolve({
          affected:
@@ -388,6 +390,119 @@ describe('StreamAgentChatJob', () => {
    );
  });

+  it('bails out without streaming when the thread no longer holds the claim for this stream', async () => {
+    const {
+      job,
+      publishedEvents,
+      threadRepository,
+      eventPublisherService,
+      agentChatStreamingService,
+    } = buildJob();
+
+    threadRepository.findOne.mockResolvedValueOnce({
+      id: 'thread-id',
+      deletedAt: null,
+      activeStreamId: 'newer-stream-id',
+    });
+
+    await job.handle(jobData);
+
+    expect(publishedEvents).toHaveLength(0);
+    expect(eventPublisherService.resetStreamState).not.toHaveBeenCalled();
+    expect(threadRepository.update).not.toHaveBeenCalled();
+    expect(
+      agentChatStreamingService.flushNextQueuedMessage,
+    ).not.toHaveBeenCalled();
+  });
+
+  it('bails out without streaming when the thread was deleted', async () => {
+    const { job, publishedEvents, threadRepository, eventPublisherService } =
+      buildJob();
+
+    threadRepository.findOne.mockResolvedValueOnce(null);
+
+    await job.handle(jobData);
+
+    expect(publishedEvents).toHaveLength(0);
+    expect(eventPublisherService.resetStreamState).not.toHaveBeenCalled();
+    expect(threadRepository.update).not.toHaveBeenCalled();
+  });
+
+  it('persists the interrupted error and publishes the terminal sequence when aborted by a worker shutdown', async () => {
+    let triggerShutdown: (() => void) | undefined;
+
+    const {
+      job,
+      publishedEvents,
+      threadRepository,
+      agentChatStreamingService,
+    } = buildJob({
+      chatStream: createFakeChatStream({
+        onFirstChunk: () => triggerShutdown?.(),
+      }),
+    });
+
+    const shutdownController = new AbortController();
+
+    triggerShutdown = () => shutdownController.abort();
+
+    await expect(
+      job.handle(jobData, { abortSignal: shutdownController.signal }),
+    ).rejects.toMatchObject({ code: AiExceptionCode.STREAM_INTERRUPTED });
+
+    const eventTypes = publishedEvents.map((event) => event.type);
+    const streamErrorIndex = eventTypes.indexOf('stream-error');
+    const queueUpdatedIndex = eventTypes.indexOf('queue-updated');
+
+    expect(publishedEvents[streamErrorIndex]).toMatchObject({
+      type: 'stream-error',
+      code: AiExceptionCode.STREAM_INTERRUPTED,
+    });
+    expect(queueUpdatedIndex).toBeGreaterThan(streamErrorIndex);
+    expect(eventTypes).not.toContain('message-persisted');
+    expect(threadRepository.update).toHaveBeenCalledWith(
+      'workspace-id',
+      { id: 'thread-id' },
+      {
+        lastStreamError: expect.objectContaining({
+          code: AiExceptionCode.STREAM_INTERRUPTED,
+        }),
+      },
+    );
+    expect(threadRepository.update).toHaveBeenCalledWith(
+      'workspace-id',
+      { id: 'thread-id', activeStreamId: 'stream-id' },
+      { activeStreamId: null },
+    );
+    expect(
+      agentChatStreamingService.flushNextQueuedMessage,
+    ).not.toHaveBeenCalled();
+  });
+
+  it('keeps user-cancel semantics when a shutdown signal is wired but never aborted', async () => {
+    let triggerUserCancel: (() => void) | undefined;
+
+    const { job, publishedEvents, agentChatStreamingService, cancelCallbacks } =
+      buildJob({
+        chatStream: createFakeChatStream({
+          onFirstChunk: () => triggerUserCancel?.(),
+        }),
+      });
+
+    triggerUserCancel = () => cancelCallbacks.forEach((callback) => callback());
+
+    const shutdownController = new AbortController();
+
+    await job.handle(jobData, { abortSignal: shutdownController.signal });
+
+    expect(publishedEvents.map((event) => event.type)).not.toContain(
+      'stream-error',
+    );
+    expect(
+      agentChatStreamingService.flushNextQueuedMessage,
+    ).not.toHaveBeenCalled();
+  });
+
  it('resolves without flushing the queue when the stream is cancelled', async () => {
    let triggerCancel: (() => void) | undefined;

@@ -12,6 +12,8 @@ import { isDefined } from 'twenty-shared/utils';
 import { Repository } from 'typeorm';
 import { v5 as uuidv5 } from 'uuid';

+import { type MessageQueueJobContext } from 'src/engine/core-modules/message-queue/interfaces/message-queue-job.interface';
+
 import { Process } from 'src/engine/core-modules/message-queue/decorators/process.decorator';
 import { Processor } from 'src/engine/core-modules/message-queue/decorators/processor.decorator';
 import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
@@ -68,7 +70,23 @@ export class StreamAgentChatJob {
  ) {}

  @Process(STREAM_AGENT_CHAT_JOB_NAME)
-  async handle(data: StreamAgentChatJobData): Promise<void> {
+  async handle(
+    data: StreamAgentChatJobData,
+    context?: MessageQueueJobContext,
+  ): Promise<void> {
+    const thread = await this.threadRepository.findOne(data.workspaceId, {
+      where: { id: data.threadId },
+      select: ['id', 'activeStreamId'],
+    });
+
+    if (thread?.activeStreamId !== data.streamId) {
+      this.logger.warn(
+        `Skipping stream ${data.streamId} for thread ${data.threadId}: the thread no longer holds this claim`,
+      );
+
+      return;
+    }
+
    await this.eventPublisherService.resetStreamState(data.threadId);

    const abortController = new AbortController();
@@ -82,6 +100,19 @@ export class StreamAgentChatJob {
      abortController.abort();
    });

+    context?.abortSignal?.addEventListener(
+      'abort',
+      () => {
+        abortController.abort(
+          new AiException(
+            'The response was interrupted before it could finish.',
+            AiExceptionCode.STREAM_INTERRUPTED,
+          ),
+        );
+      },
+      { once: true },
+    );
+
    try {
      const workspace = await this.workspaceRepository.findOne({
        where: { id: data.workspaceId },
@@ -268,7 +299,15 @@ export class StreamAgentChatJob {
      abortSignal.addEventListener(
        'abort',
        () => {
-          void streamFinishedPromise.then(() => resolve());
+          const reason = abortSignal.reason;
+
+          void streamFinishedPromise.then(() => {
+            if (reason instanceof AiException) {
+              reject(reason);
+            } else {
+              resolve();
+            }
+          });
        },
        { once: true },
      );
@@ -100,29 +100,23 @@ export class AgentChatSubscriptionResolver {
    });
  }

-  // Reaping from the keep-alive loop means a user watching a stream whose
-  // worker died sees the interrupted state without having to interact;
-  // reapDeadStream publishes the stream-error event they are subscribed to.
  private async reapWatchedStreamIfDead(
    workspaceId: string,
    threadId: string,
  ): Promise<void> {
-    try {
-      const thread = await this.threadRepository.findOne(workspaceId, {
+    const thread = await this.threadRepository
+      .findOne(workspaceId, {
        where: { id: threadId },
        select: ['id', 'activeStreamId'],
-      });
+      })
+      .catch(() => null);

-      if (!isDefined(thread) || !isDefined(thread.activeStreamId)) {
-        return;
-      }
-
-      await this.agentChatStreamingService.reapDeadStream({
-        thread,
-        workspaceId,
-      });
-    } catch {
-      // The keep-alive tick must never die with the check
+    if (!isDefined(thread) || !isDefined(thread.activeStreamId)) {
+      return;
    }
+
+    await this.agentChatStreamingService
+      .reapDeadStream({ thread, workspaceId })
+      .catch(() => {});
  }
 }
@@ -1,5 +1,5 @@
 import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
-import { STREAM_INTERRUPTED_CODE } from 'src/engine/metadata-modules/ai/ai-chat/constants/stream-interrupted-code.constant';
+import { AiExceptionCode } from 'src/engine/metadata-modules/ai/ai.exception';
 import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';

 describe('AgentChatStreamingService claim & reap', () => {
@@ -175,7 +175,7 @@ describe('AgentChatStreamingService claim & reap', () => {
      });

      expect(reaped).toEqual(
-        expect.objectContaining({ code: STREAM_INTERRUPTED_CODE }),
+        expect.objectContaining({ code: AiExceptionCode.STREAM_INTERRUPTED }),
      );
      expect(threadRepository.update).toHaveBeenCalledWith(
        'workspace-id',
@@ -183,7 +183,7 @@ describe('AgentChatStreamingService claim & reap', () => {
        expect.objectContaining({
          activeStreamId: null,
          lastStreamError: expect.objectContaining({
-            code: STREAM_INTERRUPTED_CODE,
+            code: AiExceptionCode.STREAM_INTERRUPTED,
          }),
        }),
      );
@@ -193,7 +193,7 @@ describe('AgentChatStreamingService claim & reap', () => {
      expect(publishedEvents).toContainEqual(
        expect.objectContaining({
          type: 'stream-error',
-          code: STREAM_INTERRUPTED_CODE,
+          code: AiExceptionCode.STREAM_INTERRUPTED,
        }),
      );
    });
@@ -16,7 +16,6 @@ import { InjectMessageQueue } from 'src/engine/core-modules/message-queue/decora
 import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
 import { MessageQueueService } from 'src/engine/core-modules/message-queue/services/message-queue.service';
 import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
-import { STREAM_INTERRUPTED_CODE } from 'src/engine/metadata-modules/ai/ai-chat/constants/stream-interrupted-code.constant';
 import {
  AgentMessageRole,
  AgentMessageStatus,
@@ -81,7 +80,7 @@ export class AgentChatStreamingService {
    }

    const interruptedError: AgentChatThreadLastStreamError = {
-      code: STREAM_INTERRUPTED_CODE,
+      code: AiExceptionCode.STREAM_INTERRUPTED,
      message: 'The response was interrupted before it could finish.',
      failedAt: new Date().toISOString(),
    };
@@ -22,6 +22,7 @@ export enum AiExceptionCode {
  ROLE_NOT_FOUND = 'ROLE_NOT_FOUND',
  ROLE_CANNOT_BE_ASSIGNED_TO_AGENTS = 'ROLE_CANNOT_BE_ASSIGNED_TO_AGENTS',
  NO_FAILED_TURN_TO_RETRY = 'NO_FAILED_TURN_TO_RETRY',
+  STREAM_INTERRUPTED = 'STREAM_INTERRUPTED',
 }

 const getAiExceptionUserFriendlyMessage = (code: AiExceptionCode) => {
@@ -60,6 +61,8 @@ const getAiExceptionUserFriendlyMessage = (code: AiExceptionCode) => {
      return msg`This role cannot be assigned to agents.`;
    case AiExceptionCode.NO_FAILED_TURN_TO_RETRY:
      return msg`There is no failed message to retry.`;
+    case AiExceptionCode.STREAM_INTERRUPTED:
+      return msg`The response was interrupted before it could finish.`;
    default:
      assertUnreachable(code);
  }
@@ -42,6 +42,7 @@ export const aiGraphqlApiExceptionHandler = (error: Error) => {
      case AiExceptionCode.AGENT_EXECUTION_FAILED:
      case AiExceptionCode.API_KEY_NOT_CONFIGURED:
      case AiExceptionCode.USER_WORKSPACE_ID_NOT_FOUND:
+      case AiExceptionCode.STREAM_INTERRUPTED:
        throw new InternalServerError(error);
      default: {
        return assertUnreachable(error.code);
				`@@ -1 +0,0 @@`
				`export const STREAM_INTERRUPTED_CODE = 'STREAM_INTERRUPTED';`