feat(ai): add a stream heartbeat and reap dead claims so a worker crash cannot brick a thread (#22482)

## Rationale If the worker process dies mid-stream (OOM, deploy, crash), nothing ever clears `activeStreamId`: `aiStreamQueue` runs with `attempts: 1`, the job's `finally` never executes, and the SSE keepalive comes from the API server — so it actively masks worker death. The thread is bricked: every send queues behind a dead claim until someone intervenes manually. This is a CONFIRMED-high from the chat-stack audit, and worker death is not hypothetical: Sentry shows an unhandled promise rejection inside the AI SDK in the worker ([TWENTY-SERVER-H7Y](https://twenty-v7.sentry.io/issues/TWENTY-SERVER-H7Y)) — unhandled rejections terminate Node by default. ## Design - **Claim-time mark**: every enqueue site marks `agent-chat-stream-alive:<streamId>` with a TTL matching the job lock horizon (600s) — covering the enqueue→pickup window where a waiting job holds no lock. - **Running refresh**: the job tightens it to **30s, refreshed every 5s**; if the process dies, the interval dies with it and the key expires. The expiry *is* the death signal. (Was 60s/15s — tightened after review: detection latency is bounded by the TTL, robustness by TTL−interval and the missed-beat tolerance; 30s/5s halves detection while tolerating *more* missed beats, 5 vs 3.) - **Read-path reap**: the send gate and the catchup query convert a heartbeat-less claim into a normal retryable `STREAM_INTERRUPTED` failed-turn state (conditional UPDATE guarded on the observed streamId, so a newer stream's claim is never touched), reset the Redis chunk state, and publish the terminal error. `isAlive` fails open on Redis errors — a liveness probe must not turn a Redis blip into a broken send path. ## Why this is the root cause, not a symptom patch The strongest alternative — BullMQ's own stalled-job detection — fails on four concrete grounds: detection latency is bounded by the deliberate 10-minute `AI_STREAM_LOCK_DURATION_MS` (long silent tool runs must not spuriously stall); the stalled checker needs a *surviving* worker in the pool; the signal fires in the worker process while the thing needing repair is a DB claim read by API-server resolvers; and a `waiting` job holds no lock at all. Reaping at the read path means recovery happens exactly when a user is looking — the moment it matters — with zero background machinery. **Relationship to the graceful-shutdown work (planned follow-ups)**: shutdown hooks + drain-then-abort will make *deploys* (cooperative SIGTERM) end streams cleanly, and disabling stalled re-runs will stop hard-killed jobs from zombie re-executing tools. This PR remains the only recovery layer for non-cooperative deaths — OOMKill is a straight SIGKILL, crashes and unhandled rejections never run shutdown hooks — and the backstop when the drain path itself fails. The two are complements, not alternatives. ## User impact Today a worker crash mid-answer bricks the thread until manual intervention; users see sends silently queue forever. With this, the next interaction (send, reload) converts it into a visible "response was interrupted" error with a working Retry, within ~30s of actual death. ## Test plan - [x] Claim spec: live stream untouched; heartbeat-less claim reaped into retryable `STREAM_INTERRUPTED` + chunk-state reset + published terminal event; no-op when the claim moved to a newer stream mid-check - [x] CI green https://claude.ai/code/session_01Lyi6zTema2FMVVh8MD6c38
2026-07-03 13:29:02 +02:00
parent 203da78b05
commit 5d3b6d05b3
11 changed files with 282 additions and 2 deletions
@@ -33,6 +33,7 @@ import { AgentChatResolver } from './resolvers/agent-chat.resolver';
 import { AgentChatSubscriptionResolver } from './resolvers/agent-chat-subscription.resolver';
 import { AgentChatCancelSubscriberService } from './services/agent-chat-cancel-subscriber.service';
 import { AgentChatEventPublisherService } from './services/agent-chat-event-publisher.service';
+import { AgentChatStreamHeartbeatService } from './services/agent-chat-stream-heartbeat.service';
 import { AgentChatStreamingService } from './services/agent-chat-streaming.service';
 import { AgentChatService } from './services/agent-chat.service';
 import { AgentTitleGenerationService } from './services/agent-title-generation.service';
@@ -69,6 +70,7 @@ import { SystemPromptBuilderService } from './services/system-prompt-builder.ser
  providers: [
    AgentChatCancelSubscriberService,
    AgentChatEventPublisherService,
+    AgentChatStreamHeartbeatService,
    AgentChatResolver,
    AgentChatSubscriptionResolver,
    AgentChatService,
@@ -0,0 +1 @@
+export const AGENT_CHAT_STREAM_REAP_CHECK_INTERVAL_MS = 10_000;
@@ -0,0 +1 @@
+export const STREAM_INTERRUPTED_CODE = 'STREAM_INTERRUPTED';
@@ -144,6 +144,11 @@ describe('StreamAgentChatJob', () => {
    const agentChatStreamingService = {
      flushNextQueuedMessage: jest.fn().mockResolvedValue(undefined),
    };
+    const streamHeartbeatService = {
+      startRunning: jest.fn().mockReturnValue(() => {}),
+      markClaimed: jest.fn().mockResolvedValue(undefined),
+      clear: jest.fn().mockResolvedValue(undefined),
+    };
    const job = new StreamAgentChatJob(
      threadRepository as never,
      workspaceRepository as never,
@@ -152,6 +157,7 @@ describe('StreamAgentChatJob', () => {
      eventPublisherService as never,
      cancelSubscriberService as never,
      agentChatStreamingService as never,
+      streamHeartbeatService as never,
    );

    return {
@@ -28,6 +28,7 @@ import {
 import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
 import { AgentChatCancelSubscriberService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-cancel-subscriber.service';
 import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service';
+import { AgentChatStreamHeartbeatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-stream-heartbeat.service';
 import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';
 import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service';
 import { ChatExecutionService } from 'src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service';
@@ -62,6 +63,7 @@ export class StreamAgentChatJob {
    private readonly eventPublisherService: AgentChatEventPublisherService,
    private readonly cancelSubscriberService: AgentChatCancelSubscriberService,
    private readonly agentChatStreamingService: AgentChatStreamingService,
+    private readonly streamHeartbeatService: AgentChatStreamHeartbeatService,
  ) {}

  @Process(STREAM_AGENT_CHAT_JOB_NAME)
@@ -71,6 +73,10 @@ export class StreamAgentChatJob {
    const abortController = new AbortController();
    const cancelChannel = getCancelChannel(data.threadId, data.streamId);

+    const stopHeartbeat = this.streamHeartbeatService.startRunning(
+      data.streamId,
+    );
+
    await this.cancelSubscriberService.subscribe(cancelChannel, () => {
      abortController.abort();
    });
@@ -132,6 +138,8 @@ export class StreamAgentChatJob {
        .catch(() => {});
      throw error;
    } finally {
+      stopHeartbeat();
+      await this.streamHeartbeatService.clear(data.streamId);
      await this.cancelSubscriberService.unsubscribe(cancelChannel);
      await this.threadRepository
        .update(
@@ -18,8 +18,10 @@ import {
 } from 'src/engine/metadata-modules/ai/ai.exception';
 import { AiGraphqlApiExceptionInterceptor } from 'src/engine/metadata-modules/ai/interceptors/ai-graphql-api-exception.interceptor';
 import { AGENT_CHAT_KEEPALIVE_INTERVAL_MS } from 'src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-keepalive-interval-ms.constant';
+import { AGENT_CHAT_STREAM_REAP_CHECK_INTERVAL_MS } from 'src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-stream-reap-check-interval-ms.constant';
 import { AgentChatEventDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/agent-chat-event.dto';
 import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
+import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';
 import { SubscriptionService } from 'src/engine/subscriptions/subscription.service';
 import { wrapAsyncIteratorWithLifecycle } from 'src/engine/subscriptions/utils/wrap-async-iterator-with-lifecycle';
 import { InjectWorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/inject-workspace-scoped-repository.decorator';
@@ -30,6 +32,7 @@ import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scope
 export class AgentChatSubscriptionResolver {
  constructor(
    private readonly subscriptionService: SubscriptionService,
+    private readonly agentChatStreamingService: AgentChatStreamingService,
    @InjectWorkspaceScopedRepository(AgentChatThreadEntity)
    private readonly threadRepository: WorkspaceScopedRepository<AgentChatThreadEntity>,
  ) {}
@@ -72,9 +75,19 @@ export class AgentChatSubscriptionResolver {
      },
    };

+    let lastReapCheckAt = 0;
+
    return wrapAsyncIteratorWithLifecycle(iterator, {
      initialValue: keepalivePayload,
      onHeartbeat: async () => {
+        if (
+          Date.now() - lastReapCheckAt >=
+          AGENT_CHAT_STREAM_REAP_CHECK_INTERVAL_MS
+        ) {
+          lastReapCheckAt = Date.now();
+          await this.reapWatchedStreamIfDead(workspace.id, threadId);
+        }
+
        await this.subscriptionService.publishToAgentChat({
          workspaceId: workspace.id,
          threadId,
@@ -86,4 +99,30 @@ export class AgentChatSubscriptionResolver {
      heartbeatIntervalMs: AGENT_CHAT_KEEPALIVE_INTERVAL_MS,
    });
  }
+
+  // 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, {
+        where: { id: threadId },
+        select: ['id', 'activeStreamId'],
+      });
+
+      if (!isDefined(thread) || !isDefined(thread.activeStreamId)) {
+        return;
+      }
+
+      await this.agentChatStreamingService.reapDeadStream({
+        thread,
+        workspaceId,
+      });
+    } catch {
+      // The keep-alive tick must never die with the check
+    }
+  }
 }
@@ -110,6 +110,17 @@ export class AgentChatResolver {
      workspaceId,
    });

+    const interruptedError =
+      await this.agentChatStreamingService.reapDeadStream({
+        thread,
+        workspaceId,
+      });
+
+    if (interruptedError) {
+      thread.activeStreamId = null;
+      thread.lastStreamError = interruptedError;
+    }
+
    const { chunks, maxSeq } =
      await this.eventPublisherService.getAccumulatedChunks(threadId);

@@ -188,6 +199,19 @@ export class AgentChatResolver {
      });
    }

+    if (isDefined(thread.activeStreamId)) {
+      const interruptedError =
+        await this.agentChatStreamingService.reapDeadStream({
+          thread,
+          workspaceId: workspace.id,
+        });
+
+      if (interruptedError) {
+        thread.activeStreamId = null;
+        thread.lastStreamError = interruptedError;
+      }
+    }
+
    if (
      isDefined(thread.activeStreamId) ||
      isDefined(thread.pendingQuestionMessageId)
@@ -1,4 +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 { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';

 describe('AgentChatStreamingService claim & reap', () => {
@@ -17,6 +18,7 @@ describe('AgentChatStreamingService claim & reap', () => {
    thread = idleThread,
    claimAffected = 1,
    queuedMessages = [] as unknown[],
+    heartbeatAlive = true,
  } = {}) => {
    const publishedEvents: Array<{ type: string }> = [];
    const threadRepository = {
@@ -47,6 +49,12 @@ describe('AgentChatStreamingService claim & reap', () => {
      }),
      resetStreamState: jest.fn().mockResolvedValue(undefined),
    };
+    const streamHeartbeatService = {
+      markClaimed: jest.fn().mockResolvedValue(undefined),
+      isAlive: jest.fn().mockResolvedValue(heartbeatAlive),
+      clear: jest.fn().mockResolvedValue(undefined),
+    };
+
    const service = new AgentChatStreamingService(
      threadRepository as never,
      { find: jest.fn().mockResolvedValue([]) } as never,
@@ -54,6 +62,7 @@ describe('AgentChatStreamingService claim & reap', () => {
      agentChatService as never,
      eventPublisherService as never,
      { signFileByIdUrl: jest.fn() } as never,
+      streamHeartbeatService as never,
    );

    return {
@@ -62,6 +71,7 @@ describe('AgentChatStreamingService claim & reap', () => {
      messageQueueService,
      agentChatService,
      eventPublisherService,
+      streamHeartbeatService,
      publishedEvents,
    };
  };
@@ -76,7 +86,8 @@ describe('AgentChatStreamingService claim & reap', () => {

  describe('streamAgentChat', () => {
    it('claims the thread conditionally before enqueueing', async () => {
-      const { service, threadRepository } = buildService();
+      const { service, threadRepository, streamHeartbeatService } =
+        buildService();

      const result = await service.streamAgentChat(sendArguments);

@@ -86,6 +97,7 @@ describe('AgentChatStreamingService claim & reap', () => {
        expect.objectContaining({ id: 'thread-id' }),
        expect.objectContaining({ lastStreamError: null }),
      );
+      expect(streamHeartbeatService.markClaimed).toHaveBeenCalled();
    });

    it('queues the message when another stream wins the claim race', async () => {
@@ -135,4 +147,71 @@ describe('AgentChatStreamingService claim & reap', () => {
      );
    });
  });
+
+  describe('reapDeadStream', () => {
+    it('leaves a live stream alone', async () => {
+      const { service, threadRepository } = buildService();
+
+      const reaped = await service.reapDeadStream({
+        thread: { id: 'thread-id', activeStreamId: 'stream-id' },
+        workspaceId: 'workspace-id',
+      });
+
+      expect(reaped).toBeNull();
+      expect(threadRepository.update).not.toHaveBeenCalled();
+    });
+
+    it('converts a heartbeat-less claim into a retryable interrupted error', async () => {
+      const {
+        service,
+        threadRepository,
+        eventPublisherService,
+        publishedEvents,
+      } = buildService({ heartbeatAlive: false });
+
+      const reaped = await service.reapDeadStream({
+        thread: { id: 'thread-id', activeStreamId: 'stream-id' },
+        workspaceId: 'workspace-id',
+      });
+
+      expect(reaped).toEqual(
+        expect.objectContaining({ code: STREAM_INTERRUPTED_CODE }),
+      );
+      expect(threadRepository.update).toHaveBeenCalledWith(
+        'workspace-id',
+        { id: 'thread-id', activeStreamId: 'stream-id' },
+        expect.objectContaining({
+          activeStreamId: null,
+          lastStreamError: expect.objectContaining({
+            code: STREAM_INTERRUPTED_CODE,
+          }),
+        }),
+      );
+      expect(eventPublisherService.resetStreamState).toHaveBeenCalledWith(
+        'thread-id',
+      );
+      expect(publishedEvents).toContainEqual(
+        expect.objectContaining({
+          type: 'stream-error',
+          code: STREAM_INTERRUPTED_CODE,
+        }),
+      );
+    });
+
+    it('does nothing when the claim moved to a newer stream mid-check', async () => {
+      const { service, publishedEvents, threadRepository } = buildService({
+        heartbeatAlive: false,
+        claimAffected: 0,
+      });
+
+      const reaped = await service.reapDeadStream({
+        thread: { id: 'thread-id', activeStreamId: 'stream-id' },
+        workspaceId: 'workspace-id',
+      });
+
+      expect(reaped).toBeNull();
+      expect(threadRepository.update).toHaveBeenCalledTimes(1);
+      expect(publishedEvents).toHaveLength(0);
+    });
+  });
 });
@@ -47,6 +47,12 @@ describe('AgentChatStreamingService.retryLastFailedTurn', () => {
      getMessagesForThread: jest.fn().mockResolvedValue(threadMessages),
    };

+    const streamHeartbeatService = {
+      markClaimed: jest.fn().mockResolvedValue(undefined),
+      isAlive: jest.fn().mockResolvedValue(true),
+      clear: jest.fn().mockResolvedValue(undefined),
+    };
+
    const service = new AgentChatStreamingService(
      threadRepository as never,
      { find: jest.fn() } as never,
@@ -54,6 +60,7 @@ describe('AgentChatStreamingService.retryLastFailedTurn', () => {
      agentChatService as never,
      { publish: jest.fn() } as never,
      { signFileByIdUrl: jest.fn() } as never,
+      streamHeartbeatService as never,
    );

    return { service, threadRepository, messageQueueService, agentChatService };
@@ -0,0 +1,55 @@
+import { Injectable } from '@nestjs/common';
+
+import { RedisClientService } from 'src/engine/core-modules/redis-client/redis-client.service';
+
+const CLAIM_TTL_SECONDS = 600;
+const RUNNING_TTL_SECONDS = 30;
+const REFRESH_INTERVAL_MS = 5_000;
+
+@Injectable()
+export class AgentChatStreamHeartbeatService {
+  constructor(private readonly redisClientService: RedisClientService) {}
+
+  private getKey(streamId: string): string {
+    return `agent-chat-stream-alive:${streamId}`;
+  }
+
+  async markClaimed(streamId: string): Promise<void> {
+    await this.redisClientService
+      .getClient()
+      .set(this.getKey(streamId), '1', 'EX', CLAIM_TTL_SECONDS);
+  }
+
+  startRunning(streamId: string): () => void {
+    const refresh = () => {
+      this.redisClientService
+        .getClient()
+        .set(this.getKey(streamId), '1', 'EX', RUNNING_TTL_SECONDS)
+        .catch(() => {});
+    };
+
+    refresh();
+    const interval = setInterval(refresh, REFRESH_INTERVAL_MS);
+
+    return () => clearInterval(interval);
+  }
+
+  async isAlive(streamId: string): Promise<boolean> {
+    try {
+      const exists = await this.redisClientService
+        .getClient()
+        .exists(this.getKey(streamId));
+
+      return exists === 1;
+    } catch {
+      return true;
+    }
+  }
+
+  async clear(streamId: string): Promise<void> {
+    await this.redisClientService
+      .getClient()
+      .del(this.getKey(streamId))
+      .catch(() => {});
+  }
+}
@@ -16,6 +16,7 @@ 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,
@@ -27,6 +28,7 @@ import { type AgentChatThreadLastStreamError } from 'src/engine/metadata-modules
 import { STREAM_AGENT_CHAT_JOB_NAME } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job-name.constant';
 import { type StreamAgentChatJobData } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job.types';
 import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service';
+import { AgentChatStreamHeartbeatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-stream-heartbeat.service';
 import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service';
 import { AiChatFileAttachment } from 'src/engine/metadata-modules/ai/ai-chat/types/ai-chat-file-attachment.type';
 import {
@@ -60,8 +62,56 @@ export class AgentChatStreamingService {
    private readonly agentChatService: AgentChatService,
    private readonly eventPublisherService: AgentChatEventPublisherService,
    private readonly fileUrlService: FileUrlService,
+    private readonly streamHeartbeatService: AgentChatStreamHeartbeatService,
  ) {}

+  async reapDeadStream({
+    thread,
+    workspaceId,
+  }: {
+    thread: Pick<AgentChatThreadEntity, 'id' | 'activeStreamId'>;
+    workspaceId: string;
+  }): Promise<AgentChatThreadLastStreamError | null> {
+    if (!isDefined(thread.activeStreamId)) {
+      return null;
+    }
+
+    if (await this.streamHeartbeatService.isAlive(thread.activeStreamId)) {
+      return null;
+    }
+
+    const interruptedError: AgentChatThreadLastStreamError = {
+      code: STREAM_INTERRUPTED_CODE,
+      message: 'The response was interrupted before it could finish.',
+      failedAt: new Date().toISOString(),
+    };
+
+    const reap = await this.threadRepository.update(
+      workspaceId,
+      { id: thread.id, activeStreamId: thread.activeStreamId },
+      { activeStreamId: null, lastStreamError: interruptedError },
+    );
+
+    if (!reap.affected) {
+      return null;
+    }
+
+    await this.eventPublisherService.resetStreamState(thread.id);
+    await this.eventPublisherService
+      .publish({
+        threadId: thread.id,
+        workspaceId,
+        event: {
+          type: 'stream-error',
+          code: interruptedError.code,
+          message: interruptedError.message,
+        },
+      })
+      .catch(() => {});
+
+    return interruptedError;
+  }
+
  private async tryClaimStream({
    threadId,
    workspaceId,
@@ -79,7 +129,13 @@ export class AgentChatStreamingService {
      { activeStreamId: streamId, lastStreamError: null },
    );

-    return claim.affected === 1;
+    if (!claim.affected) {
+      return false;
+    }
+
+    await this.streamHeartbeatService.markClaimed(streamId);
+
+    return true;
  }

  async streamAgentChat({
@@ -343,6 +399,8 @@ export class AgentChatStreamingService {
      workspace.id,
    );

+    await this.streamHeartbeatService.markClaimed(streamId);
+
    await this.messageQueueService.add<StreamAgentChatJobData>(
      STREAM_AGENT_CHAT_JOB_NAME,
      {
				`@@ -0,0 +1 @@`
				`export const AGENT_CHAT_STREAM_REAP_CHECK_INTERVAL_MS = 10_000;`
				`@@ -0,0 +1 @@`
				`export const STREAM_INTERRUPTED_CODE = 'STREAM_INTERRUPTED';`