From 9be2b21e516a28f2c3a79678d637d6a8ce99946a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?F=C3=A9lix=20Malfait?= Date: Fri, 3 Jul 2026 18:50:10 +0200 Subject: [PATCH] fix(ai): gate chat thread totals on stream ownership to prevent usage under-counting (#22534) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. Review in cubic --- .../__tests__/stream-agent-chat.job.spec.ts | 61 ++++++++++++------- .../ai/ai-chat/jobs/stream-agent-chat.job.ts | 34 +++++------ 2 files changed, 55 insertions(+), 40 deletions(-) diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts index 34e48ac5b1..238d8efafb 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/__tests__/stream-agent-chat.job.spec.ts @@ -87,14 +87,14 @@ describe('StreamAgentChatJob', () => { streamChatRejection, addMessageRejection, assistantPersistRejection, - assistantMessageExistedAtStart = false, + totalsUpdateAffected = 1, }: { workspaceFound?: boolean; chatStream?: ReturnType; 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(); }); }); diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts index 84aef08ec8..c43188dd72 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job.ts @@ -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((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({ 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; 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,