Restructure agent chat messages with parts-based architecture (#14749)

Co-authored-by: Félix Malfait <felix@twenty.com>

packages/twenty-server/src/engine/metadata-modules/agent/agent-chat-message-part.entity.ts

@@ -0,0 +1,98 @@
+import { JSONValue } from 'ai';
+import {
+  Column,
+  CreateDateColumn,
+  Entity,
+  Index,
+  JoinColumn,
+  ManyToOne,
+  PrimaryGeneratedColumn,
+  Relation,
+} from 'typeorm';
+
+import { AgentChatMessageEntity } from './agent-chat-message.entity';
+
+@Entity('agentChatMessagePart')
+export class AgentChatMessagePartEntity {
+  @PrimaryGeneratedColumn('uuid')
+  id: string;
+
+  @Column('uuid')
+  @Index()
+  messageId: string;
+
+  @ManyToOne(() => AgentChatMessageEntity, (message) => message.parts, {
+    onDelete: 'CASCADE',
+  })
+  @JoinColumn({ name: 'messageId' })
+  message: Relation<AgentChatMessageEntity>;
+
+  @Column({ type: 'int' })
+  orderIndex: number;
+
+  @Column({ type: 'varchar' })
+  type: string;
+
+  @Column({ type: 'text', nullable: true })
+  textContent: string | null;
+
+  @Column({ type: 'text', nullable: true })
+  reasoningContent: string | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  toolName: string | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  toolCallId: string | null;
+
+  @Column({ type: 'jsonb', nullable: true })
+  toolInput: unknown | null;
+
+  @Column({ type: 'jsonb', nullable: true })
+  toolOutput: unknown | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  state: string | null;
+
+  @Column({ type: 'text', nullable: true })
+  errorMessage: string | null;
+
+  @Column({ type: 'jsonb', nullable: true })
+  errorDetails: Record<string, unknown> | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  sourceUrlSourceId: string | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  sourceUrlUrl: string | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  sourceUrlTitle: string | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  sourceDocumentSourceId: string | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  sourceDocumentMediaType: string | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  sourceDocumentTitle: string | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  sourceDocumentFilename: string | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  fileMediaType: string | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  fileFilename: string | null;
+
+  @Column({ type: 'varchar', nullable: true })
+  fileUrl: string | null;
+
+  @Column({ type: 'jsonb', nullable: true })
+  providerMetadata: Record<string, Record<string, JSONValue>> | null;
+
+  @CreateDateColumn()
+  createdAt: Date;
+}

packages/twenty-server/src/engine/metadata-modules/agent/agent-chat-message.entity.ts

@@ -13,6 +13,8 @@ import {
 import { FileEntity } from 'src/engine/core-modules/file/entities/file.entity';
 import { AgentChatThreadEntity } from 'src/engine/metadata-modules/agent/agent-chat-thread.entity';
 
+import { AgentChatMessagePartEntity } from './agent-chat-message-part.entity';
+
 export enum AgentChatMessageRole {
   USER = 'user',
   ASSISTANT = 'assistant',
@@ -36,8 +38,8 @@ export class AgentChatMessageEntity {
   @Column({ type: 'enum', enum: AgentChatMessageRole })
   role: AgentChatMessageRole;
 
-  @Column({ type: 'text', nullable: true })
-  rawContent: string | null;
+  @OneToMany(() => AgentChatMessagePartEntity, (part) => part.message)
+  parts: Relation<AgentChatMessagePartEntity[]>;
 
   @OneToMany(() => FileEntity, (file) => file.message)
   files: Relation<FileEntity[]>;

packages/twenty-server/src/engine/metadata-modules/agent/agent-chat.controller.ts

@@ -9,14 +9,15 @@ import {
   UseGuards,
 } from '@nestjs/common';
 
+import { UIDataTypes, UIMessage, UITools } from 'ai';
 import { Response } from 'express';
 
 import { RestApiExceptionFilter } from 'src/engine/api/rest/rest-api-exception.filter';
+import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
 import { AuthUserWorkspaceId } from 'src/engine/decorators/auth/auth-user-workspace-id.decorator';
+import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
 import { JwtAuthGuard } from 'src/engine/guards/jwt-auth.guard';
 import { WorkspaceAuthGuard } from 'src/engine/guards/workspace-auth.guard';
-import { AuthWorkspace } from 'src/engine/decorators/auth/auth-workspace.decorator';
-import { Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
 import { type RecordIdsByObjectMetadataNameSingularType } from 'src/engine/metadata-modules/agent/types/recordIdsByObjectMetadataNameSingular.type';
 
 import { AgentChatService } from './agent-chat.service';
@@ -63,45 +64,21 @@ export class AgentChatController {
     @Body()
     body: {
       threadId: string;
-      userMessage: string;
-      fileIds?: string[];
+      messages: UIMessage<unknown, UIDataTypes, UITools>[];
       recordIdsByObjectMetadataNameSingular?: RecordIdsByObjectMetadataNameSingularType;
     },
     @AuthUserWorkspaceId() userWorkspaceId: string,
     @AuthWorkspace() workspace: Workspace,
-    @Res() res: Response,
+    @Res() response: Response,
   ) {
-    try {
-      await this.agentStreamingService.streamAgentChat({
-        threadId: body.threadId,
-        userMessage: body.userMessage,
-        userWorkspaceId,
-        workspace,
-        fileIds: body.fileIds || [],
-        recordIdsByObjectMetadataNameSingular:
-          body.recordIdsByObjectMetadataNameSingular || [],
-        res,
-      });
-    } catch (error) {
-      // Handle errors at controller level for streaming responses
-      // since the RestApiExceptionFilter interferes with our streaming error handling
-      const errorMessage =
-        error instanceof Error ? error.message : 'Unknown error occurred';
-
-      if (!res.headersSent) {
-        res.setHeader('Content-Type', 'text/plain; charset=utf-8');
-        res.setHeader('Transfer-Encoding', 'chunked');
-        res.setHeader('Cache-Control', 'no-cache');
-      }
-
-      res.write(
-        JSON.stringify({
-          type: 'error',
-          message: errorMessage,
-        }) + '\n',
-      );
-
-      res.end();
-    }
+    this.agentStreamingService.streamAgentChat({
+      threadId: body.threadId,
+      messages: body.messages,
+      userWorkspaceId,
+      workspace,
+      recordIdsByObjectMetadataNameSingular:
+        body.recordIdsByObjectMetadataNameSingular || [],
+      response,
+    });
   }
 }

packages/twenty-server/src/engine/metadata-modules/agent/agent-chat.service.ts

@@ -3,16 +3,19 @@ import { InjectRepository } from '@nestjs/typeorm';
 
 import { Repository } from 'typeorm';
 
-import { FileEntity } from 'src/engine/core-modules/file/entities/file.entity';
+import type { UIDataTypes, UIMessage, UIMessagePart, UITools } from 'ai';
+
+import { AgentChatMessagePartEntity } from 'src/engine/metadata-modules/agent/agent-chat-message-part.entity';
 import {
   AgentChatMessageEntity,
-  type AgentChatMessageRole,
+  AgentChatMessageRole,
 } from 'src/engine/metadata-modules/agent/agent-chat-message.entity';
 import { AgentChatThreadEntity } from 'src/engine/metadata-modules/agent/agent-chat-thread.entity';
 import {
   AgentException,
   AgentExceptionCode,
 } from 'src/engine/metadata-modules/agent/agent.exception';
+import { mapUIMessagePartsToDBParts } from 'src/engine/metadata-modules/agent/utils/mapUIMessagePartsToDBParts';
 
 import { AgentTitleGenerationService } from './agent-title-generation.service';
 
@@ -23,8 +26,8 @@ export class AgentChatService {
     private readonly threadRepository: Repository<AgentChatThreadEntity>,
     @InjectRepository(AgentChatMessageEntity)
     private readonly messageRepository: Repository<AgentChatMessageEntity>,
-    @InjectRepository(FileEntity)
-    private readonly fileRepository: Repository<FileEntity>,
+    @InjectRepository(AgentChatMessagePartEntity)
+    private readonly messagePartRepository: Repository<AgentChatMessagePartEntity>,
     private readonly titleGenerationService: AgentTitleGenerationService,
   ) {}
 
@@ -67,32 +70,32 @@ export class AgentChatService {
 
   async addMessage({
     threadId,
-    role,
-    rawContent,
-    fileIds,
+    uiMessage,
   }: {
     threadId: string;
-    role: AgentChatMessageRole;
-    rawContent: string | null;
-    fileIds?: string[];
+    uiMessage: Omit<UIMessage<unknown, UIDataTypes, UITools>, 'id'>;
+    uiMessageParts?: UIMessagePart<UIDataTypes, UITools>[];
   }) {
     const message = this.messageRepository.create({
       threadId,
-      role,
-      rawContent,
+      role: uiMessage.role as AgentChatMessageRole,
     });
 
     const savedMessage = await this.messageRepository.save(message);
 
-    if (fileIds && fileIds.length > 0) {
-      for (const fileId of fileIds) {
-        await this.fileRepository.update(fileId, {
-          messageId: savedMessage.id,
-        });
-      }
+    if (uiMessage.parts && uiMessage.parts.length > 0) {
+      const dbParts = mapUIMessagePartsToDBParts(
+        uiMessage.parts,
+        savedMessage.id,
+      );
+
+      await this.messagePartRepository.save(dbParts);
     }
 
-    this.generateTitleIfNeeded(threadId, rawContent);
+    this.generateTitleIfNeeded(
+      threadId,
+      uiMessage.parts.find((part) => part.type === 'text')?.text,
+    );
 
     return savedMessage;
   }
@@ -115,13 +118,13 @@ export class AgentChatService {
     return this.messageRepository.find({
       where: { threadId },
       order: { createdAt: 'ASC' },
-      relations: ['files'],
+      relations: ['parts', 'files'],
     });
   }
 
   private async generateTitleIfNeeded(
     threadId: string,
-    messageContent: string | null,
+    messageContent?: string | null,
   ) {
     const thread = await this.threadRepository.findOne({
       where: { id: threadId },

packages/twenty-server/src/engine/metadata-modules/agent/agent-execution.service.ts

@@ -2,14 +2,14 @@ import { Injectable, Logger } from '@nestjs/common';
 import { InjectRepository } from '@nestjs/typeorm';
 
 import {
-  type FilePart,
-  type ImagePart,
+  convertToModelMessages,
   LanguageModelUsage,
-  type ModelMessage,
+  stepCountIs,
   streamText,
   ToolSet,
-  type UserContent,
-  UserModelMessage,
+  UIDataTypes,
+  UIMessage,
+  UITools,
 } from 'ai';
 import { AppPath } from 'twenty-shared/types';
 import { getAppPath } from 'twenty-shared/utils';
@@ -20,20 +20,13 @@ import { AiModelRegistryService } from 'src/engine/core-modules/ai/services/ai-m
 import { DomainManagerService } from 'src/engine/core-modules/domain-manager/services/domain-manager.service';
 import { FileEntity } from 'src/engine/core-modules/file/entities/file.entity';
 import { FileService } from 'src/engine/core-modules/file/services/file.service';
-import { extractFolderPathAndFilename } from 'src/engine/core-modules/file/utils/extract-folderpath-and-filename.utils';
 import { type Workspace } from 'src/engine/core-modules/workspace/workspace.entity';
-import {
-  type AgentChatMessageEntity,
-  AgentChatMessageRole,
-} from 'src/engine/metadata-modules/agent/agent-chat-message.entity';
 import { AgentHandoffToolService } from 'src/engine/metadata-modules/agent/agent-handoff-tool.service';
 import { AGENT_CONFIG } from 'src/engine/metadata-modules/agent/constants/agent-config.const';
 import { AGENT_SYSTEM_PROMPTS } from 'src/engine/metadata-modules/agent/constants/agent-system-prompts.const';
 import { type RecordIdsByObjectMetadataNameSingularType } from 'src/engine/metadata-modules/agent/types/recordIdsByObjectMetadataNameSingular.type';
-import { constructAssistantMessageContentFromStream } from 'src/engine/metadata-modules/agent/utils/constructAssistantMessageContentFromStream';
 import { WorkspacePermissionsCacheService } from 'src/engine/metadata-modules/workspace-permissions-cache/workspace-permissions-cache.service';
 import { TwentyORMGlobalManager } from 'src/engine/twenty-orm/twenty-orm-global.manager';
-import { streamToBuffer } from 'src/utils/stream-to-buffer';
 
 import { AgentToolGeneratorService } from './agent-tool-generator.service';
 import { AgentEntity } from './agent.entity';
@@ -70,7 +63,7 @@ export class AgentExecutionService {
   }: {
     system: string;
     agent: AgentEntity | null;
-    messages: ModelMessage[];
+    messages: UIMessage<unknown, UIDataTypes, UITools>[];
   }) {
     try {
       if (agent) {
@@ -106,8 +99,8 @@ export class AgentExecutionService {
         system,
         tools,
         model: registeredModel.model,
-        messages,
-        maxSteps: AGENT_CONFIG.MAX_STEPS,
+        messages: convertToModelMessages(messages),
+        stopWhen: stepCountIs(AGENT_CONFIG.MAX_STEPS),
         ...(registeredModel.doesSupportThinking && {
           providerOptions: {
             anthropic: {
@@ -128,39 +121,6 @@ export class AgentExecutionService {
     }
   }
 
-  private async buildUserMessageWithFiles(
-    fileIds: string[],
-  ): Promise<(ImagePart | FilePart)[]> {
-    const files = await this.fileRepository.find({
-      where: {
-        id: In(fileIds),
-      },
-    });
-
-    return await Promise.all(files.map((file) => this.createFilePart(file)));
-  }
-
-  private async buildUserMessage(
-    userMessage: string,
-    fileIds: string[],
-  ): Promise<UserModelMessage> {
-    const content: Exclude<UserContent, string> = [
-      {
-        type: 'text',
-        text: userMessage,
-      },
-    ];
-
-    if (fileIds.length !== 0) {
-      content.push(...(await this.buildUserMessageWithFiles(fileIds)));
-    }
-
-    return {
-      role: AgentChatMessageRole.USER,
-      content,
-    };
-  }
-
   private async getContextForSystemPrompt(
     workspace: Workspace,
     recordIdsByObjectMetadataNameSingular: RecordIdsByObjectMetadataNameSingularType,
@@ -225,119 +185,72 @@ export class AgentExecutionService {
     return JSON.stringify(contextObject);
   }
 
-  private async createFilePart(
-    file: FileEntity,
-  ): Promise<ImagePart | FilePart> {
-    const { folderPath, filename } = extractFolderPathAndFilename(
-      file.fullPath,
-    );
-    const fileStream = await this.fileService.getFileStream(
-      folderPath,
-      filename,
-      file.workspaceId,
-    );
-    const fileBuffer = await streamToBuffer(fileStream);
-
-    if (file.type.startsWith('image')) {
-      return {
-        type: 'image',
-        image: fileBuffer,
-        mediaType: file.type,
-      };
-    }
-
-    return {
-      type: 'file',
-      data: fileBuffer,
-      mediaType: file.type,
-    };
-  }
-
-  private mapMessagesToCoreMessages(
-    messages: AgentChatMessageEntity[],
-  ): ModelMessage[] {
-    return messages
-      .map(({ role, rawContent }): ModelMessage => {
-        if (role === AgentChatMessageRole.USER) {
-          return {
-            role: 'user',
-            content: rawContent ?? '',
-          };
-        }
-
-        return {
-          role: 'assistant',
-          content: constructAssistantMessageContentFromStream(rawContent ?? ''),
-        };
-      })
-      .filter((message) => message.content.length > 0);
-  }
-
   async streamChatResponse({
     workspace,
     userWorkspaceId,
     agentId,
-    userMessage,
     messages,
-    fileIds,
     recordIdsByObjectMetadataNameSingular,
   }: {
     workspace: Workspace;
     userWorkspaceId: string;
     agentId: string;
-    userMessage: string;
-    messages: AgentChatMessageEntity[];
-    fileIds: string[];
+    messages: UIMessage<unknown, UIDataTypes, UITools>[];
     recordIdsByObjectMetadataNameSingular: RecordIdsByObjectMetadataNameSingularType;
   }) {
-    const agent = await this.agentRepository.findOneOrFail({
-      where: { id: agentId },
-    });
+    try {
+      const agent = await this.agentRepository.findOneOrFail({
+        where: { id: agentId },
+      });
 
-    const llmMessages: ModelMessage[] =
-      this.mapMessagesToCoreMessages(messages);
+      let contextString = '';
 
-    let contextString = '';
+      if (recordIdsByObjectMetadataNameSingular.length > 0) {
+        const contextPart = await this.getContextForSystemPrompt(
+          workspace,
+          recordIdsByObjectMetadataNameSingular,
+          userWorkspaceId,
+        );
 
-    if (recordIdsByObjectMetadataNameSingular.length > 0) {
-      const contextPart = await this.getContextForSystemPrompt(
-        workspace,
-        recordIdsByObjectMetadataNameSingular,
-        userWorkspaceId,
+        contextString = `\n\nCONTEXT:\n${contextPart}`;
+      }
+
+      const aiRequestConfig = await this.prepareAIRequestConfig({
+        system: `${AGENT_SYSTEM_PROMPTS.AGENT_CHAT}\n\n${agent.prompt}${contextString}`,
+        agent,
+        messages,
+      });
+
+      this.logger.log(
+        `Sending request to AI model with ${messages.length} messages`,
       );
 
-      contextString = `\n\nCONTEXT:\n${contextPart}`;
+      const model =
+        await this.aiModelRegistryService.resolveModelForAgent(agent);
+
+      const stream = streamText(aiRequestConfig);
+
+      stream.usage
+        .then((usage) => {
+          this.aiBillingService.calculateAndBillUsage(
+            model.modelId,
+            usage,
+            workspace.id,
+          );
+        })
+        .catch((usageError) => {
+          this.logger.error('Failed to get usage information:', usageError);
+        });
+
+      return stream;
+    } catch (error) {
+      this.logger.error('Error in streamChatResponse:', error);
+      throw new AgentException(
+        error instanceof Error
+          ? error.message
+          : 'Failed to stream chat response',
+        AgentExceptionCode.AGENT_EXECUTION_FAILED,
+      );
     }
-
-    const userMessageWithFiles = await this.buildUserMessage(
-      userMessage,
-      fileIds,
-    );
-
-    llmMessages.push(userMessageWithFiles);
-
-    const aiRequestConfig = await this.prepareAIRequestConfig({
-      system: `${AGENT_SYSTEM_PROMPTS.AGENT_CHAT}\n\n${agent.prompt}${contextString}`,
-      agent,
-      messages: llmMessages,
-    });
-
-    this.logger.log(
-      `Sending request to AI model with ${llmMessages.length} messages`,
-    );
-
-    const model = await this.aiModelRegistryService.resolveModelForAgent(agent);
-
-    const stream = streamText(aiRequestConfig);
-
-    stream.usage.then((usage) => {
-      this.aiBillingService.calculateAndBillUsage(
-        model.modelId,
-        usage,
-        workspace.id,
-      );
-    });
-
-    return stream;
   }
 }

packages/twenty-server/src/engine/metadata-modules/agent/agent-streaming.service.ts

@@ -1,6 +1,13 @@
 import { Injectable, Logger } from '@nestjs/common';
 import { InjectRepository } from '@nestjs/typeorm';
 
+import {
+  createUIMessageStream,
+  pipeUIMessageStreamToResponse,
+  UIDataTypes,
+  UIMessage,
+  UITools,
+} from 'ai';
 import { type Response } from 'express';
 import { Repository } from 'typeorm';
 
@@ -17,24 +24,13 @@ import { type RecordIdsByObjectMetadataNameSingularType } from 'src/engine/metad
 
 export type StreamAgentChatOptions = {
   threadId: string;
-  userMessage: string;
   userWorkspaceId: string;
   workspace: Workspace;
-  fileIds: string[];
   recordIdsByObjectMetadataNameSingular: RecordIdsByObjectMetadataNameSingularType;
-  res: Response;
+  response: Response;
+  messages: UIMessage<unknown, UIDataTypes, UITools>[];
 };
 
-const CLIENT_FORWARDED_EVENT_TYPES = [
-  'text-delta',
-  'reasoning',
-  'reasoning-delta',
-  'tool-call',
-  'tool-input-delta',
-  'tool-result',
-  'error',
-];
-
 @Injectable()
 export class AgentStreamingService {
   private readonly logger = new Logger(AgentStreamingService.name);
@@ -48,15 +44,12 @@ export class AgentStreamingService {
 
   async streamAgentChat({
     threadId,
-    userMessage,
     userWorkspaceId,
     workspace,
-    fileIds,
+    messages,
     recordIdsByObjectMetadataNameSingular,
-    res,
+    response,
   }: StreamAgentChatOptions) {
-    let rawStreamString = '';
-
     try {
       const thread = await this.threadRepository.findOne({
         where: {
@@ -73,74 +66,56 @@ export class AgentStreamingService {
         );
       }
 
-      this.setupStreamingHeaders(res);
+      const stream = createUIMessageStream({
+        execute: async ({ writer }) => {
+          const result = await this.agentExecutionService.streamChatResponse({
+            workspace,
+            agentId: thread.agent.id,
+            userWorkspaceId,
+            messages,
+            recordIdsByObjectMetadataNameSingular,
+          });
 
-      const { fullStream } =
-        await this.agentExecutionService.streamChatResponse({
-          workspace,
-          agentId: thread.agent.id,
-          userWorkspaceId,
-          userMessage,
-          messages: thread.messages,
-          fileIds,
-          recordIdsByObjectMetadataNameSingular,
-        });
+          writer.merge(
+            result.toUIMessageStream({
+              onError: (error) => {
+                return error instanceof Error ? error.message : String(error);
+              },
+              onFinish: async ({ responseMessage }) => {
+                if (responseMessage.parts.length === 0) {
+                  return;
+                }
 
-      for await (const chunk of fullStream) {
-        rawStreamString += JSON.stringify(chunk) + '\n';
+                await this.agentChatService.addMessage({
+                  threadId,
+                  uiMessage: {
+                    role: AgentChatMessageRole.USER,
+                    parts: [
+                      {
+                        type: 'text',
+                        text:
+                          messages[messages.length - 1].parts.find(
+                            (part) => part.type === 'text',
+                          )?.text ?? '',
+                      },
+                    ],
+                  },
+                });
+                await this.agentChatService.addMessage({
+                  threadId,
+                  uiMessage: responseMessage,
+                });
+              },
+              sendReasoning: true,
+            }),
+          );
+        },
+      });
 
-        this.sendStreamEvent(
-          res,
-          CLIENT_FORWARDED_EVENT_TYPES.includes(chunk.type)
-            ? chunk
-            : { type: chunk.type },
-        );
-      }
+      pipeUIMessageStreamToResponse({ stream, response });
     } catch (error) {
-      const errorMessage =
-        error instanceof Error ? error.message : 'Unknown error occurred';
-
-      if (error instanceof AgentException) {
-        this.logger.error(`Agent Exception Code: ${error.code}`);
-      }
-
-      if (!res.headersSent) {
-        this.setupStreamingHeaders(res);
-      }
-
-      const errorChunk = {
-        type: 'error',
-        message: errorMessage,
-      };
-
-      rawStreamString += JSON.stringify(errorChunk) + '\n';
-
-      this.sendStreamEvent(res, errorChunk);
+      this.logger.error(error.message);
+      response.end();
     }
-
-    await this.agentChatService.addMessage({
-      threadId,
-      role: AgentChatMessageRole.USER,
-      rawContent: userMessage,
-      fileIds,
-    });
-
-    await this.agentChatService.addMessage({
-      threadId,
-      role: AgentChatMessageRole.ASSISTANT,
-      rawContent: rawStreamString.trim() || null,
-    });
-
-    res.end();
-  }
-
-  private sendStreamEvent(res: Response, event: object): void {
-    res.write(JSON.stringify(event) + '\n');
-  }
-
-  private setupStreamingHeaders(res: Response): void {
-    res.setHeader('Content-Type', 'text/plain; charset=utf-8');
-    res.setHeader('Transfer-Encoding', 'chunked');
-    res.setHeader('Cache-Control', 'no-cache');
   }
 }

packages/twenty-server/src/engine/metadata-modules/agent/agent.module.ts

@@ -21,6 +21,7 @@ import { WorkspacePermissionsCacheModule } from 'src/engine/metadata-modules/wor
 import { WorkspaceCacheStorageModule } from 'src/engine/workspace-cache-storage/workspace-cache-storage.module';
 import { WorkflowToolsModule } from 'src/modules/workflow/workflow-tools/workflow-tools.module';
 
+import { AgentChatMessagePartEntity } from './agent-chat-message-part.entity';
 import { AgentChatMessageEntity } from './agent-chat-message.entity';
 import { AgentChatThreadEntity } from './agent-chat-thread.entity';
 import { AgentChatResolver } from './agent-chat.resolver';
@@ -45,6 +46,7 @@ import { AgentService } from './agent.service';
       RoleEntity,
       RoleTargetsEntity,
       AgentChatMessageEntity,
+      AgentChatMessagePartEntity,
       AgentChatThreadEntity,
       FileEntity,
       UserWorkspace,
@@ -89,6 +91,7 @@ import { AgentService } from './agent.service';
     TypeOrmModule.forFeature([
       AgentEntity,
       AgentChatMessageEntity,
+      AgentChatMessagePartEntity,
       AgentChatThreadEntity,
     ]),
     AgentHandoffExecutorService,

packages/twenty-server/src/engine/metadata-modules/agent/constants/agent-handoff-schema.const.ts

@@ -76,7 +76,18 @@ export const AGENT_HANDOFF_SCHEMA = z.object({
           }),
           z.object({
             role: z.literal('tool'),
-            content: z.string(),
+            content: z.union([
+              z.string(),
+              z.array(
+                z.object({
+                  type: z.literal('tool-result'),
+                  toolCallId: z.string(),
+                  toolName: z.string(),
+                  result: z.unknown(),
+                  isError: z.boolean().optional(),
+                }),
+              ),
+            ]),
             toolCallId: z.string(),
           }),
         ]),

packages/twenty-server/src/engine/metadata-modules/agent/dtos/agent-chat-message-part.dto.ts

@@ -0,0 +1,84 @@
+import { Field, Int, ObjectType } from '@nestjs/graphql';
+
+import { JSONValue } from 'ai';
+import GraphQLJSON from 'graphql-type-json';
+
+import { UUIDScalarType } from 'src/engine/api/graphql/workspace-schema-builder/graphql-types/scalars';
+
+@ObjectType('AgentChatMessagePart')
+export class AgentChatMessagePartDTO {
+  @Field(() => UUIDScalarType)
+  id: string;
+
+  @Field(() => UUIDScalarType)
+  messageId: string;
+
+  @Field(() => Int)
+  orderIndex: number;
+
+  @Field()
+  type: string;
+
+  @Field({ nullable: true })
+  textContent?: string;
+
+  @Field({ nullable: true })
+  reasoningContent?: string;
+
+  @Field({ nullable: true })
+  toolName?: string;
+
+  @Field({ nullable: true })
+  toolCallId?: string;
+
+  @Field(() => GraphQLJSON, { nullable: true })
+  toolInput?: Record<string, unknown>;
+
+  @Field(() => GraphQLJSON, { nullable: true })
+  toolOutput?: Record<string, unknown>;
+
+  @Field({ nullable: true })
+  state?: string;
+
+  @Field({ nullable: true })
+  errorMessage?: string;
+
+  @Field(() => GraphQLJSON, { nullable: true })
+  errorDetails?: Record<string, unknown>;
+
+  @Field({ nullable: true })
+  sourceUrlSourceId?: string;
+
+  @Field({ nullable: true })
+  sourceUrlUrl?: string;
+
+  @Field({ nullable: true })
+  sourceUrlTitle?: string;
+
+  @Field({ nullable: true })
+  sourceDocumentSourceId?: string;
+
+  @Field({ nullable: true })
+  sourceDocumentMediaType?: string;
+
+  @Field({ nullable: true })
+  sourceDocumentTitle?: string;
+
+  @Field({ nullable: true })
+  sourceDocumentFilename?: string;
+
+  @Field({ nullable: true })
+  fileMediaType?: string;
+
+  @Field({ nullable: true })
+  fileFilename?: string;
+
+  @Field({ nullable: true })
+  fileUrl?: string;
+
+  @Field(() => GraphQLJSON, { nullable: true })
+  providerMetadata?: Record<string, Record<string, JSONValue>>;
+
+  @Field()
+  createdAt: Date;
+}

packages/twenty-server/src/engine/metadata-modules/agent/dtos/agent-chat-message.dto.ts

@@ -3,6 +3,8 @@ import { Field, ObjectType } from '@nestjs/graphql';
 import { UUIDScalarType } from 'src/engine/api/graphql/workspace-schema-builder/graphql-types/scalars';
 import { FileDTO } from 'src/engine/core-modules/file/dtos/file.dto';
 
+import { AgentChatMessagePartDTO } from './agent-chat-message-part.dto';
+
 @ObjectType('AgentChatMessage')
 export class AgentChatMessageDTO {
   @Field(() => UUIDScalarType)
@@ -14,8 +16,8 @@ export class AgentChatMessageDTO {
   @Field()
   role: 'user' | 'assistant';
 
-  @Field({ nullable: true })
-  rawContent: string;
+  @Field(() => [AgentChatMessagePartDTO])
+  parts: AgentChatMessagePartDTO[];
 
   @Field(() => [FileDTO])
   files: FileDTO[];

packages/twenty-server/src/engine/metadata-modules/agent/utils/constructAssistantMessageContentFromStream.ts

@@ -1,96 +0,0 @@
-import { type ReasoningPart } from '@ai-sdk/provider-utils';
-import { type TextPart } from 'ai';
-import {
-  parseStreamLine,
-  splitStreamIntoLines,
-  type TextBlock,
-} from 'twenty-shared/ai';
-
-export const constructAssistantMessageContentFromStream = (
-  rawContent: string,
-) => {
-  const lines = splitStreamIntoLines(rawContent);
-
-  const output: Array<TextPart | ReasoningPart> = [];
-  let currentTextBlock: TextBlock = null;
-
-  const flushTextBlock = () => {
-    if (currentTextBlock) {
-      if (currentTextBlock.type === 'reasoning') {
-        output.push({
-          type: 'reasoning',
-          text: currentTextBlock.content,
-        });
-      } else {
-        output.push({
-          type: 'text',
-          text: currentTextBlock.content,
-        });
-      }
-      currentTextBlock = null;
-    }
-  };
-
-  for (const line of lines) {
-    const event = parseStreamLine(line);
-
-    if (!event) {
-      continue;
-    }
-
-    switch (event.type) {
-      case 'reasoning-start':
-        flushTextBlock();
-        currentTextBlock = {
-          type: 'reasoning',
-          content: '',
-          isThinking: true,
-        };
-        break;
-
-      case 'reasoning-delta':
-        if (!currentTextBlock || currentTextBlock.type !== 'reasoning') {
-          flushTextBlock();
-          currentTextBlock = {
-            type: 'reasoning',
-            content: '',
-            isThinking: true,
-          };
-        }
-        currentTextBlock.content += event.text || '';
-        break;
-
-      case 'reasoning-end':
-        if (currentTextBlock?.type === 'reasoning') {
-          currentTextBlock.isThinking = false;
-        }
-        break;
-
-      case 'text-delta':
-        if (!currentTextBlock || currentTextBlock.type !== 'text') {
-          flushTextBlock();
-          currentTextBlock = { type: 'text', content: '' };
-        }
-        currentTextBlock.content += event.text || '';
-        break;
-
-      case 'step-finish':
-        if (currentTextBlock?.type === 'reasoning') {
-          currentTextBlock.isThinking = false;
-        }
-        flushTextBlock();
-        break;
-
-      case 'error':
-        flushTextBlock();
-        break;
-
-      default:
-        break;
-    }
-  }
-
-  flushTextBlock();
-
-  return output;
-};

packages/twenty-server/src/engine/metadata-modules/agent/utils/mapUIMessagePartsToDBParts.ts

@@ -0,0 +1,81 @@
+import {
+  type ToolUIPart,
+  type UIDataTypes,
+  type UIMessagePart,
+  type UITools,
+} from 'ai';
+
+import { type AgentChatMessagePartEntity } from 'src/engine/metadata-modules/agent/agent-chat-message-part.entity';
+
+const isToolPart = (
+  part: UIMessagePart<UIDataTypes, UITools>,
+): part is ToolUIPart => {
+  return part.type.includes('tool-') && 'toolCallId' in part;
+};
+
+export const mapUIMessagePartsToDBParts = (
+  uiMessageParts: UIMessagePart<UIDataTypes, UITools>[],
+  messageId: string,
+): Partial<AgentChatMessagePartEntity>[] => {
+  return uiMessageParts.map((part, index) => {
+    const basePart: Partial<AgentChatMessagePartEntity> = {
+      messageId,
+      orderIndex: index,
+      type: part.type,
+    };
+
+    switch (part.type) {
+      case 'text':
+        return {
+          ...basePart,
+          textContent: part.text,
+        };
+      case 'reasoning':
+        return {
+          ...basePart,
+          reasoningContent: part.text,
+        };
+      case 'file':
+        return {
+          ...basePart,
+          fileMediaType: part.mediaType,
+          fileFilename: part.filename,
+          fileUrl: part.url,
+        };
+      case 'source-url':
+        return {
+          ...basePart,
+          sourceUrlSourceId: part.sourceId,
+          sourceUrlUrl: part.url,
+          sourceUrlTitle: part.title,
+          providerMetadata: part.providerMetadata ?? null,
+        };
+      case 'source-document':
+        return {
+          ...basePart,
+          sourceDocumentSourceId: part.sourceId,
+          sourceDocumentMediaType: part.mediaType,
+          sourceDocumentTitle: part.title,
+          sourceDocumentFilename: part.filename,
+          providerMetadata: part.providerMetadata ?? null,
+        };
+      case 'step-start':
+        return basePart;
+      default:
+        {
+          if (isToolPart(part)) {
+            const { toolCallId, input, output, errorText } = part;
+
+            return {
+              ...basePart,
+              toolCallId: toolCallId,
+              toolInput: input,
+              toolOutput: output,
+              errorMessage: errorText,
+            };
+          }
+        }
+        throw new Error(`Unsupported part type: ${part.type}`);
+    }
+  });
+};

packages/twenty-server/src/engine/metadata-modules/remote-server/remote-server.service.ts

@@ -3,6 +3,7 @@ import { InjectDataSource, InjectRepository } from '@nestjs/typeorm';
 
 import isEmpty from 'lodash.isempty';
 import { DataSource, type EntityManager, Repository } from 'typeorm';
+import { type DeepPartial } from 'typeorm/common/DeepPartial';
 import { v4 } from 'uuid';
 
 import { ForeignDataWrapperServerQueryFactory } from 'src/engine/api/graphql/workspace-query-builder/factories/foreign-data-wrapper-server-query.factory';

packages/twenty-server/src/engine/metadata-modules/remote-server/utils/build-update-remote-server-raw-query.utils.ts

@@ -1,3 +1,5 @@
+import { type DeepPartial } from 'typeorm';
+
 import {
   type ForeignDataWrapperOptions,
   type RemoteServerEntity,
@@ -9,10 +11,6 @@ import {
 } from 'src/engine/metadata-modules/remote-server/remote-server.exception';
 import { type UserMappingOptions } from 'src/engine/metadata-modules/remote-server/types/user-mapping-options';
 
-export type DeepPartial<T> = {
-  [P in keyof T]?: DeepPartial<T[P]>;
-};
-
 export const buildUpdateRemoteServerRawQuery = (
   remoteServerToUpdate: DeepPartial<RemoteServerEntity<RemoteServerType>> &
     Pick<RemoteServerEntity<RemoteServerType>, 'workspaceId' | 'id'>,