diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/ai-chat.module.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/ai-chat.module.ts index 2e8f2248cd..7cf1f6dac8 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/ai-chat.module.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/ai-chat.module.ts @@ -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, diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-stream-reap-check-interval-ms.constant.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-stream-reap-check-interval-ms.constant.ts new file mode 100644 index 0000000000..6374738a62 --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-stream-reap-check-interval-ms.constant.ts @@ -0,0 +1 @@ +export const AGENT_CHAT_STREAM_REAP_CHECK_INTERVAL_MS = 10_000; diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/stream-interrupted-code.constant.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/stream-interrupted-code.constant.ts new file mode 100644 index 0000000000..98be7902d1 --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/stream-interrupted-code.constant.ts @@ -0,0 +1 @@ +export const STREAM_INTERRUPTED_CODE = 'STREAM_INTERRUPTED'; 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 755a9c098f..f443b8014b 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 @@ -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 { 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 ee6871d9c9..2ab12c30da 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 @@ -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( diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat-subscription.resolver.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat-subscription.resolver.ts index 4e37f7a7dd..6a378d9204 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat-subscription.resolver.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat-subscription.resolver.ts @@ -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, ) {} @@ -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 { + 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 + } + } } diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts index 1310c267a0..6841384edf 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/resolvers/agent-chat.resolver.ts @@ -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) diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.claim.spec.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.claim.spec.ts index a63668f935..00c9e01cd8 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.claim.spec.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.claim.spec.ts @@ -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); + }); + }); }); diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.retry.spec.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.retry.spec.ts index cb1fa2bf35..791eef09ec 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.retry.spec.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/__tests__/agent-chat-streaming.service.retry.spec.ts @@ -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 }; diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-stream-heartbeat.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-stream-heartbeat.service.ts new file mode 100644 index 0000000000..3bef3438e0 --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-stream-heartbeat.service.ts @@ -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 { + 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 { + try { + const exists = await this.redisClientService + .getClient() + .exists(this.getKey(streamId)); + + return exists === 1; + } catch { + return true; + } + } + + async clear(streamId: string): Promise { + await this.redisClientService + .getClient() + .del(this.getKey(streamId)) + .catch(() => {}); + } +} diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts index 19ca567a61..0d00945a2e 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service.ts @@ -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; + workspaceId: string; + }): Promise { + 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( STREAM_AGENT_CHAT_JOB_NAME, {