Compare commits
6 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| b4f5e4b55e | |||
| fc546f90ff | |||
| 79d9e75b6e | |||
| 2c90edaa59 | |||
| 480b78c55f | |||
| fbd6c9daea |
@@ -52,3 +52,4 @@ mcp.json
|
||||
/.junie/
|
||||
TRANSLATION_QA_REPORT.md
|
||||
.playwright-mcp/
|
||||
.claude/
|
||||
|
||||
+125
@@ -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)
|
||||
+16
@@ -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 };
|
||||
};
|
||||
-10
@@ -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,
|
||||
});
|
||||
-16
@@ -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];
|
||||
};
|
||||
@@ -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",
|
||||
|
||||
+19
@@ -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"`,
|
||||
);
|
||||
}
|
||||
}
|
||||
+53
@@ -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,
|
||||
|
||||
+67
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
+3
@@ -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[]>;
|
||||
|
||||
|
||||
+77
@@ -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);
|
||||
}
|
||||
}
|
||||
+50
-11
@@ -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) {
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user