Compare commits

...

6 Commits

Author SHA1 Message Date
Lucas Bordeau b4f5e4b55e Added handoff 2026-03-27 17:36:33 +01:00
Lucas Bordeau fc546f90ff Added debug logs 2026-03-27 17:32:45 +01:00
Lucas Bordeau 79d9e75b6e WIP 2026-03-27 14:55:26 +01:00
Lucas Bordeau 2c90edaa59 Removed claude config from versionning 2026-03-27 14:54:41 +01:00
Lucas Bordeau 480b78c55f WIP 2026-03-27 14:54:41 +01:00
Lucas Bordeau fbd6c9daea WIP 2026-03-27 14:54:36 +01:00
19 changed files with 664 additions and 147 deletions
+1
View File
@@ -52,3 +52,4 @@ mcp.json
/.junie/
TRANSLATION_QA_REPORT.md
.playwright-mcp/
.claude/
+125
View File
@@ -0,0 +1,125 @@
# Handoff: AI Chat Streaming — Alternating Messages Bug
## Branch
`fix/ai-chat-thread-switching`
## Problem
During streaming, the UI flickers: the last assistant message alternates between two versions with different `parts` arrays. This continues until streaming finishes, at which point the message stabilizes.
## Root Cause Analysis
The bug comes from having **two independent sources writing messages to the Chat instance**, creating a tug-of-war during streaming.
### Source 1: AI SDK streaming (internal)
`useChat` subscribes to the Chat instance's messages via `useSyncExternalStore`. During streaming, the SDK calls `Chat.replaceMessage()` on every chunk, which:
1. Calls `structuredClone(message)` on the updated last message
2. Replaces it in the internal `#messages` array
3. Fires all registered callbacks (throttled at 100ms via `experimental_throttle`)
### Source 2: GraphQL fetch (external mutation)
`AgentChatMessagesFetchEffect.handleDataLoaded` (line 76) directly mutates the Chat instance:
```ts
chatInstance.messages = uiMessages;
```
This triggers the Chat's `set messages()` setter, which creates `[...newMessages]` and fires ALL registered callbacks — including the throttled one that `useSyncExternalStore` uses.
### The guard that doesn't fully work
There IS a guard in `handleDataLoaded` (lines 68-74):
```ts
const isStreaming = chatInstance.status === 'streaming' || chatInstance.status === 'submitted';
if (isStreaming) { return; }
```
But this guard only protects against the **initial fetch** and **explicit refetch** during streaming. The problem occurs because:
1. `onStreamingComplete` dispatches `AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME`
2. This triggers `refetchAgentChatMessages()` — a GraphQL network call
3. The response arrives asynchronously, potentially AFTER the SDK has already updated internal state but BEFORE React has committed the final render
4. More critically: Apollo's `useQuery` can update its cache and trigger `onDataLoaded` from cache operations (optimistic updates, cache writes from other queries) even during streaming
### The alternation mechanism
When both sources write to the Chat instance in rapid succession:
1. SDK streaming updates `#messages` with the latest chunk → snapshot A (has new parts)
2. Apollo cache/refetch triggers `handleDataLoaded``chatInstance.messages = uiMessages` → snapshot B (DB version, missing latest streamed parts)
3. `useSyncExternalStore` detects the snapshot changed → re-render with B
4. Next throttle tick, SDK has already updated internally → snapshot A again
5. React renders with A, then B arrives again → alternation
The user sees the last message's `parts` flickering between the streaming version (more content) and the DB version (less content).
## Architecture Context
### How useChat works internally (`@ai-sdk/react`)
```
useChat({ chat: instance, experimental_throttle: 100 })
└── chatRef = useRef(instance)
└── useSyncExternalStore(
subscribe: throttle(onChange, 100ms), // only fires every 100ms
getSnapshot: () => chatRef.current.messages // but reads LATEST on every render
)
```
Key detail: `getSnapshot` always reads the latest `chatRef.current.messages`. React can call this during ANY render — not just when the subscribe callback fires. If something else triggers a render (Jotai state update, parent re-render), React will read the current snapshot and may detect it differs from the last committed value, causing an additional re-render.
### Data flow (current)
```
AI SDK streaming ──replaceMessage()──► Chat#messages ──useSyncExternalStore──► useChat.messages
▲ │
GraphQL fetch ──handleDataLoaded──► chatInstance.messages = uiMessages │
setAgentChatMessages (Jotai)
agentChatMessagesComponentFamilyState
UI renders
```
The problem is the two arrows writing to `Chat#messages`.
### Key files
| File | Role |
|---|---|
| `useAgentChat.ts` | Calls `useChat({ chat: agentChatInstanceForThread })`, returns messages/status |
| `useAgentChatInstanceForThread.ts` | Creates/caches Chat instances per thread in Jotai atom family |
| `AgentChatAiSdkStreamEffect.tsx` | Bridges useChat → Jotai: `setAgentChatMessages(chatState.messages)` |
| `AgentChatMessagesFetchEffect.tsx` | Fetches historical messages from GraphQL, writes to Chat instance |
| `agentChatInstanceByThreadIdFamilyState.ts` | Jotai atom family holding Chat instances keyed by `{ threadId }` |
### AI SDK internals (relevant)
| Method | What it does |
|---|---|
| `Chat.replaceMessage(index, message)` | `structuredClone`s the message, replaces at index, fires callbacks |
| `Chat.messages` setter | Creates `[...newMessages]`, fires callbacks |
| `Chat.messages` getter | Returns internal `#messages` reference |
| `~registerMessagesCallback(onChange, throttleMs)` | Wraps onChange with `throttleit` if throttleMs provided |
## Logging
There are `console.log` statements in `AgentChatAiSdkStreamEffect.tsx` and a `lastMessageMap` (module-level Map) that captures every message snapshot with a timestamp. These were added for debugging this exact issue and will be kept as-is in the PR.
## Possible Fix Directions
1. **Make Chat instance the single source of truth during streaming**: Remove the `chatInstance.messages = uiMessages` path entirely. Instead, only seed messages at Chat construction time (before streaming starts). After streaming ends, refetch and set.
2. **Use setMessages from useChat instead of direct mutation**: The AI SDK's `useChat` returns a `setMessages` function that properly coordinates with the internal state. Replace `chatInstance.messages = uiMessages` with the SDK's own `setMessages`. This ensures the SDK controls the timing.
3. **Strengthen the guard**: Instead of checking `chatInstance.status`, use a dedicated atom that tracks whether the SDK is actively streaming for this thread. Set it before the stream starts, clear it after the effect processes the final message.
4. **Debounce/skip the fetch effect during streaming**: Don't refetch at all while streaming. Only fetch historical messages when switching TO a thread that isn't streaming.
## What's Working
- Thread creation (draft → first send → thread created)
- Message sending and receiving (when not observing the alternation)
- Thread switching (historical messages load correctly for non-streaming threads)
- Backend resumable streams (Redis infrastructure, `GET /agent-chat/:threadId/stream`)
- `onFinish` callback (title update, usage tracking)
@@ -2,10 +2,15 @@ import { AIChatMessage } from '@/ai/components/AIChatMessage';
import { agentChatErrorState } from '@/ai/states/agentChatErrorState';
import { agentChatIsStreamingState } from '@/ai/states/agentChatIsStreamingState';
import { agentChatLastMessageIdComponentSelector } from '@/ai/states/agentChatLastMessageIdComponentSelector';
import { agentChatMessageComponentFamilySelector } from '@/ai/states/agentChatMessageComponentFamilySelector';
import { useAtomComponentFamilySelectorValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilySelectorValue';
import { useAtomComponentSelectorValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentSelectorValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
import { isDefined } from 'twenty-shared/utils';
const messageMap = new Map<string, ExtendedUIMessage>();
export const AIChatLastMessageWithStreamingState = () => {
const lastMessageId = useAtomComponentSelectorValue(
agentChatLastMessageIdComponentSelector,
@@ -14,6 +19,17 @@ export const AIChatLastMessageWithStreamingState = () => {
const agentChatIsStreaming = useAtomStateValue(agentChatIsStreamingState);
const agentChatError = useAtomStateValue(agentChatErrorState);
const agentChatMessage = useAtomComponentFamilySelectorValue(
agentChatMessageComponentFamilySelector,
{ messageId: lastMessageId },
);
messageMap.set(new Date().toISOString(), agentChatMessage!);
console.log({
messageMap,
});
if (!isDefined(lastMessageId)) {
return null;
}
@@ -1,35 +1,31 @@
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 { 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 { agentChatIsInitialScrollPendingOnThreadChangeState } from '@/ai/states/agentChatIsInitialScrollPendingOnThreadChangeState';
import { agentChatIsLoadingState } from '@/ai/states/agentChatIsLoadingState';
import { agentChatIsStreamingState } from '@/ai/states/agentChatIsStreamingState';
import { normalizeAiSdkError } from '@/ai/utils/normalizeAiSdkError';
import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMessagesComponentFamilyState';
import { agentChatMessagesLoadingState } from '@/ai/states/agentChatMessagesLoadingState';
import { agentChatThreadsLoadingState } from '@/ai/states/agentChatThreadsLoadingState';
import { agentChatDisplayedThreadState } from '@/ai/states/agentChatDisplayedThreadState';
import { agentChatFetchedMessagesComponentFamilyState } from '@/ai/states/agentChatFetchedMessagesComponentFamilyState';
import { agentChatIsInitialScrollPendingOnThreadChangeState } from '@/ai/states/agentChatIsInitialScrollPendingOnThreadChangeState';
import { mergeAgentChatFetchedAndStreamingMessages } from '@/ai/utils/mergeAgentChatFetchedAndStreamingMessages';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { normalizeAiSdkError } from '@/ai/utils/normalizeAiSdkError';
import { useListenToBrowserEvent } from '@/browser-event/hooks/useListenToBrowserEvent';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
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 { type ExtendedUIMessage } from 'twenty-shared/ai';
const lastMessageMap = new Map<string, ExtendedUIMessage>();
export const AgentChatAiSdkStreamEffect = () => {
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const agentChatFetchedMessages = useAtomComponentFamilyStateValue(
agentChatFetchedMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
const { createChatThread } = useCreateAgentChatThread();
@@ -48,11 +44,7 @@ export const AgentChatAiSdkStreamEffect = () => {
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
}, []);
const chatState = useAgentChat(
agentChatFetchedMessages,
ensureThreadIdForSend,
onStreamingComplete,
);
const chatState = useAgentChat(ensureThreadIdForSend, onStreamingComplete);
const setAgentChatMessages = useSetAtomComponentFamilyState(
agentChatMessagesComponentFamilyState,
@@ -72,20 +64,24 @@ export const AgentChatAiSdkStreamEffect = () => {
);
useEffect(() => {
const mergedMessages = mergeAgentChatFetchedAndStreamingMessages(
agentChatFetchedMessages,
chatState.messages,
);
setAgentChatMessages(mergedMessages);
setAgentChatMessages(chatState.messages);
const lastMessage = chatState.messages.at(-1);
lastMessageMap.set(new Date().toISOString(), lastMessage!);
console.log('AgentChatAiSdkStreamEffect - Updated messages', {
messages: chatState.messages,
lastMessage,
lastMessageMap,
});
if (currentAIChatThread !== agentChatDisplayedThread) {
if (mergedMessages.length > 0) {
if (chatState.messages.length > 0) {
setAgentChatIsInitialScrollPendingOnThreadChange(true);
}
setAgentChatDisplayedThread(currentAIChatThread);
}
}, [
agentChatFetchedMessages,
chatState.messages,
chatState.status,
setAgentChatMessages,
@@ -1,10 +1,11 @@
import { useCallback, useMemo } from 'react';
import { useStore } from 'jotai';
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 { AGENT_CHAT_NEW_THREAD_DRAFT_KEY } from '@/ai/states/agentChatDraftsByThreadIdState';
import { agentChatFetchedMessagesComponentFamilyState } from '@/ai/states/agentChatFetchedMessagesComponentFamilyState';
import { agentChatInstanceByThreadIdFamilyState } from '@/ai/states/agentChatInstanceByThreadIdFamilyState';
import { agentChatMessagesLoadingState } from '@/ai/states/agentChatMessagesLoadingState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { skipMessagesSkeletonUntilLoadedState } from '@/ai/states/skipMessagesSkeletonUntilLoadedState';
@@ -12,7 +13,6 @@ import { mapDBMessagesToUIMessages } from '@/ai/utils/mapDBMessagesToUIMessages'
import { useQueryWithCallbacks } from '@/apollo/hooks/useQueryWithCallbacks';
import { useListenToBrowserEvent } from '@/browser-event/hooks/useListenToBrowserEvent';
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 {
GetChatMessagesDocument,
@@ -21,6 +21,7 @@ import {
export const AgentChatMessagesFetchEffect = () => {
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const store = useStore();
const isNewThread = useMemo(
() =>
@@ -37,11 +38,6 @@ export const AgentChatMessagesFetchEffect = () => {
skipMessagesSkeletonUntilLoadedState,
);
const setAgentChatFetchedMessages = useSetAtomComponentFamilyState(
agentChatFetchedMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
const handleFirstLoad = useCallback(
(_data: GetChatMessagesQuery) => {
setSkipMessagesSkeletonUntilLoaded(false);
@@ -52,9 +48,34 @@ export const AgentChatMessagesFetchEffect = () => {
const handleDataLoaded = useCallback(
(data: GetChatMessagesQuery) => {
const uiMessages = mapDBMessagesToUIMessages(data.chatMessages ?? []);
setAgentChatFetchedMessages(uiMessages);
const threadId = store.get(currentAIChatThreadState.atom);
if (!isDefined(threadId)) {
return;
}
const threadAtom = agentChatInstanceByThreadIdFamilyState.atomFamily({
threadId,
});
const chatInstance = store.get(threadAtom);
if (chatInstance === null) {
return;
}
const isStreaming =
chatInstance.status === 'streaming' ||
chatInstance.status === 'submitted';
if (isStreaming) {
return;
}
chatInstance.messages = uiMessages;
},
[setAgentChatFetchedMessages],
[store],
);
const handleLoadingChange = useCallback(
@@ -10,13 +10,14 @@ import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTi
import { AGENT_CHAT_RETRY_EVENT_NAME } from '@/ai/constants/AgentChatRetryEventName';
import { AGENT_CHAT_STOP_EVENT_NAME } from '@/ai/constants/AgentChatStopEventName';
import { AGENT_CHAT_UNKNOWN_THREAD_ID } from '@/ai/constants/AgentChatUnknownThreadId';
import { useAgentChatInstanceForThread } from '@/ai/hooks/useAgentChatInstanceForThread';
import { useAgentChatModelId } from '@/ai/hooks/useAgentChatModelId';
import {
AGENT_CHAT_NEW_THREAD_DRAFT_KEY,
agentChatDraftsByThreadIdState,
} from '@/ai/states/agentChatDraftsByThreadIdState';
import { agentChatInputState } from '@/ai/states/agentChatInputState';
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';
@@ -25,14 +26,14 @@ 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, useState } from 'react';
import { useCallback, useEffect, useState } from 'react';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
import { isDefined } from 'twenty-shared/utils';
import { REACT_APP_SERVER_BASE_URL } from '~/config';
import { cookieStorage } from '~/utils/cookie-storage';
export const useAgentChat = (
uiMessages: ExtendedUIMessage[],
ensureThreadIdForSend: () => Promise<string | null>,
onStreamingComplete?: () => void,
) => {
@@ -63,86 +64,53 @@ export const useAgentChat = (
agentChatDraftsByThreadIdState,
);
const retryFetchWithRenewedToken = async (
input: RequestInfo | URL,
init?: RequestInit,
) => {
const tokenPair = getTokenPair();
const retryFetchWithRenewedToken = useCallback(
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);
if (!isDefined(tokenPair)) {
return null;
}
const renewedAccessToken =
renewedTokens.accessOrWorkspaceAgnosticToken?.token;
try {
const renewedTokens = await renewToken(
`${REACT_APP_SERVER_BASE_URL}/metadata`,
tokenPair,
);
if (!isDefined(renewedAccessToken)) {
if (!isDefined(renewedTokens)) {
setTokenPair(null);
return null;
}
const renewedAccessToken =
renewedTokens.accessOrWorkspaceAgnosticToken?.token;
if (!isDefined(renewedAccessToken)) {
setTokenPair(null);
return null;
}
cookieStorage.setItem('tokenPair', JSON.stringify(renewedTokens));
setTokenPair(renewedTokens);
const updatedHeaders = new Headers(init?.headers ?? {});
updatedHeaders.set('Authorization', `Bearer ${renewedAccessToken}`);
return fetch(input, {
...init,
headers: updatedHeaders,
});
} catch {
setTokenPair(null);
return null;
}
},
[setTokenPair],
);
setTokenPair(renewedTokens);
const updatedHeaders = new Headers(init?.headers ?? {});
updatedHeaders.set('Authorization', `Bearer ${renewedAccessToken}`);
return fetch(input, {
...init,
headers: updatedHeaders,
});
} catch {
setTokenPair(null);
return null;
}
};
const { sendMessage, messages, status, error, regenerate, stop } = useChat({
transport: new DefaultChatTransport({
api: `${REST_API_BASE_URL}/agent-chat/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;
}
// For non-2xx responses, parse the error body and throw with the code
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;
},
}),
messages: uiMessages,
id: `${currentAIChatThread}-${uiMessages.length}`,
experimental_throttle: 100,
onFinish: ({ message }) => {
const handleOnFinish = useCallback(
({ message }: { message: ExtendedUIMessage }) => {
type UsageMetadata = {
inputTokens: number;
outputTokens: number;
@@ -183,7 +151,8 @@ export const useAgentChat = (
);
setPendingThreadIdAfterFirstSend((pendingId) => {
const threadIdForTitle = pendingId ?? currentAIChatThread;
const threadIdForTitle =
pendingId ?? store.get(currentAIChatThreadState.atom);
if (isDefined(titlePart) && titlePart.type === 'data-thread-title') {
setCurrentAIChatThreadTitle(titlePart.data.title);
if (isDefined(threadIdForTitle)) {
@@ -209,8 +178,47 @@ export const useAgentChat = (
onStreamingComplete?.();
},
[
setAgentChatUsage,
store,
apolloClient.cache,
setCurrentAIChatThreadTitle,
setCurrentAIChatThread,
onStreamingComplete,
],
);
const { agentChatInstanceForThread } = useAgentChatInstanceForThread({
threadId: currentAIChatThread ?? AGENT_CHAT_UNKNOWN_THREAD_ID,
onFinish: handleOnFinish,
retryFetchWithRenewedToken,
});
const {
sendMessage,
messages,
status,
error,
regenerate,
stop,
resumeStream,
} = useChat({
chat: agentChatInstanceForThread,
resume: true,
experimental_throttle: 100,
});
useEffect(() => {
const isRealThreadId =
isDefined(currentAIChatThread) &&
currentAIChatThread !== AGENT_CHAT_UNKNOWN_THREAD_ID &&
currentAIChatThread !== AGENT_CHAT_NEW_THREAD_DRAFT_KEY;
if (isRealThreadId) {
resumeStream();
}
}, [currentAIChatThread, resumeStream]);
const isStreaming = status === 'streaming';
const isLoading = isStreaming || agentChatSelectedFiles.length > 0;
@@ -0,0 +1,95 @@
import { Chat } from '@ai-sdk/react';
import { DefaultChatTransport } from 'ai';
import { useStore } from 'jotai';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
import { agentChatInstanceByThreadIdFamilyState } from '@/ai/states/agentChatInstanceByThreadIdFamilyState';
import { REST_API_BASE_URL } from '@/apollo/constant/rest-api-base-url';
import { getTokenPair } from '@/apollo/utils/getTokenPair';
type UseAgentChatInstanceForThreadOptions = {
threadId: string;
onFinish: (options: { message: ExtendedUIMessage }) => void;
retryFetchWithRenewedToken: (
input: RequestInfo | URL,
init?: RequestInit,
) => Promise<Response | null>;
};
const createAgentChatInstanceForThread = ({
threadId,
onFinish,
retryFetchWithRenewedToken,
}: UseAgentChatInstanceForThreadOptions): Chat<ExtendedUIMessage> => {
return new Chat<ExtendedUIMessage>({
id: threadId,
messages: [],
transport: 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 (errorBody.code !== undefined) {
error.code = errorBody.code;
}
throw error;
}
return response;
},
}),
onFinish,
});
};
export const useAgentChatInstanceForThread = ({
threadId,
onFinish,
retryFetchWithRenewedToken,
}: UseAgentChatInstanceForThreadOptions): {
agentChatInstanceForThread: Chat<ExtendedUIMessage>;
} => {
const store = useStore();
const threadAtom = agentChatInstanceByThreadIdFamilyState.atomFamily({
threadId,
});
const existingInstance = store.get(threadAtom);
if (existingInstance !== null) {
return { agentChatInstanceForThread: existingInstance };
}
const newInstance = createAgentChatInstanceForThread({
threadId,
onFinish,
retryFetchWithRenewedToken,
});
store.set(threadAtom, newInstance);
return { agentChatInstanceForThread: newInstance };
};
@@ -1,10 +0,0 @@
import { AgentChatComponentInstanceContext } from '@/ai/states/AgentChatComponentInstanceContext';
import { createAtomComponentFamilyState } from '@/ui/utilities/state/jotai/utils/createAtomComponentFamilyState';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
export const agentChatFetchedMessagesComponentFamilyState =
createAtomComponentFamilyState<ExtendedUIMessage[], { threadId: string }>({
key: 'agentChatFetchedMessagesComponentFamilyState',
defaultValue: [],
componentInstanceContext: AgentChatComponentInstanceContext,
});
@@ -0,0 +1,12 @@
import { type Chat } from '@ai-sdk/react';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
import { createAtomFamilyState } from '@/ui/utilities/state/jotai/utils/createAtomFamilyState';
export const agentChatInstanceByThreadIdFamilyState = createAtomFamilyState<
Chat<ExtendedUIMessage> | null,
{ threadId: string }
>({
key: 'agentChatInstanceByThreadIdFamilyState',
defaultValue: null,
});
@@ -1,16 +0,0 @@
import { type ExtendedUIMessage } from 'twenty-shared/ai';
export const mergeAgentChatFetchedAndStreamingMessages = (
fetchedMessages: ExtendedUIMessage[],
streamingMessages: ExtendedUIMessage[],
): ExtendedUIMessage[] => {
const fetchedMessageIds = new Set(
fetchedMessages.map((message) => message.id),
);
const streamingOnlyMessages = streamingMessages.filter(
(message) => !fetchedMessageIds.has(message.id),
);
return [...fetchedMessages, ...streamingOnlyMessages];
};
+1
View File
@@ -168,6 +168,7 @@
"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,19 @@
import { MigrationInterface, QueryRunner } from 'typeorm';
export class AddActiveStreamIdToAgentChatThread1774003611071
implements MigrationInterface
{
name = 'AddActiveStreamIdToAgentChatThread1774003611071';
public async up(queryRunner: QueryRunner): Promise<void> {
await queryRunner.query(
`ALTER TABLE "core"."agentChatThread" ADD "activeStreamId" character varying`,
);
}
public async down(queryRunner: QueryRunner): Promise<void> {
await queryRunner.query(
`ALTER TABLE "core"."agentChatThread" DROP COLUMN "activeStreamId"`,
);
}
}
@@ -0,0 +1,53 @@
import { type ToolUIPart } from 'ai';
import { type ExtendedUIMessage, type ExtendedUIMessagePart } from 'twenty-shared/ai';
import { type AgentMessagePartEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message-part.entity';
import { type AgentMessageEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
const mapAgentMessagePartEntityToUIMessagePart = (
part: AgentMessagePartEntity,
): ExtendedUIMessagePart | null => {
switch (part.type) {
case 'text':
return { type: 'text', text: part.textContent ?? '' };
case 'reasoning':
return { type: 'reasoning', text: part.reasoningContent ?? '', state: 'done' as const };
case 'step-start':
return { type: 'step-start' };
case 'source-url':
return {
type: 'source-url',
sourceId: part.sourceUrlSourceId ?? '',
url: part.sourceUrlUrl ?? '',
title: part.sourceUrlTitle ?? undefined,
providerMetadata: part.providerMetadata ?? undefined,
};
default: {
if (part.type.includes('tool-') && part.toolCallId) {
return {
type: part.type as `tool-${string}`,
toolCallId: part.toolCallId,
input: part.toolInput ?? {},
output: part.toolOutput ?? undefined,
errorText: part.errorMessage ?? undefined,
state: part.state ?? undefined,
} as ToolUIPart;
}
return null;
}
}
};
export const mapAgentMessageEntitiesToUIMessages = (
messageEntities: AgentMessageEntity[],
): ExtendedUIMessage[] => {
return messageEntities.map((message) => ({
id: message.id,
role: message.role as ExtendedUIMessage['role'],
parts: (message.parts ?? [])
.sort((partA, partB) => partA.orderIndex - partB.orderIndex)
.map(mapAgentMessagePartEntityToUIMessagePart)
.filter((part): part is ExtendedUIMessagePart => part !== null),
metadata: { createdAt: message.createdAt.toISOString() },
}));
};
@@ -36,6 +36,7 @@ import { AgentChatController } from './controllers/agent-chat.controller';
import { AgentChatThreadDTO } from './dtos/agent-chat-thread.dto';
import { AgentChatThreadEntity } from './entities/agent-chat-thread.entity';
import { AgentChatResolver } from './resolvers/agent-chat.resolver';
import { AgentChatResumableStreamService } from './services/agent-chat-resumable-stream.service';
import { AgentChatStreamingService } from './services/agent-chat-streaming.service';
import { AgentChatService } from './services/agent-chat.service';
import { AgentTitleGenerationService } from './services/agent-title-generation.service';
@@ -99,6 +100,7 @@ import { SystemPromptBuilderService } from './services/system-prompt-builder.ser
controllers: [AgentChatController],
providers: [
AgentChatResolver,
AgentChatResumableStreamService,
AgentChatService,
AgentChatStreamingService,
AgentTitleGenerationService,
@@ -1,6 +1,9 @@
import {
Body,
Controller,
Get,
Logger,
Param,
Post,
Res,
UseFilters,
@@ -9,8 +12,12 @@ import {
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 {
@@ -33,6 +40,8 @@ import {
} 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 { 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';
@@ -44,11 +53,16 @@ import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models
BillingRestApiExceptionFilter,
)
export class AgentChatController {
private readonly logger = new Logger(AgentChatController.name);
constructor(
private readonly agentStreamingService: AgentChatStreamingService,
private readonly resumableStreamService: AgentChatResumableStreamService,
private readonly billingService: BillingService,
private readonly twentyConfigService: TwentyConfigService,
private readonly aiModelRegistryService: AiModelRegistryService,
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
) {}
@Post('stream')
@@ -103,4 +117,57 @@ export class AgentChatController {
response,
});
}
@Get(':threadId/stream')
@UseGuards(SettingsPermissionGuard(PermissionFlagType.AI))
async resumeAgentChatStream(
@Param('threadId') threadId: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@Res() response: Response,
) {
console.log('Received request to resume stream for threadId:', threadId);
const uuidRegex =
/^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i;
if (!uuidRegex.test(threadId)) {
response.status(204).end();
return;
}
const thread = await this.threadRepository.findOne({
where: { id: threadId, userWorkspaceId },
});
this.logger.log(
`resume requested — thread found: ${!!thread} | activeStreamId: ${thread?.activeStreamId}`,
);
if (!isDefined(thread) || !isDefined(thread.activeStreamId)) {
response.status(204).end();
return;
}
const resumedNodeReadable =
await this.resumableStreamService.resumeExistingStreamAsNodeReadable(
thread.activeStreamId,
);
if (!isDefined(resumedNodeReadable)) {
this.logger.log(
'resumeExistingStream returned null — Redis stream already done or expired',
);
response.status(204).end();
return;
}
this.logger.log('resuming stream from Redis');
response.writeHead(200, UI_MESSAGE_STREAM_HEADERS);
resumedNodeReadable.pipe(response);
}
}
@@ -51,6 +51,9 @@ export class AgentChatThreadEntity {
@Column({ type: 'bigint', default: 0 })
totalOutputCredits: number;
@Column({ type: 'varchar', nullable: true })
activeStreamId: string | null;
@OneToMany(() => AgentTurnEntity, (turn) => turn.thread)
turns: EntityRelation<AgentTurnEntity[]>;
@@ -0,0 +1,77 @@
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 publisher: Redis;
private subscriber: Redis;
constructor(private readonly redisClientService: RedisClientService) {
const redisClient = this.redisClientService.getClient();
this.publisher = redisClient.duplicate();
this.subscriber = redisClient.duplicate();
this.streamContext = createResumableStreamContext({
waitUntil: (callback) => {
void Promise.resolve(callback);
},
publisher: this.publisher,
subscriber: this.subscriber,
});
}
async onModuleDestroy() {
await this.publisher.quit();
await this.subscriber.quit();
}
async createResumableStream(
streamId: string,
streamFactory: () => ReadableStream<string>,
) {
const resumableStream = await this.streamContext.createNewResumableStream(
streamId,
streamFactory,
);
if (!resumableStream) {
return;
}
const reader = resumableStream.getReader();
void (async () => {
try {
while (true) {
const { done } = await reader.read();
if (done) {
break;
}
}
} catch {
// ignore errors from background stream consumption
}
})();
}
async resumeExistingStreamAsNodeReadable(
streamId: string,
): Promise<Readable | null> {
const webStream = await this.streamContext.resumeExistingStream(streamId);
if (!webStream) {
return null;
}
return Readable.fromWeb(webStream as NodeWebReadableStream);
}
}
@@ -1,7 +1,11 @@
import { Injectable } from '@nestjs/common';
import { Injectable, Logger } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { createUIMessageStream, pipeUIMessageStreamToResponse } from 'ai';
import {
createUIMessageStream,
generateId,
pipeUIMessageStreamToResponse,
} from 'ai';
import { type Response } from 'express';
import {
type CodeExecutionData,
@@ -11,6 +15,7 @@ import { type Repository } from 'typeorm';
import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
import { AgentMessageRole } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
import { mapAgentMessageEntitiesToUIMessages } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapAgentMessageEntitiesToUIMessages';
import {
AgentException,
AgentExceptionCode,
@@ -19,10 +24,11 @@ import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agen
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 { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.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 { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util';
import { AIModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type';
import { AgentChatResumableStreamService } from './agent-chat-resumable-stream.service';
import { AgentChatService } from './agent-chat.service';
import { ChatExecutionService } from './chat-execution.service';
@@ -38,11 +44,14 @@ export type StreamAgentChatOptions = {
@Injectable()
export class AgentChatStreamingService {
private readonly logger = new Logger(AgentChatStreamingService.name);
constructor(
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
private readonly agentChatService: AgentChatService,
private readonly chatExecutionService: ChatExecutionService,
private readonly resumableStreamService: AgentChatResumableStreamService,
) {}
async streamAgentChat({
@@ -68,12 +77,25 @@ export class AgentChatStreamingService {
);
}
// Fire user-message save without awaiting to avoid delaying time-to-first-letter.
// The promise is awaited inside onFinish where we need the turnId.
const lastUserMessage = messages[messages.length - 1];
const lastUserText =
lastUserMessage?.parts.find((part) => part.type === 'text')?.text ?? '';
const historicalMessageEntities =
await this.agentChatService.getMessagesForThread(
thread.id,
userWorkspaceId,
);
const historicalUIMessages = mapAgentMessageEntitiesToUIMessages(
historicalMessageEntities,
);
const messagesWithHistory: ExtendedUIMessage[] = [
...historicalUIMessages,
...(lastUserMessage ? [lastUserMessage] : []),
];
const userMessagePromise = this.agentChatService.addMessage({
threadId: thread.id,
uiMessage: {
@@ -113,7 +135,7 @@ export class AgentChatStreamingService {
await this.chatExecutionService.streamChat({
workspace,
userWorkspaceId,
messages,
messages: messagesWithHistory,
browsingContext,
onCodeExecutionUpdate,
modelId,
@@ -212,6 +234,7 @@ export class AgentChatStreamingService {
`"totalOutputCredits" + ${streamUsage.outputCredits}`,
contextWindowTokens: modelConfig.contextWindowTokens,
conversationSize: lastStepConversationSize,
activeStreamId: null,
});
const generatedTitle = await titlePromise;
@@ -233,10 +256,26 @@ export class AgentChatStreamingService {
pipeUIMessageStreamToResponse({
stream: uiStream,
response,
// Consume the stream independently so onFinish fires even if
// the client disconnects (e.g., page refresh mid-stream)
consumeSseStream: ({ stream }) => {
stream.pipeTo(new WritableStream()).catch(() => {});
consumeSseStream: async ({ stream }) => {
try {
const streamId = generateId();
this.logger.log(`consumeSseStream called, streamId: ${streamId}`);
await this.resumableStreamService.createResumableStream(
streamId,
() => stream,
);
this.logger.log(
'createResumableStream done, writing activeStreamId to DB',
);
await this.threadRepository.update(thread.id, {
activeStreamId: streamId,
});
this.logger.log(`activeStreamId written to DB: ${streamId}`);
} catch (error) {
this.logger.error('consumeSseStream error:', error);
}
},
});
} catch (error) {
+8
View File
@@ -55145,6 +55145,13 @@ __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"
@@ -59932,6 +59939,7 @@ __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"