diff --git a/packages/twenty-front/src/modules/ai/components/AgentChatRuntimeEffects.tsx b/packages/twenty-front/src/modules/ai/components/AgentChatRuntimeEffects.tsx index df4e0b0b46..dc351447c1 100644 --- a/packages/twenty-front/src/modules/ai/components/AgentChatRuntimeEffects.tsx +++ b/packages/twenty-front/src/modules/ai/components/AgentChatRuntimeEffects.tsx @@ -1,5 +1,6 @@ import { AgentChatMessagesFetchEffect } from '@/ai/components/AgentChatMessagesFetchEffect'; import { AgentChatSessionStartTimeEffect } from '@/ai/components/AgentChatSessionStartTimeEffect'; +import { AgentChatStreamKeepAliveEffect } from '@/ai/components/AgentChatStreamKeepAliveEffect'; import { AgentChatStreamSubscriptionEffect } from '@/ai/components/AgentChatStreamSubscriptionEffect'; import { AgentChatStreamingAutoScrollEffect } from '@/ai/components/AgentChatStreamingAutoScrollEffect'; import { AgentChatStreamingPartsDiffSyncEffect } from '@/ai/components/AgentChatStreamingPartsDiffSyncEffect'; @@ -36,6 +37,7 @@ export const AgentChatRuntimeEffects = () => { <> + {isAgentChatOpen && ( <> diff --git a/packages/twenty-front/src/modules/ai/components/AgentChatStreamKeepAliveEffect.tsx b/packages/twenty-front/src/modules/ai/components/AgentChatStreamKeepAliveEffect.tsx new file mode 100644 index 0000000000..a71f03f59d --- /dev/null +++ b/packages/twenty-front/src/modules/ai/components/AgentChatStreamKeepAliveEffect.tsx @@ -0,0 +1,75 @@ +import { useStore } from 'jotai'; +import { useCallback, useEffect } from 'react'; +import { isDefined, isValidUuid } from 'twenty-shared/utils'; + +import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName'; +import { AGENT_CHAT_STREAM_LIVENESS_CHECK_INTERVAL_IN_MS } from '@/ai/constants/AgentChatStreamLivenessCheckIntervalInMs'; +import { AGENT_CHAT_STREAM_LIVENESS_TIMEOUT_IN_MS } from '@/ai/constants/AgentChatStreamLivenessTimeoutInMs'; +import { agentChatIsStreamingComponentFamilyState } from '@/ai/states/agentChatIsStreamingComponentFamilyState'; +import { agentChatStreamLastEventTimestampState } from '@/ai/states/agentChatStreamLastEventTimestampState'; +import { agentChatStreamResubscribeNonceState } from '@/ai/states/agentChatStreamResubscribeNonceState'; +import { currentAiChatThreadState } from '@/ai/states/currentAiChatThreadState'; +import { useListenToBrowserEvent } from '@/browser-event/hooks/useListenToBrowserEvent'; +import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent'; +import { SSE_CLIENT_RECONNECTED_EVENT_NAME } from '@/sse-db-event/constants/SseClientReconnectedEventName'; +import { useAtomComponentFamilyStateCallbackState } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateCallbackState'; +import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue'; + +export const AgentChatStreamKeepAliveEffect = () => { + const store = useStore(); + const currentAiChatThread = useAtomStateValue(currentAiChatThreadState); + + const isStreamingFamilyCallback = useAtomComponentFamilyStateCallbackState( + agentChatIsStreamingComponentFamilyState, + ); + + const hasActiveSubscription = + isDefined(currentAiChatThread) && isValidUuid(currentAiChatThread); + + const recoverStreamIfStreaming = useCallback(() => { + const isStreaming = store.get( + isStreamingFamilyCallback({ threadId: currentAiChatThread }), + ); + + if (!isStreaming) { + return; + } + + store.set(agentChatStreamLastEventTimestampState.atom, Date.now()); + store.set(agentChatStreamResubscribeNonceState.atom, (nonce) => nonce + 1); + dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME); + }, [store, isStreamingFamilyCallback, currentAiChatThread]); + + useEffect(() => { + if (!hasActiveSubscription) { + return; + } + + const interval = setInterval(() => { + const lastTimestamp = store.get( + agentChatStreamLastEventTimestampState.atom, + ); + + if (!isDefined(lastTimestamp)) { + return; + } + + const timeSinceLastEventInMs = Date.now() - lastTimestamp; + + if (timeSinceLastEventInMs <= AGENT_CHAT_STREAM_LIVENESS_TIMEOUT_IN_MS) { + return; + } + + recoverStreamIfStreaming(); + }, AGENT_CHAT_STREAM_LIVENESS_CHECK_INTERVAL_IN_MS); + + return () => clearInterval(interval); + }, [hasActiveSubscription, store, recoverStreamIfStreaming]); + + useListenToBrowserEvent({ + eventName: SSE_CLIENT_RECONNECTED_EVENT_NAME, + onBrowserEvent: recoverStreamIfStreaming, + }); + + return null; +}; diff --git a/packages/twenty-front/src/modules/ai/constants/AgentChatStreamLivenessCheckIntervalInMs.ts b/packages/twenty-front/src/modules/ai/constants/AgentChatStreamLivenessCheckIntervalInMs.ts new file mode 100644 index 0000000000..c22e600053 --- /dev/null +++ b/packages/twenty-front/src/modules/ai/constants/AgentChatStreamLivenessCheckIntervalInMs.ts @@ -0,0 +1 @@ +export const AGENT_CHAT_STREAM_LIVENESS_CHECK_INTERVAL_IN_MS = 2_000; diff --git a/packages/twenty-front/src/modules/ai/constants/AgentChatStreamLivenessTimeoutInMs.ts b/packages/twenty-front/src/modules/ai/constants/AgentChatStreamLivenessTimeoutInMs.ts new file mode 100644 index 0000000000..b37666d770 --- /dev/null +++ b/packages/twenty-front/src/modules/ai/constants/AgentChatStreamLivenessTimeoutInMs.ts @@ -0,0 +1 @@ +export const AGENT_CHAT_STREAM_LIVENESS_TIMEOUT_IN_MS = 5_000; diff --git a/packages/twenty-front/src/modules/ai/hooks/useAgentChatSubscription.ts b/packages/twenty-front/src/modules/ai/hooks/useAgentChatSubscription.ts index 1b770c3992..4d07381010 100644 --- a/packages/twenty-front/src/modules/ai/hooks/useAgentChatSubscription.ts +++ b/packages/twenty-front/src/modules/ai/hooks/useAgentChatSubscription.ts @@ -18,6 +18,8 @@ import { agentChatHandleEventCallbackComponentFamilyState } from '@/ai/states/ag import { agentChatIsAwaitingPersistedRefetchComponentFamilyState } from '@/ai/states/agentChatIsAwaitingPersistedRefetchComponentFamilyState'; import { agentChatIsStreamingComponentFamilyState } from '@/ai/states/agentChatIsStreamingComponentFamilyState'; import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMessagesComponentFamilyState'; +import { agentChatStreamLastEventTimestampState } from '@/ai/states/agentChatStreamLastEventTimestampState'; +import { agentChatStreamResubscribeNonceState } from '@/ai/states/agentChatStreamResubscribeNonceState'; import { agentChatUsageComponentFamilyState } from '@/ai/states/agentChatUsageComponentFamilyState'; import { currentAiChatThreadTitleComponentFamilyState } from '@/ai/states/currentAiChatThreadTitleComponentFamilyState'; import { AiChatErrorCode } from '@/ai/utils/aiChatErrorCode'; @@ -99,6 +101,9 @@ type AgentChatEventPayload = { export const useAgentChatSubscription = (threadId: string | null) => { const store = useStore(); const sseClient = useAtomStateValue(sseClientState); + const agentChatStreamResubscribeNonce = useAtomStateValue( + agentChatStreamResubscribeNonceState, + ); const errorFamilyCallback = useAtomComponentFamilyStateCallbackState( agentChatErrorComponentFamilyState, @@ -152,6 +157,7 @@ export const useAgentChatSubscription = (threadId: string | null) => { let disposed = false; store.set(firstLiveSeqAtom, null); + store.set(agentChatStreamLastEventTimestampState.atom, Date.now()); const closeWriter = () => { if (isDefined(writer)) { @@ -320,6 +326,10 @@ export const useAgentChatSubscription = (threadId: string | null) => { break; } + case 'keepalive': { + break; + } + case 'stream-error': { const streamError = new Error(event.message) as Error & { code?: string; @@ -388,6 +398,8 @@ export const useAgentChatSubscription = (threadId: string | null) => { }, { next: (value: ExecutionResult) => { + store.set(agentChatStreamLastEventTimestampState.atom, Date.now()); + if (isDefined(value.data?.onAgentChatEvent?.event)) { handleEvent( value.data.onAgentChatEvent.event as AgentChatSubscriptionEvent, @@ -417,6 +429,7 @@ export const useAgentChatSubscription = (threadId: string | null) => { }, [ threadId, sseClient, + agentChatStreamResubscribeNonce, store, errorFamilyCallback, isStreamingFamilyCallback, diff --git a/packages/twenty-front/src/modules/ai/states/agentChatStreamLastEventTimestampState.ts b/packages/twenty-front/src/modules/ai/states/agentChatStreamLastEventTimestampState.ts new file mode 100644 index 0000000000..2d9d2b967e --- /dev/null +++ b/packages/twenty-front/src/modules/ai/states/agentChatStreamLastEventTimestampState.ts @@ -0,0 +1,8 @@ +import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState'; + +export const agentChatStreamLastEventTimestampState = createAtomState< + number | null +>({ + key: 'agentChatStreamLastEventTimestampState', + defaultValue: null, +}); diff --git a/packages/twenty-front/src/modules/ai/states/agentChatStreamResubscribeNonceState.ts b/packages/twenty-front/src/modules/ai/states/agentChatStreamResubscribeNonceState.ts new file mode 100644 index 0000000000..f5322287ba --- /dev/null +++ b/packages/twenty-front/src/modules/ai/states/agentChatStreamResubscribeNonceState.ts @@ -0,0 +1,6 @@ +import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState'; + +export const agentChatStreamResubscribeNonceState = createAtomState({ + key: 'agentChatStreamResubscribeNonceState', + defaultValue: 0, +}); diff --git a/packages/twenty-front/src/modules/sse-db-event/components/SSEClientEffect.tsx b/packages/twenty-front/src/modules/sse-db-event/components/SSEClientEffect.tsx index 3ac7e0deb0..85421c2f71 100644 --- a/packages/twenty-front/src/modules/sse-db-event/components/SSEClientEffect.tsx +++ b/packages/twenty-front/src/modules/sse-db-event/components/SSEClientEffect.tsx @@ -1,5 +1,7 @@ import { useHasAccessTokenPair } from '@/auth/hooks/useHasAccessTokenPair'; import { tokenPairState } from '@/auth/states/tokenPairState'; +import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent'; +import { SSE_CLIENT_RECONNECTED_EVENT_NAME } from '@/sse-db-event/constants/SseClientReconnectedEventName'; import { useHandleSseClientConnectionRetry } from '@/sse-db-event/hooks/useHandleSseClientConnectionRetry'; import { activeQueryListenersState } from '@/sse-db-event/states/activeQueryListenersState'; import { sseClientState } from '@/sse-db-event/states/sseClientState'; @@ -18,15 +20,22 @@ export const SSEClientEffect = () => { const [sseClient, setSseClient] = useAtomState(sseClientState); const tokenPair = useAtomStateValue(tokenPairState); - const handleSSEClientConnected = useCallback(() => { - const currentActiveQueryListeners = store.get( - activeQueryListenersState.atom, - ); + const handleSSEClientConnected = useCallback( + (reconnected: boolean) => { + const currentActiveQueryListeners = store.get( + activeQueryListenersState.atom, + ); - if (isNonEmptyArray(currentActiveQueryListeners)) { - store.set(activeQueryListenersState.atom, []); - } - }, [store]); + if (isNonEmptyArray(currentActiveQueryListeners)) { + store.set(activeQueryListenersState.atom, []); + } + + if (reconnected) { + dispatchBrowserEvent(SSE_CLIENT_RECONNECTED_EVENT_NAME); + } + }, + [store], + ); const { handleSseClientConnectionRetry } = useHandleSseClientConnectionRetry(); diff --git a/packages/twenty-front/src/modules/sse-db-event/constants/SseClientReconnectedEventName.ts b/packages/twenty-front/src/modules/sse-db-event/constants/SseClientReconnectedEventName.ts new file mode 100644 index 0000000000..72c8fbe682 --- /dev/null +++ b/packages/twenty-front/src/modules/sse-db-event/constants/SseClientReconnectedEventName.ts @@ -0,0 +1,2 @@ +export const SSE_CLIENT_RECONNECTED_EVENT_NAME = + 'sse-client-reconnected' as const; diff --git a/packages/twenty-server/src/engine/core-modules/message-queue/constants/ai-stream-lock-duration.constant.ts b/packages/twenty-server/src/engine/core-modules/message-queue/constants/ai-stream-lock-duration.constant.ts new file mode 100644 index 0000000000..c739aa8943 --- /dev/null +++ b/packages/twenty-server/src/engine/core-modules/message-queue/constants/ai-stream-lock-duration.constant.ts @@ -0,0 +1 @@ +export const AI_STREAM_LOCK_DURATION_MS = 10 * 60 * 1000; diff --git a/packages/twenty-server/src/engine/core-modules/message-queue/drivers/bullmq.driver.ts b/packages/twenty-server/src/engine/core-modules/message-queue/drivers/bullmq.driver.ts index 5daf3e0134..258760ba67 100644 --- a/packages/twenty-server/src/engine/core-modules/message-queue/drivers/bullmq.driver.ts +++ b/packages/twenty-server/src/engine/core-modules/message-queue/drivers/bullmq.driver.ts @@ -108,6 +108,9 @@ export class BullMQDriver ...(isDefined(options?.concurrency) ? { concurrency: options.concurrency } : {}), + ...(isDefined(options?.lockDuration) + ? { lockDuration: options.lockDuration } + : {}), metrics: { maxDataPoints: MetricsTime.ONE_WEEK, collectInterval: 60000, diff --git a/packages/twenty-server/src/engine/core-modules/message-queue/interfaces/message-queue-worker-options.interface.ts b/packages/twenty-server/src/engine/core-modules/message-queue/interfaces/message-queue-worker-options.interface.ts index c4def7121e..777d122050 100644 --- a/packages/twenty-server/src/engine/core-modules/message-queue/interfaces/message-queue-worker-options.interface.ts +++ b/packages/twenty-server/src/engine/core-modules/message-queue/interfaces/message-queue-worker-options.interface.ts @@ -1,3 +1,4 @@ export interface MessageQueueWorkerOptions { concurrency?: number; + lockDuration?: number; } diff --git a/packages/twenty-server/src/engine/core-modules/message-queue/message-queue-worker-options.constant.ts b/packages/twenty-server/src/engine/core-modules/message-queue/message-queue-worker-options.constant.ts index 004bb703ab..24b8e9b5be 100644 --- a/packages/twenty-server/src/engine/core-modules/message-queue/message-queue-worker-options.constant.ts +++ b/packages/twenty-server/src/engine/core-modules/message-queue/message-queue-worker-options.constant.ts @@ -1,8 +1,13 @@ +import { AI_STREAM_LOCK_DURATION_MS } from 'src/engine/core-modules/message-queue/constants/ai-stream-lock-duration.constant'; +import { type MessageQueueWorkerOptions } from 'src/engine/core-modules/message-queue/interfaces/message-queue-worker-options.interface'; import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants'; export const QUEUE_WORKER_OPTIONS: Partial< - Record + Record > = { - [MessageQueue.aiStreamQueue]: { concurrency: 20 }, + [MessageQueue.aiStreamQueue]: { + concurrency: 20, + lockDuration: AI_STREAM_LOCK_DURATION_MS, + }, [MessageQueue.logicFunctionQueue]: { concurrency: 10 }, }; diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-keepalive-interval-ms.constant.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-keepalive-interval-ms.constant.ts new file mode 100644 index 0000000000..34673133b7 --- /dev/null +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/constants/agent-chat-keepalive-interval-ms.constant.ts @@ -0,0 +1 @@ +export const AGENT_CHAT_KEEPALIVE_INTERVAL_MS = 2_000; 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 c35ebf57b8..15d52a6e63 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 @@ -509,6 +509,16 @@ export class StreamAgentChatJob { const userMessage = await userMessagePromise; + if ( + isDefined(userMessage.turnId) && + (await this.agentChatService.hasAssistantMessageForTurn({ + turnId: userMessage.turnId, + workspaceId, + })) + ) { + return; + } + await this.agentChatService.addMessage({ threadId, uiMessage: responseMessage, 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 77221edc7f..4e37f7a7dd 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 @@ -17,9 +17,11 @@ import { AiExceptionCode, } 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 { 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 { 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'; import { WorkspaceScopedRepository } from 'src/engine/twenty-orm/workspace-scoped-repository/workspace-scoped-repository'; @MetadataResolver() @@ -58,9 +60,30 @@ export class AgentChatSubscriptionResolver { ); } - return this.subscriptionService.subscribeToAgentChat({ + const iterator = await this.subscriptionService.subscribeToAgentChat({ workspaceId: workspace.id, threadId, }); + + const keepalivePayload = { + onAgentChatEvent: { + threadId, + event: { type: 'keepalive' as const }, + }, + }; + + return wrapAsyncIteratorWithLifecycle(iterator, { + initialValue: keepalivePayload, + onHeartbeat: async () => { + await this.subscriptionService.publishToAgentChat({ + workspaceId: workspace.id, + threadId, + payload: keepalivePayload, + }); + + return true; + }, + heartbeatIntervalMs: AGENT_CHAT_KEEPALIVE_INTERVAL_MS, + }); } } diff --git a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts index 6750de2a47..634cc83ae1 100644 --- a/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts +++ b/packages/twenty-server/src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service.ts @@ -1,6 +1,7 @@ import { Injectable, Logger } from '@nestjs/common'; import { ExtendedUIMessage } from 'twenty-shared/ai'; +import { isDefined } from 'twenty-shared/utils'; import { In, IsNull, Not } from 'typeorm'; import type { QueryDeepPartialEntity } from 'typeorm/query-builder/QueryPartialEntity'; @@ -259,6 +260,21 @@ export class AgentChatService { } as AgentMessageEntity; } + async hasAssistantMessageForTurn({ + turnId, + workspaceId, + }: { + turnId: string; + workspaceId: string; + }): Promise { + const existingMessage = await this.messageRepository.findOne(workspaceId, { + where: { turnId, role: AgentMessageRole.ASSISTANT }, + select: ['id'], + }); + + return isDefined(existingMessage); + } + async getMessagesForThread({ threadId, userWorkspaceId, diff --git a/packages/twenty-shared/src/ai/types/AgentChatSubscriptionEvent.ts b/packages/twenty-shared/src/ai/types/AgentChatSubscriptionEvent.ts index 7ec44a20b5..a175453c40 100644 --- a/packages/twenty-shared/src/ai/types/AgentChatSubscriptionEvent.ts +++ b/packages/twenty-shared/src/ai/types/AgentChatSubscriptionEvent.ts @@ -6,4 +6,5 @@ export type AgentChatSubscriptionEvent = | { type: 'message-persisted'; messageId: string } | { type: 'queue-updated' } | { type: 'stream-error'; code: string; message: string } - | { type: 'credits-exhausted' }; + | { type: 'credits-exhausted' } + | { type: 'keepalive' };