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
+80
-70
@@ -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;
|
||||
};
|
||||
|
||||
Reference in New Issue
Block a user