feat: queue messages + replace AI SDK with GraphQL SSE subscription (#19203)

## Summary

- **Queue messages while streaming**: Messages sent during active AI
streaming are queued server-side and auto-flushed when the current
stream completes. Frontend renders queued messages optimistically in a
dedicated queue UI.
- **Drop `@ai-sdk/react` + `resumable-stream`**: Replace the dual HTTP
SSE + AI SDK client architecture with a single GraphQL SSE subscription
per thread. All events (token chunks, message persistence, queue
updates, errors) flow through Redis PubSub → GraphQL subscription.
- **Server-driven architecture**: The server decides whether to queue or
stream (via `POST /:threadId/message`). The frontend mirrors this
decision for optimistic rendering but defers to the server response.
- **Reuse AI SDK accumulation logic**: `readUIMessageStream` from the
`ai` package handles chunk-to-message accumulation on the frontend,
avoiding a custom 780-line accumulator.

## Key files

**Backend:**
- `agent-chat-event-publisher.service.ts` — publishes events to Redis
PubSub
- `agent-chat-subscription.resolver.ts` — GraphQL subscription resolver
- `stream-agent-chat.job.ts` — publishes chunks via PubSub instead of
resumable-stream
- `agent-chat.controller.ts` — unified `POST /:threadId/message`
endpoint

**Frontend:**
- `useAgentChatSubscription.ts` — subscribes to `onAgentChatEvent`,
bridges to `readUIMessageStream`
- `useAgentChat.ts` — send/stop/optimistic rendering (no more AI SDK)
- `AgentChatStreamSubscriptionEffect.tsx` — replaces
`AgentChatAiSdkStreamEffect.tsx`

## Test plan

- [ ] Send message on new thread → optimistic render, streaming response
appears
- [ ] Send message while streaming → queued instantly (no flash in main
thread)
- [ ] Queued message auto-flushes after current stream completes
- [ ] Remove queued message via queue UI
- [ ] Stop streaming mid-response
- [ ] Leave chat idle for several minutes → streaming still works after
(SSE client recycling)
- [ ] Token refresh during session → requests succeed (authenticated
fetch)
- [ ] Switch threads while streaming → clean subscription handoff


Made with [Cursor](https://cursor.com)

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Félix Malfait
2026-04-02 10:10:13 +02:00
committed by GitHub
parent 3733a5a763
commit fd7387928c
79 changed files with 2617 additions and 1428 deletions
@@ -2787,10 +2787,12 @@ type AgentChatThread {
type AgentMessage {
id: UUID!
threadId: UUID!
turnId: UUID!
turnId: UUID
agentId: UUID
role: String!
status: String!
parts: [AgentMessagePart!]!
processedAt: DateTime
createdAt: DateTime!
}
@@ -2805,6 +2807,22 @@ type AISystemPromptPreview {
estimatedTokenCount: Int!
}
type ChatStreamCatchupChunks {
chunks: [JSON!]!
maxSeq: Int!
}
type SendChatMessageResult {
messageId: String!
queued: Boolean!
streamId: String
}
type AgentChatEvent {
threadId: String!
event: JSON!
}
type AgentChatThreadEdge {
"""The node containing the AgentChatThread"""
node: AgentChatThread!
@@ -3152,6 +3170,7 @@ type Query {
minimalMetadata: MinimalMetadata!
chatThread(id: UUID!): AgentChatThread!
chatMessages(threadId: UUID!): [AgentMessage!]!
chatStreamCatchupChunks(threadId: UUID!): ChatStreamCatchupChunks!
getAISystemPromptPreview: AISystemPromptPreview!
skills: [Skill!]!
skill(id: UUID!): Skill
@@ -3452,6 +3471,9 @@ type Mutation {
updateWebhook(input: UpdateWebhookInput!): Webhook!
deleteWebhook(id: UUID!): Webhook!
createChatThread: AgentChatThread!
sendChatMessage(threadId: UUID!, text: String!, messageId: UUID!, browsingContext: JSON, modelId: String): SendChatMessageResult!
stopAgentChatStream(threadId: UUID!): Boolean!
deleteQueuedChatMessage(messageId: UUID!): Boolean!
createSkill(input: CreateSkillInput!): Skill!
updateSkill(input: UpdateSkillInput!): Skill!
deleteSkill(id: UUID!): Skill!
@@ -4541,6 +4563,7 @@ enum FileFolder {
type Subscription {
onEventSubscription(eventStreamId: String!): EventSubscription
logicFunctionLogs(input: LogicFunctionLogsInput!): LogicFunctionLogs!
onAgentChatEvent(threadId: UUID!): AgentChatEvent!
}
input LogicFunctionLogsInput {
@@ -2448,10 +2448,12 @@ export interface AgentChatThread {
export interface AgentMessage {
id: Scalars['UUID']
threadId: Scalars['UUID']
turnId: Scalars['UUID']
turnId?: Scalars['UUID']
agentId?: Scalars['UUID']
role: Scalars['String']
status: Scalars['String']
parts: AgentMessagePart[]
processedAt?: Scalars['DateTime']
createdAt: Scalars['DateTime']
__typename: 'AgentMessage'
}
@@ -2469,6 +2471,25 @@ export interface AISystemPromptPreview {
__typename: 'AISystemPromptPreview'
}
export interface ChatStreamCatchupChunks {
chunks: Scalars['JSON'][]
maxSeq: Scalars['Int']
__typename: 'ChatStreamCatchupChunks'
}
export interface SendChatMessageResult {
messageId: Scalars['String']
queued: Scalars['Boolean']
streamId?: Scalars['String']
__typename: 'SendChatMessageResult'
}
export interface AgentChatEvent {
threadId: Scalars['String']
event: Scalars['JSON']
__typename: 'AgentChatEvent'
}
export interface AgentChatThreadEdge {
/** The node containing the AgentChatThread */
node: AgentChatThread
@@ -2713,6 +2734,7 @@ export interface Query {
minimalMetadata: MinimalMetadata
chatThread: AgentChatThread
chatMessages: AgentMessage[]
chatStreamCatchupChunks: ChatStreamCatchupChunks
getAISystemPromptPreview: AISystemPromptPreview
skills: Skill[]
skill?: Skill
@@ -2899,6 +2921,9 @@ export interface Mutation {
updateWebhook: Webhook
deleteWebhook: Webhook
createChatThread: AgentChatThread
sendChatMessage: SendChatMessageResult
stopAgentChatStream: Scalars['Boolean']
deleteQueuedChatMessage: Scalars['Boolean']
createSkill: Skill
updateSkill: Skill
deleteSkill: Skill
@@ -3001,6 +3026,7 @@ export type FileFolder = 'ProfilePicture' | 'WorkspaceLogo' | 'Attachment' | 'Pe
export interface Subscription {
onEventSubscription?: EventSubscription
logicFunctionLogs: LogicFunctionLogs
onAgentChatEvent: AgentChatEvent
__typename: 'Subscription'
}
@@ -5598,7 +5624,9 @@ export interface AgentMessageGenqlSelection{
turnId?: boolean | number
agentId?: boolean | number
role?: boolean | number
status?: boolean | number
parts?: AgentMessagePartGenqlSelection
processedAt?: boolean | number
createdAt?: boolean | number
__typename?: boolean | number
__scalar?: boolean | number
@@ -5619,6 +5647,28 @@ export interface AISystemPromptPreviewGenqlSelection{
__scalar?: boolean | number
}
export interface ChatStreamCatchupChunksGenqlSelection{
chunks?: boolean | number
maxSeq?: boolean | number
__typename?: boolean | number
__scalar?: boolean | number
}
export interface SendChatMessageResultGenqlSelection{
messageId?: boolean | number
queued?: boolean | number
streamId?: boolean | number
__typename?: boolean | number
__scalar?: boolean | number
}
export interface AgentChatEventGenqlSelection{
threadId?: boolean | number
event?: boolean | number
__typename?: boolean | number
__scalar?: boolean | number
}
export interface AgentChatThreadEdgeGenqlSelection{
/** The node containing the AgentChatThread */
node?: AgentChatThreadGenqlSelection
@@ -5868,6 +5918,7 @@ export interface QueryGenqlSelection{
minimalMetadata?: MinimalMetadataGenqlSelection
chatThread?: (AgentChatThreadGenqlSelection & { __args: {id: Scalars['UUID']} })
chatMessages?: (AgentMessageGenqlSelection & { __args: {threadId: Scalars['UUID']} })
chatStreamCatchupChunks?: (ChatStreamCatchupChunksGenqlSelection & { __args: {threadId: Scalars['UUID']} })
getAISystemPromptPreview?: AISystemPromptPreviewGenqlSelection
skills?: SkillGenqlSelection
skill?: (SkillGenqlSelection & { __args: {id: Scalars['UUID']} })
@@ -6079,6 +6130,9 @@ export interface MutationGenqlSelection{
updateWebhook?: (WebhookGenqlSelection & { __args: {input: UpdateWebhookInput} })
deleteWebhook?: (WebhookGenqlSelection & { __args: {id: Scalars['UUID']} })
createChatThread?: AgentChatThreadGenqlSelection
sendChatMessage?: (SendChatMessageResultGenqlSelection & { __args: {threadId: Scalars['UUID'], text: Scalars['String'], messageId: Scalars['UUID'], browsingContext?: (Scalars['JSON'] | null), modelId?: (Scalars['String'] | null)} })
stopAgentChatStream?: { __args: {threadId: Scalars['UUID']} }
deleteQueuedChatMessage?: { __args: {messageId: Scalars['UUID']} }
createSkill?: (SkillGenqlSelection & { __args: {input: CreateSkillInput} })
updateSkill?: (SkillGenqlSelection & { __args: {input: UpdateSkillInput} })
deleteSkill?: (SkillGenqlSelection & { __args: {id: Scalars['UUID']} })
@@ -6496,6 +6550,7 @@ export interface WorkspaceMigrationDeleteActionInput {type: WorkspaceMigrationAc
export interface SubscriptionGenqlSelection{
onEventSubscription?: (EventSubscriptionGenqlSelection & { __args: {eventStreamId: Scalars['String']} })
logicFunctionLogs?: (LogicFunctionLogsGenqlSelection & { __args: {input: LogicFunctionLogsInput} })
onAgentChatEvent?: (AgentChatEventGenqlSelection & { __args: {threadId: Scalars['UUID']} })
__typename?: boolean | number
__scalar?: boolean | number
}
@@ -8423,6 +8478,30 @@ export interface LogicFunctionLogsInput {applicationId?: (Scalars['UUID'] | null
const ChatStreamCatchupChunks_possibleTypes: string[] = ['ChatStreamCatchupChunks']
export const isChatStreamCatchupChunks = (obj?: { __typename?: any } | null): obj is ChatStreamCatchupChunks => {
if (!obj?.__typename) throw new Error('__typename is missing in "isChatStreamCatchupChunks"')
return ChatStreamCatchupChunks_possibleTypes.includes(obj.__typename)
}
const SendChatMessageResult_possibleTypes: string[] = ['SendChatMessageResult']
export const isSendChatMessageResult = (obj?: { __typename?: any } | null): obj is SendChatMessageResult => {
if (!obj?.__typename) throw new Error('__typename is missing in "isSendChatMessageResult"')
return SendChatMessageResult_possibleTypes.includes(obj.__typename)
}
const AgentChatEvent_possibleTypes: string[] = ['AgentChatEvent']
export const isAgentChatEvent = (obj?: { __typename?: any } | null): obj is AgentChatEvent => {
if (!obj?.__typename) throw new Error('__typename is missing in "isAgentChatEvent"')
return AgentChatEvent_possibleTypes.includes(obj.__typename)
}
const AgentChatThreadEdge_possibleTypes: string[] = ['AgentChatThreadEdge']
export const isAgentChatThreadEdge = (obj?: { __typename?: any } | null): obj is AgentChatThreadEdge => {
if (!obj?.__typename) throw new Error('__typename is missing in "isAgentChatThreadEdge"')
File diff suppressed because it is too large Load Diff
-1
View File
@@ -29,7 +29,6 @@
"workerDirectory": "public"
},
"dependencies": {
"@ai-sdk/react": "3.0.99",
"@apollo/client": "^4.0.0",
"@blocknote/mantine": "^0.47.1",
"@blocknote/react": "^0.47.1",
File diff suppressed because one or more lines are too long
@@ -1,11 +1,6 @@
import { AGENT_CHAT_RETRY_EVENT_NAME } from '@/ai/constants/AgentChatRetryEventName';
import { agentChatIsStreamingState } from '@/ai/states/agentChatIsStreamingState';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { styled } from '@linaria/react';
import { t } from '@lingui/core/macro';
import { IconAlertCircle, IconRefresh } from 'twenty-ui/display';
import { Button } from 'twenty-ui/input';
import { IconAlertCircle } from 'twenty-ui/display';
import { useContext } from 'react';
import { ThemeContext, themeCssVariables } from 'twenty-ui/theme-constants';
@@ -49,12 +44,6 @@ type AIChatErrorMessageProps = {
};
export const AIChatErrorMessage = ({ error }: AIChatErrorMessageProps) => {
const agentChatIsStreaming = useAtomStateValue(agentChatIsStreamingState);
const handleRetryClick = () => {
dispatchBrowserEvent(AGENT_CHAT_RETRY_EVENT_NAME);
};
const { theme } = useContext(ThemeContext);
return (
@@ -68,14 +57,6 @@ export const AIChatErrorMessage = ({ error }: AIChatErrorMessageProps) => {
{error.message || t`An error occurred while processing your message`}
</StyledErrorMessage>
</StyledErrorContent>
<Button
variant="secondary"
size="small"
Icon={IconRefresh}
onClick={handleRetryClick}
disabled={agentChatIsStreaming}
title={t`Retry`}
/>
</StyledErrorContainer>
);
};
@@ -189,23 +189,22 @@ export const AIChatMessage = ({
<AIChatErrorRenderer error={error} />
)}
</StyledMessageContainer>
{agentChatMessage.parts.length > 0 &&
agentChatMessage.metadata?.createdAt && (
<StyledMessageFooter className="message-footer">
<StyledMessageTimestamp>
{beautifyPastDateRelativeToNow(
agentChatMessage.metadata?.createdAt,
localeCatalog,
)}
</StyledMessageTimestamp>
<LightCopyIconButton
copyText={
agentChatMessage.parts.find((part) => part.type === 'text')
?.text ?? ''
}
/>
</StyledMessageFooter>
)}
{agentChatMessage.parts.length > 0 && (
<StyledMessageFooter className="message-footer">
<StyledMessageTimestamp>
{beautifyPastDateRelativeToNow(
agentChatMessage.metadata?.createdAt ?? new Date(),
localeCatalog,
)}
</StyledMessageTimestamp>
<LightCopyIconButton
copyText={
agentChatMessage.parts.find((part) => part.type === 'text')
?.text ?? ''
}
/>
</StyledMessageFooter>
)}
</StyledMessageBubble>
);
};
@@ -0,0 +1,78 @@
import { styled } from '@linaria/react';
import { agentChatQueuedMessagesComponentFamilyState } from '@/ai/states/agentChatQueuedMessagesComponentFamilyState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { useDeleteQueuedMessage } from '@/ai/hooks/useDeleteQueuedMessage';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { isDefined } from 'twenty-shared/utils';
import { IconX } from 'twenty-ui/display';
import { LightIconButton } from 'twenty-ui/input';
import { themeCssVariables } from 'twenty-ui/theme-constants';
const StyledQueueContainer = styled.div`
display: flex;
flex-direction: column;
gap: ${themeCssVariables.spacing[1]};
padding: 0 ${themeCssVariables.spacing[3]};
`;
const StyledQueueLabel = styled.div`
color: ${themeCssVariables.font.color.light};
font-size: ${themeCssVariables.font.size.xs};
padding-left: ${themeCssVariables.spacing[1]};
`;
const StyledQueuedItem = styled.div`
align-items: center;
background: ${themeCssVariables.background.tertiary};
border-radius: ${themeCssVariables.border.radius.sm};
color: ${themeCssVariables.font.color.secondary};
display: flex;
font-size: ${themeCssVariables.font.size.md};
gap: ${themeCssVariables.spacing[2]};
justify-content: space-between;
padding: ${themeCssVariables.spacing[1]} ${themeCssVariables.spacing[2]};
`;
const StyledQueuedText = styled.span`
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
`;
export const AIChatQueuedMessages = () => {
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const agentChatQueuedMessages = useAtomComponentFamilyStateValue(
agentChatQueuedMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
const { deleteQueuedMessage } = useDeleteQueuedMessage();
if (!isDefined(currentAIChatThread) || agentChatQueuedMessages.length === 0) {
return null;
}
return (
<StyledQueueContainer>
<StyledQueueLabel>
{agentChatQueuedMessages.length} Queued
</StyledQueueLabel>
{agentChatQueuedMessages.map((message) => {
const textPart = message.parts?.find((part) => part.type === 'text');
const displayText = textPart && 'text' in textPart ? textPart.text : '';
return (
<StyledQueuedItem key={message.id}>
<StyledQueuedText>{displayText}</StyledQueuedText>
<LightIconButton
Icon={IconX}
onClick={() => deleteQueuedMessage(message.id)}
size="small"
/>
</StyledQueuedItem>
);
})}
</StyledQueueContainer>
);
};
@@ -10,6 +10,7 @@ import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { threadIdCreatedFromDraftState } from '@/ai/states/threadIdCreatedFromDraftState';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { AIChatQueuedMessages } from '@/ai/components/AIChatQueuedMessages';
import { AIChatTabMessageList } from '@/ai/components/AIChatTabMessageList';
const StyledContainer = styled.div<{ isDraggingFile: boolean }>`
@@ -51,6 +52,7 @@ export const AIChatTab = () => {
{!isDraggingFile && (
<>
<AIChatTabMessageList />
<AIChatQueuedMessages />
<AIChatEditorSection key={editorSectionKey} />
</>
)}
@@ -1,11 +1,16 @@
import { useCallback, useMemo } from 'react';
import { useStore } from 'jotai';
import { type AgentChatSubscriptionEvent } from 'twenty-shared/ai';
import { isDefined } from 'twenty-shared/utils';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { AGENT_CHAT_UNKNOWN_THREAD_ID } from '@/ai/constants/AgentChatUnknownThreadId';
import { agentChatFirstLiveSeqState } from '@/ai/states/agentChatFirstLiveSeqState';
import { agentChatHandleEventCallbackState } from '@/ai/states/agentChatHandleEventCallbackState';
import { AGENT_CHAT_NEW_THREAD_DRAFT_KEY } from '@/ai/states/agentChatDraftsByThreadIdState';
import { agentChatFetchedMessagesComponentFamilyState } from '@/ai/states/agentChatFetchedMessagesComponentFamilyState';
import { agentChatMessagesLoadingState } from '@/ai/states/agentChatMessagesLoadingState';
import { agentChatQueuedMessagesComponentFamilyState } from '@/ai/states/agentChatQueuedMessagesComponentFamilyState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { skipMessagesSkeletonUntilLoadedState } from '@/ai/states/skipMessagesSkeletonUntilLoadedState';
import { mapDBMessagesToUIMessages } from '@/ai/utils/mapDBMessagesToUIMessages';
@@ -20,6 +25,7 @@ import {
} from '~/generated-metadata/graphql';
export const AgentChatMessagesFetchEffect = () => {
const store = useStore();
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const isNewThread = useMemo(
@@ -42,6 +48,11 @@ export const AgentChatMessagesFetchEffect = () => {
{ threadId: currentAIChatThread },
);
const setAgentChatQueuedMessages = useSetAtomComponentFamilyState(
agentChatQueuedMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
const handleFirstLoad = useCallback(
(_data: GetChatMessagesQuery) => {
setSkipMessagesSkeletonUntilLoaded(false);
@@ -52,9 +63,42 @@ export const AgentChatMessagesFetchEffect = () => {
const handleDataLoaded = useCallback(
(data: GetChatMessagesQuery) => {
const uiMessages = mapDBMessagesToUIMessages(data.chatMessages ?? []);
setAgentChatFetchedMessages(uiMessages);
setAgentChatFetchedMessages(
uiMessages.filter((message) => message.status !== 'queued'),
);
setAgentChatQueuedMessages(
uiMessages.filter((message) => message.status === 'queued'),
);
const catchup = data.chatStreamCatchupChunks;
if (!isDefined(catchup) || catchup.chunks.length === 0) {
return;
}
const handleEvent = store.get(agentChatHandleEventCallbackState.atom);
if (!isDefined(handleEvent)) {
return;
}
const firstLiveSeq = store.get(agentChatFirstLiveSeqState.atom);
for (let index = 0; index < catchup.chunks.length; index++) {
const chunkSeq = index + 1;
if (firstLiveSeq !== null && chunkSeq >= firstLiveSeq) {
break;
}
handleEvent({
type: 'stream-chunk',
chunk: catchup.chunks[index],
seq: chunkSeq,
} as AgentChatSubscriptionEvent);
}
},
[setAgentChatFetchedMessages],
[setAgentChatFetchedMessages, setAgentChatQueuedMessages, store],
);
const handleLoadingChange = useCallback(
@@ -1,4 +1,4 @@
import { AgentChatAiSdkStreamEffect } from '@/ai/components/AgentChatAiSdkStreamEffect';
import { AgentChatStreamSubscriptionEffect } from '@/ai/components/AgentChatStreamSubscriptionEffect';
import { AgentChatMessagesFetchEffect } from '@/ai/components/AgentChatMessagesFetchEffect';
import { AgentChatSessionStartTimeEffect } from '@/ai/components/AgentChatSessionStartTimeEffect';
@@ -20,7 +20,7 @@ export const AgentChatProviderContent = ({
>
<AgentChatThreadInitializationEffect />
<AgentChatMessagesFetchEffect />
<AgentChatAiSdkStreamEffect />
<AgentChatStreamSubscriptionEffect />
<AgentChatStreamingPartsDiffSyncEffect />
<AgentChatSessionStartTimeEffect />
<AgentChatStreamingAutoScrollEffect />
@@ -1,11 +1,13 @@
import { useCallback, useEffect } from 'react';
import { isValidUuid } from 'twenty-shared/utils';
import { AGENT_CHAT_ENSURE_THREAD_FOR_DRAFT_EVENT_NAME } from '@/ai/constants/AgentChatEnsureThreadForDraftEventName';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { useAgentChat } from '@/ai/hooks/useAgentChat';
import { useAgentChatSubscription } from '@/ai/hooks/useAgentChatSubscription';
import { useCreateAgentChatThread } from '@/ai/hooks/useCreateAgentChatThread';
import { useEnsureAgentChatThreadExistsForDraft } from '@/ai/hooks/useEnsureAgentChatThreadExistsForDraft';
import { useEnsureAgentChatThreadIdForSend } from '@/ai/hooks/useEnsureAgentChatThreadIdForSend';
import { agentChatDisplayedThreadState } from '@/ai/states/agentChatDisplayedThreadState';
import { agentChatErrorState } from '@/ai/states/agentChatErrorState';
import { agentChatFetchedMessagesComponentFamilyState } from '@/ai/states/agentChatFetchedMessagesComponentFamilyState';
import { agentChatIsInitialScrollPendingOnThreadChangeState } from '@/ai/states/agentChatIsInitialScrollPendingOnThreadChangeState';
import { agentChatIsLoadingState } from '@/ai/states/agentChatIsLoadingState';
@@ -14,23 +16,14 @@ import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMess
import { agentChatMessagesLoadingState } from '@/ai/states/agentChatMessagesLoadingState';
import { agentChatThreadsLoadingState } from '@/ai/states/agentChatThreadsLoadingState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { mergeAgentChatFetchedAndStreamingMessages } from '@/ai/utils/mergeAgentChatFetchedAndStreamingMessages';
import { normalizeAiSdkError } from '@/ai/utils/normalizeAiSdkError';
import { useListenToBrowserEvent } from '@/browser-event/hooks/useListenToBrowserEvent';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { useSetAtomComponentFamilyState } from '@/ui/utilities/state/jotai/hooks/useSetAtomComponentFamilyState';
import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
import { useCallback, useEffect } from 'react';
import { isValidUuid } from 'twenty-shared/utils';
export const AgentChatAiSdkStreamEffect = () => {
export const AgentChatStreamSubscriptionEffect = () => {
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const agentChatFetchedMessages = useAtomComponentFamilyStateValue(
agentChatFetchedMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
const { createChatThread } = useCreateAgentChatThread();
@@ -45,47 +38,27 @@ export const AgentChatAiSdkStreamEffect = () => {
onBrowserEvent: ensureThreadExistsForDraft,
});
const onStreamingComplete = useCallback(() => {
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
}, []);
useAgentChat(ensureThreadIdForSend);
const chatState = useAgentChat(
agentChatFetchedMessages,
ensureThreadIdForSend,
onStreamingComplete,
const subscriptionThreadId =
currentAIChatThread !== null && isValidUuid(currentAIChatThread)
? currentAIChatThread
: null;
useAgentChatSubscription(subscriptionThreadId);
const agentChatFetchedMessages = useAtomComponentFamilyStateValue(
agentChatFetchedMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
// Attempt to resume an active stream when navigating to an existing
// thread. We call resumeStream() manually instead of using useChat's
// resume:true option so that the stop button can coexist with
// resumption (resume:true is incompatible with abort signals).
// Only resume when the thread already has fetched messages — this
// avoids resuming on newly created threads where the thread ID
// transitions from a placeholder to a real UUID mid-conversation.
useEffect(() => {
if (
currentAIChatThread === null ||
!isValidUuid(currentAIChatThread) ||
agentChatFetchedMessages.length === 0 ||
chatState.status === 'streaming' ||
chatState.status === 'submitted'
) {
return;
}
chatState.resumeStream();
// We intentionally omit chatState.resumeStream and status from deps
// to avoid resume loops. We do include agentChatFetchedMessages.length
// so that resume fires once messages are fetched (they may arrive
// after the thread ID is set).
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [currentAIChatThread, agentChatFetchedMessages.length]);
const setAgentChatMessages = useSetAtomComponentFamilyState(
agentChatMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
const agentChatIsStreaming = useAtomStateValue(agentChatIsStreamingState);
const agentChatDisplayedThread = useAtomStateValue(
agentChatDisplayedThreadState,
);
@@ -99,22 +72,21 @@ export const AgentChatAiSdkStreamEffect = () => {
);
useEffect(() => {
const mergedMessages = mergeAgentChatFetchedAndStreamingMessages(
agentChatFetchedMessages,
chatState.messages,
);
if (agentChatIsStreaming) {
return;
}
setAgentChatMessages(mergedMessages);
setAgentChatMessages(agentChatFetchedMessages);
if (currentAIChatThread !== agentChatDisplayedThread) {
if (mergedMessages.length > 0) {
if (agentChatFetchedMessages.length > 0) {
setAgentChatIsInitialScrollPendingOnThreadChange(true);
}
setAgentChatDisplayedThread(currentAIChatThread);
}
}, [
agentChatFetchedMessages,
chatState.messages,
agentChatIsStreaming,
setAgentChatMessages,
currentAIChatThread,
agentChatDisplayedThread,
@@ -130,31 +102,20 @@ export const AgentChatAiSdkStreamEffect = () => {
agentChatMessagesLoadingState,
);
useEffect(() => {
const handleLoadingChange = useCallback(() => {
const combinedIsLoading =
chatState.isLoading ||
agentChatMessagesLoading ||
agentChatThreadsLoading;
agentChatMessagesLoading || agentChatThreadsLoading;
setAgentChatIsLoading(combinedIsLoading);
}, [
chatState.isLoading,
agentChatMessagesLoading,
agentChatThreadsLoading,
setAgentChatIsLoading,
]);
const setAgentChatError = useSetAtomState(agentChatErrorState);
useEffect(() => {
setAgentChatError(normalizeAiSdkError(chatState.error));
}, [chatState.error, setAgentChatError]);
const setAgentChatIsStreaming = useSetAtomState(agentChatIsStreamingState);
useEffect(() => {
setAgentChatIsStreaming(chatState.status === 'streaming');
}, [chatState.status, setAgentChatIsStreaming]);
handleLoadingChange();
}, [handleLoadingChange]);
return null;
};
@@ -1,25 +1,24 @@
import { useStore } from 'jotai';
import { useCallback } from 'react';
import { useAtomValue, useStore } from 'jotai';
import { useEffect } from 'react';
import { isDefined } from 'twenty-shared/utils';
import { CHAT_THREADS_PAGE_SIZE } from '@/ai/constants/ChatThreads';
import {
AGENT_CHAT_NEW_THREAD_DRAFT_KEY,
agentChatDraftsByThreadIdState,
} from '@/ai/states/agentChatDraftsByThreadIdState';
import { agentChatInputState } from '@/ai/states/agentChatInputState';
import { agentChatThreadsLoadingState } from '@/ai/states/agentChatThreadsLoadingState';
import { agentChatThreadsSelector } from '@/ai/states/agentChatThreadsSelector';
import { agentChatUsageState } from '@/ai/states/agentChatUsageState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTitleState';
import { hasInitializedAgentChatThreadsState } from '@/ai/states/hasInitializedAgentChatThreadsState';
import { hasTriggeredCreateForDraftState } from '@/ai/states/hasTriggeredCreateForDraftState';
import { useQueryWithCallbacks } from '@/apollo/hooks/useQueryWithCallbacks';
import { metadataStoreState } from '@/metadata-store/states/metadataStoreState';
import { type FlatAgentChatThread } from '@/metadata-store/types/FlatAgentChatThread';
import { useAtomState } from '@/ui/utilities/state/jotai/hooks/useAtomState';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
import {
type GetChatThreadsQuery,
GetChatThreadsDocument,
} from '~/generated-metadata/graphql';
export const AgentChatThreadInitializationEffect = () => {
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
@@ -33,71 +32,82 @@ export const AgentChatThreadInitializationEffect = () => {
agentChatThreadsLoadingState,
);
const store = useStore();
const handleFirstLoad = useCallback(
(data: GetChatThreadsQuery) => {
const threads = data.chatThreads.edges.map((edge) => edge.node);
if (threads.length > 0) {
const firstThread = threads[0];
const draftForThread =
store.get(agentChatDraftsByThreadIdState.atom)[firstThread.id] ?? '';
setCurrentAIChatThread(firstThread.id);
setAgentChatInput(draftForThread);
setCurrentAIChatThreadTitle(firstThread.title ?? null);
const hasUsageData =
(firstThread.conversationSize ?? 0) > 0 &&
isDefined(firstThread.contextWindowTokens);
setAgentChatUsage(
hasUsageData
? {
lastMessage: null,
conversationSize: firstThread.conversationSize ?? 0,
contextWindowTokens: firstThread.contextWindowTokens ?? 0,
inputTokens: firstThread.totalInputTokens,
outputTokens: firstThread.totalOutputTokens,
inputCredits: firstThread.totalInputCredits,
outputCredits: firstThread.totalOutputCredits,
}
: null,
);
} else {
store.set(hasTriggeredCreateForDraftState.atom, false);
setCurrentAIChatThread(AGENT_CHAT_NEW_THREAD_DRAFT_KEY);
setAgentChatInput(
store.get(agentChatDraftsByThreadIdState.atom)[
AGENT_CHAT_NEW_THREAD_DRAFT_KEY
] ?? '',
);
setCurrentAIChatThreadTitle(null);
setAgentChatUsage(null);
}
},
[
setCurrentAIChatThread,
setAgentChatInput,
setCurrentAIChatThreadTitle,
setAgentChatUsage,
store,
],
const agentChatThreads = useAtomStateValue(agentChatThreadsSelector);
const storeEntry = useAtomValue(
metadataStoreState.atomFamily('agentChatThreads'),
);
const [hasInitializedAgentChatThreads, setHasInitializedAgentChatThreads] =
useAtomState(hasInitializedAgentChatThreadsState);
const handleLoadingChange = useCallback(
(loading: boolean) => {
setAgentChatThreadsLoading(loading);
},
[setAgentChatThreadsLoading],
);
useEffect(() => {
setAgentChatThreadsLoading(storeEntry.status === 'empty');
}, [storeEntry.status, setAgentChatThreadsLoading]);
useQueryWithCallbacks(GetChatThreadsDocument, {
variables: { paging: { first: CHAT_THREADS_PAGE_SIZE } },
skip: isDefined(currentAIChatThread),
onFirstLoad: handleFirstLoad,
onLoadingChange: handleLoadingChange,
});
useEffect(() => {
if (hasInitializedAgentChatThreads || isDefined(currentAIChatThread)) {
return;
}
if (storeEntry.status === 'empty') {
return;
}
setHasInitializedAgentChatThreads(true);
const sortedThreads = agentChatThreads.toSorted(
(a: FlatAgentChatThread, b: FlatAgentChatThread) =>
new Date(b.updatedAt).getTime() - new Date(a.updatedAt).getTime(),
);
if (sortedThreads.length > 0) {
const firstThread = sortedThreads[0];
const draftForThread =
store.get(agentChatDraftsByThreadIdState.atom)[firstThread.id] ?? '';
setCurrentAIChatThread(firstThread.id);
setAgentChatInput(draftForThread);
setCurrentAIChatThreadTitle(firstThread.title ?? null);
const hasUsageData =
(firstThread.conversationSize ?? 0) > 0 &&
isDefined(firstThread.contextWindowTokens);
setAgentChatUsage(
hasUsageData
? {
lastMessage: null,
conversationSize: firstThread.conversationSize ?? 0,
contextWindowTokens: firstThread.contextWindowTokens ?? 0,
inputTokens: firstThread.totalInputTokens,
outputTokens: firstThread.totalOutputTokens,
inputCredits: firstThread.totalInputCredits,
outputCredits: firstThread.totalOutputCredits,
}
: null,
);
} else {
store.set(hasTriggeredCreateForDraftState.atom, false);
setCurrentAIChatThread(AGENT_CHAT_NEW_THREAD_DRAFT_KEY);
setAgentChatInput(
store.get(agentChatDraftsByThreadIdState.atom)[
AGENT_CHAT_NEW_THREAD_DRAFT_KEY
] ?? '',
);
setCurrentAIChatThreadTitle(null);
setAgentChatUsage(null);
}
}, [
agentChatThreads,
currentAIChatThread,
hasInitializedAgentChatThreads,
setHasInitializedAgentChatThreads,
storeEntry.status,
setCurrentAIChatThread,
setAgentChatInput,
setCurrentAIChatThreadTitle,
setAgentChatUsage,
store,
]);
return null;
};
@@ -0,0 +1 @@
export const AGENT_CHAT_INSTANCE_ID = 'agentChatComponentInstance';
@@ -1 +0,0 @@
export const AGENT_CHAT_RETRY_EVENT_NAME = 'agent-chat-retry' as const;
@@ -0,0 +1,7 @@
import { gql } from '@apollo/client';
export const DELETE_QUEUED_CHAT_MESSAGE = gql`
mutation DeleteQueuedChatMessage($messageId: UUID!) {
deleteQueuedChatMessage(messageId: $messageId)
}
`;
@@ -0,0 +1,23 @@
import { gql } from '@apollo/client';
export const SEND_CHAT_MESSAGE = gql`
mutation SendChatMessage(
$threadId: UUID!
$text: String!
$messageId: UUID!
$browsingContext: JSON
$modelId: String
) {
sendChatMessage(
threadId: $threadId
text: $text
messageId: $messageId
browsingContext: $browsingContext
modelId: $modelId
) {
messageId
queued
streamId
}
}
`;
@@ -0,0 +1,7 @@
import { gql } from '@apollo/client';
export const STOP_AGENT_CHAT_STREAM = gql`
mutation StopAgentChatStream($threadId: UUID!) {
stopAgentChatStream(threadId: $threadId)
}
`;
@@ -7,6 +7,7 @@ export const GET_CHAT_MESSAGES = gql`
threadId
turnId
role
status
createdAt
parts {
id
@@ -37,5 +38,9 @@ export const GET_CHAT_MESSAGES = gql`
createdAt
}
}
chatStreamCatchupChunks(threadId: $threadId) {
chunks
maxSeq
}
}
`;
@@ -0,0 +1,10 @@
import { gql } from '@apollo/client';
export const ON_AGENT_CHAT_EVENT = gql`
subscription OnAgentChatEvent($threadId: UUID!) {
onAgentChatEvent(threadId: $threadId) {
threadId
event
}
}
`;
@@ -1,58 +1,44 @@
import { AGENT_CHAT_SEND_MESSAGE_EVENT_NAME } from '@/ai/constants/AgentChatSendMessageEventName';
import { useApolloClient } from '@apollo/client/react';
import { useStore } from 'jotai';
import { useCallback, useState } from 'react';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
import { isDefined, isValidUuid } from 'twenty-shared/utils';
import { v4 } from 'uuid';
import { useGetBrowsingContext } from '@/ai/hooks/useBrowsingContext';
import { agentChatSelectedFilesState } from '@/ai/states/agentChatSelectedFilesState';
import { agentChatUploadedFilesState } from '@/ai/states/agentChatUploadedFilesState';
import { agentChatUsageState } from '@/ai/states/agentChatUsageState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTitleState';
import { AGENT_CHAT_RETRY_EVENT_NAME } from '@/ai/constants/AgentChatRetryEventName';
import { AGENT_CHAT_INSTANCE_ID } from '@/ai/constants/AgentChatInstanceId';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { AGENT_CHAT_SEND_MESSAGE_EVENT_NAME } from '@/ai/constants/AgentChatSendMessageEventName';
import { AGENT_CHAT_STOP_EVENT_NAME } from '@/ai/constants/AgentChatStopEventName';
import { SEND_CHAT_MESSAGE } from '@/ai/graphql/mutations/sendChatMessage';
import { STOP_AGENT_CHAT_STREAM } from '@/ai/graphql/mutations/stopAgentChatStream';
import {
AGENT_CHAT_NEW_THREAD_DRAFT_KEY,
agentChatDraftsByThreadIdState,
} from '@/ai/states/agentChatDraftsByThreadIdState';
import { agentChatInputState } from '@/ai/states/agentChatInputState';
import { agentChatSelectedFilesState } from '@/ai/states/agentChatSelectedFilesState';
import { agentChatUploadedFilesState } from '@/ai/states/agentChatUploadedFilesState';
import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMessagesComponentFamilyState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { useGetBrowsingContext } from '@/ai/hooks/useBrowsingContext';
import { useAgentChatModelId } from '@/ai/hooks/useAgentChatModelId';
import { REST_API_BASE_URL } from '@/apollo/constant/rest-api-base-url';
import { getTokenPair } from '@/apollo/utils/getTokenPair';
import { renewToken } from '@/auth/services/AuthService';
import { tokenPairState } from '@/auth/states/tokenPairState';
import { useListenToBrowserEvent } from '@/browser-event/hooks/useListenToBrowserEvent';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { useAtomState } from '@/ui/utilities/state/jotai/hooks/useAtomState';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
import { useChat } from '@ai-sdk/react';
import { DefaultChatTransport } from 'ai';
import { useStore } from 'jotai';
import { useCallback, useMemo, useState } from 'react';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
import { isDefined, isValidUuid } from 'twenty-shared/utils';
import { REACT_APP_SERVER_BASE_URL } from '~/config';
export const useAgentChat = (
uiMessages: ExtendedUIMessage[],
ensureThreadIdForSend: () => Promise<string | null>,
onStreamingComplete?: () => void,
) => {
const setTokenPair = useSetAtomState(tokenPairState);
const setAgentChatUsage = useSetAtomState(agentChatUsageState);
const { modelIdForRequest } = useAgentChatModelId();
const { getBrowsingContext } = useGetBrowsingContext();
const setCurrentAIChatThreadTitle = useSetAtomState(
currentAIChatThreadTitleState,
);
const setCurrentAIChatThread = useSetAtomState(currentAIChatThreadState);
const apolloClient = useApolloClient();
const setCurrentAIChatThread = useSetAtomState(currentAIChatThreadState);
const store = useStore();
const agentChatSelectedFiles = useAtomStateValue(agentChatSelectedFilesState);
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const [, setPendingThreadIdAfterFirstSend] = useState<string | null>(null);
const [agentChatUploadedFiles, setAgentChatUploadedFiles] = useAtomState(
@@ -64,183 +50,6 @@ export const useAgentChat = (
agentChatDraftsByThreadIdState,
);
const retryFetchWithRenewedToken = async (
input: RequestInfo | URL,
init?: RequestInit,
) => {
const tokenPair = getTokenPair();
if (!isDefined(tokenPair)) {
return null;
}
try {
const renewedTokens = await renewToken(
`${REACT_APP_SERVER_BASE_URL}/metadata`,
tokenPair,
);
if (!isDefined(renewedTokens)) {
setTokenPair(null);
return null;
}
const renewedAccessToken =
renewedTokens.accessOrWorkspaceAgnosticToken?.token;
if (!isDefined(renewedAccessToken)) {
setTokenPair(null);
return null;
}
setTokenPair(renewedTokens);
const updatedHeaders = new Headers(init?.headers ?? {});
updatedHeaders.set('Authorization', `Bearer ${renewedAccessToken}`);
return fetch(input, {
...init,
headers: updatedHeaders,
});
} catch {
setTokenPair(null);
return null;
}
};
// eslint-disable-next-line react-hooks/exhaustive-deps
const transport = useMemo(
() =>
new DefaultChatTransport({
api: `${REST_API_BASE_URL}/agent-chat/stream`,
headers: () => ({
Authorization: `Bearer ${getTokenPair()?.accessOrWorkspaceAgnosticToken.token}`,
}),
prepareReconnectToStreamRequest: ({ id }) => ({
api: `${REST_API_BASE_URL}/agent-chat/${id}/stream`,
headers: {
Authorization: `Bearer ${getTokenPair()?.accessOrWorkspaceAgnosticToken.token}`,
},
}),
fetch: async (input, init) => {
const response = await fetch(input, init);
if (response.status === 401) {
const retriedResponse = await retryFetchWithRenewedToken(
input,
init,
);
return retriedResponse ?? response;
}
if (!response.ok) {
const errorBody = await response.json().catch(() => ({}));
const error = new Error(
errorBody.messages?.[0] ||
`Request failed with status ${response.status}`,
) as Error & { code?: string };
if (isDefined(errorBody.code)) {
error.code = errorBody.code;
}
throw error;
}
return response;
},
}),
// Intentionally created once — closures inside (getTokenPair, etc.)
// read fresh values via function references, not stale captures.
[],
);
const {
sendMessage,
messages,
status,
error,
regenerate,
stop,
resumeStream,
} = useChat({
transport,
messages: uiMessages,
id: currentAIChatThread ?? undefined,
experimental_throttle: 100,
onFinish: ({ message }) => {
type UsageMetadata = {
inputTokens: number;
outputTokens: number;
cachedInputTokens: number;
inputCredits: number;
outputCredits: number;
conversationSize: number;
};
type ModelMetadata = {
contextWindowTokens: number;
};
const metadata = message.metadata as
| { usage?: UsageMetadata; model?: ModelMetadata }
| undefined;
const usage = metadata?.usage;
const model = metadata?.model;
if (isDefined(usage) && isDefined(model)) {
setAgentChatUsage((prev) => ({
lastMessage: {
inputTokens: usage.inputTokens,
outputTokens: usage.outputTokens,
cachedInputTokens: usage.cachedInputTokens,
inputCredits: usage.inputCredits,
outputCredits: usage.outputCredits,
},
conversationSize: usage.conversationSize,
contextWindowTokens: model.contextWindowTokens,
inputTokens: (prev?.inputTokens ?? 0) + usage.inputTokens,
outputTokens: (prev?.outputTokens ?? 0) + usage.outputTokens,
inputCredits: (prev?.inputCredits ?? 0) + usage.inputCredits,
outputCredits: (prev?.outputCredits ?? 0) + usage.outputCredits,
}));
}
const titlePart = message.parts.find(
(part) => part.type === 'data-thread-title',
);
setPendingThreadIdAfterFirstSend((pendingId) => {
const threadIdForTitle =
pendingId ?? store.get(currentAIChatThreadState.atom);
if (isDefined(titlePart) && titlePart.type === 'data-thread-title') {
setCurrentAIChatThreadTitle(titlePart.data.title);
if (isDefined(threadIdForTitle)) {
const threadRef = apolloClient.cache.identify({
__typename: 'AgentChatThread',
id: threadIdForTitle,
});
if (isDefined(threadRef)) {
apolloClient.cache.modify({
id: threadRef,
fields: {
title: () => titlePart.data.title,
},
});
}
}
}
if (isDefined(pendingId)) {
setCurrentAIChatThread(pendingId);
}
return null;
});
onStreamingComplete?.();
},
});
const isStreaming = status === 'streaming';
const isLoading = isStreaming || agentChatSelectedFiles.length > 0;
const handleSendMessage = useCallback(async () => {
const draftKey =
store.get(currentAIChatThreadState.atom) ??
@@ -254,12 +63,19 @@ export const useAgentChat = (
).trim()
: store.get(agentChatInputState.atom).trim();
if (contentToSend === '' || isLoading) {
if (contentToSend === '') {
return;
}
const isLoading = agentChatSelectedFiles.length > 0;
if (isLoading) {
return;
}
const threadId = await ensureThreadIdForSend();
if (!threadId) {
if (!isDefined(threadId)) {
return;
}
@@ -274,35 +90,95 @@ export const useAgentChat = (
}));
const browsingContext = getBrowsingContext();
const messageId = v4();
sendMessage(
{
text: contentToSend,
files: agentChatUploadedFiles,
const optimisticUserMessage: ExtendedUIMessage = {
id: messageId,
role: 'user',
parts: [
{ type: 'text' as const, text: contentToSend },
...agentChatUploadedFiles,
],
metadata: {
createdAt: new Date().toISOString(),
},
{
body: {
threadId,
browsingContext,
...(isDefined(modelIdForRequest) && {
modelId: modelIdForRequest,
}),
},
},
);
status: 'sent',
};
const messagesAtom = agentChatMessagesComponentFamilyState.atomFamily({
instanceId: AGENT_CHAT_INSTANCE_ID,
familyKey: { threadId },
});
const currentMessages = store.get(messagesAtom);
store.set(messagesAtom, [...currentMessages, optimisticUserMessage]);
setAgentChatUploadedFiles([]);
try {
const { data } = await apolloClient.mutate<{
sendChatMessage: {
messageId: string;
queued: boolean;
streamId?: string;
};
}>({
mutation: SEND_CHAT_MESSAGE,
variables: {
threadId,
text: contentToSend,
messageId,
browsingContext: browsingContext ?? null,
modelId: modelIdForRequest ?? undefined,
},
});
if (data?.sendChatMessage?.queued) {
const latestMessages = store.get(messagesAtom);
store.set(
messagesAtom,
latestMessages.filter((message) => message.id !== messageId),
);
}
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
setPendingThreadIdAfterFirstSend((pendingId) => {
if (isDefined(pendingId)) {
setCurrentAIChatThread(pendingId);
}
return null;
});
} catch {
setAgentChatInput(contentToSend);
setAgentChatDraftsByThreadId((prev) => ({
...prev,
[draftKey]: contentToSend,
}));
const latestMessages = store.get(messagesAtom);
store.set(
messagesAtom,
latestMessages.filter((message) => message.id !== messageId),
);
}
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [
store,
isLoading,
agentChatSelectedFiles,
ensureThreadIdForSend,
setAgentChatInput,
getBrowsingContext,
sendMessage,
agentChatUploadedFiles,
setAgentChatUploadedFiles,
setAgentChatDraftsByThreadId,
modelIdForRequest,
setCurrentAIChatThread,
apolloClient,
]);
useListenToBrowserEvent({
@@ -311,45 +187,27 @@ export const useAgentChat = (
});
const handleStop = useCallback(async () => {
stop();
const threadId = store.get(currentAIChatThreadState.atom);
if (!isDefined(threadId) || !isValidUuid(threadId)) {
return;
}
const tokenPair = getTokenPair();
if (!isDefined(tokenPair)) {
return;
}
fetch(`${REST_API_BASE_URL}/agent-chat/${threadId}/stream`, {
method: 'DELETE',
headers: {
Authorization: `Bearer ${tokenPair.accessOrWorkspaceAgnosticToken.token}`,
},
}).catch(() => {});
}, [stop, store]);
apolloClient
.mutate({
mutation: STOP_AGENT_CHAT_STREAM,
variables: { threadId },
})
.catch(() => {});
}, [store, apolloClient]);
useListenToBrowserEvent({
eventName: AGENT_CHAT_STOP_EVENT_NAME,
onBrowserEvent: handleStop,
});
useListenToBrowserEvent({
eventName: AGENT_CHAT_RETRY_EVENT_NAME,
onBrowserEvent: regenerate,
});
return {
messages,
handleSendMessage,
handleStop,
resumeStream,
isLoading,
error,
status,
};
};
@@ -0,0 +1,358 @@
import { useCallback, useEffect } from 'react';
import { readUIMessageStream, type UIMessageChunk } from 'ai';
import { print, type ExecutionResult } from 'graphql';
import { useStore } from 'jotai';
import {
type AgentChatSubscriptionEvent,
type ExtendedUIMessage,
} from 'twenty-shared/ai';
import { isDefined } from 'twenty-shared/utils';
import { v4 } from 'uuid';
import { AGENT_CHAT_INSTANCE_ID } from '@/ai/constants/AgentChatInstanceId';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { ON_AGENT_CHAT_EVENT } from '@/ai/graphql/subscriptions/OnAgentChatEvent';
import { agentChatErrorState } from '@/ai/states/agentChatErrorState';
import { agentChatFirstLiveSeqState } from '@/ai/states/agentChatFirstLiveSeqState';
import { agentChatHandleEventCallbackState } from '@/ai/states/agentChatHandleEventCallbackState';
import { agentChatIsStreamingState } from '@/ai/states/agentChatIsStreamingState';
import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMessagesComponentFamilyState';
import { agentChatStreamWriterState } from '@/ai/states/agentChatStreamWriterState';
import { agentChatSubscriptionDisposeState } from '@/ai/states/agentChatSubscriptionDisposeState';
import { agentChatUsageState } from '@/ai/states/agentChatUsageState';
import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTitleState';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { sseClientState } from '@/sse-db-event/states/sseClientState';
const THROTTLE_MS = 100;
// readUIMessageStream requires initialization chunks (start, start-step,
// text-start) before content chunks. When reconnecting to a thread mid-stream,
// those chunks were already sent before we subscribed. This adapter injects
// synthetic initialization chunks so the reader can process mid-stream content.
const createMidStreamAdapter = () => {
let hasSeenStart = false;
const knownTextPartIds = new Set<string>();
const knownReasoningPartIds = new Set<string>();
const knownToolCallIds = new Set<string>();
return new TransformStream<UIMessageChunk, UIMessageChunk>({
transform(chunk, controller) {
if (!hasSeenStart) {
hasSeenStart = true;
if (chunk.type !== 'start') {
controller.enqueue({ type: 'start', messageId: v4() });
controller.enqueue({ type: 'start-step' });
}
}
if (chunk.type === 'text-start') {
knownTextPartIds.add(chunk.id);
} else if (
(chunk.type === 'text-delta' || chunk.type === 'text-end') &&
!knownTextPartIds.has(chunk.id)
) {
controller.enqueue({ type: 'text-start', id: chunk.id });
knownTextPartIds.add(chunk.id);
}
if (chunk.type === 'reasoning-start') {
knownReasoningPartIds.add(chunk.id);
} else if (
(chunk.type === 'reasoning-delta' || chunk.type === 'reasoning-end') &&
!knownReasoningPartIds.has(chunk.id)
) {
controller.enqueue({ type: 'reasoning-start', id: chunk.id });
knownReasoningPartIds.add(chunk.id);
}
if (chunk.type === 'tool-input-start') {
knownToolCallIds.add(chunk.toolCallId);
} else if (
chunk.type === 'tool-input-delta' &&
!knownToolCallIds.has(chunk.toolCallId)
) {
controller.enqueue({
type: 'tool-input-start',
toolCallId: chunk.toolCallId,
toolName: 'unknown',
});
knownToolCallIds.add(chunk.toolCallId);
}
controller.enqueue(chunk);
},
});
};
type AgentChatEventPayload = {
onAgentChatEvent: {
threadId: string;
event: AgentChatSubscriptionEvent;
};
};
export const useAgentChatSubscription = (threadId: string | null) => {
const store = useStore();
const sseClient = useAtomStateValue(sseClientState);
const cleanup = useCallback(() => {
const writer = store.get(agentChatStreamWriterState.atom);
if (isDefined(writer)) {
writer.close().catch(() => {});
store.set(agentChatStreamWriterState.atom, null);
}
const dispose = store.get(agentChatSubscriptionDisposeState.atom);
if (isDefined(dispose)) {
dispose();
store.set(agentChatSubscriptionDisposeState.atom, null);
}
if (store.get(agentChatIsStreamingState.atom)) {
store.set(agentChatIsStreamingState.atom, false);
}
}, [store]);
useEffect(() => {
if (!isDefined(threadId)) {
cleanup();
return;
}
if (!isDefined(sseClient)) {
return;
}
let bridge: TransformStream<UIMessageChunk> | null = null;
let throttleTimer: ReturnType<typeof setTimeout> | null = null;
let latestMessage: ExtendedUIMessage | null = null;
store.set(agentChatFirstLiveSeqState.atom, null);
const flushToAtom = () => {
const messageToFlush = latestMessage;
if (!isDefined(messageToFlush) || !isDefined(threadId)) {
return;
}
const atomKey = {
instanceId: AGENT_CHAT_INSTANCE_ID,
familyKey: { threadId },
};
const currentMessages = store.get(
agentChatMessagesComponentFamilyState.atomFamily(atomKey),
);
const streamingMsgIndex = currentMessages.findIndex(
(message) => message.id === messageToFlush.id,
);
if (streamingMsgIndex >= 0) {
const updatedMessages = [...currentMessages];
updatedMessages[streamingMsgIndex] = messageToFlush;
store.set(
agentChatMessagesComponentFamilyState.atomFamily(atomKey),
updatedMessages,
);
} else {
store.set(agentChatMessagesComponentFamilyState.atomFamily(atomKey), [
...currentMessages,
messageToFlush,
]);
}
};
const scheduleAtomUpdate = (message: ExtendedUIMessage) => {
latestMessage = message;
if (!isDefined(throttleTimer)) {
flushToAtom();
throttleTimer = setTimeout(() => {
throttleTimer = null;
flushToAtom();
}, THROTTLE_MS);
}
};
const startReadLoop = async (readable: ReadableStream<UIMessageChunk>) => {
const messageStream = readUIMessageStream({ stream: readable });
for await (const message of messageStream) {
const extendedMessage = message as ExtendedUIMessage;
const titlePart = extendedMessage.parts.find(
(part) => part.type === 'data-thread-title',
);
if (isDefined(titlePart) && titlePart.type === 'data-thread-title') {
store.set(currentAIChatThreadTitleState.atom, titlePart.data.title);
}
const metadata = extendedMessage.metadata as
| {
usage?: {
inputTokens: number;
outputTokens: number;
cachedInputTokens: number;
inputCredits: number;
outputCredits: number;
conversationSize: number;
};
model?: {
contextWindowTokens: number;
};
}
| undefined;
if (isDefined(metadata?.usage) && isDefined(metadata?.model)) {
const usage = metadata.usage;
const model = metadata.model;
store.set(agentChatUsageState.atom, (prev) => ({
lastMessage: {
inputTokens: usage.inputTokens,
outputTokens: usage.outputTokens,
cachedInputTokens: usage.cachedInputTokens,
inputCredits: usage.inputCredits,
outputCredits: usage.outputCredits,
},
conversationSize: usage.conversationSize,
contextWindowTokens: model.contextWindowTokens,
inputTokens: (prev?.inputTokens ?? 0) + usage.inputTokens,
outputTokens: (prev?.outputTokens ?? 0) + usage.outputTokens,
inputCredits: (prev?.inputCredits ?? 0) + usage.inputCredits,
outputCredits: (prev?.outputCredits ?? 0) + usage.outputCredits,
}));
}
scheduleAtomUpdate(extendedMessage);
}
if (isDefined(throttleTimer)) {
clearTimeout(throttleTimer);
throttleTimer = null;
}
flushToAtom();
store.set(agentChatIsStreamingState.atom, false);
};
const handleEvent = (event: AgentChatSubscriptionEvent) => {
switch (event.type) {
case 'stream-chunk': {
if (
isDefined(event.seq) &&
store.get(agentChatFirstLiveSeqState.atom) === null
) {
store.set(agentChatFirstLiveSeqState.atom, event.seq);
}
if (!store.get(agentChatIsStreamingState.atom)) {
store.set(agentChatIsStreamingState.atom, true);
bridge = new TransformStream<UIMessageChunk>();
store.set(
agentChatStreamWriterState.atom,
bridge.writable.getWriter(),
);
const adaptedReadable = bridge.readable.pipeThrough(
createMidStreamAdapter(),
);
startReadLoop(adaptedReadable).catch(() => {
store.set(agentChatIsStreamingState.atom, false);
});
}
const writer = store.get(agentChatStreamWriterState.atom);
if (isDefined(writer)) {
writer.write(event.chunk as UIMessageChunk).catch(() => {});
}
break;
}
case 'message-persisted': {
const writer = store.get(agentChatStreamWriterState.atom);
if (isDefined(writer)) {
writer.close().catch(() => {});
store.set(agentChatStreamWriterState.atom, null);
}
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
break;
}
case 'queue-updated': {
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
break;
}
case 'stream-error': {
const streamError = new Error(event.message) as Error & {
code?: string;
};
streamError.code = event.code;
store.set(agentChatErrorState.atom, streamError);
const writer = store.get(agentChatStreamWriterState.atom);
if (isDefined(writer)) {
writer.close().catch(() => {});
store.set(agentChatStreamWriterState.atom, null);
}
store.set(agentChatIsStreamingState.atom, false);
break;
}
}
};
store.set(agentChatHandleEventCallbackState.atom, () => handleEvent);
const dispose = sseClient.subscribe<AgentChatEventPayload>(
{
query: print(ON_AGENT_CHAT_EVENT),
variables: { threadId },
},
{
next: (value: ExecutionResult<AgentChatEventPayload>) => {
if (isDefined(value.data?.onAgentChatEvent?.event)) {
handleEvent(
value.data.onAgentChatEvent.event as AgentChatSubscriptionEvent,
);
}
},
error: () => {
// graphql-sse handles reconnection automatically
},
complete: () => {
cleanup();
},
},
);
store.set(agentChatSubscriptionDisposeState.atom, () => dispose);
return () => {
store.set(agentChatHandleEventCallbackState.atom, null);
if (isDefined(throttleTimer)) {
clearTimeout(throttleTimer);
}
cleanup();
};
}, [threadId, sseClient, store, cleanup]);
};
@@ -1,86 +1,29 @@
import { useCallback, useEffect, useState } from 'react';
import { useInView } from 'react-intersection-observer';
import { useAtomValue } from 'jotai';
import { useMemo } from 'react';
import { isDefined } from 'twenty-shared/utils';
import { CHAT_THREADS_PAGE_SIZE } from '@/ai/constants/ChatThreads';
import { useQuery } from '@apollo/client/react';
import { GetChatThreadsDocument } from '~/generated-metadata/graphql';
const FETCH_MORE_ROOT_MARGIN = '200px';
import { agentChatThreadsSelector } from '@/ai/states/agentChatThreadsSelector';
import { metadataStoreState } from '@/metadata-store/states/metadataStoreState';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
export const useChatThreads = () => {
const [shouldFetchMore, setShouldFetchMore] = useState(false);
const { data, loading, fetchMore } = useQuery(GetChatThreadsDocument, {
variables: {
paging: { first: CHAT_THREADS_PAGE_SIZE },
},
});
const agentChatThreads = useAtomStateValue(agentChatThreadsSelector);
const storeEntry = useAtomValue(
metadataStoreState.atomFamily('agentChatThreads'),
);
useEffect(() => {
if (data) {
setShouldFetchMore(false);
}
}, [data]);
const edges = data?.chatThreads?.edges ?? [];
const threads = edges.map((edge) => edge.node);
const pageInfo = data?.chatThreads?.pageInfo;
const endCursor = pageInfo?.endCursor ?? undefined;
const hasNextPage = pageInfo?.hasNextPage ?? false;
const { ref: fetchMoreRef, inView } = useInView({
rootMargin: FETCH_MORE_ROOT_MARGIN,
});
const loadMore = useCallback(() => {
if (!hasNextPage || loading || !endCursor) {
return;
}
return fetchMore({
variables: {
paging: {
first: CHAT_THREADS_PAGE_SIZE,
after: endCursor,
},
},
updateQuery: (previousResult, { fetchMoreResult }) => {
const newEdges = fetchMoreResult?.chatThreads?.edges ?? [];
if (newEdges.length === 0) {
return previousResult;
}
return {
chatThreads: {
...fetchMoreResult.chatThreads,
edges: [...(previousResult.chatThreads?.edges ?? []), ...newEdges],
pageInfo:
fetchMoreResult.chatThreads?.pageInfo ??
previousResult.chatThreads?.pageInfo,
},
};
},
});
}, [hasNextPage, loading, endCursor, fetchMore]);
useEffect(() => {
if (inView && hasNextPage && !loading && !shouldFetchMore) {
setShouldFetchMore(true);
const promise = loadMore();
if (isDefined(promise)) {
promise.finally(() => setShouldFetchMore(false));
} else {
setShouldFetchMore(false);
}
}
}, [inView, hasNextPage, loading, shouldFetchMore, loadMore]);
const threads = useMemo(
() =>
[...agentChatThreads].sort(
(a, b) =>
new Date(b.updatedAt).getTime() - new Date(a.updatedAt).getTime(),
),
[agentChatThreads],
);
return {
threads,
hasNextPage,
loading,
fetchMoreRef,
hasNextPage: false,
loading: storeEntry.status === 'empty',
fetchMoreRef: undefined,
};
};
@@ -14,14 +14,12 @@ import { isCreatingChatThreadState } from '@/ai/states/isCreatingChatThreadState
import { isCreatingForFirstSendState } from '@/ai/states/isCreatingForFirstSendState';
import { skipMessagesSkeletonUntilLoadedState } from '@/ai/states/skipMessagesSkeletonUntilLoadedState';
import { threadIdCreatedFromDraftState } from '@/ai/states/threadIdCreatedFromDraftState';
import { useUpdateMetadataStoreDraft } from '@/metadata-store/hooks/useUpdateMetadataStoreDraft';
import { type FlatAgentChatThread } from '@/metadata-store/types/FlatAgentChatThread';
import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
import { useMutation } from '@apollo/client/react';
import {
CreateChatThreadDocument,
GetChatThreadsDocument,
} from '~/generated-metadata/graphql';
import { getOperationName } from '~/utils/getOperationName';
import { CreateChatThreadDocument } from '~/generated-metadata/graphql';
export const useCreateAgentChatThread = () => {
const setCurrentAIChatThread = useSetAtomState(currentAIChatThreadState);
@@ -35,9 +33,26 @@ export const useCreateAgentChatThread = () => {
agentChatDraftsByThreadIdState,
);
const store = useStore();
const { addToDraft, applyChanges } = useUpdateMetadataStoreDraft();
const [createChatThread] = useMutation(CreateChatThreadDocument, {
onCompleted: (data) => {
const newThread: FlatAgentChatThread = {
id: data.createChatThread.id,
title: data.createChatThread.title ?? null,
createdAt: data.createChatThread.createdAt,
updatedAt: data.createChatThread.updatedAt,
conversationSize: 0,
contextWindowTokens: null,
totalInputTokens: 0,
totalOutputTokens: 0,
totalInputCredits: 0,
totalOutputCredits: 0,
};
addToDraft({ key: 'agentChatThreads', items: [newThread] });
applyChanges();
if (store.get(isCreatingForFirstSendState.atom)) {
store.set(isCreatingForFirstSendState.atom, false);
setIsCreatingChatThread(false);
@@ -80,9 +95,6 @@ export const useCreateAgentChatThread = () => {
store.set(isCreatingForFirstSendState.atom, false);
store.set(hasTriggeredCreateForDraftState.atom, false);
},
refetchQueries: [
getOperationName(GetChatThreadsDocument) ?? 'GetChatThreads',
],
});
return { createChatThread };
@@ -0,0 +1,28 @@
import { useApolloClient } from '@apollo/client/react';
import { useCallback } from 'react';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { DELETE_QUEUED_CHAT_MESSAGE } from '@/ai/graphql/mutations/deleteQueuedChatMessage';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
export const useDeleteQueuedMessage = () => {
const apolloClient = useApolloClient();
const deleteQueuedMessage = useCallback(
async (messageId: string) => {
const { data } = await apolloClient.mutate<{
deleteQueuedChatMessage: boolean;
}>({
mutation: DELETE_QUEUED_CHAT_MESSAGE,
variables: { messageId },
});
if (data?.deleteQueuedChatMessage) {
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
}
},
[apolloClient],
);
return { deleteQueuedMessage };
};
@@ -0,0 +1,6 @@
import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
export const agentChatFirstLiveSeqState = createAtomState<number | null>({
key: 'agentChatFirstLiveSeqState',
defaultValue: null,
});
@@ -0,0 +1,10 @@
import { type AgentChatSubscriptionEvent } from 'twenty-shared/ai';
import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
export const agentChatHandleEventCallbackState = createAtomState<
((event: AgentChatSubscriptionEvent) => void) | null
>({
key: 'agentChatHandleEventCallbackState',
defaultValue: null,
});
@@ -0,0 +1,10 @@
import { AgentChatComponentInstanceContext } from '@/ai/states/AgentChatComponentInstanceContext';
import { createAtomComponentFamilyState } from '@/ui/utilities/state/jotai/utils/createAtomComponentFamilyState';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
export const agentChatQueuedMessagesComponentFamilyState =
createAtomComponentFamilyState<ExtendedUIMessage[], { threadId: string }>({
key: 'agentChatQueuedMessagesComponentFamilyState',
defaultValue: [],
componentInstanceContext: AgentChatComponentInstanceContext,
});
@@ -0,0 +1,9 @@
import { type UIMessageChunk } from 'ai';
import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
export const agentChatStreamWriterState =
createAtomState<WritableStreamDefaultWriter<UIMessageChunk> | null>({
key: 'agentChatStreamWriterState',
defaultValue: null,
});
@@ -0,0 +1,8 @@
import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
export const agentChatSubscriptionDisposeState = createAtomState<
(() => void) | null
>({
key: 'agentChatSubscriptionDisposeState',
defaultValue: null,
});
@@ -0,0 +1,14 @@
import { metadataStoreState } from '@/metadata-store/states/metadataStoreState';
import { type FlatAgentChatThread } from '@/metadata-store/types/FlatAgentChatThread';
import { createAtomSelector } from '@/ui/utilities/state/jotai/utils/createAtomSelector';
export const agentChatThreadsSelector = createAtomSelector<
FlatAgentChatThread[]
>({
key: 'agentChatThreadsSelector',
get: ({ get }) => {
const storeItem = get(metadataStoreState, 'agentChatThreads');
return storeItem.current as FlatAgentChatThread[];
},
});
@@ -0,0 +1,6 @@
import { createAtomState } from '@/ui/utilities/state/jotai/utils/createAtomState';
export const hasInitializedAgentChatThreadsState = createAtomState<boolean>({
key: 'hasInitializedAgentChatThreadsState',
defaultValue: false,
});
@@ -1,60 +0,0 @@
import { normalizeAiSdkError } from '@/ai/utils/normalizeAiSdkError';
describe('normalizeAiSdkError', () => {
it('should return undefined for undefined input', () => {
expect(normalizeAiSdkError(undefined)).toBeUndefined();
});
it('should return the error unchanged if it already has a code', () => {
const error = new Error('test') as Error & { code: string };
error.code = 'EXISTING_CODE';
const result = normalizeAiSdkError(error);
expect(result).toBe(error);
expect((result as Error & { code: string }).code).toBe('EXISTING_CODE');
});
it('should parse JSON message and attach code from response body', () => {
const error = new Error(
'{"statusCode":402,"error":"Error","messages":["Credits exhausted"],"code":"BILLING_CREDITS_EXHAUSTED"}',
);
const result = normalizeAiSdkError(error);
expect(result).not.toBe(error);
expect((result as Error & { code: string }).code).toBe(
'BILLING_CREDITS_EXHAUSTED',
);
expect(result?.message).toBe(error.message);
});
it('should return original error for non-JSON message', () => {
const error = new Error('Something went wrong');
const result = normalizeAiSdkError(error);
expect(result).toBe(error);
});
it('should return original error for JSON message without code', () => {
const error = new Error(
'{"statusCode":500,"error":"Internal Server Error"}',
);
const result = normalizeAiSdkError(error);
expect(result).toBe(error);
});
it('should preserve the original stack trace', () => {
const error = new Error(
'{"statusCode":402,"code":"BILLING_CREDITS_EXHAUSTED"}',
);
const originalStack = error.stack;
const result = normalizeAiSdkError(error);
expect(result?.stack).toBe(originalStack);
});
});
@@ -8,6 +8,7 @@ export const mapDBMessagesToUIMessages = (
return dbMessages.map((dbMessage) => ({
id: dbMessage.id,
role: dbMessage.role as ExtendedUIMessage['role'],
status: dbMessage.status as 'queued' | 'sent',
parts: dbMessage.parts.map(mapDBPartToUIMessagePart),
metadata: {
createdAt: dbMessage.createdAt,
@@ -5,6 +5,10 @@ import {
} from 'twenty-shared/ai';
import { type AgentMessagePart } from '~/generated-metadata/graphql';
// Maps GraphQL DTO fields to UI message parts.
// A parallel mapping for TypeORM entities exists in the server at:
// packages/twenty-server/src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartsToUIMessageParts.ts
export const mapDBPartToUIMessagePart = (
part: AgentMessagePart,
): ExtendedUIMessagePart => {
@@ -1,15 +0,0 @@
import { type ExtendedUIMessage } from 'twenty-shared/ai';
// The SDK's messages array is always [...seedMessages, ...newMessages]
// where seedMessages are the fetchedMessages we passed to useChat.
// SDK-generated IDs differ from server-persisted IDs, so comparing
// by ID doesn't work. Instead, slice by position: everything beyond
// the fetched count is a new streaming message.
export const mergeAgentChatFetchedAndStreamingMessages = (
fetchedMessages: ExtendedUIMessage[],
streamingMessages: ExtendedUIMessage[],
): ExtendedUIMessage[] => {
const newStreamingMessages = streamingMessages.slice(fetchedMessages.length);
return [...fetchedMessages, ...newStreamingMessages];
};
@@ -1,43 +0,0 @@
import { isDefined } from 'twenty-shared/utils';
// The Vercel AI SDK wraps non-200 responses into an Error whose message
// is the raw JSON response body, losing any custom properties (like `code`).
// This function parses that JSON and re-attaches the code so downstream
// consumers (extractErrorCode) can find it without knowing about the SDK.
export const normalizeAiSdkError = (
error: Error | undefined,
): Error | undefined => {
if (!isDefined(error) || !(error instanceof Error)) {
return error;
}
if (
'code' in error &&
typeof (error as Error & { code: string }).code === 'string'
) {
return error;
}
try {
const parsed: unknown = JSON.parse(error.message);
if (
isDefined(parsed) &&
typeof parsed === 'object' &&
'code' in parsed &&
typeof (parsed as { code: unknown }).code === 'string'
) {
const normalizedError = new Error(error.message) as Error & {
code: string;
};
normalizedError.code = (parsed as { code: string }).code;
normalizedError.stack = error.stack;
return normalizedError;
}
} catch {
// message is not JSON — nothing to normalize
}
return error;
};
@@ -1,8 +1,8 @@
import { METADATA_OPERATION_BROWSER_EVENT_NAME } from '@/browser-event/constants/MetadataOperationBrowserEventName';
import { type BroadcastEntityName } from '@/browser-event/types/BroadcastEntityName';
import { type MetadataOperation } from '@/browser-event/types/MetadataOperation';
import { type MetadataOperationBrowserEventDetail } from '@/browser-event/types/MetadataOperationBrowserEventDetail';
import { useEffect } from 'react';
import { type AllMetadataName } from 'twenty-shared/metadata';
import { isDefined, isNonEmptyArray } from 'twenty-shared/utils';
export const useListenToMetadataOperationBrowserEvent = <
@@ -15,7 +15,7 @@ export const useListenToMetadataOperationBrowserEvent = <
onMetadataOperationBrowserEvent: (
detail: MetadataOperationBrowserEventDetail<T>,
) => void;
metadataName?: AllMetadataName;
metadataName?: BroadcastEntityName;
operationTypes?: MetadataOperation<T>['type'][];
}) => {
useEffect(() => {
@@ -0,0 +1,15 @@
import { type AllMetadataName } from 'twenty-shared/metadata';
// Non-syncable entities that broadcast real-time events through the SSE
// event stream via WorkspaceEventBroadcaster, but are NOT part of the
// workspace migration / syncable entity system (ALL_METADATA_NAME).
const ALL_NON_SYNCABLE_BROADCAST_ENTITY_NAME = {
agentChatThread: 'agentChatThread',
} as const;
type NonSyncableBroadcastEntityName =
keyof typeof ALL_NON_SYNCABLE_BROADCAST_ENTITY_NAME;
export type BroadcastEntityName =
| AllMetadataName
| NonSyncableBroadcastEntityName;
@@ -1,10 +1,19 @@
import { type MetadataOperation } from '@/browser-event/types/MetadataOperation';
import { type AllMetadataName } from 'twenty-shared/metadata';
import { type BroadcastEntityName } from '@/browser-event/types/BroadcastEntityName';
// TODO: The browser event layer currently uses "Metadata" naming throughout
// (MetadataOperation, MetadataOperationBrowserEventDetail, etc.) because it
// originally only carried events for syncable metadata entities (views, fields,
// etc.). With WorkspaceEventBroadcaster, non-syncable entities like
// agentChatThread also broadcast through the same SSE stream.
// BroadcastEntityName captures this distinction at the type level, but the
// naming of the surrounding types/hooks/files still says "Metadata." A future
// rename of this layer (e.g. Metadata* → Entity* or Broadcast*) would make
// the naming consistent with the actual scope.
export type MetadataOperationBrowserEventDetail<
T extends Record<string, unknown>,
> = {
metadataName: AllMetadataName;
metadataName: BroadcastEntityName;
operation: MetadataOperation<T>;
updatedCollectionHash?: string;
};
@@ -6,17 +6,20 @@ import { splitViewWithRelated } from '@/metadata-store/utils/splitViewWithRelate
import { FIND_MANY_OBJECT_METADATA_ITEMS } from '@/object-metadata/graphql/queries';
import { transformPageLayout } from '@/page-layout/utils/transformPageLayout';
import { logicFunctionsState } from '@/settings/logic-functions/states/logicFunctionsState';
import { useIsFeatureEnabled } from '@/workspace/hooks/useIsFeatureEnabled';
import { useApolloClient } from '@apollo/client/react';
import { useStore } from 'jotai';
import { useCallback } from 'react';
import { isDefined } from 'twenty-shared/utils';
import {
FeatureFlagKey,
FindAllViewsDocument,
FindManyCommandMenuItemsDocument,
FindAllRecordPageLayoutsDocument,
FindFieldsWidgetViewsDocument,
FindManyLogicFunctionsDocument,
FindManyNavigationMenuItemsDocument,
GetChatThreadsDocument,
type ObjectMetadataItemsQuery,
ViewType,
} from '~/generated-metadata/graphql';
@@ -55,6 +58,7 @@ export const useLoadStaleMetadataEntities = () => {
const client = useApolloClient();
const store = useStore();
const { replaceDraft, applyChanges } = useUpdateMetadataStoreDraft();
const isAiEnabled = useIsFeatureEnabled(FeatureFlagKey.IS_AI_ENABLED);
const loadStaleMetadataEntities = useCallback(
async (staleEntityKeys: MetadataEntityKey[]) => {
@@ -211,10 +215,32 @@ export const useLoadStaleMetadataEntities = () => {
);
}
if (staleEntityKeys.includes('agentChatThreads') && isAiEnabled) {
fetchPromises.push(
client
.query({
query: GetChatThreadsDocument,
variables: { paging: { first: 500 } },
fetchPolicy: 'network-only',
})
.then((result) => {
if (!isDefined(result.data?.chatThreads?.edges)) {
return;
}
const threads = result.data.chatThreads.edges.map(
(edge) => edge.node,
);
replaceDraft('agentChatThreads', threads);
}),
);
}
await Promise.all(fetchPromises);
applyChanges();
},
[client, store, replaceDraft, applyChanges],
[client, store, replaceDraft, applyChanges, isAiEnabled],
);
return { loadStaleMetadataEntities };
@@ -30,6 +30,9 @@ export const ALL_METADATA_ENTITY_KEYS = [
'skills',
'rowLevelPermissionPredicates',
'rowLevelPermissionPredicateGroups',
// TODO: clarify what really is metadata (syncable entity?)
// vs 'core engine entity' or 'broadcastable entity'
'agentChatThreads',
] as const;
export type MetadataEntityKey = (typeof ALL_METADATA_ENTITY_KEYS)[number];
@@ -0,0 +1,3 @@
import { type AgentChatThread } from '~/generated-metadata/graphql';
export type FlatAgentChatThread = AgentChatThread;
@@ -1,4 +1,5 @@
import { type FlatAgent } from '@/metadata-store/types/FlatAgent';
import { type FlatAgentChatThread } from '@/metadata-store/types/FlatAgentChatThread';
import { type FlatCommandMenuItem } from '@/metadata-store/types/FlatCommandMenuItem';
import { type FlatFieldMetadataItem } from '@/metadata-store/types/FlatFieldMetadataItem';
import { type FlatFrontComponent } from '@/metadata-store/types/FlatFrontComponent';
@@ -48,4 +49,5 @@ export type MetadataEntityTypeMap = {
skills: FlatSkill;
rowLevelPermissionPredicates: FlatRowLevelPermissionPredicate;
rowLevelPermissionPredicateGroups: FlatRowLevelPermissionPredicateGroup;
agentChatThreads: FlatAgentChatThread;
};
@@ -25,6 +25,7 @@ const METADATA_NAME_TO_ENTITY_KEY: Record<string, MetadataEntityKey> = {
navigationMenuItem: 'navigationMenuItems',
frontComponent: 'frontComponents',
webhook: 'webhooks',
agentChatThread: 'agentChatThreads',
};
export const mapAllMetadataNameToEntityKey = (
@@ -1,25 +1,23 @@
import { type BroadcastEntityName } from '@/browser-event/types/BroadcastEntityName';
import { dispatchMetadataOperationBrowserEvent } from '@/browser-event/utils/dispatchMetadataOperationBrowserEvent';
import { turnSseMetadataEventsToMetadataOperationBrowserEvents } from '@/sse-db-event/utils/turnSseMetadataEventsToMetadataOperationBrowserEvents';
import { useCallback } from 'react';
import {
type AllMetadataName,
type MetadataEvent,
} from '~/generated-metadata/graphql';
import { type MetadataEvent } from '~/generated-metadata/graphql';
const groupSseMetadataEventsByMetadataName = (
const groupSseMetadataEventsByEntityName = (
sseMetadataEvents: MetadataEvent[],
): Map<AllMetadataName, MetadataEvent[]> => {
const eventsByMetadataName = new Map<AllMetadataName, MetadataEvent[]>();
): Map<BroadcastEntityName, MetadataEvent[]> => {
const eventsByEntityName = new Map<BroadcastEntityName, MetadataEvent[]>();
for (const event of sseMetadataEvents) {
const metadataName = event.metadataName as AllMetadataName;
const entityName = event.metadataName as BroadcastEntityName;
const existing = eventsByMetadataName.get(metadataName) ?? [];
const existing = eventsByEntityName.get(entityName) ?? [];
eventsByMetadataName.set(metadataName, [...existing, event]);
eventsByEntityName.set(entityName, [...existing, event]);
}
return eventsByMetadataName;
return eventsByEntityName;
};
export const useDispatchMetadataEventsFromSseToBrowserEvents = <
@@ -27,10 +25,10 @@ export const useDispatchMetadataEventsFromSseToBrowserEvents = <
>() => {
const dispatchMetadataEventsFromSseToBrowserEvents = useCallback(
(metadataEvents: MetadataEvent[]) => {
const eventsByMetadataName =
groupSseMetadataEventsByMetadataName(metadataEvents);
const eventsByEntityName =
groupSseMetadataEventsByEntityName(metadataEvents);
for (const [metadataName, events] of eventsByMetadataName) {
for (const [metadataName, events] of eventsByEntityName) {
const metadataOperationBrowserEvents =
turnSseMetadataEventsToMetadataOperationBrowserEvents<T>({
metadataName,
@@ -1,8 +1,8 @@
import { type BroadcastEntityName } from '@/browser-event/types/BroadcastEntityName';
import { type MetadataOperationBrowserEventDetail } from '@/browser-event/types/MetadataOperationBrowserEventDetail';
import { isDefined } from 'twenty-shared/utils';
import {
MetadataEventAction,
type AllMetadataName,
type MetadataEvent,
} from '~/generated-metadata/graphql';
@@ -12,7 +12,7 @@ export const turnSseMetadataEventsToMetadataOperationBrowserEvents = <
metadataName,
sseMetadataEvents,
}: {
metadataName: AllMetadataName;
metadataName: BroadcastEntityName;
sseMetadataEvents: MetadataEvent[];
}): MetadataOperationBrowserEventDetail<T>[] => {
return sseMetadataEvents
-1
View File
@@ -168,7 +168,6 @@
"react-dom": "18.3.1",
"redis": "^4.7.0",
"reflect-metadata": "0.2.2",
"resumable-stream": "^2.2.12",
"rxjs": "7.8.1",
"semver": "7.6.3",
"sharp": "0.32.6",
@@ -0,0 +1,41 @@
import { MigrationInterface, QueryRunner } from 'typeorm';
export class AddStatusToAgentMessage1775001600000
implements MigrationInterface
{
public async up(queryRunner: QueryRunner): Promise<void> {
await queryRunner.query(
`CREATE TYPE "core"."agentMessage_status_enum" AS ENUM ('queued', 'sent')`,
);
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" ADD COLUMN "status" "core"."agentMessage_status_enum" NOT NULL DEFAULT 'sent'`,
);
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" ALTER COLUMN "turnId" DROP NOT NULL`,
);
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" ADD COLUMN "processedAt" TIMESTAMPTZ`,
);
await queryRunner.query(
`UPDATE "core"."agentMessage" SET "processedAt" = "createdAt" WHERE "status" = 'sent'`,
);
}
public async down(queryRunner: QueryRunner): Promise<void> {
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" DROP COLUMN "processedAt"`,
);
// Queued messages have turnId=NULL. They must be deleted before
// restoring the NOT NULL constraint on turnId.
await queryRunner.query(
`DELETE FROM "core"."agentMessage" WHERE "turnId" IS NULL`,
);
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" ALTER COLUMN "turnId" SET NOT NULL`,
);
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" DROP COLUMN "status"`,
);
await queryRunner.query(`DROP TYPE "core"."agentMessage_status_enum"`);
}
}
@@ -13,8 +13,8 @@ export class AgentMessageDTO {
@Field(() => UUIDScalarType)
threadId: string;
@Field(() => UUIDScalarType)
turnId: string;
@Field(() => UUIDScalarType, { nullable: true })
turnId: string | null;
@Field(() => UUIDScalarType, { nullable: true })
agentId: string | null;
@@ -22,9 +22,16 @@ export class AgentMessageDTO {
@Field()
role: 'system' | 'user' | 'assistant';
@Field()
status: 'queued' | 'sent';
@Field(() => [AgentMessagePartDTO])
parts: AgentMessagePartDTO[];
@IsDateString()
@Field(() => Date, { nullable: true })
processedAt: Date | null;
@IsDateString()
@Field()
createdAt: Date;
@@ -20,6 +20,11 @@ export enum AgentMessageRole {
ASSISTANT = 'assistant',
}
export enum AgentMessageStatus {
QUEUED = 'queued',
SENT = 'sent',
}
@Entity('agentMessage')
export class AgentMessageEntity {
@PrimaryGeneratedColumn('uuid')
@@ -35,15 +40,16 @@ export class AgentMessageEntity {
@JoinColumn({ name: 'threadId' })
thread: Relation<AgentChatThreadEntity>;
@Column('uuid')
@Column({ type: 'uuid', nullable: true })
@Index()
turnId: string;
turnId: string | null;
@ManyToOne(() => AgentTurnEntity, {
onDelete: 'CASCADE',
nullable: true,
})
@JoinColumn({ name: 'turnId' })
turn: Relation<AgentTurnEntity>;
turn: Relation<AgentTurnEntity> | null;
@Column({ type: 'uuid', nullable: true })
@Index()
@@ -52,9 +58,19 @@ export class AgentMessageEntity {
@Column({ type: 'enum', enum: AgentMessageRole })
role: AgentMessageRole;
@Column({
type: 'enum',
enum: AgentMessageStatus,
default: AgentMessageStatus.SENT,
})
status: AgentMessageStatus;
@OneToMany(() => AgentMessagePartEntity, (part) => part.message)
parts: Relation<AgentMessagePartEntity[]>;
@Column({ type: 'timestamptz', nullable: true })
processedAt: Date | null;
@CreateDateColumn({ type: 'timestamptz' })
createdAt: Date;
}
@@ -0,0 +1,71 @@
import { type ExtendedUIMessagePart } from 'twenty-shared/ai';
import { type AgentMessagePartEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message-part.entity';
// Maps TypeORM entity fields to UI message parts.
// A parallel mapping for GraphQL DTOs exists in the frontend at:
// packages/twenty-front/src/modules/ai/utils/mapDBPartToUIMessagePart.ts
export const mapDBPartToUIMessagePart = (
part: AgentMessagePartEntity,
): ExtendedUIMessagePart | null => {
switch (part.type) {
case 'text':
return {
type: 'text',
text: part.textContent ?? '',
};
case 'reasoning':
return {
type: 'reasoning',
text: part.reasoningContent ?? '',
state: (part.state as 'streaming' | 'done') ?? 'done',
};
case 'file':
return {
type: 'file',
mediaType: part.fileFilename?.endsWith('.png')
? 'image/png'
: 'application/octet-stream',
filename: part.fileFilename ?? '',
url: '',
};
case 'source-url':
return {
type: 'source-url',
sourceId: part.sourceUrlSourceId ?? '',
url: part.sourceUrlUrl ?? '',
title: part.sourceUrlTitle ?? '',
providerMetadata: part.providerMetadata ?? undefined,
};
case 'source-document':
return {
type: 'source-document',
sourceId: part.sourceDocumentSourceId ?? '',
mediaType: part.sourceDocumentMediaType ?? '',
title: part.sourceDocumentTitle ?? '',
filename: part.sourceDocumentFilename ?? '',
providerMetadata: part.providerMetadata ?? undefined,
};
case 'step-start':
return {
type: 'step-start',
};
case 'data-routing-status':
return null;
default: {
if (part.type.includes('tool-') && part.toolCallId) {
return {
type: part.type,
toolCallId: part.toolCallId,
input: part.toolInput ?? {},
output: part.toolOutput,
errorText: part.errorMessage ?? '',
state: part.state,
} as ExtendedUIMessagePart;
}
return null;
}
}
};
@@ -0,0 +1,13 @@
import { type ExtendedUIMessagePart } from 'twenty-shared/ai';
import { type AgentMessagePartEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message-part.entity';
import { mapDBPartToUIMessagePart } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartToUIMessagePart';
export const mapDBPartsToUIMessageParts = (
parts: AgentMessagePartEntity[],
): ExtendedUIMessagePart[] => {
return parts
.sort((a, b) => a.orderIndex - b.orderIndex)
.map(mapDBPartToUIMessagePart)
.filter((part): part is ExtendedUIMessagePart => part !== null);
};
@@ -33,13 +33,13 @@ import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache
import { DashboardToolsModule } from 'src/modules/dashboard/tools/dashboard-tools.module';
import { WorkflowToolsModule } from 'src/modules/workflow/workflow-tools/workflow-tools.module';
import { AgentChatController } from './controllers/agent-chat.controller';
import { AgentChatThreadDTO } from './dtos/agent-chat-thread.dto';
import { AgentChatThreadEntity } from './entities/agent-chat-thread.entity';
import { StreamAgentChatJob } from './jobs/stream-agent-chat.job';
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 { AgentChatResumableStreamService } from './services/agent-chat-resumable-stream.service';
import { AgentChatEventPublisherService } from './services/agent-chat-event-publisher.service';
import { AgentChatStreamingService } from './services/agent-chat-streaming.service';
import { AgentChatService } from './services/agent-chat.service';
import { AgentTitleGenerationService } from './services/agent-title-generation.service';
@@ -101,11 +101,11 @@ import { SystemPromptBuilderService } from './services/system-prompt-builder.ser
DashboardToolsModule,
WorkflowToolsModule,
],
controllers: [AgentChatController],
providers: [
AgentChatCancelSubscriberService,
AgentChatEventPublisherService,
AgentChatResolver,
AgentChatResumableStreamService,
AgentChatSubscriptionResolver,
AgentChatService,
AgentChatStreamingService,
AgentTitleGenerationService,
@@ -1,183 +0,0 @@
import {
Body,
Controller,
Delete,
Get,
Param,
Post,
Res,
UseFilters,
UseGuards,
} from '@nestjs/common';
import { PermissionFlagType } from 'twenty-shared/constants';
import { InjectRepository } from '@nestjs/typeorm';
import { UI_MESSAGE_STREAM_HEADERS } from 'ai';
import type { Response } from 'express';
import type { ExtendedUIMessage } from 'twenty-shared/ai';
import { isDefined } from 'twenty-shared/utils';
import type { Repository } from 'typeorm';
import { RestApiExceptionFilter } from 'src/engine/api/rest/rest-api-exception.filter';
import {
BillingException,
BillingExceptionCode,
} from 'src/engine/core-modules/billing/billing.exception';
import { BillingProductKey } from 'src/engine/core-modules/billing/enums/billing-product-key.enum';
import { BillingRestApiExceptionFilter } from 'src/engine/core-modules/billing/filters/billing-api-exception.filter';
import { BillingService } from 'src/engine/core-modules/billing/services/billing.service';
import { RedisClientService } from 'src/engine/core-modules/redis-client/redis-client.service';
import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';
import type { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-workspace-id.decorator';
import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
import { JwtAuthGuard } from 'src/engine/guards/jwt-auth.guard';
import { SettingsPermissionGuard } from 'src/engine/guards/settings-permission.guard';
import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard';
import {
AgentException,
AgentExceptionCode,
} from 'src/engine/metadata-modules/ai/ai-agent/agent.exception';
import { AgentRestApiExceptionFilter } from 'src/engine/metadata-modules/ai/ai-agent/filters/agent-api-exception.filter';
import type { BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util';
import { AgentChatResumableStreamService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-resumable-stream.service';
import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';
import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service';
@Controller('rest/agent-chat')
@UseGuards(JwtAuthGuard, WorkspaceAuthGuard)
@UseFilters(
RestApiExceptionFilter,
AgentRestApiExceptionFilter,
BillingRestApiExceptionFilter,
)
export class AgentChatController {
constructor(
private readonly agentStreamingService: AgentChatStreamingService,
private readonly resumableStreamService: AgentChatResumableStreamService,
private readonly billingService: BillingService,
private readonly twentyConfigService: TwentyConfigService,
private readonly aiModelRegistryService: AiModelRegistryService,
private readonly redisClientService: RedisClientService,
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
) {}
@Post('stream')
@UseGuards(SettingsPermissionGuard(PermissionFlagType.AI))
async streamAgentChat(
@Body()
body: {
threadId: string;
messages: ExtendedUIMessage[];
browsingContext?: BrowsingContextType | null;
modelId?: string;
},
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() workspace: WorkspaceEntity,
@Res() response: Response,
) {
if (this.aiModelRegistryService.getAvailableModels().length === 0) {
throw new AgentException(
'No AI models are available. Configure at least one AI provider.',
AgentExceptionCode.API_KEY_NOT_CONFIGURED,
);
}
const resolvedModelId = body.modelId ?? workspace.smartModel;
this.aiModelRegistryService.validateModelAvailability(
resolvedModelId,
workspace,
);
if (this.twentyConfigService.get('IS_BILLING_ENABLED')) {
const canBill = await this.billingService.canBillMeteredProduct(
workspace.id,
BillingProductKey.WORKFLOW_NODE_EXECUTION,
);
if (!canBill) {
throw new BillingException(
'Credits exhausted',
BillingExceptionCode.BILLING_CREDITS_EXHAUSTED,
);
}
}
return this.agentStreamingService.streamAgentChat({
threadId: body.threadId,
messages: body.messages,
browsingContext: body.browsingContext ?? null,
modelId: body.modelId,
userWorkspaceId,
workspace,
response,
});
}
@Get(':threadId/stream')
@UseGuards(SettingsPermissionGuard(PermissionFlagType.AI))
async resumeAgentChatStream(
@Param('threadId') threadId: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@Res() response: Response,
) {
const thread = await this.threadRepository.findOne({
where: { id: threadId, userWorkspaceId },
});
if (!isDefined(thread) || !isDefined(thread.activeStreamId)) {
response.status(204).end();
return;
}
const resumedNodeReadable =
await this.resumableStreamService.resumeExistingStreamAsNodeReadable(
thread.activeStreamId,
);
if (!isDefined(resumedNodeReadable)) {
response.status(204).end();
return;
}
response.writeHead(200, UI_MESSAGE_STREAM_HEADERS);
resumedNodeReadable.pipe(response);
}
@Delete(':threadId/stream')
@UseGuards(SettingsPermissionGuard(PermissionFlagType.AI))
async stopAgentChatStream(
@Param('threadId') threadId: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
) {
const thread = await this.threadRepository.findOne({
where: { id: threadId, userWorkspaceId },
});
if (!isDefined(thread) || !isDefined(thread.activeStreamId)) {
return { success: true };
}
// Publish a cancel signal via Redis pub/sub. The BullMQ worker
// processing this thread's stream subscribes to this channel and
// will abort the LLM connection when the message arrives — stopping
// token generation and billing immediately.
const redis = this.redisClientService.getClient();
await redis.publish(getCancelChannel(threadId), 'cancel');
await this.threadRepository.update(
{ id: threadId, userWorkspaceId },
{ activeStreamId: null },
);
return { success: true };
}
}
@@ -0,0 +1,14 @@
import { Field, ObjectType } from '@nestjs/graphql';
import GraphQLJSON from 'graphql-type-json';
// Typed as JSON because the payload is AgentChatSubscriptionEvent
// (a discriminated union defined in twenty-shared).
@ObjectType('AgentChatEvent')
export class AgentChatEventDTO {
@Field(() => String)
threadId: string;
@Field(() => GraphQLJSON)
event: Record<string, unknown>;
}
@@ -0,0 +1,12 @@
import { Field, Int, ObjectType } from '@nestjs/graphql';
import GraphQLJSON from 'graphql-type-json';
@ObjectType('ChatStreamCatchupChunks')
export class ChatStreamCatchupChunksDTO {
@Field(() => [GraphQLJSON])
chunks: Record<string, unknown>[];
@Field(() => Int)
maxSeq: number;
}
@@ -0,0 +1,13 @@
import { Field, ObjectType } from '@nestjs/graphql';
@ObjectType('SendChatMessageResult')
export class SendChatMessageResultDTO {
@Field(() => String)
messageId: string;
@Field(() => Boolean)
queued: boolean;
@Field(() => String, { nullable: true })
streamId?: string;
}
@@ -0,0 +1 @@
export const STREAM_AGENT_CHAT_JOB_NAME = 'StreamAgentChatJob';
@@ -0,0 +1,20 @@
import type {
ExtendedUIMessage,
ExtendedUIMessagePart,
} from 'twenty-shared/ai';
import type { BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
export type StreamAgentChatJobData = {
threadId: string;
streamId: string;
userWorkspaceId: string;
workspaceId: string;
messages: ExtendedUIMessage[];
browsingContext: BrowsingContextType | null;
modelId?: string;
lastUserMessageText: string;
lastUserMessageParts: ExtendedUIMessagePart[];
hasTitle: boolean;
existingTurnId?: string;
};
@@ -1,7 +1,7 @@
import { Logger, Scope } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { createUIMessageStream, JsonToSseTransformStream } from 'ai';
import { createUIMessageStream } from 'ai';
import type {
CodeExecutionData,
ExtendedUIMessage,
@@ -10,37 +10,27 @@ import type {
import { Repository } from 'typeorm';
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 { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util';
import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
import { Process } from 'src/engine/core-modules/message-queue/decorators/process.decorator';
import { Processor } from 'src/engine/core-modules/message-queue/decorators/processor.decorator';
import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
import { AgentMessageRole } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
import type { BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';
import { computeCostBreakdown } from 'src/engine/metadata-modules/ai/ai-billing/utils/compute-cost-breakdown.util';
import { convertDollarsToBillingCredits } from 'src/engine/metadata-modules/ai/ai-billing/utils/convert-dollars-to-billing-credits.util';
import { extractCacheCreationTokens } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util';
import type { AIModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import { AgentChatResumableStreamService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-resumable-stream.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';
import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util';
export const STREAM_AGENT_CHAT_JOB_NAME = 'StreamAgentChatJob';
import { STREAM_AGENT_CHAT_JOB_NAME } from './stream-agent-chat-job-name.constant';
import { type StreamAgentChatJobData } from './stream-agent-chat-job.types';
export type StreamAgentChatJobData = {
threadId: string;
streamId: string;
userWorkspaceId: string;
workspaceId: string;
messages: ExtendedUIMessage[];
browsingContext: BrowsingContextType | null;
modelId?: string;
lastUserMessageText: string;
lastUserMessageParts: ExtendedUIMessagePart[];
hasTitle: boolean;
};
export { STREAM_AGENT_CHAT_JOB_NAME, type StreamAgentChatJobData };
@Processor({ queueName: MessageQueue.aiStreamQueue, scope: Scope.REQUEST })
export class StreamAgentChatJob {
@@ -53,8 +43,9 @@ export class StreamAgentChatJob {
private readonly workspaceRepository: Repository<WorkspaceEntity>,
private readonly agentChatService: AgentChatService,
private readonly chatExecutionService: ChatExecutionService,
private readonly resumableStreamService: AgentChatResumableStreamService,
private readonly eventPublisherService: AgentChatEventPublisherService,
private readonly cancelSubscriberService: AgentChatCancelSubscriberService,
private readonly agentChatStreamingService: AgentChatStreamingService,
) {}
@Process(STREAM_AGENT_CHAT_JOB_NAME)
@@ -65,9 +56,14 @@ export class StreamAgentChatJob {
if (!workspace) {
this.logger.error(`Workspace ${data.workspaceId} not found`);
await this.resumableStreamService.writeStreamError(data.streamId, {
code: 'WORKSPACE_NOT_FOUND',
message: `Workspace ${data.workspaceId} not found`,
await this.eventPublisherService.publish({
threadId: data.threadId,
workspaceId: data.workspaceId,
event: {
type: 'stream-error',
code: 'WORKSPACE_NOT_FOUND',
message: `Workspace ${data.workspaceId} not found`,
},
});
return;
@@ -86,11 +82,18 @@ export class StreamAgentChatJob {
this.logger.error(
`Stream ${data.streamId} failed: ${error instanceof Error ? error.message : String(error)}`,
);
await this.resumableStreamService
.writeStreamError(data.streamId, {
code: 'STREAM_EXECUTION_FAILED',
message:
error instanceof Error ? error.message : 'Stream execution failed',
await this.eventPublisherService
.publish({
threadId: data.threadId,
workspaceId: data.workspaceId,
event: {
type: 'stream-error',
code: 'STREAM_EXECUTION_FAILED',
message:
error instanceof Error
? error.message
: 'Stream execution failed',
},
})
.catch(() => {});
} finally {
@@ -105,6 +108,21 @@ export class StreamAgentChatJob {
})
.execute()
.catch(() => {});
if (!abortController.signal.aborted) {
await this.agentChatStreamingService
.flushNextQueuedMessage(
data.threadId,
data.userWorkspaceId,
data.workspaceId,
data.hasTitle,
)
.catch((error) => {
this.logger.error(
`Failed to flush queued message for thread ${data.threadId}: ${error instanceof Error ? error.message : String(error)}`,
);
});
}
}
}
@@ -113,26 +131,34 @@ export class StreamAgentChatJob {
workspace: WorkspaceEntity,
abortSignal: AbortSignal,
): Promise<void> {
const userMessagePromise = this.agentChatService.addMessage({
threadId: data.threadId,
uiMessage: {
role: AgentMessageRole.USER,
parts: data.lastUserMessageParts.filter(
(part): part is ExtendedUIMessagePart =>
part.type === 'text' || part.type === 'file',
),
},
});
// When processing a promoted queued message, the user message already
// exists in the DB with a turn — skip persisting it again.
const userMessagePromise = data.existingTurnId
? Promise.resolve({ turnId: data.existingTurnId })
: this.agentChatService.addMessage({
threadId: data.threadId,
uiMessage: {
role: AgentMessageRole.USER,
parts: data.lastUserMessageParts.filter(
(part): part is ExtendedUIMessagePart =>
part.type === 'text' || part.type === 'file',
),
},
});
userMessagePromise.catch(() => {});
const titlePromise = data.hasTitle
? Promise.resolve(null)
: this.agentChatService
.generateTitleIfNeeded(data.threadId, data.lastUserMessageText)
.generateTitleIfNeeded({
threadId: data.threadId,
messageContent: data.lastUserMessageText,
workspaceId: data.workspaceId,
})
.catch(() => null);
await this.buildAndPipeStream({
await this.buildAndPublishStream({
workspace,
data,
userMessagePromise,
@@ -141,7 +167,7 @@ export class StreamAgentChatJob {
});
}
private async buildAndPipeStream({
private async buildAndPublishStream({
workspace,
data,
userMessagePromise,
@@ -150,7 +176,7 @@ export class StreamAgentChatJob {
}: {
workspace: WorkspaceEntity;
data: StreamAgentChatJobData;
userMessagePromise: Promise<{ turnId: string }>;
userMessagePromise: Promise<{ turnId: string | null }>;
titlePromise: Promise<string | null>;
abortSignal: AbortSignal;
}): Promise<void> {
@@ -164,6 +190,14 @@ export class StreamAgentChatJob {
let lastStepConversationSize = 0;
let totalCacheCreationTokens = 0;
// onFinish fires before the uiStream is fully drained. We use this
// promise to coordinate: the IIFE waits for DB persist to complete
// before publishing message-persisted (after all chunks).
let resolveStreamFinished: () => void;
const streamFinishedPromise = new Promise<void>((res) => {
resolveStreamFinished = res;
});
abortSignal.addEventListener('abort', () => resolve(), { once: true });
const uiStream = createUIMessageStream<ExtendedUIMessage>({
@@ -233,7 +267,7 @@ export class StreamAgentChatJob {
userMessagePromise,
});
await titleWritePromise;
resolve();
resolveStreamFinished();
} catch (error) {
reject(error);
}
@@ -244,11 +278,34 @@ export class StreamAgentChatJob {
},
});
const sseStream = uiStream.pipeThrough(new JsonToSseTransformStream());
// Publish all chunks first, then signal completion. This guarantees
// message-persisted arrives after every stream-chunk on the client.
(async () => {
try {
for await (const chunk of uiStream) {
await this.eventPublisherService.publish({
threadId: data.threadId,
workspaceId: data.workspaceId,
event: {
type: 'stream-chunk',
chunk: chunk as Record<string, unknown>,
},
});
}
this.resumableStreamService
.createResumableStream(data.streamId, () => sseStream)
.catch(reject);
await streamFinishedPromise;
await this.eventPublisherService.publish({
threadId: data.threadId,
workspaceId: data.workspaceId,
event: { type: 'message-persisted', messageId: data.threadId },
});
resolve();
} catch (error) {
reject(error);
}
})();
});
}
@@ -358,7 +415,7 @@ export class StreamAgentChatJob {
};
lastStepConversationSize: number;
modelConfig: AIModelConfig;
userMessagePromise: Promise<{ turnId: string }>;
userMessagePromise: Promise<{ turnId: string | null }>;
}): Promise<void> {
if (responseMessage.parts.length === 0) {
return;
@@ -369,7 +426,7 @@ export class StreamAgentChatJob {
await this.agentChatService.addMessage({
threadId,
uiMessage: responseMessage,
turnId: userMessage.turnId,
turnId: userMessage.turnId ?? undefined,
});
await this.threadRepository.update(threadId, {
@@ -0,0 +1,68 @@
import { UseGuards } from '@nestjs/common';
import { Args, Subscription } from '@nestjs/graphql';
import { InjectRepository } from '@nestjs/typeorm';
import { PermissionFlagType } from 'twenty-shared/constants';
import { isDefined } from 'twenty-shared/utils';
import { Repository } from 'typeorm';
import { MetadataResolver } from 'src/engine/api/graphql/graphql-config/decorators/metadata-resolver.decorator';
import { UUIDScalarType } from 'src/engine/api/graphql/workspace-schema-builder/graphql-types/scalars';
import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-workspace-id.decorator';
import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
import { RequireFeatureFlag } from 'src/engine/guards/feature-flag.guard';
import { SettingsPermissionGuard } from 'src/engine/guards/settings-permission.guard';
import { UserAuthGuard } from 'src/engine/guards/user-auth.guard';
import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard';
import {
AgentException,
AgentExceptionCode,
} from 'src/engine/metadata-modules/ai/ai-agent/agent.exception';
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 { FeatureFlagKey } from 'twenty-shared/types';
@MetadataResolver()
@UseGuards(WorkspaceAuthGuard, UserAuthGuard)
export class AgentChatSubscriptionResolver {
constructor(
private readonly subscriptionService: SubscriptionService,
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
) {}
@Subscription(() => AgentChatEventDTO, {
filter: (
payload: { onAgentChatEvent: AgentChatEventDTO },
variables: { threadId: string },
) => {
return payload.onAgentChatEvent.threadId === variables.threadId;
},
})
@RequireFeatureFlag(FeatureFlagKey.IS_AI_ENABLED)
@UseGuards(SettingsPermissionGuard(PermissionFlagType.AI))
async onAgentChatEvent(
@Args('threadId', { type: () => UUIDScalarType }) threadId: string,
@AuthWorkspace() workspace: WorkspaceEntity,
@AuthUserWorkspaceId() userWorkspaceId: string,
) {
const thread = await this.threadRepository.findOne({
where: { id: threadId, userWorkspaceId },
select: ['id'],
});
if (!isDefined(thread)) {
throw new AgentException(
'Thread not found',
AgentExceptionCode.AGENT_EXECUTION_FAILED,
);
}
return this.subscriptionService.subscribeToAgentChat({
workspaceId: workspace.id,
threadId,
});
}
}
@@ -8,11 +8,23 @@ import {
ResolveField,
} from '@nestjs/graphql';
import { InjectRepository } from '@nestjs/typeorm';
import { PermissionFlagType } from 'twenty-shared/constants';
import { FeatureFlagKey } from 'twenty-shared/types';
import { isDefined } from 'twenty-shared/utils';
import { Repository } from 'typeorm';
import GraphQLJSON from 'graphql-type-json';
import { MetadataResolver } from 'src/engine/api/graphql/graphql-config/decorators/metadata-resolver.decorator';
import { UUIDScalarType } from 'src/engine/api/graphql/workspace-schema-builder/graphql-types/scalars';
import {
BillingException,
BillingExceptionCode,
} from 'src/engine/core-modules/billing/billing.exception';
import { BillingProductKey } from 'src/engine/core-modules/billing/enums/billing-product-key.enum';
import { BillingService } from 'src/engine/core-modules/billing/services/billing.service';
import { RedisClientService } from 'src/engine/core-modules/redis-client/redis-client.service';
import { TwentyConfigService } from 'src/engine/core-modules/twenty-config/twenty-config.service';
import { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util';
import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-workspace-id.decorator';
@@ -23,12 +35,23 @@ import {
} from 'src/engine/guards/feature-flag.guard';
import { SettingsPermissionGuard } from 'src/engine/guards/settings-permission.guard';
import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard';
import {
AgentException,
AgentExceptionCode,
} from 'src/engine/metadata-modules/ai/ai-agent/agent.exception';
import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
import { AgentMessageDTO } from 'src/engine/metadata-modules/ai/ai-agent-execution/dtos/agent-message.dto';
import { AgentChatThreadDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/agent-chat-thread.dto';
import { AISystemPromptPreviewDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/ai-system-prompt-preview.dto';
import { type AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import { ChatStreamCatchupChunksDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/chat-stream-catchup-chunks.dto';
import { SendChatMessageResultDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/send-chat-message-result.dto';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.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 { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util';
import { SystemPromptBuilderService } from 'src/engine/metadata-modules/ai/ai-chat/services/system-prompt-builder.service';
import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service';
@UseGuards(
WorkspaceAuthGuard,
@@ -39,7 +62,15 @@ import { SystemPromptBuilderService } from 'src/engine/metadata-modules/ai/ai-ch
export class AgentChatResolver {
constructor(
private readonly agentChatService: AgentChatService,
private readonly agentChatStreamingService: AgentChatStreamingService,
private readonly eventPublisherService: AgentChatEventPublisherService,
private readonly systemPromptBuilderService: SystemPromptBuilderService,
private readonly billingService: BillingService,
private readonly twentyConfigService: TwentyConfigService,
private readonly aiModelRegistryService: AiModelRegistryService,
private readonly redisClientService: RedisClientService,
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
) {}
@Query(() => AgentChatThreadDTO)
@@ -63,10 +94,178 @@ export class AgentChatResolver {
);
}
@Query(() => ChatStreamCatchupChunksDTO)
@RequireFeatureFlag(FeatureFlagKey.IS_AI_ENABLED)
async chatStreamCatchupChunks(
@Args('threadId', { type: () => UUIDScalarType }) threadId: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
) {
await this.agentChatService.getThreadById(threadId, userWorkspaceId);
return this.eventPublisherService.getAccumulatedChunks(threadId);
}
@Mutation(() => AgentChatThreadDTO)
@RequireFeatureFlag(FeatureFlagKey.IS_AI_ENABLED)
async createChatThread(@AuthUserWorkspaceId() userWorkspaceId: string) {
return this.agentChatService.createThread(userWorkspaceId);
async createChatThread(
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() workspace: WorkspaceEntity,
) {
return this.agentChatService.createThread({
userWorkspaceId,
workspaceId: workspace.id,
});
}
@Mutation(() => SendChatMessageResultDTO)
@RequireFeatureFlag(FeatureFlagKey.IS_AI_ENABLED)
async sendChatMessage(
@Args('threadId', { type: () => UUIDScalarType }) threadId: string,
@Args('text') text: string,
@Args('messageId', { type: () => UUIDScalarType }) messageId: string,
@Args('browsingContext', { type: () => GraphQLJSON, nullable: true })
browsingContext: BrowsingContextType | null,
@Args('modelId', { type: () => String, nullable: true })
modelId: string | undefined,
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() workspace: WorkspaceEntity,
): Promise<SendChatMessageResultDTO> {
if (this.aiModelRegistryService.getAvailableModels().length === 0) {
throw new AgentException(
'No AI models are available. Configure at least one AI provider.',
AgentExceptionCode.API_KEY_NOT_CONFIGURED,
);
}
const resolvedModelId = modelId ?? workspace.smartModel;
this.aiModelRegistryService.validateModelAvailability(
resolvedModelId,
workspace,
);
if (this.twentyConfigService.get('IS_BILLING_ENABLED')) {
const canBill = await this.billingService.canBillMeteredProduct(
workspace.id,
BillingProductKey.WORKFLOW_NODE_EXECUTION,
);
if (!canBill) {
throw new BillingException(
'Credits exhausted',
BillingExceptionCode.BILLING_CREDITS_EXHAUSTED,
);
}
}
const thread = await this.threadRepository.findOne({
where: { id: threadId, userWorkspaceId },
});
if (!isDefined(thread)) {
throw new AgentException(
'Thread not found',
AgentExceptionCode.AGENT_EXECUTION_FAILED,
);
}
if (isDefined(thread.activeStreamId)) {
const queuedMessage = await this.agentChatService.queueMessage({
threadId,
text,
id: messageId,
});
await this.eventPublisherService.publish({
threadId,
workspaceId: workspace.id,
event: { type: 'queue-updated' },
});
return { messageId: queuedMessage.id, queued: true };
}
const result = await this.agentChatStreamingService.streamAgentChat({
threadId,
browsingContext: browsingContext ?? null,
modelId,
userWorkspaceId,
workspace,
text,
messageId,
});
return {
messageId: result.messageId,
queued: false,
streamId: result.streamId,
};
}
@Mutation(() => Boolean)
@RequireFeatureFlag(FeatureFlagKey.IS_AI_ENABLED)
async stopAgentChatStream(
@Args('threadId', { type: () => UUIDScalarType }) threadId: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
): Promise<boolean> {
const thread = await this.threadRepository.findOne({
where: { id: threadId, userWorkspaceId },
});
if (!isDefined(thread) || !isDefined(thread.activeStreamId)) {
return true;
}
const redis = this.redisClientService.getClient();
await redis.publish(getCancelChannel(threadId), 'cancel');
await this.threadRepository.update(
{ id: threadId, userWorkspaceId },
{ activeStreamId: null },
);
return true;
}
@Mutation(() => Boolean)
@RequireFeatureFlag(FeatureFlagKey.IS_AI_ENABLED)
async deleteQueuedChatMessage(
@Args('messageId', { type: () => UUIDScalarType }) messageId: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() workspace: WorkspaceEntity,
): Promise<boolean> {
const message = await this.agentChatService.findQueuedMessage(messageId);
if (!isDefined(message)) {
throw new AgentException(
'Queued message not found',
AgentExceptionCode.AGENT_EXECUTION_FAILED,
);
}
const thread = await this.threadRepository.findOne({
where: { id: message.threadId, userWorkspaceId },
});
if (!isDefined(thread)) {
throw new AgentException(
'Thread not found',
AgentExceptionCode.AGENT_EXECUTION_FAILED,
);
}
const deleted = await this.agentChatService.deleteQueuedMessage(messageId);
if (deleted) {
await this.eventPublisherService.publish({
threadId: message.threadId,
workspaceId: workspace.id,
event: { type: 'queue-updated' },
});
}
return deleted;
}
@Query(() => AISystemPromptPreviewDTO)
@@ -0,0 +1,74 @@
import { Injectable } from '@nestjs/common';
import { type AgentChatSubscriptionEvent } from 'twenty-shared/ai';
import { RedisClientService } from 'src/engine/core-modules/redis-client/redis-client.service';
import { SubscriptionService } from 'src/engine/subscriptions/subscription.service';
const STREAM_CHUNKS_TTL_SECONDS = 3600;
@Injectable()
export class AgentChatEventPublisherService {
constructor(
private readonly subscriptionService: SubscriptionService,
private readonly redisClientService: RedisClientService,
) {}
private getStreamChunksKey(threadId: string): string {
return `agent-chat-stream-chunks:${threadId}`;
}
async publish({
threadId,
workspaceId,
event,
}: {
threadId: string;
workspaceId: string;
event: AgentChatSubscriptionEvent;
}): Promise<void> {
let publishedEvent = event;
if (event.type === 'stream-chunk') {
const redis = this.redisClientService.getClient();
const key = this.getStreamChunksKey(threadId);
// RPUSH returns the new list length — use it as a 1-based sequence number
const seq = await redis.rpush(key, JSON.stringify(event.chunk));
await redis.expire(key, STREAM_CHUNKS_TTL_SECONDS);
publishedEvent = { ...event, seq };
} else if (event.type === 'message-persisted') {
const redis = this.redisClientService.getClient();
await redis.del(this.getStreamChunksKey(threadId));
}
await this.subscriptionService.publishToAgentChat({
workspaceId,
threadId,
payload: {
onAgentChatEvent: {
threadId,
event: publishedEvent,
},
},
});
}
async getAccumulatedChunks(threadId: string): Promise<{
chunks: Record<string, unknown>[];
maxSeq: number;
}> {
const redis = this.redisClientService.getClient();
const rawChunks = await redis.lrange(
this.getStreamChunksKey(threadId),
0,
-1,
);
return {
chunks: rawChunks.map((raw) => JSON.parse(raw)),
maxSeq: rawChunks.length,
};
}
}
@@ -1,104 +0,0 @@
import { Injectable, OnModuleDestroy } from '@nestjs/common';
import { Readable } from 'stream';
import { type ReadableStream as NodeWebReadableStream } from 'stream/web';
import type { Redis } from 'ioredis';
import { createResumableStreamContext } from 'resumable-stream/ioredis';
import { RedisClientService } from 'src/engine/core-modules/redis-client/redis-client.service';
@Injectable()
export class AgentChatResumableStreamService implements OnModuleDestroy {
private streamContext: ReturnType<typeof createResumableStreamContext>;
private streamPublisher: Redis;
private streamSubscriber: Redis;
private redisClient: Redis;
constructor(private readonly redisClientService: RedisClientService) {
const baseClient = this.redisClientService.getClient();
this.streamPublisher = baseClient.duplicate();
this.streamSubscriber = baseClient.duplicate();
this.redisClient = baseClient.duplicate();
this.streamContext = createResumableStreamContext({
waitUntil: () => {},
publisher: this.streamPublisher,
subscriber: this.streamSubscriber,
});
}
async onModuleDestroy() {
await this.streamPublisher.quit();
await this.streamSubscriber.quit();
await this.redisClient.quit();
}
async createResumableStream(
streamId: string,
streamFactory: () => ReadableStream<string>,
) {
const resumableStream = await this.streamContext.createNewResumableStream(
streamId,
streamFactory,
);
if (!resumableStream) {
return;
}
// Read the stream to completion in the background so chunks are
// published to Redis and available for later resume consumers.
const reader = resumableStream.getReader();
void (async () => {
try {
while (true) {
const { done } = await reader.read();
if (done) {
break;
}
}
} catch {
// Stream interrupted — chunks already published are still resumable.
}
})();
}
async resumeExistingStreamAsNodeReadable(
streamId: string,
): Promise<Readable | null> {
const webStream = await this.streamContext.resumeExistingStream(streamId);
if (!webStream) {
return null;
}
return Readable.fromWeb(webStream as NodeWebReadableStream);
}
async writeStreamError(
streamId: string,
error: { code: string; message: string },
): Promise<void> {
await this.redisClient.set(
`ai-stream:error:${streamId}`,
JSON.stringify(error),
'EX',
60,
);
}
async readStreamError(
streamId: string,
): Promise<{ code: string; message: string } | null> {
const raw = await this.redisClient.get(`ai-stream:error:${streamId}`);
if (!raw) {
return null;
}
return JSON.parse(raw);
}
}
@@ -1,10 +1,7 @@
import { Injectable, Logger } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { type Readable } from 'stream';
import { generateId, UI_MESSAGE_STREAM_HEADERS } from 'ai';
import { type Response } from 'express';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
import { generateId } from 'ai';
import { type Repository } from 'typeorm';
import { InjectMessageQueue } from 'src/engine/core-modules/message-queue/decorators/message-queue.decorator';
@@ -16,27 +13,27 @@ import {
AgentExceptionCode,
} from 'src/engine/metadata-modules/ai/ai-agent/agent.exception';
import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import {
STREAM_AGENT_CHAT_JOB_NAME,
type StreamAgentChatJobData,
} from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job';
import { AgentChatResumableStreamService } from './agent-chat-resumable-stream.service';
AgentMessageRole,
AgentMessageStatus,
} from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
import { mapDBPartsToUIMessageParts } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartsToUIMessageParts';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
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 { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service';
export type StreamAgentChatOptions = {
threadId: string;
userWorkspaceId: string;
workspace: WorkspaceEntity;
response: Response;
messages: ExtendedUIMessage[];
text: string;
browsingContext: BrowsingContextType | null;
modelId?: string;
messageId?: string;
};
const STREAM_READY_TIMEOUT_MS = 5_000;
const STREAM_READY_POLL_INTERVAL_MS = 50;
@Injectable()
export class AgentChatStreamingService {
private readonly logger = new Logger(AgentChatStreamingService.name);
@@ -46,18 +43,19 @@ export class AgentChatStreamingService {
private readonly threadRepository: Repository<AgentChatThreadEntity>,
@InjectMessageQueue(MessageQueue.aiStreamQueue)
private readonly messageQueueService: MessageQueueService,
private readonly resumableStreamService: AgentChatResumableStreamService,
private readonly agentChatService: AgentChatService,
private readonly eventPublisherService: AgentChatEventPublisherService,
) {}
async streamAgentChat({
threadId,
userWorkspaceId,
workspace,
messages,
text,
browsingContext,
response,
modelId,
}: StreamAgentChatOptions) {
messageId,
}: StreamAgentChatOptions): Promise<{ streamId: string; messageId: string }> {
const thread = await this.threadRepository.findOne({
where: {
id: threadId,
@@ -72,10 +70,21 @@ export class AgentChatStreamingService {
);
}
const savedUserMessage = await this.agentChatService.addMessage({
threadId,
id: messageId,
uiMessage: {
role: AgentMessageRole.USER,
parts: [{ type: 'text' as const, text }],
},
});
const previousMessages = await this.loadMessagesFromDB(
threadId,
userWorkspaceId,
);
const streamId = generateId();
const lastUserMessage = messages[messages.length - 1];
const lastUserText =
lastUserMessage?.parts.find((part) => part.type === 'text')?.text ?? '';
await this.messageQueueService.add<StreamAgentChatJobData>(
STREAM_AGENT_CHAT_JOB_NAME,
@@ -84,12 +93,13 @@ export class AgentChatStreamingService {
streamId,
userWorkspaceId,
workspaceId: workspace.id,
messages,
messages: previousMessages,
browsingContext,
modelId,
lastUserMessageText: lastUserText,
lastUserMessageParts: lastUserMessage?.parts ?? [],
lastUserMessageText: text,
lastUserMessageParts: [{ type: 'text', text }],
hasTitle: !!thread.title,
existingTurnId: savedUserMessage.turnId ?? undefined,
},
);
@@ -97,67 +107,92 @@ export class AgentChatStreamingService {
activeStreamId: streamId,
});
try {
const result = await this.waitForResumableStream(streamId);
if ('error' in result) {
response.status(500).json(result.error);
return;
}
if (!result.readable) {
this.logger.error(
`Stream ${streamId} did not become available within timeout`,
);
response
.status(500)
.json({ code: 'WORKER_UNREACHABLE', message: 'Stream timed out' });
return;
}
response.writeHead(200, UI_MESSAGE_STREAM_HEADERS);
result.readable.pipe(response);
} catch (error) {
response.end();
throw error;
}
return { streamId, messageId: savedUserMessage.id };
}
private async waitForResumableStream(
streamId: string,
): Promise<
| { readable: Readable }
| { error: { code: string; message: string } }
| { readable: null }
> {
const maxAttempts = Math.ceil(
STREAM_READY_TIMEOUT_MS / STREAM_READY_POLL_INTERVAL_MS,
async flushNextQueuedMessage(
threadId: string,
userWorkspaceId: string,
workspaceId: string,
hasTitle: boolean,
): Promise<void> {
const queuedMessages =
await this.agentChatService.getQueuedMessages(threadId);
const nextQueued = queuedMessages[0];
if (!nextQueued) {
return;
}
const textPart = nextQueued.parts?.find((part) => part.type === 'text');
const messageText = textPart?.textContent ?? '';
if (messageText === '') {
await this.agentChatService.deleteQueuedMessage(nextQueued.id);
return;
}
const turnId = await this.agentChatService.promoteQueuedMessage(
nextQueued.id,
threadId,
);
for (let attempt = 0; attempt < maxAttempts; attempt++) {
const streamError =
await this.resumableStreamService.readStreamError(streamId);
if (streamError) {
return { error: streamError };
}
const readable =
await this.resumableStreamService.resumeExistingStreamAsNodeReadable(
streamId,
);
if (readable) {
return { readable };
}
await new Promise((resolve) =>
setTimeout(resolve, STREAM_READY_POLL_INTERVAL_MS),
);
if (turnId === null) {
return;
}
return { readable: null };
await this.eventPublisherService.publish({
threadId,
workspaceId,
event: { type: 'queue-updated' },
});
await this.eventPublisherService.publish({
threadId,
workspaceId,
event: { type: 'message-persisted', messageId: nextQueued.id },
});
const uiMessages = await this.loadMessagesFromDB(threadId, userWorkspaceId);
const streamId = generateId();
await this.messageQueueService.add<StreamAgentChatJobData>(
STREAM_AGENT_CHAT_JOB_NAME,
{
threadId,
streamId,
userWorkspaceId,
workspaceId,
messages: uiMessages,
browsingContext: null,
lastUserMessageText: messageText,
lastUserMessageParts: [{ type: 'text', text: messageText }],
hasTitle,
existingTurnId: turnId,
},
);
await this.threadRepository.update(threadId, {
activeStreamId: streamId,
});
}
private async loadMessagesFromDB(threadId: string, userWorkspaceId: string) {
const allMessages = await this.agentChatService.getMessagesForThread(
threadId,
userWorkspaceId,
);
return allMessages
.filter((message) => message.status !== AgentMessageStatus.QUEUED)
.map((message) => ({
id: message.id,
role: message.role as 'user' | 'assistant' | 'system',
parts: mapDBPartsToUIMessageParts(message.parts ?? []),
createdAt: message.createdAt,
}));
}
}
@@ -10,6 +10,7 @@ import { AgentMessagePartEntity } from 'src/engine/metadata-modules/ai/ai-agent-
import {
AgentMessageEntity,
AgentMessageRole,
AgentMessageStatus,
} from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
import { AgentTurnEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-turn.entity';
import { mapUIMessagePartsToDBParts } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapUIMessagePartsToDBParts';
@@ -18,9 +19,23 @@ import {
AgentExceptionCode,
} from 'src/engine/metadata-modules/ai/ai-agent/agent.exception';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import { WorkspaceEventBroadcaster } from 'src/engine/subscriptions/workspace-event-broadcaster/workspace-event-broadcaster.service';
import { AgentTitleGenerationService } from './agent-title-generation.service';
const serializeThreadForBroadcast = (thread: AgentChatThreadEntity) => ({
id: thread.id,
title: thread.title,
totalInputTokens: thread.totalInputTokens,
totalOutputTokens: thread.totalOutputTokens,
contextWindowTokens: thread.contextWindowTokens,
conversationSize: thread.conversationSize,
totalInputCredits: thread.totalInputCredits,
totalOutputCredits: thread.totalOutputCredits,
createdAt: thread.createdAt.toISOString(),
updatedAt: thread.updatedAt.toISOString(),
});
@Injectable()
export class AgentChatService {
constructor(
@@ -33,14 +48,37 @@ export class AgentChatService {
@InjectRepository(AgentMessagePartEntity)
private readonly messagePartRepository: Repository<AgentMessagePartEntity>,
private readonly titleGenerationService: AgentTitleGenerationService,
private readonly workspaceEventBroadcaster: WorkspaceEventBroadcaster,
) {}
async createThread(userWorkspaceId: string) {
async createThread({
userWorkspaceId,
workspaceId,
}: {
userWorkspaceId: string;
workspaceId: string;
}) {
const thread = this.threadRepository.create({
userWorkspaceId,
});
return this.threadRepository.save(thread);
const savedThread = await this.threadRepository.save(thread);
await this.workspaceEventBroadcaster.broadcast({
workspaceId,
events: [
{
type: 'created',
entityName: 'agentChatThread',
recordId: savedThread.id,
properties: {
after: serializeThreadForBroadcast(savedThread),
},
},
],
});
return savedThread;
}
async getThreadById(threadId: string, userWorkspaceId: string) {
@@ -66,12 +104,14 @@ export class AgentChatService {
uiMessage,
agentId,
turnId,
id,
}: {
threadId: string;
uiMessage: Omit<ExtendedUIMessage, 'id'>;
uiMessageParts?: UIMessagePart<UIDataTypes, UITools>[];
agentId?: string;
turnId?: string;
id?: string;
}) {
let actualTurnId = turnId;
@@ -87,10 +127,12 @@ export class AgentChatService {
}
const message = this.messageRepository.create({
...(id ? { id } : {}),
threadId,
turnId: actualTurnId,
role: uiMessage.role as AgentMessageRole,
agentId: agentId ?? null,
processedAt: new Date(),
});
const savedMessage = await this.messageRepository.save(message);
@@ -124,18 +166,111 @@ export class AgentChatService {
return this.messageRepository.find({
where: { threadId },
order: { createdAt: 'ASC' },
order: { processedAt: { direction: 'ASC', nulls: 'LAST' } },
relations: ['parts', 'parts.file'],
});
}
async generateTitleIfNeeded(
async queueMessage({
threadId,
text,
id,
}: {
threadId: string;
text: string;
id?: string;
}): Promise<AgentMessageEntity> {
const message = this.messageRepository.create({
...(id ? { id } : {}),
threadId,
turnId: null,
role: AgentMessageRole.USER,
agentId: null,
status: AgentMessageStatus.QUEUED,
});
const savedMessage = await this.messageRepository.save(message);
const part = this.messagePartRepository.create({
messageId: savedMessage.id,
orderIndex: 0,
type: 'text',
textContent: text,
});
await this.messagePartRepository.save(part);
return savedMessage;
}
async getQueuedMessages(threadId: string): Promise<AgentMessageEntity[]> {
return this.messageRepository.find({
where: {
threadId,
status: AgentMessageStatus.QUEUED,
},
order: { createdAt: 'ASC' },
relations: ['parts'],
});
}
async findQueuedMessage(
messageId: string,
): Promise<AgentMessageEntity | null> {
return this.messageRepository.findOne({
where: { id: messageId, status: AgentMessageStatus.QUEUED },
});
}
async deleteQueuedMessage(messageId: string): Promise<boolean> {
const result = await this.messageRepository.delete({
id: messageId,
status: AgentMessageStatus.QUEUED,
});
return (result.affected ?? 0) > 0;
}
async promoteQueuedMessage(
messageId: string,
threadId: string,
messageContent: string,
): Promise<string | null> {
const turn = this.turnRepository.create({
threadId,
agentId: null,
});
const savedTurn = await this.turnRepository.save(turn);
const result = await this.messageRepository.update(
{ id: messageId, threadId, status: AgentMessageStatus.QUEUED },
{
status: AgentMessageStatus.SENT,
processedAt: new Date(),
turnId: savedTurn.id,
},
);
if ((result.affected ?? 0) === 0) {
await this.turnRepository.delete(savedTurn.id);
return null;
}
return savedTurn.id;
}
async generateTitleIfNeeded({
threadId,
messageContent,
workspaceId,
}: {
threadId: string;
messageContent: string;
workspaceId: string;
}): Promise<string | null> {
const thread = await this.threadRepository.findOne({
where: { id: threadId },
select: ['id', 'title'],
});
if (!thread || thread.title || !messageContent) {
@@ -147,6 +282,21 @@ export class AgentChatService {
await this.threadRepository.update(threadId, { title });
await this.workspaceEventBroadcaster.broadcast({
workspaceId,
events: [
{
type: 'updated',
entityName: 'agentChatThread',
recordId: threadId,
properties: {
updatedFields: ['title'],
after: serializeThreadForBroadcast({ ...thread, title }),
},
},
],
});
return title;
}
}
@@ -1,4 +1,5 @@
export enum SubscriptionChannel {
LOGIC_FUNCTION_LOGS_CHANNEL = 'LOGIC_FUNCTION_LOGS_CHANNEL',
EVENT_STREAM_CHANNEL = 'EVENT_STREAM_CHANNEL',
AGENT_CHAT_CHANNEL = 'AGENT_CHAT_CHANNEL',
}
@@ -6,56 +6,35 @@ import { OBJECT_METADATA_STANDARD_OVERRIDES_PROPERTIES } from 'src/engine/metada
import { WorkspaceManyOrAllFlatEntityMapsCacheService } from 'src/engine/metadata-modules/flat-entity/services/workspace-many-or-all-flat-entity-maps-cache.service';
import { NavigationMenuItemRecordIdentifierService } from 'src/engine/metadata-modules/navigation-menu-item/services/navigation-menu-item-record-identifier.service';
import { type MetadataEventBatch } from 'src/engine/subscriptions/metadata-event/types/metadata-event-batch.type';
import { type EventStreamPayload } from 'src/engine/subscriptions/types/event-stream-payload.type';
import { EventStreamService } from 'src/engine/subscriptions/event-stream.service';
import { SubscriptionService } from 'src/engine/subscriptions/subscription.service';
import { WorkspaceEventBroadcaster } from 'src/engine/subscriptions/workspace-event-broadcaster/workspace-event-broadcaster.service';
import { enrichFieldMetadataEventWithRelations } from 'src/engine/subscriptions/metadata-event/utils/enrich-field-metadata-event-with-relations.util';
@Injectable()
export class MetadataEventPublisher {
constructor(
private readonly subscriptionService: SubscriptionService,
private readonly eventStreamService: EventStreamService,
private readonly workspaceEventBroadcaster: WorkspaceEventBroadcaster,
private readonly workspaceManyOrAllFlatEntityMapsCacheService: WorkspaceManyOrAllFlatEntityMapsCacheService,
private readonly navigationMenuItemRecordIdentifierService: NavigationMenuItemRecordIdentifierService,
) {}
async publish(metadataEventBatch: MetadataEventBatch): Promise<void> {
const workspaceId = metadataEventBatch.workspaceId;
const activeStreamIds =
await this.eventStreamService.getActiveStreamIds(workspaceId);
if (activeStreamIds.length === 0) {
if (!isNonEmptyArray(metadataEventBatch.events)) {
return;
}
const streamsData = await this.eventStreamService.getStreamsData(
workspaceId,
activeStreamIds,
);
const enrichedBatch =
await this.enrichMetadataEventBatch(metadataEventBatch);
const streamIdsToRemove: string[] = [];
for (const [streamChannelId, streamData] of streamsData) {
if (!isDefined(streamData)) {
streamIdsToRemove.push(streamChannelId);
continue;
}
await this.publishToStream({
streamChannelId,
metadataEventBatch: enrichedBatch,
});
}
await this.eventStreamService.removeFromActiveStreams(
workspaceId,
streamIdsToRemove,
);
await this.workspaceEventBroadcaster.broadcast({
workspaceId: enrichedBatch.workspaceId,
updatedCollectionHash: enrichedBatch.updatedCollectionHash,
events: enrichedBatch.events.map((event) => ({
type: event.type,
entityName: event.metadataName,
recordId: event.recordId,
properties: event.properties as Record<string, unknown>,
})),
});
}
private async enrichMetadataEventBatch(
@@ -79,34 +58,6 @@ export class MetadataEventPublisher {
}
}
private async publishToStream({
streamChannelId,
metadataEventBatch,
}: {
streamChannelId: string;
metadataEventBatch: MetadataEventBatch;
}): Promise<void> {
if (!isNonEmptyArray(metadataEventBatch.events)) {
return;
}
const metadataEvents = metadataEventBatch.events.map((metadataEvent) => ({
...metadataEvent,
updatedCollectionHash: metadataEventBatch.updatedCollectionHash,
}));
const payload: EventStreamPayload = {
objectRecordEventsWithQueryIds: [],
metadataEvents,
};
await this.subscriptionService.publishToEventStream({
workspaceId: metadataEventBatch.workspaceId,
eventStreamChannelId: streamChannelId,
payload,
});
}
private async enrichFieldMetadataEventsWithRelations(
metadataEventBatch: MetadataEventBatch<'fieldMetadata'>,
): Promise<MetadataEventBatch<'fieldMetadata'>> {
@@ -88,4 +88,45 @@ export class SubscriptionService {
payload,
);
}
private getAgentChatChannel({
workspaceId,
threadId,
}: {
workspaceId: string;
threadId: string;
}) {
return `${SubscriptionChannel.AGENT_CHAT_CHANNEL}:${workspaceId}:${threadId}`;
}
async subscribeToAgentChat({
workspaceId,
threadId,
}: {
workspaceId: string;
threadId: string;
}) {
const client = this.redisClient.getPubSubClient();
return client.asyncIterator(
this.getAgentChatChannel({ workspaceId, threadId }),
);
}
async publishToAgentChat<T>({
workspaceId,
threadId,
payload,
}: {
workspaceId: string;
threadId: string;
payload: T;
}): Promise<void> {
const client = this.redisClient.getPubSubClient();
await client.publish(
this.getAgentChatChannel({ workspaceId, threadId }),
payload,
);
}
}
@@ -18,6 +18,7 @@ import { MetadataEventPublisher } from 'src/engine/subscriptions/metadata-event/
import { MetadataEventsToDbListener } from 'src/engine/subscriptions/metadata-event/metadata-events-to-db.listener';
import { ObjectRecordEventPublisher } from 'src/engine/subscriptions/object-record-event/object-record-event-publisher';
import { SubscriptionService } from 'src/engine/subscriptions/subscription.service';
import { WorkspaceEventBroadcaster } from 'src/engine/subscriptions/workspace-event-broadcaster/workspace-event-broadcaster.service';
import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache.module';
@Global()
@@ -40,6 +41,7 @@ import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache
MetadataEventPublisher,
MetadataEventEmitter,
MetadataEventsToDbListener,
WorkspaceEventBroadcaster,
ProcessNestedRelationsHelper,
ProcessNestedRelationsV2Helper,
CommonSelectFieldsHelper,
@@ -48,6 +50,7 @@ import { WorkspaceCacheModule } from 'src/engine/workspace-cache/workspace-cache
SubscriptionService,
ObjectRecordEventPublisher,
MetadataEventEmitter,
WorkspaceEventBroadcaster,
],
})
export class SubscriptionsModule {}
@@ -0,0 +1,12 @@
export type EventStreamMetadataEvent = {
metadataName: string;
type: 'created' | 'updated' | 'deleted';
recordId: string;
properties: {
updatedFields?: string[];
before?: Record<string, unknown>;
after?: Record<string, unknown>;
diff?: Record<string, unknown>;
};
updatedCollectionHash?: string;
};
@@ -1,10 +1,10 @@
import { type EventStreamMetadataEvent } from 'src/engine/subscriptions/types/event-stream-metadata-event.type';
import { type ObjectRecordSubscriptionEvent } from 'src/engine/subscriptions/types/object-record-subscription-event.type';
import { type MetadataEvent } from 'src/engine/workspace-manager/workspace-migration/workspace-migration-runner/types/metadata-event';
export type EventStreamPayload = {
objectRecordEventsWithQueryIds: {
queryIds: string[];
objectRecordEvent: ObjectRecordSubscriptionEvent;
}[];
metadataEvents: MetadataEvent[];
metadataEvents: EventStreamMetadataEvent[];
};
@@ -0,0 +1,11 @@
export type WorkspaceBroadcastEvent = {
type: 'created' | 'updated' | 'deleted';
entityName: string;
recordId: string;
properties: {
updatedFields?: string[];
before?: Record<string, unknown>;
after?: Record<string, unknown>;
diff?: Record<string, unknown>;
};
};
@@ -0,0 +1,73 @@
import { Injectable } from '@nestjs/common';
import { isDefined } from 'twenty-shared/utils';
import { EventStreamService } from 'src/engine/subscriptions/event-stream.service';
import { SubscriptionService } from 'src/engine/subscriptions/subscription.service';
import { type EventStreamPayload } from 'src/engine/subscriptions/types/event-stream-payload.type';
import { type WorkspaceBroadcastEvent } from 'src/engine/subscriptions/workspace-event-broadcaster/types/workspace-broadcast-event.type';
@Injectable()
export class WorkspaceEventBroadcaster {
constructor(
private readonly eventStreamService: EventStreamService,
private readonly subscriptionService: SubscriptionService,
) {}
async broadcast({
workspaceId,
events,
updatedCollectionHash,
}: {
workspaceId: string;
events: WorkspaceBroadcastEvent[];
updatedCollectionHash?: string;
}): Promise<void> {
if (events.length === 0) {
return;
}
const activeStreamIds =
await this.eventStreamService.getActiveStreamIds(workspaceId);
if (activeStreamIds.length === 0) {
return;
}
const streamsData = await this.eventStreamService.getStreamsData(
workspaceId,
activeStreamIds,
);
const streamIdsToRemove: string[] = [];
const payload: EventStreamPayload = {
objectRecordEventsWithQueryIds: [],
metadataEvents: events.map((event) => ({
metadataName: event.entityName,
type: event.type,
recordId: event.recordId,
properties: event.properties,
updatedCollectionHash,
})),
};
for (const [streamChannelId, streamData] of streamsData) {
if (!isDefined(streamData)) {
streamIdsToRemove.push(streamChannelId);
continue;
}
await this.subscriptionService.publishToEventStream({
workspaceId,
eventStreamChannelId: streamChannelId,
payload,
});
}
await this.eventStreamService.removeFromActiveStreams(
workspaceId,
streamIdsToRemove,
);
}
}
+1
View File
@@ -18,6 +18,7 @@ export type {
AgentResponseFieldType,
AgentResponseSchema,
} from './types/agent-response-schema.type';
export type { AgentChatSubscriptionEvent } from './types/AgentChatSubscriptionEvent';
export type {
CodeExecutionFile,
ExtendedFileUIPart,
@@ -0,0 +1,8 @@
// `chunk` is typed as Record<string, unknown> rather than UIMessageChunk
// to avoid importing the `ai` package in twenty-shared. The frontend casts
// it back to UIMessageChunk when feeding it into readUIMessageStream.
export type AgentChatSubscriptionEvent =
| { type: 'stream-chunk'; chunk: Record<string, unknown>; seq?: number }
| { type: 'message-persisted'; messageId: string }
| { type: 'queue-updated' }
| { type: 'stream-error'; code: string; message: string };
@@ -23,4 +23,5 @@ type Metadata = {
export type ExtendedUIMessage = UIMessage<Metadata, DataMessagePart> & {
threadId?: Nullable<string>;
status?: 'queued' | 'sent';
};
-42
View File
@@ -195,20 +195,6 @@ __metadata:
languageName: node
linkType: hard
"@ai-sdk/react@npm:3.0.99":
version: 3.0.99
resolution: "@ai-sdk/react@npm:3.0.99"
dependencies:
"@ai-sdk/provider-utils": "npm:4.0.15"
ai: "npm:6.0.97"
swr: "npm:^2.2.5"
throttleit: "npm:2.1.0"
peerDependencies:
react: ^18 || ~19.0.1 || ~19.1.2 || ^19.2.1
checksum: 10c0/e8ee12df235e13ef238ffd4a78c34b53c989ab5803cfe26453458bb7ecf001e15e9afb303f294d8f1c189f31260db4726a34615657ac94ffd7e058172549fd43
languageName: node
linkType: hard
"@ai-sdk/xai@npm:^3.0.57":
version: 3.0.74
resolution: "@ai-sdk/xai@npm:3.0.74"
@@ -55145,13 +55131,6 @@ __metadata:
languageName: node
linkType: hard
"resumable-stream@npm:^2.2.12":
version: 2.2.12
resolution: "resumable-stream@npm:2.2.12"
checksum: 10c0/6fa3ddb31d7d88e720c8313532799e7e913dec31890fd1132d4a324718fe55eeb82782c4a571c3be3f6ce8e0396253eb5b775983f035a6a88d01a886154420c2
languageName: node
linkType: hard
"retext-latin@npm:^4.0.0":
version: 4.0.0
resolution: "retext-latin@npm:4.0.0"
@@ -58280,18 +58259,6 @@ __metadata:
languageName: node
linkType: hard
"swr@npm:^2.2.5":
version: 2.3.6
resolution: "swr@npm:2.3.6"
dependencies:
dequal: "npm:^2.0.3"
use-sync-external-store: "npm:^1.4.0"
peerDependencies:
react: ^16.11.0 || ^17.0.0 || ^18.0.0 || ^19.0.0
checksum: 10c0/9534f350982e36a3ae0a13da8c0f7da7011fc979e77f306e60c4e5db0f9b84f17172c44f973441ba56bb684b69b0d9838ab40011a6b6b3e32d0cd7f3d5405f99
languageName: node
linkType: hard
"symbol-observable@npm:4.0.0":
version: 4.0.0
resolution: "symbol-observable@npm:4.0.0"
@@ -58692,13 +58659,6 @@ __metadata:
languageName: node
linkType: hard
"throttleit@npm:2.1.0":
version: 2.1.0
resolution: "throttleit@npm:2.1.0"
checksum: 10c0/1696ae849522cea6ba4f4f3beac1f6655d335e51b42d99215e196a718adced0069e48deaaf77f7e89f526ab31de5b5c91016027da182438e6f9280be2f3d5265
languageName: node
linkType: hard
"through2@npm:4.0.2, through2@npm:^4.0.2":
version: 4.0.2
resolution: "through2@npm:4.0.2"
@@ -59634,7 +59594,6 @@ __metadata:
version: 0.0.0-use.local
resolution: "twenty-front@workspace:packages/twenty-front"
dependencies:
"@ai-sdk/react": "npm:3.0.99"
"@apollo/client": "npm:^4.0.0"
"@blocknote/mantine": "npm:^0.47.1"
"@blocknote/react": "npm:^0.47.1"
@@ -59991,7 +59950,6 @@ __metadata:
react-dom: "npm:18.3.1"
redis: "npm:^4.7.0"
reflect-metadata: "npm:0.2.2"
resumable-stream: "npm:^2.2.12"
rimraf: "npm:^5.0.5"
rxjs: "npm:7.8.1"
semver: "npm:7.6.3"