fix(ai): gate chat thread totals on stream ownership to prevent usage under-counting (#22534)

Fixes a usage under-counting bug introduced by #22524 (progressive assistant-message persistence), flagged by @cubic-dev-ai and confirmed by @FelixMalfait. ## Root cause #22524 gated the thread-totals update (tokens, credits, `conversationSize`) on `assistantMessageExistedAtStart` — an existence check on the deterministic `uuidv5(streamId)` message id captured at job start. That was a sound idempotency signal *before* #22524, when the message only ever existed if a prior run had completed and applied totals. Progressive checkpoints broke that assumption: a checkpoint creates the message row ~2s into the stream, without applying totals. So if the worker is SIGKILLed after a checkpoint but before `handleStreamFinish`, and the job is re-delivered (BullMQ's stalled re-run — `aiStreamQueue` has no `maxStalledCount: 0` yet, that's #22518 — or an admin `retryJobs`), the re-run sees `assistantMessageExistedAtStart === true` and returns before the totals update. The turn's usage is lost permanently. cubic's P2 (the non-transactional `delete`+`insert` in `upsertAssistantMessage`) is the same root cause: its partless window only mattered because it tripped the same existence-based gate. ## Fix Stop inferring "totals already applied" from message existence. Gate the totals update on **still owning the stream** — a conditional `UPDATE ... WHERE id = :threadId AND activeStreamId = :streamId`, and only `notifyThreadUsageUpdated` when it affects a row. This is the same claim pattern the stream already uses (#22481), and it's idempotent by construction: - The run that completes while holding the claim → `affected = 1` → totals applied exactly once. This holds **even when a checkpoint already created the message**, which is precisely the bug. - A duplicate/zombie run after another run completed (and its `finally` cleared `activeStreamId`) → `affected = 0` → skipped, no double-count. - A superseded run whose thread has moved to a newer stream → `affected = 0` → skipped (defense-in-depth, aligns with #22518's ownership pre-check). The `assistantMessageExistedAtStart` flag and its start-of-stream `hasMessageById` query are removed entirely — the message write is already idempotent via the deterministic id + `upsert`, so it needs no gate. This subsumes cubic's P2: the totals are no longer lost regardless of the `delete`+`insert` window, so no transaction is required for correctness (the residual window is a benign sub-millisecond transient for an actively-streaming message; happy to add a workspace-datasource transaction as separate hardening if you'd prefer). ## Validation `stream-agent-chat.job.spec.ts` (9 green): - New: totals update returns `affected: 0` → `notifyThreadUsageUpdated` **not** called (prior completion not double-counted), message still upserted. - New: message already exists from a checkpoint but claim still held (`affected: 1`) → totals **are** applied — the exact regression #22524 caused. - Existing success/error/cancel/abort flows updated for the conditional criteria and still green.  <a href="https://cubic.dev/pr/twentyhq/twenty/pull/22534?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-03 18:50:10 +02:00
parent e5fc1b3702
commit 9be2b21e51
2 changed files with 55 additions and 40 deletions
@@ -87,14 +87,14 @@ describe('StreamAgentChatJob', () => {
    streamChatRejection,
    addMessageRejection,
    assistantPersistRejection,
-    assistantMessageExistedAtStart = false,
+    totalsUpdateAffected = 1,
  }: {
    workspaceFound?: boolean;
    chatStream?: ReturnType<typeof createFakeChatStream>;
    streamChatRejection?: Error;
    addMessageRejection?: Error;
    assistantPersistRejection?: Error;
-    assistantMessageExistedAtStart?: boolean;
+    totalsUpdateAffected?: number;
  } = {}) => {
    const publishedEvents: PublishedEvent[] = [];

@@ -102,7 +102,14 @@ describe('StreamAgentChatJob', () => {
      findOne: jest
        .fn()
        .mockResolvedValue({ id: 'thread-id', deletedAt: null }),
-      update: jest.fn().mockResolvedValue({ affected: 1 }),
+      update: jest.fn().mockImplementation((_workspaceId, _criteria, values) =>
+        Promise.resolve({
+          affected:
+            values && typeof values.totalInputTokens === 'function'
+              ? totalsUpdateAffected
+              : 1,
+        }),
+      ),
    };
    const workspaceRepository = {
      findOne: jest.fn().mockResolvedValue(workspaceFound ? workspace : null),
@@ -114,9 +121,6 @@ describe('StreamAgentChatJob', () => {
      upsertAssistantMessage: assistantPersistRejection
        ? jest.fn().mockRejectedValue(assistantPersistRejection)
        : jest.fn().mockResolvedValue(undefined),
-      hasMessageById: jest
-        .fn()
-        .mockResolvedValue(assistantMessageExistedAtStart),
      generateTitleIfNeeded: jest.fn().mockResolvedValue(null),
      notifyThreadUsageUpdated: jest.fn().mockResolvedValue(undefined),
    };
@@ -210,7 +214,7 @@ describe('StreamAgentChatJob', () => {
    );
    expect(threadRepository.update).toHaveBeenCalledWith(
      'workspace-id',
-      { id: 'thread-id' },
+      { id: 'thread-id', activeStreamId: 'stream-id' },
      expect.objectContaining({ lastStreamError: null }),
    );
    expect(threadRepository.update).toHaveBeenCalledWith(
@@ -221,9 +225,9 @@ describe('StreamAgentChatJob', () => {
    expect(agentChatStreamingService.flushNextQueuedMessage).toHaveBeenCalled();
  });

-  it('persists the assistant message but does not re-apply thread totals when a prior execution already persisted it', async () => {
+  it('gates the thread totals on still owning the stream so a prior completion is not double-counted', async () => {
    const { job, agentChatService, threadRepository } = buildJob({
-      assistantMessageExistedAtStart: true,
+      totalsUpdateAffected: 0,
    });

    await job.handle(jobData);
@@ -231,18 +235,23 @@ describe('StreamAgentChatJob', () => {
    expect(agentChatService.upsertAssistantMessage).toHaveBeenCalledWith(
      expect.objectContaining({ turnId: 'turn-id' }),
    );
-    // The thread-totals accumulation must not run twice for the same stream.
-    const totalsUpdate = threadRepository.update.mock.calls.find(
-      ([, criteria]) =>
-        criteria &&
-        typeof criteria === 'object' &&
-        !('activeStreamId' in criteria),
+    expect(threadRepository.update).toHaveBeenCalledWith(
+      'workspace-id',
+      { id: 'thread-id', activeStreamId: 'stream-id' },
+      expect.objectContaining({ lastStreamError: null }),
    );
-
-    expect(totalsUpdate).toBeUndefined();
    expect(agentChatService.notifyThreadUsageUpdated).not.toHaveBeenCalled();
  });

+  it('applies thread totals when the claim is still held even if the message already exists from a checkpoint', async () => {
+    const { job, agentChatService } = buildJob({ totalsUpdateAffected: 1 });
+
+    await job.handle(jobData);
+
+    expect(agentChatService.upsertAssistantMessage).toHaveBeenCalled();
+    expect(agentChatService.notifyThreadUsageUpdated).toHaveBeenCalled();
+  });
+
  it('never publishes the opaque error chunk to subscribers', async () => {
    const { job, publishedEvents } = buildJob({
      chatStream: createFakeChatStream({
@@ -382,12 +391,17 @@ describe('StreamAgentChatJob', () => {
  it('resolves without flushing the queue when the stream is cancelled', async () => {
    let triggerCancel: (() => void) | undefined;

-    const { job, publishedEvents, agentChatStreamingService, cancelCallbacks } =
-      buildJob({
-        chatStream: createFakeChatStream({
-          onFirstChunk: () => triggerCancel?.(),
-        }),
-      });
+    const {
+      job,
+      publishedEvents,
+      agentChatService,
+      agentChatStreamingService,
+      cancelCallbacks,
+    } = buildJob({
+      chatStream: createFakeChatStream({
+        onFirstChunk: () => triggerCancel?.(),
+      }),
+    });

    triggerCancel = () => cancelCallbacks.forEach((callback) => callback());

@@ -399,5 +413,6 @@ describe('StreamAgentChatJob', () => {
    expect(
      agentChatStreamingService.flushNextQueuedMessage,
    ).not.toHaveBeenCalled();
+    expect(agentChatService.notifyThreadUsageUpdated).toHaveBeenCalled();
  });
 });
@@ -227,12 +227,6 @@ export class StreamAgentChatJob {
      ASSISTANT_MESSAGE_ID_NAMESPACE,
    );

-    const assistantMessageExistedAtStart =
-      await this.agentChatService.hasMessageById({
-        id: assistantMessageId,
-        workspaceId: data.workspaceId,
-      });
-
    return new Promise<void>((resolve, reject) => {
      let streamUsage = {
        inputTokens: 0,
@@ -271,7 +265,13 @@ export class StreamAgentChatJob {
        resolveStreamFinished = res;
      });

-      abortSignal.addEventListener('abort', () => resolve(), { once: true });
+      abortSignal.addEventListener(
+        'abort',
+        () => {
+          void streamFinishedPromise.then(() => resolve());
+        },
+        { once: true },
+      );

      const uiStream = createUIMessageStream<ExtendedUIMessage>({
        execute: async ({ writer }) => {
@@ -352,7 +352,7 @@ export class StreamAgentChatJob {
                  await persistChain;
                  await this.handleStreamFinish({
                    assistantMessageId,
-                    assistantMessageExistedAtStart,
+                    streamId: data.streamId,
                    responseMessage,
                    isAborted,
                    streamError,
@@ -572,7 +572,7 @@ export class StreamAgentChatJob {

  private async handleStreamFinish({
    assistantMessageId,
-    assistantMessageExistedAtStart,
+    streamId,
    responseMessage,
    isAborted,
    streamError,
@@ -587,7 +587,7 @@ export class StreamAgentChatJob {
    userMessagePromise,
  }: {
    assistantMessageId: string;
-    assistantMessageExistedAtStart: boolean;
+    streamId: string;
    responseMessage: Omit<ExtendedUIMessage, 'id'>;
    isAborted: boolean;
    streamError: unknown;
@@ -649,7 +649,7 @@ export class StreamAgentChatJob {
        parts: responseMessage.parts,
        workspaceId,
      });
-    } else if (!assistantMessageExistedAtStart) {
+    } else {
      await this.agentChatService.addMessage({
        threadId,
        uiMessage: responseMessage,
@@ -658,13 +658,9 @@ export class StreamAgentChatJob {
      });
    }

-    if (assistantMessageExistedAtStart) {
-      return;
-    }
-
-    await this.threadRepository.update(
+    const totalsUpdate = await this.threadRepository.update(
      workspaceId,
-      { id: threadId },
+      { id: threadId, activeStreamId: streamId },
      {
        totalInputTokens: () =>
          `"totalInputTokens" + ${streamUsage.inputTokens}`,
@@ -687,6 +683,10 @@ export class StreamAgentChatJob {
      },
    );

+    if (!totalsUpdate.affected) {
+      return;
+    }
+
    await this.agentChatService.notifyThreadUsageUpdated({
      threadId,
      userWorkspaceId,