Compare commits

...

13 Commits

Author SHA1 Message Date
Claude 571f79978b fix: address PR review comments on agent chat queue and SSE refactor
- Remove .claude/launch.json (avoid new root-level folder)
- Extract useAuthenticatedFetch hook from useAgentChat for reuse
- Use authenticatedFetch in AIChatQueuedMessages for proper token renewal
- Replace timestamp-based optimistic IDs with uuid v4
- Remove redundant comments, keep only useful ones
- Split mapDBPartsToUIMessageParts into one export per file
- Split STREAM_AGENT_CHAT_JOB_NAME constant into its own file
- Simplify delete queue endpoint to use only messageId (derive threadId)
- Add thread ownership check in subscription resolver

https://claude.ai/code/session_01NGv3c62sYGMww7iuVsWzcJ
2026-04-01 18:27:45 +00:00
Félix Malfait 6aedeb1cb2 Merge branch 'main' into feat/agent-chat-queue-and-sse-refactor 2026-04-01 19:28:30 +02:00
Félix Malfait eefc5c1914 feat: replace AI SDK client with GraphQL SSE subscription for agent chat
Drop @ai-sdk/react and resumable-stream in favor of a unified GraphQL SSE
subscription for all agent chat events (token streaming, message persistence,
queue updates). This makes the architecture server-driven and eliminates
the class of bugs caused by dual-state (AI SDK internal state vs Jotai atoms).

Key changes:
- Backend: publish UIMessageChunks via Redis PubSub instead of HTTP SSE +
  resumable-stream. New AgentChatEventPublisherService, subscription resolver,
  and unified POST /:threadId/message endpoint (server decides queue vs stream).
- Frontend: new useAgentChatSubscription hook bridges graphql-sse events into
  readUIMessageStream (from ai package) for chunk accumulation. Optimistic
  rendering preserves user messages during streaming. SSE client recycling
  triggers automatic resubscription.
- Queue: messages sent during active streaming are placed directly in queue
  atoms (no flash). Server auto-flushes next queued message after stream ends.

Made-with: Cursor
2026-04-01 13:18:58 +02:00
Félix Malfait f5b2c71943 Fix 2026-03-31 18:28:16 +02:00
Félix Malfait 74e6e20921 Merge branch 'feat/agent-chat-message-queue' of github.com:twentyhq/twenty into pr/19118 2026-03-31 17:57:07 +02:00
Félix Malfait c3ee312d26 Fix client SDK generated types: make turnId nullable
The DTO exposes turnId as nullable for queued messages, update
generated schema to match.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-31 17:45:52 +02:00
Félix Malfait 6cdca784fd Merge branch 'main' into feat/agent-chat-message-queue 2026-03-31 17:31:36 +02:00
Félix Malfait 9573e11879 Fix front typecheck and update client SDK generated types
- Cast status to union type in mapDBMessagesToUIMessages
- Update client SDK generated schema/types to include status field

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-31 09:57:51 +02:00
Félix Malfait 717989d36d Fix circular dependency and improve code quality
- Extract STREAM_AGENT_CHAT_JOB_NAME and StreamAgentChatJobData to
  separate types file to break circular dependency between
  stream-agent-chat.job.ts and agent-chat-streaming.service.ts
- Optimize flushNextQueuedMessage to check queued messages first
  (lightweight query) before loading full conversation
- Add missing state field to reasoning parts in server-side mapper
- Use entity status value in controller response instead of string literal
- Narrow ExtendedUIMessage.status to 'queued' | 'sent' union type

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-31 09:45:09 +02:00
Félix Malfait e6cfa46998 fix: typecheck errors for status field and reasoning part
- Add optional status field to ExtendedUIMessage in twenty-shared
- Remove invalid providerOptions from reasoning part mapper

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-31 09:34:03 +02:00
Félix Malfait d4acf77338 fix: migration schema prefix, review fixes
- Add "core" schema prefix to migration SQL (fixes CI)
- Remove unused modelId param from queueMessage
- Scope promoteQueuedMessage to QUEUED status only
- Use AgentMessageStatus enum instead of string literals
- Create proper mapDBPartsToUIMessageParts server-side mapper
  (fixes lossy mapping that dropped tool calls, reasoning, etc.)
- Merge getQueuedMessages + getMessagesForThread into single query
- Pass hasTitle from job data to avoid redundant thread lookup
- Remove modelId from frontend queue POST body

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-31 09:28:18 +02:00
Félix Malfait 220b32b19b feat: server-side message queue for agent chat
Replace localStorage-based frontend queue with server-side persistence.
When a user sends a message while streaming, the message is stored in
the database with status='queued'. When streaming finishes, the worker
automatically picks up the next queued message and starts processing it.

- Add `status` column (queued/sent) and make `turnId` nullable on AgentMessage
- Add REST endpoints POST/:threadId/queue and DELETE/:threadId/queue/:messageId
- Add auto-flush logic in stream job's finally block
- Frontend reads queued messages from GraphQL instead of localStorage
- Delete queued messages via REST DELETE endpoint

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-30 22:40:50 +02:00
Félix Malfait f8d92f9626 feat: queue messages sent while AI is streaming
When the user sends a message while the AI is still streaming a
response, instead of silently dropping it, queue it per-thread.
The queue is persisted to localStorage so it survives page refreshes.
When streaming completes, the next queued message is automatically
sent.

