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
@@ -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';`