- New Jotai state: agentChatQueuedMessagesByThreadIdState (localStorage-backed)
- New UI component: AIChatQueuedMessages showing queued items with remove buttons
- Modified handleSendMessage to push to queue when streaming is active
- Added flushNextQueuedMessage to auto-send after stream completion

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
2026-03-30 22:28:16 +02:00
44 changed files with 1388 additions and 732 deletions
@@ -2787,10 +2787,12 @@ type AgentChatThread {
type AgentMessage {
id: UUID!
threadId: UUID!
turnId: UUID!
turnId: UUID
agentId: UUID
role: String!
status: String!
parts: [AgentMessagePart!]!
processedAt: DateTime
createdAt: DateTime!
}
@@ -2448,10 +2448,12 @@ export interface AgentChatThread {
export interface AgentMessage {
id: Scalars['UUID']
threadId: Scalars['UUID']
turnId: Scalars['UUID']
turnId?: Scalars['UUID']
agentId?: Scalars['UUID']
role: Scalars['String']
status: Scalars['String']
parts: AgentMessagePart[]
processedAt?: Scalars['DateTime']
createdAt: Scalars['DateTime']
__typename: 'AgentMessage'
}
@@ -5598,7 +5600,9 @@ export interface AgentMessageGenqlSelection{
turnId?: boolean | number
agentId?: boolean | number
role?: boolean | number
status?: boolean | number
parts?: AgentMessagePartGenqlSelection
processedAt?: boolean | number
createdAt?: boolean | number
__typename?: boolean | number
__scalar?: boolean | number
@@ -5543,9 +5543,15 @@ export default {
"role": [
1
],
"status": [
1
],
"parts": [
296
],
"processedAt": [
4
],
"createdAt": [
4
],
-1
View File
@@ -29,7 +29,6 @@
"workerDirectory": "public"
},
"dependencies": {
"@ai-sdk/react": "3.0.99",
"@apollo/client": "^4.0.0",
"@blocknote/mantine": "^0.47.1",
"@blocknote/react": "^0.47.1",
@@ -171,9 +171,11 @@ export type AgentMessage = {
createdAt: Scalars['DateTime'];
id: Scalars['UUID'];
parts: Array<AgentMessagePart>;
processedAt?: Maybe<Scalars['DateTime']>;
role: Scalars['String'];
status: Scalars['String'];
threadId: Scalars['UUID'];
turnId: Scalars['UUID'];
turnId?: Maybe<Scalars['UUID']>;
};
export type AgentMessagePart = {
@@ -6283,7 +6285,7 @@ export type GetChatMessagesQueryVariables = Exact<{
}>;
export type GetChatMessagesQuery = { __typename?: 'Query', chatMessages: Array<{ __typename?: 'AgentMessage', id: string, threadId: string, turnId: string, role: string, createdAt: string, parts: Array<{ __typename?: 'AgentMessagePart', id: string, messageId: string, orderIndex: number, type: string, textContent?: string | null, reasoningContent?: string | null, toolName?: string | null, toolCallId?: string | null, toolInput?: any | null, toolOutput?: any | null, state?: string | null, errorMessage?: string | null, errorDetails?: any | null, sourceUrlSourceId?: string | null, sourceUrlUrl?: string | null, sourceUrlTitle?: string | null, sourceDocumentSourceId?: string | null, sourceDocumentMediaType?: string | null, sourceDocumentTitle?: string | null, sourceDocumentFilename?: string | null, fileMediaType?: string | null, fileFilename?: string | null, fileUrl?: string | null, fileId?: string | null, providerMetadata?: any | null, createdAt: string }> }> };
export type GetChatMessagesQuery = { __typename?: 'Query', chatMessages: Array<{ __typename?: 'AgentMessage', id: string, threadId: string, turnId?: string | null, role: string, status: string, createdAt: string, parts: Array<{ __typename?: 'AgentMessagePart', id: string, messageId: string, orderIndex: number, type: string, textContent?: string | null, reasoningContent?: string | null, toolName?: string | null, toolCallId?: string | null, toolInput?: any | null, toolOutput?: any | null, state?: string | null, errorMessage?: string | null, errorDetails?: any | null, sourceUrlSourceId?: string | null, sourceUrlUrl?: string | null, sourceUrlTitle?: string | null, sourceDocumentSourceId?: string | null, sourceDocumentMediaType?: string | null, sourceDocumentTitle?: string | null, sourceDocumentFilename?: string | null, fileMediaType?: string | null, fileFilename?: string | null, fileUrl?: string | null, fileId?: string | null, providerMetadata?: any | null, createdAt: string }> }> };
export type GetChatThreadsQueryVariables = Exact<{
paging?: InputMaybe<CursorPaging>;
@@ -8121,7 +8123,7 @@ export const FindManySkillsDocument = {"kind":"Document","definitions":[{"kind":
export const FindOneAgentDocument = {"kind":"Document","definitions":[{"kind":"OperationDefinition","operation":"query","name":{"kind":"Name","value":"FindOneAgent"},"variableDefinitions":[{"kind":"VariableDefinition","variable":{"kind":"Variable","name":{"kind":"Name","value":"id"}},"type":{"kind":"NonNullType","type":{"kind":"NamedType","name":{"kind":"Name","value":"UUID"}}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"findOneAgent"},"arguments":[{"kind":"Argument","name":{"kind":"Name","value":"input"},"value":{"kind":"ObjectValue","fields":[{"kind":"ObjectField","name":{"kind":"Name","value":"id"},"value":{"kind":"Variable","name":{"kind":"Name","value":"id"}}}]}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"FragmentSpread","name":{"kind":"Name","value":"AgentFields"}}]}}]}},{"kind":"FragmentDefinition","name":{"kind":"Name","value":"AgentFields"},"typeCondition":{"kind":"NamedType","name":{"kind":"Name","value":"Agent"}},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"id"}},{"kind":"Field","name":{"kind":"Name","value":"name"}},{"kind":"Field","name":{"kind":"Name","value":"label"}},{"kind":"Field","name":{"kind":"Name","value":"description"}},{"kind":"Field","name":{"kind":"Name","value":"icon"}},{"kind":"Field","name":{"kind":"Name","value":"prompt"}},{"kind":"Field","name":{"kind":"Name","value":"modelId"}},{"kind":"Field","name":{"kind":"Name","value":"responseFormat"}},{"kind":"Field","name":{"kind":"Name","value":"roleId"}},{"kind":"Field","name":{"kind":"Name","value":"isCustom"}},{"kind":"Field","name":{"kind":"Name","value":"modelConfiguration"}},{"kind":"Field","name":{"kind":"Name","value":"evaluationInputs"}},{"kind":"Field","name":{"kind":"Name","value":"applicationId"}},{"kind":"Field","name":{"kind":"Name","value":"createdAt"}},{"kind":"Field","name":{"kind":"Name","value":"updatedAt"}}]}}]} as unknown as DocumentNode<FindOneAgentQuery, FindOneAgentQueryVariables>;
export const FindOneSkillDocument = {"kind":"Document","definitions":[{"kind":"OperationDefinition","operation":"query","name":{"kind":"Name","value":"FindOneSkill"},"variableDefinitions":[{"kind":"VariableDefinition","variable":{"kind":"Variable","name":{"kind":"Name","value":"id"}},"type":{"kind":"NonNullType","type":{"kind":"NamedType","name":{"kind":"Name","value":"UUID"}}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"skill"},"arguments":[{"kind":"Argument","name":{"kind":"Name","value":"id"},"value":{"kind":"Variable","name":{"kind":"Name","value":"id"}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"FragmentSpread","name":{"kind":"Name","value":"SkillFields"}}]}}]}},{"kind":"FragmentDefinition","name":{"kind":"Name","value":"SkillFields"},"typeCondition":{"kind":"NamedType","name":{"kind":"Name","value":"Skill"}},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"id"}},{"kind":"Field","name":{"kind":"Name","value":"name"}},{"kind":"Field","name":{"kind":"Name","value":"label"}},{"kind":"Field","name":{"kind":"Name","value":"description"}},{"kind":"Field","name":{"kind":"Name","value":"icon"}},{"kind":"Field","name":{"kind":"Name","value":"content"}},{"kind":"Field","name":{"kind":"Name","value":"isCustom"}},{"kind":"Field","name":{"kind":"Name","value":"isActive"}},{"kind":"Field","name":{"kind":"Name","value":"createdAt"}},{"kind":"Field","name":{"kind":"Name","value":"updatedAt"}}]}}]} as unknown as DocumentNode<FindOneSkillQuery, FindOneSkillQueryVariables>;
export const GetAgentTurnsDocument = {"kind":"Document","definitions":[{"kind":"OperationDefinition","operation":"query","name":{"kind":"Name","value":"GetAgentTurns"},"variableDefinitions":[{"kind":"VariableDefinition","variable":{"kind":"Variable","name":{"kind":"Name","value":"agentId"}},"type":{"kind":"NonNullType","type":{"kind":"NamedType","name":{"kind":"Name","value":"UUID"}}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"agentTurns"},"arguments":[{"kind":"Argument","name":{"kind":"Name","value":"agentId"},"value":{"kind":"Variable","name":{"kind":"Name","value":"agentId"}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"id"}},{"kind":"Field","name":{"kind":"Name","value":"threadId"}},{"kind":"Field","name":{"kind":"Name","value":"agentId"}},{"kind":"Field","name":{"kind":"Name","value":"createdAt"}},{"kind":"Field","name":{"kind":"Name","value":"evaluations"},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"id"}},{"kind":"Field","name":{"kind":"Name","value":"score"}},{"kind":"Field","name":{"kind":"Name","value":"comment"}},{"kind":"Field","name":{"kind":"Name","value":"createdAt"}}]}},{"kind":"Field","name":{"kind":"Name","value":"messages"},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"id"}},{"kind":"Field","name":{"kind":"Name","value":"role"}},{"kind":"Field","name":{"kind":"Name","value":"createdAt"}},{"kind":"Field","name":{"kind":"Name","value":"parts"},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"id"}},{"kind":"Field","name":{"kind":"Name","value":"type"}},{"kind":"Field","name":{"kind":"Name","value":"textContent"}},{"kind":"Field","name":{"kind":"Name","value":"reasoningContent"}},{"kind":"Field","name":{"kind":"Name","value":"toolName"}},{"kind":"Field","name":{"kind":"Name","value":"toolCallId"}},{"kind":"Field","name":{"kind":"Name","value":"toolInput"}},{"kind":"Field","name":{"kind":"Name","value":"toolOutput"}},{"kind":"Field","name":{"kind":"Name","value":"errorMessage"}},{"kind":"Field","name":{"kind":"Name","value":"state"}},{"kind":"Field","name":{"kind":"Name","value":"errorDetails"}},{"kind":"Field","name":{"kind":"Name","value":"sourceUrlSourceId"}},{"kind":"Field","name":{"kind":"Name","value":"sourceUrlUrl"}},{"kind":"Field","name":{"kind":"Name","value":"sourceUrlTitle"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentSourceId"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentMediaType"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentTitle"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentFilename"}},{"kind":"Field","name":{"kind":"Name","value":"fileMediaType"}},{"kind":"Field","name":{"kind":"Name","value":"fileFilename"}},{"kind":"Field","name":{"kind":"Name","value":"fileUrl"}},{"kind":"Field","name":{"kind":"Name","value":"providerMetadata"}}]}}]}}]}}]}}]} as unknown as DocumentNode<GetAgentTurnsQuery, GetAgentTurnsQueryVariables>;
export const GetChatMessagesDocument = {"kind":"Document","definitions":[{"kind":"OperationDefinition","operation":"query","name":{"kind":"Name","value":"GetChatMessages"},"variableDefinitions":[{"kind":"VariableDefinition","variable":{"kind":"Variable","name":{"kind":"Name","value":"threadId"}},"type":{"kind":"NonNullType","type":{"kind":"NamedType","name":{"kind":"Name","value":"UUID"}}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"chatMessages"},"arguments":[{"kind":"Argument","name":{"kind":"Name","value":"threadId"},"value":{"kind":"Variable","name":{"kind":"Name","value":"threadId"}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"id"}},{"kind":"Field","name":{"kind":"Name","value":"threadId"}},{"kind":"Field","name":{"kind":"Name","value":"turnId"}},{"kind":"Field","name":{"kind":"Name","value":"role"}},{"kind":"Field","name":{"kind":"Name","value":"createdAt"}},{"kind":"Field","name":{"kind":"Name","value":"parts"},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"id"}},{"kind":"Field","name":{"kind":"Name","value":"messageId"}},{"kind":"Field","name":{"kind":"Name","value":"orderIndex"}},{"kind":"Field","name":{"kind":"Name","value":"type"}},{"kind":"Field","name":{"kind":"Name","value":"textContent"}},{"kind":"Field","name":{"kind":"Name","value":"reasoningContent"}},{"kind":"Field","name":{"kind":"Name","value":"toolName"}},{"kind":"Field","name":{"kind":"Name","value":"toolCallId"}},{"kind":"Field","name":{"kind":"Name","value":"toolInput"}},{"kind":"Field","name":{"kind":"Name","value":"toolOutput"}},{"kind":"Field","name":{"kind":"Name","value":"state"}},{"kind":"Field","name":{"kind":"Name","value":"errorMessage"}},{"kind":"Field","name":{"kind":"Name","value":"errorDetails"}},{"kind":"Field","name":{"kind":"Name","value":"sourceUrlSourceId"}},{"kind":"Field","name":{"kind":"Name","value":"sourceUrlUrl"}},{"kind":"Field","name":{"kind":"Name","value":"sourceUrlTitle"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentSourceId"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentMediaType"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentTitle"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentFilename"}},{"kind":"Field","name":{"kind":"Name","value":"fileMediaType"}},{"kind":"Field","name":{"kind":"Name","value":"fileFilename"}},{"kind":"Field","name":{"kind":"Name","value":"fileUrl"}},{"kind":"Field","name":{"kind":"Name","value":"fileId"}},{"kind":"Field","name":{"kind":"Name","value":"providerMetadata"}},{"kind":"Field","name":{"kind":"Name","value":"createdAt"}}]}}]}}]}}]} as unknown as DocumentNode<GetChatMessagesQuery, GetChatMessagesQueryVariables>;
export const GetChatMessagesDocument = {"kind":"Document","definitions":[{"kind":"OperationDefinition","operation":"query","name":{"kind":"Name","value":"GetChatMessages"},"variableDefinitions":[{"kind":"VariableDefinition","variable":{"kind":"Variable","name":{"kind":"Name","value":"threadId"}},"type":{"kind":"NonNullType","type":{"kind":"NamedType","name":{"kind":"Name","value":"UUID"}}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"chatMessages"},"arguments":[{"kind":"Argument","name":{"kind":"Name","value":"threadId"},"value":{"kind":"Variable","name":{"kind":"Name","value":"threadId"}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"id"}},{"kind":"Field","name":{"kind":"Name","value":"threadId"}},{"kind":"Field","name":{"kind":"Name","value":"turnId"}},{"kind":"Field","name":{"kind":"Name","value":"role"}},{"kind":"Field","name":{"kind":"Name","value":"status"}},{"kind":"Field","name":{"kind":"Name","value":"createdAt"}},{"kind":"Field","name":{"kind":"Name","value":"parts"},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"id"}},{"kind":"Field","name":{"kind":"Name","value":"messageId"}},{"kind":"Field","name":{"kind":"Name","value":"orderIndex"}},{"kind":"Field","name":{"kind":"Name","value":"type"}},{"kind":"Field","name":{"kind":"Name","value":"textContent"}},{"kind":"Field","name":{"kind":"Name","value":"reasoningContent"}},{"kind":"Field","name":{"kind":"Name","value":"toolName"}},{"kind":"Field","name":{"kind":"Name","value":"toolCallId"}},{"kind":"Field","name":{"kind":"Name","value":"toolInput"}},{"kind":"Field","name":{"kind":"Name","value":"toolOutput"}},{"kind":"Field","name":{"kind":"Name","value":"state"}},{"kind":"Field","name":{"kind":"Name","value":"errorMessage"}},{"kind":"Field","name":{"kind":"Name","value":"errorDetails"}},{"kind":"Field","name":{"kind":"Name","value":"sourceUrlSourceId"}},{"kind":"Field","name":{"kind":"Name","value":"sourceUrlUrl"}},{"kind":"Field","name":{"kind":"Name","value":"sourceUrlTitle"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentSourceId"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentMediaType"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentTitle"}},{"kind":"Field","name":{"kind":"Name","value":"sourceDocumentFilename"}},{"kind":"Field","name":{"kind":"Name","value":"fileMediaType"}},{"kind":"Field","name":{"kind":"Name","value":"fileFilename"}},{"kind":"Field","name":{"kind":"Name","value":"fileUrl"}},{"kind":"Field","name":{"kind":"Name","value":"fileId"}},{"kind":"Field","name":{"kind":"Name","value":"providerMetadata"}},{"kind":"Field","name":{"kind":"Name","value":"createdAt"}}]}}]}}]}}]} as unknown as DocumentNode<GetChatMessagesQuery, GetChatMessagesQueryVariables>;
export const GetChatThreadsDocument = {"kind":"Document","definitions":[{"kind":"OperationDefinition","operation":"query","name":{"kind":"Name","value":"GetChatThreads"},"variableDefinitions":[{"kind":"VariableDefinition","variable":{"kind":"Variable","name":{"kind":"Name","value":"paging"}},"type":{"kind":"NamedType","name":{"kind":"Name","value":"CursorPaging"}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"chatThreads"},"arguments":[{"kind":"Argument","name":{"kind":"Name","value":"paging"},"value":{"kind":"Variable","name":{"kind":"Name","value":"paging"}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"edges"},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"node"},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"id"}},{"kind":"Field","name":{"kind":"Name","value":"title"}},{"kind":"Field","name":{"kind":"Name","value":"totalInputTokens"}},{"kind":"Field","name":{"kind":"Name","value":"totalOutputTokens"}},{"kind":"Field","name":{"kind":"Name","value":"contextWindowTokens"}},{"kind":"Field","name":{"kind":"Name","value":"conversationSize"}},{"kind":"Field","name":{"kind":"Name","value":"totalInputCredits"}},{"kind":"Field","name":{"kind":"Name","value":"totalOutputCredits"}},{"kind":"Field","name":{"kind":"Name","value":"createdAt"}},{"kind":"Field","name":{"kind":"Name","value":"updatedAt"}}]}},{"kind":"Field","name":{"kind":"Name","value":"cursor"}}]}},{"kind":"Field","name":{"kind":"Name","value":"pageInfo"},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"endCursor"}},{"kind":"Field","name":{"kind":"Name","value":"hasNextPage"}}]}}]}}]}}]} as unknown as DocumentNode<GetChatThreadsQuery, GetChatThreadsQueryVariables>;
export const GetToolIndexDocument = {"kind":"Document","definitions":[{"kind":"OperationDefinition","operation":"query","name":{"kind":"Name","value":"GetToolIndex"},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"getToolIndex"},"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"name"}},{"kind":"Field","name":{"kind":"Name","value":"description"}},{"kind":"Field","name":{"kind":"Name","value":"category"}},{"kind":"Field","name":{"kind":"Name","value":"objectName"}}]}}]}}]} as unknown as DocumentNode<GetToolIndexQuery, GetToolIndexQueryVariables>;
export const GetToolInputSchemaDocument = {"kind":"Document","definitions":[{"kind":"OperationDefinition","operation":"query","name":{"kind":"Name","value":"GetToolInputSchema"},"variableDefinitions":[{"kind":"VariableDefinition","variable":{"kind":"Variable","name":{"kind":"Name","value":"toolName"}},"type":{"kind":"NonNullType","type":{"kind":"NamedType","name":{"kind":"Name","value":"String"}}}}],"selectionSet":{"kind":"SelectionSet","selections":[{"kind":"Field","name":{"kind":"Name","value":"getToolInputSchema"},"arguments":[{"kind":"Argument","name":{"kind":"Name","value":"toolName"},"value":{"kind":"Variable","name":{"kind":"Name","value":"toolName"}}}]}]}}]} as unknown as DocumentNode<GetToolInputSchemaQuery, GetToolInputSchemaQueryVariables>;
@@ -189,23 +189,22 @@ export const AIChatMessage = ({
<AIChatErrorRenderer error={error} />
)}
</StyledMessageContainer>
{agentChatMessage.parts.length > 0 &&
agentChatMessage.metadata?.createdAt && (
<StyledMessageFooter className="message-footer">
<StyledMessageTimestamp>
{beautifyPastDateRelativeToNow(
agentChatMessage.metadata?.createdAt,
localeCatalog,
)}
</StyledMessageTimestamp>
<LightCopyIconButton
copyText={
agentChatMessage.parts.find((part) => part.type === 'text')
?.text ?? ''
}
/>
</StyledMessageFooter>
)}
{agentChatMessage.parts.length > 0 && (
<StyledMessageFooter className="message-footer">
<StyledMessageTimestamp>
{beautifyPastDateRelativeToNow(
agentChatMessage.metadata?.createdAt ?? new Date(),
localeCatalog,
)}
</StyledMessageTimestamp>
<LightCopyIconButton
copyText={
agentChatMessage.parts.find((part) => part.type === 'text')
?.text ?? ''
}
/>
</StyledMessageFooter>
)}
</StyledMessageBubble>
);
};
@@ -0,0 +1,101 @@
import { styled } from '@linaria/react';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { useAuthenticatedFetch } from '@/ai/hooks/useAuthenticatedFetch';
import { agentChatQueuedMessagesComponentFamilyState } from '@/ai/states/agentChatQueuedMessagesComponentFamilyState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { REST_API_BASE_URL } from '@/apollo/constant/rest-api-base-url';
import { getTokenPair } from '@/apollo/utils/getTokenPair';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { isDefined } from 'twenty-shared/utils';
import { IconX } from 'twenty-ui/display';
import { LightIconButton } from 'twenty-ui/input';
import { themeCssVariables } from 'twenty-ui/theme-constants';
const StyledQueueContainer = styled.div`
display: flex;
flex-direction: column;
gap: ${themeCssVariables.spacing[1]};
padding: 0 ${themeCssVariables.spacing[3]};
`;
const StyledQueueLabel = styled.div`
color: ${themeCssVariables.font.color.light};
font-size: ${themeCssVariables.font.size.xs};
padding-left: ${themeCssVariables.spacing[1]};
`;
const StyledQueuedItem = styled.div`
align-items: center;
background: ${themeCssVariables.background.tertiary};
border-radius: ${themeCssVariables.border.radius.sm};
color: ${themeCssVariables.font.color.secondary};
display: flex;
font-size: ${themeCssVariables.font.size.md};
gap: ${themeCssVariables.spacing[2]};
justify-content: space-between;
padding: ${themeCssVariables.spacing[1]} ${themeCssVariables.spacing[2]};
`;
const StyledQueuedText = styled.span`
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
`;
export const AIChatQueuedMessages = () => {
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const { authenticatedFetch } = useAuthenticatedFetch();
const agentChatQueuedMessages = useAtomComponentFamilyStateValue(
agentChatQueuedMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
if (!isDefined(currentAIChatThread) || agentChatQueuedMessages.length === 0) {
return null;
}
const handleRemove = (messageId: string) => {
const tokenPair = getTokenPair();
if (!isDefined(tokenPair)) {
return;
}
authenticatedFetch(`${REST_API_BASE_URL}/agent-chat/queue/${messageId}`, {
method: 'DELETE',
headers: {
Authorization: `Bearer ${tokenPair.accessOrWorkspaceAgnosticToken.token}`,
},
})
.then(() => {
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
})
.catch(() => {});
};
return (
<StyledQueueContainer>
<StyledQueueLabel>
{agentChatQueuedMessages.length} Queued
</StyledQueueLabel>
{agentChatQueuedMessages.map((message) => {
const textPart = message.parts?.find((part) => part.type === 'text');
const displayText = textPart && 'text' in textPart ? textPart.text : '';
return (
<StyledQueuedItem key={message.id}>
<StyledQueuedText>{displayText}</StyledQueuedText>
<LightIconButton
Icon={IconX}
onClick={() => handleRemove(message.id)}
size="small"
/>
</StyledQueuedItem>
);
})}
</StyledQueueContainer>
);
};
@@ -10,6 +10,7 @@ import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { threadIdCreatedFromDraftState } from '@/ai/states/threadIdCreatedFromDraftState';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { AIChatQueuedMessages } from '@/ai/components/AIChatQueuedMessages';
import { AIChatTabMessageList } from '@/ai/components/AIChatTabMessageList';
const StyledContainer = styled.div<{ isDraggingFile: boolean }>`
@@ -51,6 +52,7 @@ export const AIChatTab = () => {
{!isDraggingFile && (
<>
<AIChatTabMessageList />
<AIChatQueuedMessages />
<AIChatEditorSection key={editorSectionKey} />
</>
)}
@@ -6,6 +6,7 @@ import { AGENT_CHAT_UNKNOWN_THREAD_ID } from '@/ai/constants/AgentChatUnknownThr
import { AGENT_CHAT_NEW_THREAD_DRAFT_KEY } from '@/ai/states/agentChatDraftsByThreadIdState';
import { agentChatFetchedMessagesComponentFamilyState } from '@/ai/states/agentChatFetchedMessagesComponentFamilyState';
import { agentChatMessagesLoadingState } from '@/ai/states/agentChatMessagesLoadingState';
import { agentChatQueuedMessagesComponentFamilyState } from '@/ai/states/agentChatQueuedMessagesComponentFamilyState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { skipMessagesSkeletonUntilLoadedState } from '@/ai/states/skipMessagesSkeletonUntilLoadedState';
import { mapDBMessagesToUIMessages } from '@/ai/utils/mapDBMessagesToUIMessages';
@@ -42,6 +43,11 @@ export const AgentChatMessagesFetchEffect = () => {
{ threadId: currentAIChatThread },
);
const setAgentChatQueuedMessages = useSetAtomComponentFamilyState(
agentChatQueuedMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
const handleFirstLoad = useCallback(
(_data: GetChatMessagesQuery) => {
setSkipMessagesSkeletonUntilLoaded(false);
@@ -52,9 +58,14 @@ export const AgentChatMessagesFetchEffect = () => {
const handleDataLoaded = useCallback(
(data: GetChatMessagesQuery) => {
const uiMessages = mapDBMessagesToUIMessages(data.chatMessages ?? []);
setAgentChatFetchedMessages(uiMessages);
setAgentChatFetchedMessages(
uiMessages.filter((message) => message.status !== 'queued'),
);
setAgentChatQueuedMessages(
uiMessages.filter((message) => message.status === 'queued'),
);
},
[setAgentChatFetchedMessages],
[setAgentChatFetchedMessages, setAgentChatQueuedMessages],
);
const handleLoadingChange = useCallback(
@@ -1,4 +1,4 @@
import { AgentChatAiSdkStreamEffect } from '@/ai/components/AgentChatAiSdkStreamEffect';
import { AgentChatStreamSubscriptionEffect } from '@/ai/components/AgentChatStreamSubscriptionEffect';
import { AgentChatMessagesFetchEffect } from '@/ai/components/AgentChatMessagesFetchEffect';
import { AgentChatSessionStartTimeEffect } from '@/ai/components/AgentChatSessionStartTimeEffect';
@@ -20,7 +20,7 @@ export const AgentChatProviderContent = ({
>
<AgentChatThreadInitializationEffect />
<AgentChatMessagesFetchEffect />
<AgentChatAiSdkStreamEffect />
<AgentChatStreamSubscriptionEffect />
<AgentChatStreamingPartsDiffSyncEffect />
<AgentChatSessionStartTimeEffect />
<AgentChatStreamingAutoScrollEffect />
@@ -1,11 +1,13 @@
import { useCallback, useEffect } from 'react';
import { isValidUuid } from 'twenty-shared/utils';
import { AGENT_CHAT_ENSURE_THREAD_FOR_DRAFT_EVENT_NAME } from '@/ai/constants/AgentChatEnsureThreadForDraftEventName';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { useAgentChat } from '@/ai/hooks/useAgentChat';
import { useAgentChatSubscription } from '@/ai/hooks/useAgentChatSubscription';
import { useCreateAgentChatThread } from '@/ai/hooks/useCreateAgentChatThread';
import { useEnsureAgentChatThreadExistsForDraft } from '@/ai/hooks/useEnsureAgentChatThreadExistsForDraft';
import { useEnsureAgentChatThreadIdForSend } from '@/ai/hooks/useEnsureAgentChatThreadIdForSend';
import { agentChatDisplayedThreadState } from '@/ai/states/agentChatDisplayedThreadState';
import { agentChatErrorState } from '@/ai/states/agentChatErrorState';
import { agentChatFetchedMessagesComponentFamilyState } from '@/ai/states/agentChatFetchedMessagesComponentFamilyState';
import { agentChatIsInitialScrollPendingOnThreadChangeState } from '@/ai/states/agentChatIsInitialScrollPendingOnThreadChangeState';
import { agentChatIsLoadingState } from '@/ai/states/agentChatIsLoadingState';
@@ -14,23 +16,14 @@ import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMess
import { agentChatMessagesLoadingState } from '@/ai/states/agentChatMessagesLoadingState';
import { agentChatThreadsLoadingState } from '@/ai/states/agentChatThreadsLoadingState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { mergeAgentChatFetchedAndStreamingMessages } from '@/ai/utils/mergeAgentChatFetchedAndStreamingMessages';
import { normalizeAiSdkError } from '@/ai/utils/normalizeAiSdkError';
import { useListenToBrowserEvent } from '@/browser-event/hooks/useListenToBrowserEvent';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { useAtomComponentFamilyStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomComponentFamilyStateValue';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { useSetAtomComponentFamilyState } from '@/ui/utilities/state/jotai/hooks/useSetAtomComponentFamilyState';
import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
import { useCallback, useEffect } from 'react';
import { isValidUuid } from 'twenty-shared/utils';
export const AgentChatAiSdkStreamEffect = () => {
export const AgentChatStreamSubscriptionEffect = () => {
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const agentChatFetchedMessages = useAtomComponentFamilyStateValue(
agentChatFetchedMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
const { createChatThread } = useCreateAgentChatThread();
@@ -45,47 +38,29 @@ export const AgentChatAiSdkStreamEffect = () => {
onBrowserEvent: ensureThreadExistsForDraft,
});
const onStreamingComplete = useCallback(() => {
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
}, []);
useAgentChat(ensureThreadIdForSend);
const chatState = useAgentChat(
agentChatFetchedMessages,
ensureThreadIdForSend,
onStreamingComplete,
// Subscribe to the thread's event stream
const subscriptionThreadId =
currentAIChatThread !== null && isValidUuid(currentAIChatThread)
? currentAIChatThread
: null;
useAgentChatSubscription(subscriptionThreadId);
// Sync fetched messages to the displayed messages atom when no stream is active
const agentChatFetchedMessages = useAtomComponentFamilyStateValue(
agentChatFetchedMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
// Attempt to resume an active stream when navigating to an existing
// thread. We call resumeStream() manually instead of using useChat's
// resume:true option so that the stop button can coexist with
// resumption (resume:true is incompatible with abort signals).
// Only resume when the thread already has fetched messages — this
// avoids resuming on newly created threads where the thread ID
// transitions from a placeholder to a real UUID mid-conversation.
useEffect(() => {
if (
currentAIChatThread === null ||
!isValidUuid(currentAIChatThread) ||
agentChatFetchedMessages.length === 0 ||
chatState.status === 'streaming' ||
chatState.status === 'submitted'
) {
return;
}
chatState.resumeStream();
// We intentionally omit chatState.resumeStream and status from deps
// to avoid resume loops. We do include agentChatFetchedMessages.length
// so that resume fires once messages are fetched (they may arrive
// after the thread ID is set).
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [currentAIChatThread, agentChatFetchedMessages.length]);
const setAgentChatMessages = useSetAtomComponentFamilyState(
agentChatMessagesComponentFamilyState,
{ threadId: currentAIChatThread },
);
const agentChatIsStreaming = useAtomStateValue(agentChatIsStreamingState);
const agentChatDisplayedThread = useAtomStateValue(
agentChatDisplayedThreadState,
);
@@ -99,22 +74,21 @@ export const AgentChatAiSdkStreamEffect = () => {
);
useEffect(() => {
const mergedMessages = mergeAgentChatFetchedAndStreamingMessages(
agentChatFetchedMessages,
chatState.messages,
);
if (agentChatIsStreaming) {
return;
}
setAgentChatMessages(mergedMessages);
setAgentChatMessages(agentChatFetchedMessages);
if (currentAIChatThread !== agentChatDisplayedThread) {
if (mergedMessages.length > 0) {
if (agentChatFetchedMessages.length > 0) {
setAgentChatIsInitialScrollPendingOnThreadChange(true);
}
setAgentChatDisplayedThread(currentAIChatThread);
}
}, [
agentChatFetchedMessages,
chatState.messages,
agentChatIsStreaming,
setAgentChatMessages,
currentAIChatThread,
agentChatDisplayedThread,
@@ -130,31 +104,16 @@ export const AgentChatAiSdkStreamEffect = () => {
agentChatMessagesLoadingState,
);
useEffect(() => {
const handleLoadingChange = useCallback(() => {
const combinedIsLoading =
chatState.isLoading ||
agentChatMessagesLoading ||
agentChatThreadsLoading;
agentChatMessagesLoading || agentChatThreadsLoading;
setAgentChatIsLoading(combinedIsLoading);
}, [
chatState.isLoading,
agentChatMessagesLoading,
agentChatThreadsLoading,
setAgentChatIsLoading,
]);
const setAgentChatError = useSetAtomState(agentChatErrorState);
}, [agentChatMessagesLoading, agentChatThreadsLoading, setAgentChatIsLoading]);
useEffect(() => {
setAgentChatError(normalizeAiSdkError(chatState.error));
}, [chatState.error, setAgentChatError]);
const setAgentChatIsStreaming = useSetAtomState(agentChatIsStreamingState);
useEffect(() => {
setAgentChatIsStreaming(chatState.status === 'streaming');
}, [chatState.status, setAgentChatIsStreaming]);
handleLoadingChange();
}, [handleLoadingChange]);
return null;
};
@@ -0,0 +1 @@
export const AGENT_CHAT_INSTANCE_ID = 'agentChatComponentInstance';
@@ -7,6 +7,7 @@ export const GET_CHAT_MESSAGES = gql`
threadId
turnId
role
status
createdAt
parts {
id
@@ -0,0 +1,10 @@
import { gql } from '@apollo/client';
export const ON_AGENT_CHAT_EVENT = gql`
subscription OnAgentChatEvent($threadId: String!) {
onAgentChatEvent(threadId: $threadId) {
threadId
event
}
}
`;
@@ -1,58 +1,47 @@
import { useStore } from 'jotai';
import { useCallback, useState } from 'react';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
import { isDefined, isValidUuid } from 'twenty-shared/utils';
import { v4 } from 'uuid';
import { AGENT_CHAT_INSTANCE_ID } from '@/ai/constants/AgentChatInstanceId';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { AGENT_CHAT_SEND_MESSAGE_EVENT_NAME } from '@/ai/constants/AgentChatSendMessageEventName';
import { useApolloClient } from '@apollo/client/react';
import { useGetBrowsingContext } from '@/ai/hooks/useBrowsingContext';
import { agentChatSelectedFilesState } from '@/ai/states/agentChatSelectedFilesState';
import { agentChatUploadedFilesState } from '@/ai/states/agentChatUploadedFilesState';
import { agentChatUsageState } from '@/ai/states/agentChatUsageState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTitleState';
import { AGENT_CHAT_RETRY_EVENT_NAME } from '@/ai/constants/AgentChatRetryEventName';
import { AGENT_CHAT_STOP_EVENT_NAME } from '@/ai/constants/AgentChatStopEventName';
import { useAgentChatModelId } from '@/ai/hooks/useAgentChatModelId';
import { useAuthenticatedFetch } from '@/ai/hooks/useAuthenticatedFetch';
import { useGetBrowsingContext } from '@/ai/hooks/useBrowsingContext';
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 { agentChatSelectedFilesState } from '@/ai/states/agentChatSelectedFilesState';
import { agentChatUploadedFilesState } from '@/ai/states/agentChatUploadedFilesState';
import { agentChatFetchedMessagesComponentFamilyState } from '@/ai/states/agentChatFetchedMessagesComponentFamilyState';
import { agentChatIsStreamingState } from '@/ai/states/agentChatIsStreamingState';
import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMessagesComponentFamilyState';
import { agentChatQueuedMessagesComponentFamilyState } from '@/ai/states/agentChatQueuedMessagesComponentFamilyState';
import { currentAIChatThreadState } from '@/ai/states/currentAIChatThreadState';
import { REST_API_BASE_URL } from '@/apollo/constant/rest-api-base-url';
import { getTokenPair } from '@/apollo/utils/getTokenPair';
import { renewToken } from '@/auth/services/AuthService';
import { tokenPairState } from '@/auth/states/tokenPairState';
import { useListenToBrowserEvent } from '@/browser-event/hooks/useListenToBrowserEvent';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { useAtomState } from '@/ui/utilities/state/jotai/hooks/useAtomState';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
import { useChat } from '@ai-sdk/react';
import { DefaultChatTransport } from 'ai';
import { useStore } from 'jotai';
import { useCallback, useMemo, useState } from 'react';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
import { isDefined, isValidUuid } from 'twenty-shared/utils';
import { REACT_APP_SERVER_BASE_URL } from '~/config';
export const useAgentChat = (
uiMessages: ExtendedUIMessage[],
ensureThreadIdForSend: () => Promise<string | null>,
onStreamingComplete?: () => void,
) => {
const setTokenPair = useSetAtomState(tokenPairState);
const setAgentChatUsage = useSetAtomState(agentChatUsageState);
const { modelIdForRequest } = useAgentChatModelId();
const { authenticatedFetch } = useAuthenticatedFetch();
const { getBrowsingContext } = useGetBrowsingContext();
const setCurrentAIChatThreadTitle = useSetAtomState(
currentAIChatThreadTitleState,
);
const setCurrentAIChatThread = useSetAtomState(currentAIChatThreadState);
const apolloClient = useApolloClient();
const store = useStore();
const agentChatSelectedFiles = useAtomStateValue(agentChatSelectedFilesState);
const currentAIChatThread = useAtomStateValue(currentAIChatThreadState);
const [, setPendingThreadIdAfterFirstSend] = useState<string | null>(null);
const [agentChatUploadedFiles, setAgentChatUploadedFiles] = useAtomState(
@@ -64,183 +53,6 @@ export const useAgentChat = (
agentChatDraftsByThreadIdState,
);
const retryFetchWithRenewedToken = async (
input: RequestInfo | URL,
init?: RequestInit,
) => {
const tokenPair = getTokenPair();
if (!isDefined(tokenPair)) {
return null;
}
try {
const renewedTokens = await renewToken(
`${REACT_APP_SERVER_BASE_URL}/metadata`,
tokenPair,
);
if (!isDefined(renewedTokens)) {
setTokenPair(null);
return null;
}
const renewedAccessToken =
renewedTokens.accessOrWorkspaceAgnosticToken?.token;
if (!isDefined(renewedAccessToken)) {
setTokenPair(null);
return null;
}
setTokenPair(renewedTokens);
const updatedHeaders = new Headers(init?.headers ?? {});
updatedHeaders.set('Authorization', `Bearer ${renewedAccessToken}`);
return fetch(input, {
...init,
headers: updatedHeaders,
});
} catch {
setTokenPair(null);
return null;
}
};
// eslint-disable-next-line react-hooks/exhaustive-deps
const transport = useMemo(
() =>
new DefaultChatTransport({
api: `${REST_API_BASE_URL}/agent-chat/stream`,
headers: () => ({
Authorization: `Bearer ${getTokenPair()?.accessOrWorkspaceAgnosticToken.token}`,
}),
prepareReconnectToStreamRequest: ({ id }) => ({
api: `${REST_API_BASE_URL}/agent-chat/${id}/stream`,
headers: {
Authorization: `Bearer ${getTokenPair()?.accessOrWorkspaceAgnosticToken.token}`,
},
}),
fetch: async (input, init) => {
const response = await fetch(input, init);
if (response.status === 401) {
const retriedResponse = await retryFetchWithRenewedToken(
input,
init,
);
return retriedResponse ?? response;
}
if (!response.ok) {
const errorBody = await response.json().catch(() => ({}));
const error = new Error(
errorBody.messages?.[0] ||
`Request failed with status ${response.status}`,
) as Error & { code?: string };
if (isDefined(errorBody.code)) {
error.code = errorBody.code;
}
throw error;
}
return response;
},
}),
// Intentionally created once — closures inside (getTokenPair, etc.)
// read fresh values via function references, not stale captures.
[],
);
const {
sendMessage,
messages,
status,
error,
regenerate,
stop,
resumeStream,
} = useChat({
transport,
messages: uiMessages,
id: currentAIChatThread ?? undefined,
experimental_throttle: 100,
onFinish: ({ message }) => {
type UsageMetadata = {
inputTokens: number;
outputTokens: number;
cachedInputTokens: number;
inputCredits: number;
outputCredits: number;
conversationSize: number;
};
type ModelMetadata = {
contextWindowTokens: number;
};
const metadata = message.metadata as
| { usage?: UsageMetadata; model?: ModelMetadata }
| undefined;
const usage = metadata?.usage;
const model = metadata?.model;
if (isDefined(usage) && isDefined(model)) {
setAgentChatUsage((prev) => ({
lastMessage: {
inputTokens: usage.inputTokens,
outputTokens: usage.outputTokens,
cachedInputTokens: usage.cachedInputTokens,
inputCredits: usage.inputCredits,
outputCredits: usage.outputCredits,
},
conversationSize: usage.conversationSize,
contextWindowTokens: model.contextWindowTokens,
inputTokens: (prev?.inputTokens ?? 0) + usage.inputTokens,
outputTokens: (prev?.outputTokens ?? 0) + usage.outputTokens,
inputCredits: (prev?.inputCredits ?? 0) + usage.inputCredits,
outputCredits: (prev?.outputCredits ?? 0) + usage.outputCredits,
}));
}
const titlePart = message.parts.find(
(part) => part.type === 'data-thread-title',
);
setPendingThreadIdAfterFirstSend((pendingId) => {
const threadIdForTitle =
pendingId ?? store.get(currentAIChatThreadState.atom);
if (isDefined(titlePart) && titlePart.type === 'data-thread-title') {
setCurrentAIChatThreadTitle(titlePart.data.title);
if (isDefined(threadIdForTitle)) {
const threadRef = apolloClient.cache.identify({
__typename: 'AgentChatThread',
id: threadIdForTitle,
});
if (isDefined(threadRef)) {
apolloClient.cache.modify({
id: threadRef,
fields: {
title: () => titlePart.data.title,
},
});
}
}
}
if (isDefined(pendingId)) {
setCurrentAIChatThread(pendingId);
}
return null;
});
onStreamingComplete?.();
},
});
const isStreaming = status === 'streaming';
const isLoading = isStreaming || agentChatSelectedFiles.length > 0;
const handleSendMessage = useCallback(async () => {
const draftKey =
store.get(currentAIChatThreadState.atom) ??
@@ -254,12 +66,19 @@ export const useAgentChat = (
).trim()
: store.get(agentChatInputState.atom).trim();
if (contentToSend === '' || isLoading) {
if (contentToSend === '') {
return;
}
const isLoading = agentChatSelectedFiles.length > 0;
if (isLoading) {
return;
}
const threadId = await ensureThreadIdForSend();
if (!threadId) {
if (!isDefined(threadId)) {
return;
}
@@ -275,34 +94,126 @@ export const useAgentChat = (
const browsingContext = getBrowsingContext();
sendMessage(
{
text: contentToSend,
files: agentChatUploadedFiles,
},
{
body: {
threadId,
browsingContext,
...(isDefined(modelIdForRequest) && {
modelId: modelIdForRequest,
}),
},
},
const fetchedMessages = store.get(
agentChatFetchedMessagesComponentFamilyState.atomFamily({
instanceId: AGENT_CHAT_INSTANCE_ID,
familyKey: { threadId },
}),
);
const optimisticUserMessage: ExtendedUIMessage = {
id: v4(),
role: 'user',
parts: [
{ type: 'text' as const, text: contentToSend },
...agentChatUploadedFiles,
],
metadata: {
createdAt: new Date().toISOString(),
},
status: 'sent',
};
const atomKey = {
instanceId: AGENT_CHAT_INSTANCE_ID,
familyKey: { threadId },
};
// If already streaming, place the optimistic message in the queue
// to mirror the server's queue decision and avoid a visible flash.
const isCurrentlyStreaming = store.get(agentChatIsStreamingState.atom);
const targetAtom = isCurrentlyStreaming
? agentChatQueuedMessagesComponentFamilyState
: agentChatMessagesComponentFamilyState;
const currentTargetMessages = store.get(targetAtom.atomFamily(atomKey));
store.set(targetAtom.atomFamily(atomKey), [
...currentTargetMessages,
optimisticUserMessage,
]);
setAgentChatUploadedFiles([]);
const tokenPair = getTokenPair();
if (!isDefined(tokenPair)) {
return;
}
const allMessages = [...fetchedMessages, optimisticUserMessage];
try {
const response = await authenticatedFetch(
`${REST_API_BASE_URL}/agent-chat/${threadId}/message`,
{
method: 'POST',
headers: {
Authorization: `Bearer ${tokenPair.accessOrWorkspaceAgnosticToken.token}`,
'Content-Type': 'application/json',
},
body: JSON.stringify({
text: contentToSend,
messages: allMessages,
browsingContext,
...(isDefined(modelIdForRequest) && {
modelId: modelIdForRequest,
}),
}),
},
);
const result = await response.json();
if (result.queued) {
if (!isCurrentlyStreaming) {
// Server queued it but we optimistically placed in main thread.
// Move it to the queue (edge case: streaming started between
// our check and the server's decision).
const latestMessages = store.get(
agentChatMessagesComponentFamilyState.atomFamily(atomKey),
);
store.set(
agentChatMessagesComponentFamilyState.atomFamily(atomKey),
latestMessages.filter(
(message) => message.id !== optimisticUserMessage.id,
),
);
}
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
}
setPendingThreadIdAfterFirstSend((pendingId) => {
if (isDefined(pendingId)) {
setCurrentAIChatThread(pendingId);
}
return null;
});
} catch {
const latestMessages = store.get(targetAtom.atomFamily(atomKey));
store.set(
targetAtom.atomFamily(atomKey),
latestMessages.filter(
(message) => message.id !== optimisticUserMessage.id,
),
);
}
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [
store,
isLoading,
agentChatSelectedFiles,
ensureThreadIdForSend,
setAgentChatInput,
getBrowsingContext,
sendMessage,
agentChatUploadedFiles,
setAgentChatUploadedFiles,
setAgentChatDraftsByThreadId,
modelIdForRequest,
setCurrentAIChatThread,
]);
useListenToBrowserEvent({
@@ -311,8 +222,6 @@ export const useAgentChat = (
});
const handleStop = useCallback(async () => {
stop();
const threadId = store.get(currentAIChatThreadState.atom);
if (!isDefined(threadId) || !isValidUuid(threadId)) {
@@ -325,31 +234,25 @@ export const useAgentChat = (
return;
}
fetch(`${REST_API_BASE_URL}/agent-chat/${threadId}/stream`, {
authenticatedFetch(`${REST_API_BASE_URL}/agent-chat/${threadId}/stream`, {
method: 'DELETE',
headers: {
Authorization: `Bearer ${tokenPair.accessOrWorkspaceAgnosticToken.token}`,
},
}).catch(() => {});
}, [stop, store]);
}, [store, authenticatedFetch]);
useListenToBrowserEvent({
eventName: AGENT_CHAT_STOP_EVENT_NAME,
onBrowserEvent: handleStop,
});
useListenToBrowserEvent({
eventName: AGENT_CHAT_RETRY_EVENT_NAME,
onBrowserEvent: regenerate,
});
// TODO: implement proper retry (re-send last user message to /message)
// The AIChatErrorMessage UI still dispatches AGENT_CHAT_RETRY_EVENT_NAME
// but no listener is wired until retry is properly implemented.
return {
messages,
handleSendMessage,
handleStop,
resumeStream,
isLoading,
error,
status,
};
};
@@ -0,0 +1,292 @@
import { useCallback, useEffect, useRef } from 'react';
import { readUIMessageStream } from 'ai';
import { type UIMessageChunk } from 'ai';
import { print, type ExecutionResult } from 'graphql';
import { useStore } from 'jotai';
import { type AgentChatSubscriptionEvent, type ExtendedUIMessage } from 'twenty-shared/ai';
import { isDefined } from 'twenty-shared/utils';
import { AGENT_CHAT_INSTANCE_ID } from '@/ai/constants/AgentChatInstanceId';
import { AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME } from '@/ai/constants/AgentChatRefetchMessagesEventName';
import { ON_AGENT_CHAT_EVENT } from '@/ai/graphql/subscriptions/OnAgentChatEvent';
import { agentChatErrorState } from '@/ai/states/agentChatErrorState';
import { agentChatIsStreamingState } from '@/ai/states/agentChatIsStreamingState';
import { agentChatMessagesComponentFamilyState } from '@/ai/states/agentChatMessagesComponentFamilyState';
import { agentChatUsageState } from '@/ai/states/agentChatUsageState';
import { currentAIChatThreadTitleState } from '@/ai/states/currentAIChatThreadTitleState';
import { dispatchBrowserEvent } from '@/browser-event/utils/dispatchBrowserEvent';
import { useAtomStateValue } from '@/ui/utilities/state/jotai/hooks/useAtomStateValue';
import { sseClientState } from '@/sse-db-event/states/sseClientState';
const THROTTLE_MS = 100;
type AgentChatEventPayload = {
onAgentChatEvent: {
threadId: string;
event: AgentChatSubscriptionEvent;
};
};
export const useAgentChatSubscription = (threadId: string | null) => {
const store = useStore();
const sseClient = useAtomStateValue(sseClientState);
const disposeRef = useRef<(() => void) | null>(null);
const writerRef = useRef<WritableStreamDefaultWriter<UIMessageChunk> | null>(
null,
);
const isStreamingRef = useRef(false);
const cleanup = useCallback(() => {
if (writerRef.current) {
writerRef.current.close().catch(() => {});
writerRef.current = null;
}
if (disposeRef.current) {
disposeRef.current();
disposeRef.current = null;
}
if (isStreamingRef.current) {
isStreamingRef.current = false;
store.set(agentChatIsStreamingState.atom, false);
}
}, [store]);
useEffect(() => {
if (!isDefined(threadId)) {
cleanup();
return;
}
if (!isDefined(sseClient)) {
return;
}
let bridge: TransformStream<UIMessageChunk> | null = null;
let throttleTimer: ReturnType<typeof setTimeout> | null = null;
let latestMessage: ExtendedUIMessage | null = null;
const flushToAtom = () => {
const messageToFlush = latestMessage;
if (!isDefined(messageToFlush) || !isDefined(threadId)) {
return;
}
const atomKey = {
instanceId: AGENT_CHAT_INSTANCE_ID,
familyKey: { threadId },
};
const currentMessages = store.get(
agentChatMessagesComponentFamilyState.atomFamily(atomKey),
);
// Find and replace the streaming assistant message by ID, or append it.
// This preserves optimistic user messages that aren't yet in fetched data.
const streamingMsgIndex = currentMessages.findIndex(
(message) => message.id === messageToFlush.id,
);
if (streamingMsgIndex >= 0) {
const updatedMessages = [...currentMessages];
updatedMessages[streamingMsgIndex] = messageToFlush;
store.set(
agentChatMessagesComponentFamilyState.atomFamily(atomKey),
updatedMessages,
);
} else {
store.set(
agentChatMessagesComponentFamilyState.atomFamily(atomKey),
[...currentMessages, messageToFlush],
);
}
};
const scheduleAtomUpdate = (message: ExtendedUIMessage) => {
latestMessage = message;
if (!isDefined(throttleTimer)) {
// Leading edge: flush immediately so the first chunk renders instantly
flushToAtom();
// Then suppress further flushes for THROTTLE_MS
throttleTimer = setTimeout(() => {
throttleTimer = null;
// Trailing edge: flush whatever accumulated during the window
flushToAtom();
}, THROTTLE_MS);
}
};
const startReadLoop = async (
readable: ReadableStream<UIMessageChunk>,
) => {
const messageStream = readUIMessageStream({ stream: readable });
for await (const message of messageStream) {
const extendedMessage = message as ExtendedUIMessage;
// Extract usage and title from data parts
const titlePart = extendedMessage.parts.find(
(part) => part.type === 'data-thread-title',
);
if (
isDefined(titlePart) &&
titlePart.type === 'data-thread-title'
) {
store.set(
currentAIChatThreadTitleState.atom,
titlePart.data.title,
);
}
const metadata = extendedMessage.metadata as
| {
usage?: {
inputTokens: number;
outputTokens: number;
cachedInputTokens: number;
inputCredits: number;
outputCredits: number;
conversationSize: number;
};
model?: {
contextWindowTokens: number;
};
}
| undefined;
if (isDefined(metadata?.usage) && isDefined(metadata?.model)) {
const usage = metadata.usage;
const model = metadata.model;
store.set(agentChatUsageState.atom, (prev) => ({
lastMessage: {
inputTokens: usage.inputTokens,
outputTokens: usage.outputTokens,
cachedInputTokens: usage.cachedInputTokens,
inputCredits: usage.inputCredits,
outputCredits: usage.outputCredits,
},
conversationSize: usage.conversationSize,
contextWindowTokens: model.contextWindowTokens,
inputTokens: (prev?.inputTokens ?? 0) + usage.inputTokens,
outputTokens: (prev?.outputTokens ?? 0) + usage.outputTokens,
inputCredits: (prev?.inputCredits ?? 0) + usage.inputCredits,
outputCredits: (prev?.outputCredits ?? 0) + usage.outputCredits,
}));
}
scheduleAtomUpdate(extendedMessage);
}
// Stream finished -- flush last message immediately
if (isDefined(throttleTimer)) {
clearTimeout(throttleTimer);
throttleTimer = null;
}
flushToAtom();
isStreamingRef.current = false;
store.set(agentChatIsStreamingState.atom, false);
};
const handleEvent = (event: AgentChatSubscriptionEvent) => {
switch (event.type) {
case 'stream-chunk': {
if (!isStreamingRef.current) {
isStreamingRef.current = true;
store.set(agentChatIsStreamingState.atom, true);
bridge = new TransformStream<UIMessageChunk>();
writerRef.current = bridge.writable.getWriter();
startReadLoop(bridge.readable).catch(() => {
isStreamingRef.current = false;
store.set(agentChatIsStreamingState.atom, false);
});
}
if (isDefined(writerRef.current)) {
writerRef.current
.write(event.chunk as UIMessageChunk)
.catch(() => {});
}
break;
}
case 'message-persisted': {
// Close the current stream writer so readUIMessageStream finishes
if (isDefined(writerRef.current)) {
writerRef.current.close().catch(() => {});
writerRef.current = null;
}
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
break;
}
case 'queue-updated': {
dispatchBrowserEvent(AGENT_CHAT_REFETCH_MESSAGES_EVENT_NAME);
break;
}
case 'stream-error': {
const streamError = new Error(event.message) as Error & {
code?: string;
};
streamError.code = event.code;
store.set(agentChatErrorState.atom, streamError);
if (isDefined(writerRef.current)) {
writerRef.current.close().catch(() => {});
writerRef.current = null;
}
isStreamingRef.current = false;
store.set(agentChatIsStreamingState.atom, false);
break;
}
}
};
const dispose = sseClient.subscribe<AgentChatEventPayload>(
{
query: print(ON_AGENT_CHAT_EVENT),
variables: { threadId },
},
{
next: (
value: ExecutionResult<AgentChatEventPayload>,
) => {
if (isDefined(value.data?.onAgentChatEvent?.event)) {
handleEvent(
value.data.onAgentChatEvent
.event as AgentChatSubscriptionEvent,
);
}
},
error: () => {
// graphql-sse handles reconnection automatically
},
complete: () => {
cleanup();
},
},
);
disposeRef.current = dispose;
return () => {
if (isDefined(throttleTimer)) {
clearTimeout(throttleTimer);
}
cleanup();
};
}, [threadId, sseClient, store, cleanup]);
};
@@ -0,0 +1,77 @@
import { useCallback } from 'react';
import { isDefined } from 'twenty-shared/utils';
import { getTokenPair } from '@/apollo/utils/getTokenPair';
import { renewToken } from '@/auth/services/AuthService';
import { tokenPairState } from '@/auth/states/tokenPairState';
import { useSetAtomState } from '@/ui/utilities/state/jotai/hooks/useSetAtomState';
import { REACT_APP_SERVER_BASE_URL } from '~/config';
export const useAuthenticatedFetch = () => {
const setTokenPair = useSetAtomState(tokenPairState);
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);
return null;
}
const renewedAccessToken =
renewedTokens.accessOrWorkspaceAgnosticToken?.token;
if (!isDefined(renewedAccessToken)) {
setTokenPair(null);
return null;
}
setTokenPair(renewedTokens);
const updatedHeaders = new Headers(init?.headers ?? {});
updatedHeaders.set('Authorization', `Bearer ${renewedAccessToken}`);
return fetch(input, {
...init,
headers: updatedHeaders,
});
} catch {
setTokenPair(null);
return null;
}
},
[setTokenPair],
);
const authenticatedFetch = useCallback(
async (input: RequestInfo | URL, init?: RequestInit) => {
const response = await fetch(input, init);
if (response.status === 401) {
const retriedResponse = await retryFetchWithRenewedToken(input, init);
return retriedResponse ?? response;
}
return response;
},
[retryFetchWithRenewedToken],
);
return { authenticatedFetch };
};
@@ -0,0 +1,10 @@
import { AgentChatComponentInstanceContext } from '@/ai/states/AgentChatComponentInstanceContext';
import { createAtomComponentFamilyState } from '@/ui/utilities/state/jotai/utils/createAtomComponentFamilyState';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
export const agentChatQueuedMessagesComponentFamilyState =
createAtomComponentFamilyState<ExtendedUIMessage[], { threadId: string }>({
key: 'agentChatQueuedMessagesComponentFamilyState',
defaultValue: [],
componentInstanceContext: AgentChatComponentInstanceContext,
});
@@ -1,60 +0,0 @@
import { normalizeAiSdkError } from '@/ai/utils/normalizeAiSdkError';
describe('normalizeAiSdkError', () => {
it('should return undefined for undefined input', () => {
expect(normalizeAiSdkError(undefined)).toBeUndefined();
});
it('should return the error unchanged if it already has a code', () => {
const error = new Error('test') as Error & { code: string };
error.code = 'EXISTING_CODE';
const result = normalizeAiSdkError(error);
expect(result).toBe(error);
expect((result as Error & { code: string }).code).toBe('EXISTING_CODE');
});
it('should parse JSON message and attach code from response body', () => {
const error = new Error(
'{"statusCode":402,"error":"Error","messages":["Credits exhausted"],"code":"BILLING_CREDITS_EXHAUSTED"}',
);
const result = normalizeAiSdkError(error);
expect(result).not.toBe(error);
expect((result as Error & { code: string }).code).toBe(
'BILLING_CREDITS_EXHAUSTED',
);
expect(result?.message).toBe(error.message);
});
it('should return original error for non-JSON message', () => {
const error = new Error('Something went wrong');
const result = normalizeAiSdkError(error);
expect(result).toBe(error);
});
it('should return original error for JSON message without code', () => {
const error = new Error(
'{"statusCode":500,"error":"Internal Server Error"}',
);
const result = normalizeAiSdkError(error);
expect(result).toBe(error);
});
it('should preserve the original stack trace', () => {
const error = new Error(
'{"statusCode":402,"code":"BILLING_CREDITS_EXHAUSTED"}',
);
const originalStack = error.stack;
const result = normalizeAiSdkError(error);
expect(result?.stack).toBe(originalStack);
});
});
@@ -8,6 +8,7 @@ export const mapDBMessagesToUIMessages = (
return dbMessages.map((dbMessage) => ({
id: dbMessage.id,
role: dbMessage.role as ExtendedUIMessage['role'],
status: dbMessage.status as 'queued' | 'sent',
parts: dbMessage.parts.map(mapDBPartToUIMessagePart),
metadata: {
createdAt: dbMessage.createdAt,
@@ -1,15 +0,0 @@
import { type ExtendedUIMessage } from 'twenty-shared/ai';
// The SDK's messages array is always [...seedMessages, ...newMessages]
// where seedMessages are the fetchedMessages we passed to useChat.
// SDK-generated IDs differ from server-persisted IDs, so comparing
// by ID doesn't work. Instead, slice by position: everything beyond
// the fetched count is a new streaming message.
export const mergeAgentChatFetchedAndStreamingMessages = (
fetchedMessages: ExtendedUIMessage[],
streamingMessages: ExtendedUIMessage[],
): ExtendedUIMessage[] => {
const newStreamingMessages = streamingMessages.slice(fetchedMessages.length);
return [...fetchedMessages, ...newStreamingMessages];
};
@@ -1,43 +0,0 @@
import { isDefined } from 'twenty-shared/utils';
// The Vercel AI SDK wraps non-200 responses into an Error whose message
// is the raw JSON response body, losing any custom properties (like `code`).
// This function parses that JSON and re-attaches the code so downstream
// consumers (extractErrorCode) can find it without knowing about the SDK.
export const normalizeAiSdkError = (
error: Error | undefined,
): Error | undefined => {
if (!isDefined(error) || !(error instanceof Error)) {
return error;
}
if (
'code' in error &&
typeof (error as Error & { code: string }).code === 'string'
) {
return error;
}
try {
const parsed: unknown = JSON.parse(error.message);
if (
isDefined(parsed) &&
typeof parsed === 'object' &&
'code' in parsed &&
typeof (parsed as { code: unknown }).code === 'string'
) {
const normalizedError = new Error(error.message) as Error & {
code: string;
};
normalizedError.code = (parsed as { code: string }).code;
normalizedError.stack = error.stack;
return normalizedError;
}
} catch {
// message is not JSON — nothing to normalize
}
return error;
};
-1
View File
@@ -168,7 +168,6 @@
"react-dom": "18.3.1",
"redis": "^4.7.0",
"reflect-metadata": "0.2.2",
"resumable-stream": "^2.2.12",
"rxjs": "7.8.1",
"semver": "7.6.3",
"sharp": "0.32.6",
@@ -0,0 +1,39 @@
import { MigrationInterface, QueryRunner } from 'typeorm';
export class AddStatusToAgentMessage1774776000000
implements MigrationInterface
{
public async up(queryRunner: QueryRunner): Promise<void> {
await queryRunner.query(
`CREATE TYPE "core"."agentMessage_status_enum" AS ENUM ('queued', 'sent')`,
);
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" ADD COLUMN "status" "core"."agentMessage_status_enum" NOT NULL DEFAULT 'sent'`,
);
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" ALTER COLUMN "turnId" DROP NOT NULL`,
);
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" ADD COLUMN "processedAt" TIMESTAMPTZ`,
);
await queryRunner.query(
`UPDATE "core"."agentMessage" SET "processedAt" = "createdAt" WHERE "status" = 'sent'`,
);
}
public async down(queryRunner: QueryRunner): Promise<void> {
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" DROP COLUMN "processedAt"`,
);
await queryRunner.query(
`DELETE FROM "core"."agentMessage" WHERE "turnId" IS NULL`,
);
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" ALTER COLUMN "turnId" SET NOT NULL`,
);
await queryRunner.query(
`ALTER TABLE "core"."agentMessage" DROP COLUMN "status"`,
);
await queryRunner.query(`DROP TYPE "core"."agentMessage_status_enum"`);
}
}
@@ -13,8 +13,8 @@ export class AgentMessageDTO {
@Field(() => UUIDScalarType)
threadId: string;
@Field(() => UUIDScalarType)
turnId: string;
@Field(() => UUIDScalarType, { nullable: true })
turnId: string | null;
@Field(() => UUIDScalarType, { nullable: true })
agentId: string | null;
@@ -22,9 +22,16 @@ export class AgentMessageDTO {
@Field()
role: 'system' | 'user' | 'assistant';
@Field()
status: 'queued' | 'sent';
@Field(() => [AgentMessagePartDTO])
parts: AgentMessagePartDTO[];
@IsDateString()
@Field(() => Date, { nullable: true })
processedAt: Date | null;
@IsDateString()
@Field()
createdAt: Date;
@@ -20,6 +20,11 @@ export enum AgentMessageRole {
ASSISTANT = 'assistant',
}
export enum AgentMessageStatus {
QUEUED = 'queued',
SENT = 'sent',
}
@Entity('agentMessage')
export class AgentMessageEntity {
@PrimaryGeneratedColumn('uuid')
@@ -35,15 +40,16 @@ export class AgentMessageEntity {
@JoinColumn({ name: 'threadId' })
thread: Relation<AgentChatThreadEntity>;
@Column('uuid')
@Column({ type: 'uuid', nullable: true })
@Index()
turnId: string;
turnId: string | null;
@ManyToOne(() => AgentTurnEntity, {
onDelete: 'CASCADE',
nullable: true,
})
@JoinColumn({ name: 'turnId' })
turn: Relation<AgentTurnEntity>;
turn: Relation<AgentTurnEntity> | null;
@Column({ type: 'uuid', nullable: true })
@Index()
@@ -52,9 +58,19 @@ export class AgentMessageEntity {
@Column({ type: 'enum', enum: AgentMessageRole })
role: AgentMessageRole;
@Column({
type: 'enum',
enum: AgentMessageStatus,
default: AgentMessageStatus.SENT,
})
status: AgentMessageStatus;
@OneToMany(() => AgentMessagePartEntity, (part) => part.message)
parts: Relation<AgentMessagePartEntity[]>;
@Column({ type: 'timestamptz', nullable: true })
processedAt: Date | null;
@CreateDateColumn({ type: 'timestamptz' })
createdAt: Date;
}
@@ -0,0 +1,67 @@
import { type ExtendedUIMessagePart } from 'twenty-shared/ai';
import { type AgentMessagePartEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message-part.entity';
export const mapDBPartToUIMessagePart = (
part: AgentMessagePartEntity,
): ExtendedUIMessagePart | null => {
switch (part.type) {
case 'text':
return {
type: 'text',
text: part.textContent ?? '',
};
case 'reasoning':
return {
type: 'reasoning',
text: part.reasoningContent ?? '',
state: (part.state as 'streaming' | 'done') ?? 'done',
};
case 'file':
return {
type: 'file',
mediaType: part.fileFilename?.endsWith('.png')
? 'image/png'
: 'application/octet-stream',
filename: part.fileFilename ?? '',
url: '',
};
case 'source-url':
return {
type: 'source-url',
sourceId: part.sourceUrlSourceId ?? '',
url: part.sourceUrlUrl ?? '',
title: part.sourceUrlTitle ?? '',
providerMetadata: part.providerMetadata ?? undefined,
};
case 'source-document':
return {
type: 'source-document',
sourceId: part.sourceDocumentSourceId ?? '',
mediaType: part.sourceDocumentMediaType ?? '',
title: part.sourceDocumentTitle ?? '',
filename: part.sourceDocumentFilename ?? '',
providerMetadata: part.providerMetadata ?? undefined,
};
case 'step-start':
return {
type: 'step-start',
};
case 'data-routing-status':
return null;
default: {
if (part.type.includes('tool-') && part.toolCallId) {
return {
type: part.type,
toolCallId: part.toolCallId,
input: part.toolInput ?? {},
output: part.toolOutput,
errorText: part.errorMessage ?? '',
state: part.state,
} as ExtendedUIMessagePart;
}
return null;
}
}
};
@@ -0,0 +1,13 @@
import { type ExtendedUIMessagePart } from 'twenty-shared/ai';
import { type AgentMessagePartEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message-part.entity';
import { mapDBPartToUIMessagePart } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartToUIMessagePart';
export const mapDBPartsToUIMessageParts = (
parts: AgentMessagePartEntity[],
): ExtendedUIMessagePart[] => {
return parts
.sort((a, b) => a.orderIndex - b.orderIndex)
.map(mapDBPartToUIMessagePart)
.filter((part): part is ExtendedUIMessagePart => part !== null);
};
@@ -38,8 +38,9 @@ import { AgentChatThreadDTO } from './dtos/agent-chat-thread.dto';
import { AgentChatThreadEntity } from './entities/agent-chat-thread.entity';
import { StreamAgentChatJob } from './jobs/stream-agent-chat.job';
import { AgentChatResolver } from './resolvers/agent-chat.resolver';
import { AgentChatSubscriptionResolver } from './resolvers/agent-chat-subscription.resolver';
import { AgentChatCancelSubscriberService } from './services/agent-chat-cancel-subscriber.service';
import { AgentChatResumableStreamService } from './services/agent-chat-resumable-stream.service';
import { AgentChatEventPublisherService } from './services/agent-chat-event-publisher.service';
import { AgentChatStreamingService } from './services/agent-chat-streaming.service';
import { AgentChatService } from './services/agent-chat.service';
import { AgentTitleGenerationService } from './services/agent-title-generation.service';
@@ -104,8 +105,9 @@ import { SystemPromptBuilderService } from './services/system-prompt-builder.ser
controllers: [AgentChatController],
providers: [
AgentChatCancelSubscriberService,
AgentChatEventPublisherService,
AgentChatResolver,
AgentChatResumableStreamService,
AgentChatSubscriptionResolver,
AgentChatService,
AgentChatStreamingService,
AgentTitleGenerationService,
@@ -2,10 +2,8 @@ import {
Body,
Controller,
Delete,
Get,
Param,
Post,
Res,
UseFilters,
UseGuards,
} from '@nestjs/common';
@@ -13,8 +11,6 @@ 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';
@@ -43,8 +39,9 @@ import { AgentRestApiExceptionFilter } from 'src/engine/metadata-modules/ai/ai-a
import type { BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util';
import { AgentChatResumableStreamService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-resumable-stream.service';
import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service';
import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';
import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service';
import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models/services/ai-model-registry.service';
@Controller('rest/agent-chat')
@@ -57,7 +54,8 @@ import { AiModelRegistryService } from 'src/engine/metadata-modules/ai/ai-models
export class AgentChatController {
constructor(
private readonly agentStreamingService: AgentChatStreamingService,
private readonly resumableStreamService: AgentChatResumableStreamService,
private readonly agentChatService: AgentChatService,
private readonly eventPublisherService: AgentChatEventPublisherService,
private readonly billingService: BillingService,
private readonly twentyConfigService: TwentyConfigService,
private readonly aiModelRegistryService: AiModelRegistryService,
@@ -66,19 +64,19 @@ export class AgentChatController {
private readonly threadRepository: Repository<AgentChatThreadEntity>,
) {}
@Post('stream')
@Post(':threadId/message')
@UseGuards(SettingsPermissionGuard(PermissionFlagType.AI))
async streamAgentChat(
async sendMessage(
@Param('threadId') threadId: string,
@Body()
body: {
threadId: string;
messages: ExtendedUIMessage[];
text: string;
messages?: ExtendedUIMessage[];
browsingContext?: BrowsingContextType | null;
modelId?: string;
},
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() workspace: WorkspaceEntity,
@Res() response: Response,
) {
if (this.aiModelRegistryService.getAvailableModels().length === 0) {
throw new AgentException(
@@ -108,47 +106,44 @@ export class AgentChatController {
}
}
return this.agentStreamingService.streamAgentChat({
threadId: body.threadId,
messages: body.messages,
browsingContext: body.browsingContext ?? null,
modelId: body.modelId,
userWorkspaceId,
workspace,
response,
});
}
@Get(':threadId/stream')
@UseGuards(SettingsPermissionGuard(PermissionFlagType.AI))
async resumeAgentChatStream(
@Param('threadId') threadId: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@Res() response: Response,
) {
const thread = await this.threadRepository.findOne({
where: { id: threadId, userWorkspaceId },
});
if (!isDefined(thread) || !isDefined(thread.activeStreamId)) {
response.status(204).end();
return;
}
const resumedNodeReadable =
await this.resumableStreamService.resumeExistingStreamAsNodeReadable(
thread.activeStreamId,
if (!isDefined(thread)) {
throw new AgentException(
'Thread not found',
AgentExceptionCode.AGENT_EXECUTION_FAILED,
);
if (!isDefined(resumedNodeReadable)) {
response.status(204).end();
return;
}
response.writeHead(200, UI_MESSAGE_STREAM_HEADERS);
resumedNodeReadable.pipe(response);
// Server decides: if the thread has an active stream, queue the message;
// otherwise start streaming immediately.
if (isDefined(thread.activeStreamId)) {
const message = await this.agentChatService.queueMessage({
threadId,
text: body.text,
});
await this.eventPublisherService.publish({
threadId,
workspaceId: workspace.id,
event: { type: 'queue-updated' },
});
return { messageId: message.id, queued: true };
}
const result = await this.agentStreamingService.streamAgentChat({
threadId,
messages: body.messages ?? [],
browsingContext: body.browsingContext ?? null,
modelId: body.modelId,
userWorkspaceId,
workspace,
});
return { messageId: threadId, queued: false, streamId: result.streamId };
}
@Delete(':threadId/stream')
@@ -165,10 +160,6 @@ export class AgentChatController {
return { success: true };
}
// Publish a cancel signal via Redis pub/sub. The BullMQ worker
// processing this thread's stream subscribes to this channel and
// will abort the LLM connection when the message arrives — stopping
// token generation and billing immediately.
const redis = this.redisClientService.getClient();
await redis.publish(getCancelChannel(threadId), 'cancel');
@@ -180,4 +171,44 @@ export class AgentChatController {
return { success: true };
}
@Delete('queue/:messageId')
@UseGuards(SettingsPermissionGuard(PermissionFlagType.AI))
async deleteQueuedMessage(
@Param('messageId') messageId: string,
@AuthUserWorkspaceId() userWorkspaceId: string,
@AuthWorkspace() workspace: WorkspaceEntity,
) {
const message = await this.agentChatService.findQueuedMessage(messageId);
if (!isDefined(message)) {
throw new AgentException(
'Queued message not found',
AgentExceptionCode.AGENT_EXECUTION_FAILED,
);
}
const thread = await this.threadRepository.findOne({
where: { id: message.threadId, userWorkspaceId },
});
if (!isDefined(thread)) {
throw new AgentException(
'Thread not found',
AgentExceptionCode.AGENT_EXECUTION_FAILED,
);
}
const deleted = await this.agentChatService.deleteQueuedMessage(messageId);
if (deleted) {
await this.eventPublisherService.publish({
threadId: message.threadId,
workspaceId: workspace.id,
event: { type: 'queue-updated' },
});
}
return { success: deleted };
}
}
@@ -0,0 +1,12 @@
import { Field, ObjectType } from '@nestjs/graphql';
import GraphQLJSON from 'graphql-type-json';
@ObjectType('AgentChatEvent')
export class AgentChatEventDTO {
@Field(() => String)
threadId: string;
@Field(() => GraphQLJSON)
event: Record<string, unknown>;
}
@@ -0,0 +1 @@
export const STREAM_AGENT_CHAT_JOB_NAME = 'StreamAgentChatJob';
@@ -0,0 +1,20 @@
import type {
ExtendedUIMessage,
ExtendedUIMessagePart,
} from 'twenty-shared/ai';
import type { BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
export type StreamAgentChatJobData = {
threadId: string;
streamId: string;
userWorkspaceId: string;
workspaceId: string;
messages: ExtendedUIMessage[];
browsingContext: BrowsingContextType | null;
modelId?: string;
lastUserMessageText: string;
lastUserMessageParts: ExtendedUIMessagePart[];
hasTitle: boolean;
existingTurnId?: string;
};
@@ -1,7 +1,7 @@
import { Logger, Scope } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { createUIMessageStream, JsonToSseTransformStream } from 'ai';
import { createUIMessageStream } from 'ai';
import type {
CodeExecutionData,
ExtendedUIMessage,
@@ -10,37 +10,27 @@ import type {
import { Repository } from 'typeorm';
import { AgentChatCancelSubscriberService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-cancel-subscriber.service';
import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service';
import { toDisplayCredits } from 'src/engine/core-modules/usage/utils/to-display-credits.util';
import { WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
import { Process } from 'src/engine/core-modules/message-queue/decorators/process.decorator';
import { Processor } from 'src/engine/core-modules/message-queue/decorators/processor.decorator';
import { MessageQueue } from 'src/engine/core-modules/message-queue/message-queue.constants';
import { AgentMessageRole } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
import type { BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
import { AgentChatStreamingService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-streaming.service';
import { computeCostBreakdown } from 'src/engine/metadata-modules/ai/ai-billing/utils/compute-cost-breakdown.util';
import { convertDollarsToBillingCredits } from 'src/engine/metadata-modules/ai/ai-billing/utils/convert-dollars-to-billing-credits.util';
import { extractCacheCreationTokens } from 'src/engine/metadata-modules/ai/ai-billing/utils/extract-cache-creation-tokens.util';
import type { AIModelConfig } from 'src/engine/metadata-modules/ai/ai-models/types/ai-model-config.type';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import { AgentChatResumableStreamService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-resumable-stream.service';
import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service';
import { ChatExecutionService } from 'src/engine/metadata-modules/ai/ai-chat/services/chat-execution.service';
import { getCancelChannel } from 'src/engine/metadata-modules/ai/ai-chat/utils/get-cancel-channel.util';
export const STREAM_AGENT_CHAT_JOB_NAME = 'StreamAgentChatJob';
import { STREAM_AGENT_CHAT_JOB_NAME } from './stream-agent-chat-job-name.constant';
import { type StreamAgentChatJobData } from './stream-agent-chat-job.types';
export type StreamAgentChatJobData = {
threadId: string;
streamId: string;
userWorkspaceId: string;
workspaceId: string;
messages: ExtendedUIMessage[];
browsingContext: BrowsingContextType | null;
modelId?: string;
lastUserMessageText: string;
lastUserMessageParts: ExtendedUIMessagePart[];
hasTitle: boolean;
};
export { STREAM_AGENT_CHAT_JOB_NAME, type StreamAgentChatJobData };
@Processor({ queueName: MessageQueue.aiStreamQueue, scope: Scope.REQUEST })
export class StreamAgentChatJob {
@@ -53,8 +43,9 @@ export class StreamAgentChatJob {
private readonly workspaceRepository: Repository<WorkspaceEntity>,
private readonly agentChatService: AgentChatService,
private readonly chatExecutionService: ChatExecutionService,
private readonly resumableStreamService: AgentChatResumableStreamService,
private readonly eventPublisherService: AgentChatEventPublisherService,
private readonly cancelSubscriberService: AgentChatCancelSubscriberService,
private readonly agentChatStreamingService: AgentChatStreamingService,
) {}
@Process(STREAM_AGENT_CHAT_JOB_NAME)
@@ -65,9 +56,14 @@ export class StreamAgentChatJob {
if (!workspace) {
this.logger.error(`Workspace ${data.workspaceId} not found`);
await this.resumableStreamService.writeStreamError(data.streamId, {
code: 'WORKSPACE_NOT_FOUND',
message: `Workspace ${data.workspaceId} not found`,
await this.eventPublisherService.publish({
threadId: data.threadId,
workspaceId: data.workspaceId,
event: {
type: 'stream-error',
code: 'WORKSPACE_NOT_FOUND',
message: `Workspace ${data.workspaceId} not found`,
},
});
return;
@@ -86,11 +82,18 @@ export class StreamAgentChatJob {
this.logger.error(
`Stream ${data.streamId} failed: ${error instanceof Error ? error.message : String(error)}`,
);
await this.resumableStreamService
.writeStreamError(data.streamId, {
code: 'STREAM_EXECUTION_FAILED',
message:
error instanceof Error ? error.message : 'Stream execution failed',
await this.eventPublisherService
.publish({
threadId: data.threadId,
workspaceId: data.workspaceId,
event: {
type: 'stream-error',
code: 'STREAM_EXECUTION_FAILED',
message:
error instanceof Error
? error.message
: 'Stream execution failed',
},
})
.catch(() => {});
} finally {
@@ -105,6 +108,20 @@ export class StreamAgentChatJob {
})
.execute()
.catch(() => {});
// Auto-flush the next queued message for this thread
await this.agentChatStreamingService
.flushNextQueuedMessage(
data.threadId,
data.userWorkspaceId,
data.workspaceId,
data.hasTitle,
)
.catch((error) => {
this.logger.error(
`Failed to flush queued message for thread ${data.threadId}: ${error instanceof Error ? error.message : String(error)}`,
);
});
}
}
@@ -113,16 +130,20 @@ export class StreamAgentChatJob {
workspace: WorkspaceEntity,
abortSignal: AbortSignal,
): Promise<void> {
const userMessagePromise = this.agentChatService.addMessage({
threadId: data.threadId,
uiMessage: {
role: AgentMessageRole.USER,
parts: data.lastUserMessageParts.filter(
(part): part is ExtendedUIMessagePart =>
part.type === 'text' || part.type === 'file',
),
},
});
// When processing a promoted queued message, the user message already
// exists in the DB with a turn — skip persisting it again.
const userMessagePromise = data.existingTurnId
? Promise.resolve({ turnId: data.existingTurnId })
: this.agentChatService.addMessage({
threadId: data.threadId,
uiMessage: {
role: AgentMessageRole.USER,
parts: data.lastUserMessageParts.filter(
(part): part is ExtendedUIMessagePart =>
part.type === 'text' || part.type === 'file',
),
},
});
userMessagePromise.catch(() => {});
@@ -132,7 +153,7 @@ export class StreamAgentChatJob {
.generateTitleIfNeeded(data.threadId, data.lastUserMessageText)
.catch(() => null);
await this.buildAndPipeStream({
await this.buildAndPublishStream({
workspace,
data,
userMessagePromise,
@@ -141,7 +162,7 @@ export class StreamAgentChatJob {
});
}
private async buildAndPipeStream({
private async buildAndPublishStream({
workspace,
data,
userMessagePromise,
@@ -150,7 +171,7 @@ export class StreamAgentChatJob {
}: {
workspace: WorkspaceEntity;
data: StreamAgentChatJobData;
userMessagePromise: Promise<{ turnId: string }>;
userMessagePromise: Promise<{ turnId: string | null }>;
titlePromise: Promise<string | null>;
abortSignal: AbortSignal;
}): Promise<void> {
@@ -164,6 +185,14 @@ export class StreamAgentChatJob {
let lastStepConversationSize = 0;
let totalCacheCreationTokens = 0;
// onFinish fires before the uiStream is fully drained. We use this
// promise to coordinate: the IIFE waits for DB persist to complete
// before publishing message-persisted (after all chunks).
let resolveStreamFinished: () => void;
const streamFinishedPromise = new Promise<void>((res) => {
resolveStreamFinished = res;
});
abortSignal.addEventListener('abort', () => resolve(), { once: true });
const uiStream = createUIMessageStream<ExtendedUIMessage>({
@@ -233,7 +262,7 @@ export class StreamAgentChatJob {
userMessagePromise,
});
await titleWritePromise;
resolve();
resolveStreamFinished();
} catch (error) {
reject(error);
}
@@ -244,11 +273,34 @@ export class StreamAgentChatJob {
},
});
const sseStream = uiStream.pipeThrough(new JsonToSseTransformStream());
// Publish all chunks first, then signal completion. This guarantees
// message-persisted arrives after every stream-chunk on the client.
(async () => {
try {
for await (const chunk of uiStream) {
await this.eventPublisherService.publish({
threadId: data.threadId,
workspaceId: data.workspaceId,
event: {
type: 'stream-chunk',
chunk: chunk as Record<string, unknown>,
},
});
}
this.resumableStreamService
.createResumableStream(data.streamId, () => sseStream)
.catch(reject);
await streamFinishedPromise;
await this.eventPublisherService.publish({
threadId: data.threadId,
workspaceId: data.workspaceId,
event: { type: 'message-persisted', messageId: data.threadId },
});
resolve();
} catch (error) {
reject(error);
}
})();
});
}
@@ -358,7 +410,7 @@ export class StreamAgentChatJob {
};
lastStepConversationSize: number;
modelConfig: AIModelConfig;
userMessagePromise: Promise<{ turnId: string }>;
userMessagePromise: Promise<{ turnId: string | null }>;
}): Promise<void> {
if (responseMessage.parts.length === 0) {
return;
@@ -369,7 +421,7 @@ export class StreamAgentChatJob {
await this.agentChatService.addMessage({
threadId,
uiMessage: responseMessage,
turnId: userMessage.turnId,
turnId: userMessage.turnId ?? undefined,
});
await this.threadRepository.update(threadId, {
@@ -0,0 +1,63 @@
import { Args, Subscription } from '@nestjs/graphql';
import { UseGuards } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { PermissionFlagType } from 'twenty-shared/constants';
import { isDefined } from 'twenty-shared/utils';
import { type Repository } from 'typeorm';
import { MetadataResolver } from 'src/engine/api/graphql/graphql-config/decorators/metadata-resolver.decorator';
import { type WorkspaceEntity } from 'src/engine/core-modules/workspace/workspace.entity';
import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-workspace-id.decorator';
import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
import { SettingsPermissionGuard } from 'src/engine/guards/settings-permission.guard';
import { UserAuthGuard } from 'src/engine/guards/user-auth.guard';
import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard';
import {
AgentException,
AgentExceptionCode,
} from 'src/engine/metadata-modules/ai/ai-agent/agent.exception';
import { AgentChatEventDTO } from 'src/engine/metadata-modules/ai/ai-chat/dtos/agent-chat-event.dto';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import { SubscriptionService } from 'src/engine/subscriptions/subscription.service';
@MetadataResolver()
@UseGuards(WorkspaceAuthGuard, UserAuthGuard)
export class AgentChatSubscriptionResolver {
constructor(
private readonly subscriptionService: SubscriptionService,
@InjectRepository(AgentChatThreadEntity)
private readonly threadRepository: Repository<AgentChatThreadEntity>,
) {}
@Subscription(() => AgentChatEventDTO, {
filter: (
payload: { onAgentChatEvent: AgentChatEventDTO },
variables: { threadId: string },
) => {
return payload.onAgentChatEvent.threadId === variables.threadId;
},
})
@UseGuards(SettingsPermissionGuard(PermissionFlagType.AI))
async onAgentChatEvent(
@Args('threadId') threadId: string,
@AuthWorkspace() workspace: WorkspaceEntity,
@AuthUserWorkspaceId() userWorkspaceId: string,
) {
const thread = await this.threadRepository.findOne({
where: { id: threadId, userWorkspaceId },
});
if (!isDefined(thread)) {
throw new AgentException(
'Thread not found',
AgentExceptionCode.AGENT_EXECUTION_FAILED,
);
}
return this.subscriptionService.subscribeToAgentChat({
workspaceId: workspace.id,
threadId,
});
}
}
@@ -0,0 +1,31 @@
import { Injectable } from '@nestjs/common';
import { type AgentChatSubscriptionEvent } from 'twenty-shared/ai';
import { SubscriptionService } from 'src/engine/subscriptions/subscription.service';
@Injectable()
export class AgentChatEventPublisherService {
constructor(private readonly subscriptionService: SubscriptionService) {}
async publish({
threadId,
workspaceId,
event,
}: {
threadId: string;
workspaceId: string;
event: AgentChatSubscriptionEvent;
}): Promise<void> {
await this.subscriptionService.publishToAgentChat({
workspaceId,
threadId,
payload: {
onAgentChatEvent: {
threadId,
event,
},
},
});
}
}
@@ -1,104 +0,0 @@
import { Injectable, OnModuleDestroy } from '@nestjs/common';
import { Readable } from 'stream';
import { type ReadableStream as NodeWebReadableStream } from 'stream/web';
import type { Redis } from 'ioredis';
import { createResumableStreamContext } from 'resumable-stream/ioredis';
import { RedisClientService } from 'src/engine/core-modules/redis-client/redis-client.service';
@Injectable()
export class AgentChatResumableStreamService implements OnModuleDestroy {
private streamContext: ReturnType<typeof createResumableStreamContext>;
private streamPublisher: Redis;
private streamSubscriber: Redis;
private redisClient: Redis;
constructor(private readonly redisClientService: RedisClientService) {
const baseClient = this.redisClientService.getClient();
this.streamPublisher = baseClient.duplicate();
this.streamSubscriber = baseClient.duplicate();
this.redisClient = baseClient.duplicate();
this.streamContext = createResumableStreamContext({
waitUntil: () => {},
publisher: this.streamPublisher,
subscriber: this.streamSubscriber,
});
}
async onModuleDestroy() {
await this.streamPublisher.quit();
await this.streamSubscriber.quit();
await this.redisClient.quit();
}
async createResumableStream(
streamId: string,
streamFactory: () => ReadableStream<string>,
) {
const resumableStream = await this.streamContext.createNewResumableStream(
streamId,
streamFactory,
);
if (!resumableStream) {
return;
}
// Read the stream to completion in the background so chunks are
// published to Redis and available for later resume consumers.
const reader = resumableStream.getReader();
void (async () => {
try {
while (true) {
const { done } = await reader.read();
if (done) {
break;
}
}
} catch {
// Stream interrupted — chunks already published are still resumable.
}
})();
}
async resumeExistingStreamAsNodeReadable(
streamId: string,
): Promise<Readable | null> {
const webStream = await this.streamContext.resumeExistingStream(streamId);
if (!webStream) {
return null;
}
return Readable.fromWeb(webStream as NodeWebReadableStream);
}
async writeStreamError(
streamId: string,
error: { code: string; message: string },
): Promise<void> {
await this.redisClient.set(
`ai-stream:error:${streamId}`,
JSON.stringify(error),
'EX',
60,
);
}
async readStreamError(
streamId: string,
): Promise<{ code: string; message: string } | null> {
const raw = await this.redisClient.get(`ai-stream:error:${streamId}`);
if (!raw) {
return null;
}
return JSON.parse(raw);
}
}
@@ -1,9 +1,7 @@
import { Injectable, Logger } from '@nestjs/common';
import { InjectRepository } from '@nestjs/typeorm';
import { type Readable } from 'stream';
import { generateId, UI_MESSAGE_STREAM_HEADERS } from 'ai';
import { type Response } from 'express';
import { generateId } from 'ai';
import { type ExtendedUIMessage } from 'twenty-shared/ai';
import { type Repository } from 'typeorm';
@@ -16,27 +14,23 @@ import {
AgentExceptionCode,
} from 'src/engine/metadata-modules/ai/ai-agent/agent.exception';
import { type BrowsingContextType } from 'src/engine/metadata-modules/ai/ai-agent/types/browsingContext.type';
import { AgentMessageStatus } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
import { mapDBPartsToUIMessageParts } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapDBPartsToUIMessageParts';
import { AgentChatThreadEntity } from 'src/engine/metadata-modules/ai/ai-chat/entities/agent-chat-thread.entity';
import {
STREAM_AGENT_CHAT_JOB_NAME,
type StreamAgentChatJobData,
} from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat.job';
import { AgentChatResumableStreamService } from './agent-chat-resumable-stream.service';
import { STREAM_AGENT_CHAT_JOB_NAME } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job-name.constant';
import { type StreamAgentChatJobData } from 'src/engine/metadata-modules/ai/ai-chat/jobs/stream-agent-chat-job.types';
import { AgentChatEventPublisherService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat-event-publisher.service';
import { AgentChatService } from 'src/engine/metadata-modules/ai/ai-chat/services/agent-chat.service';
export type StreamAgentChatOptions = {
threadId: string;
userWorkspaceId: string;
workspace: WorkspaceEntity;
response: Response;
messages: ExtendedUIMessage[];
browsingContext: BrowsingContextType | null;
modelId?: string;
};
const STREAM_READY_TIMEOUT_MS = 5_000;
const STREAM_READY_POLL_INTERVAL_MS = 50;
@Injectable()
export class AgentChatStreamingService {
private readonly logger = new Logger(AgentChatStreamingService.name);
@@ -46,7 +40,8 @@ export class AgentChatStreamingService {
private readonly threadRepository: Repository<AgentChatThreadEntity>,
@InjectMessageQueue(MessageQueue.aiStreamQueue)
private readonly messageQueueService: MessageQueueService,
private readonly resumableStreamService: AgentChatResumableStreamService,
private readonly agentChatService: AgentChatService,
private readonly eventPublisherService: AgentChatEventPublisherService,
) {}
async streamAgentChat({
@@ -55,9 +50,8 @@ export class AgentChatStreamingService {
workspace,
messages,
browsingContext,
response,
modelId,
}: StreamAgentChatOptions) {
}: StreamAgentChatOptions): Promise<{ streamId: string }> {
const thread = await this.threadRepository.findOne({
where: {
id: threadId,
@@ -97,67 +91,84 @@ export class AgentChatStreamingService {
activeStreamId: streamId,
});
try {
const result = await this.waitForResumableStream(streamId);
if ('error' in result) {
response.status(500).json(result.error);
return;
}
if (!result.readable) {
this.logger.error(
`Stream ${streamId} did not become available within timeout`,
);
response
.status(500)
.json({ code: 'WORKER_UNREACHABLE', message: 'Stream timed out' });
return;
}
response.writeHead(200, UI_MESSAGE_STREAM_HEADERS);
result.readable.pipe(response);
} catch (error) {
response.end();
throw error;
}
return { streamId };
}
private async waitForResumableStream(
streamId: string,
): Promise<
| { readable: Readable }
| { error: { code: string; message: string } }
| { readable: null }
> {
const maxAttempts = Math.ceil(
STREAM_READY_TIMEOUT_MS / STREAM_READY_POLL_INTERVAL_MS,
async flushNextQueuedMessage(
threadId: string,
userWorkspaceId: string,
workspaceId: string,
hasTitle: boolean,
): Promise<void> {
const queuedMessages =
await this.agentChatService.getQueuedMessages(threadId);
const nextQueued = queuedMessages[0];
if (!nextQueued) {
return;
}
const textPart = nextQueued.parts?.find((part) => part.type === 'text');
const messageText = textPart?.textContent ?? '';
if (messageText === '') {
await this.agentChatService.deleteQueuedMessage(nextQueued.id);
return;
}
const turnId = await this.agentChatService.promoteQueuedMessage(
nextQueued.id,
threadId,
);
for (let attempt = 0; attempt < maxAttempts; attempt++) {
const streamError =
await this.resumableStreamService.readStreamError(streamId);
await this.eventPublisherService.publish({
threadId,
workspaceId,
event: { type: 'queue-updated' },
});
if (streamError) {
return { error: streamError };
}
await this.eventPublisherService.publish({
threadId,
workspaceId,
event: { type: 'message-persisted', messageId: nextQueued.id },
});
const readable =
await this.resumableStreamService.resumeExistingStreamAsNodeReadable(
streamId,
);
const allMessages = await this.agentChatService.getMessagesForThread(
threadId,
userWorkspaceId,
);
if (readable) {
return { readable };
}
const uiMessages = allMessages
.filter((message) => message.status !== AgentMessageStatus.QUEUED)
.map((message) => ({
id: message.id,
role: message.role as 'user' | 'assistant' | 'system',
parts: mapDBPartsToUIMessageParts(message.parts ?? []),
createdAt: message.createdAt,
}));
await new Promise((resolve) =>
setTimeout(resolve, STREAM_READY_POLL_INTERVAL_MS),
);
}
const streamId = generateId();
return { readable: null };
await this.messageQueueService.add<StreamAgentChatJobData>(
STREAM_AGENT_CHAT_JOB_NAME,
{
threadId,
streamId,
userWorkspaceId,
workspaceId,
messages: uiMessages,
browsingContext: null,
lastUserMessageText: messageText,
lastUserMessageParts: [{ type: 'text', text: messageText }],
hasTitle,
existingTurnId: turnId,
},
);
await this.threadRepository.update(threadId, {
activeStreamId: streamId,
});
}
}
@@ -10,6 +10,7 @@ import { AgentMessagePartEntity } from 'src/engine/metadata-modules/ai/ai-agent-
import {
AgentMessageEntity,
AgentMessageRole,
AgentMessageStatus,
} from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-message.entity';
import { AgentTurnEntity } from 'src/engine/metadata-modules/ai/ai-agent-execution/entities/agent-turn.entity';
import { mapUIMessagePartsToDBParts } from 'src/engine/metadata-modules/ai/ai-agent-execution/utils/mapUIMessagePartsToDBParts';
@@ -91,6 +92,7 @@ export class AgentChatService {
turnId: actualTurnId,
role: uiMessage.role as AgentMessageRole,
agentId: agentId ?? null,
processedAt: new Date(),
});
const savedMessage = await this.messageRepository.save(message);
@@ -124,11 +126,91 @@ export class AgentChatService {
return this.messageRepository.find({
where: { threadId },
order: { createdAt: 'ASC' },
order: { processedAt: { direction: 'ASC', nulls: 'LAST' } },
relations: ['parts', 'parts.file'],
});
}
async queueMessage({
threadId,
text,
}: {
threadId: string;
text: string;
}): Promise<AgentMessageEntity> {
const message = this.messageRepository.create({
threadId,
turnId: null,
role: AgentMessageRole.USER,
agentId: null,
status: AgentMessageStatus.QUEUED,
});
const savedMessage = await this.messageRepository.save(message);
const part = this.messagePartRepository.create({
messageId: savedMessage.id,
orderIndex: 0,
type: 'text',
textContent: text,
});
await this.messagePartRepository.save(part);
return savedMessage;
}
async getQueuedMessages(threadId: string): Promise<AgentMessageEntity[]> {
return this.messageRepository.find({
where: {
threadId,
status: AgentMessageStatus.QUEUED,
},
order: { createdAt: 'ASC' },
relations: ['parts'],
});
}
async findQueuedMessage(
messageId: string,
): Promise<AgentMessageEntity | null> {
return this.messageRepository.findOne({
where: { id: messageId, status: AgentMessageStatus.QUEUED },
});
}
async deleteQueuedMessage(messageId: string): Promise<boolean> {
const result = await this.messageRepository.delete({
id: messageId,
status: AgentMessageStatus.QUEUED,
});
return (result.affected ?? 0) > 0;
}
async promoteQueuedMessage(
messageId: string,
threadId: string,
): Promise<string> {
const turn = this.turnRepository.create({
threadId,
agentId: null,
});
const savedTurn = await this.turnRepository.save(turn);
await this.messageRepository.update(
{ id: messageId, threadId, status: AgentMessageStatus.QUEUED },
{
status: AgentMessageStatus.SENT,
processedAt: new Date(),
turnId: savedTurn.id,
},
);
return savedTurn.id;
}
async generateTitleIfNeeded(
threadId: string,
messageContent: string,
@@ -1,4 +1,5 @@
export enum SubscriptionChannel {
LOGIC_FUNCTION_LOGS_CHANNEL = 'LOGIC_FUNCTION_LOGS_CHANNEL',
EVENT_STREAM_CHANNEL = 'EVENT_STREAM_CHANNEL',
AGENT_CHAT_CHANNEL = 'AGENT_CHAT_CHANNEL',
}
@@ -88,4 +88,45 @@ export class SubscriptionService {
payload,
);
}
private getAgentChatChannel({
workspaceId,
threadId,
}: {
workspaceId: string;
threadId: string;
}) {
return `${SubscriptionChannel.AGENT_CHAT_CHANNEL}:${workspaceId}:${threadId}`;
}
async subscribeToAgentChat({
workspaceId,
threadId,
}: {
workspaceId: string;
threadId: string;
}) {
const client = this.redisClient.getPubSubClient();
return client.asyncIterator(
this.getAgentChatChannel({ workspaceId, threadId }),
);
}
async publishToAgentChat<T>({
workspaceId,
threadId,
payload,
}: {
workspaceId: string;
threadId: string;
payload: T;
}): Promise<void> {
const client = this.redisClient.getPubSubClient();
await client.publish(
this.getAgentChatChannel({ workspaceId, threadId }),
payload,
);
}
}
+1
View File
@@ -18,6 +18,7 @@ export type {
AgentResponseFieldType,
AgentResponseSchema,
} from './types/agent-response-schema.type';
export type { AgentChatSubscriptionEvent } from './types/AgentChatSubscriptionEvent';
export type {
CodeExecutionFile,
ExtendedFileUIPart,
@@ -0,0 +1,8 @@
// `chunk` is typed as Record<string, unknown> rather than UIMessageChunk
// to avoid importing the `ai` package in twenty-shared. The frontend casts
// it back to UIMessageChunk when feeding it into readUIMessageStream.
export type AgentChatSubscriptionEvent =
| { type: 'stream-chunk'; chunk: Record<string, unknown> }
| { type: 'message-persisted'; messageId: string }
| { type: 'queue-updated' }
| { type: 'stream-error'; code: string; message: string };
@@ -23,4 +23,5 @@ type Metadata = {
export type ExtendedUIMessage = UIMessage<Metadata, DataMessagePart> & {
threadId?: Nullable<string>;
status?: 'queued' | 'sent';
};