fix(ui): message sources (#998)

PetrBulanek · web-flow · commit 32a3361f7c8e · 2025-08-04T09:52:24.000+02:00
diff --git a/apps/beeai-ui/src/api/a2a/client.ts b/apps/beeai-ui/src/api/a2a/client.ts
@@ -8,9 +8,8 @@ import { A2AClient } from '@a2a-js/sdk/client';
 import { Subject } from 'rxjs';
 import { match } from 'ts-pattern';
 
-import type { FileEntity } from '#modules/files/types.ts';
-import type { UIMessagePart } from '#modules/messages/types.ts';
-import type { ContextId } from '#modules/tasks/api/types.ts';
+import type { UIMessagePart, UIUserMessage } from '#modules/messages/types.ts';
+import type { ContextId, TaskId } from '#modules/tasks/api/types.ts';
 import { getBaseUrl } from '#utils/api/getBaseUrl.ts';
 import { isNotNull } from '#utils/helpers.ts';
 
@@ -56,26 +55,24 @@ export const buildA2AClient = (providerId: string) => {
   const agentUrl = `${getBaseUrl()}/api/v1/a2a/${providerId}`;
   const client = new A2AClient(agentUrl);
 
-  const chat = ({ text, files, contextId }: { text: string; files: FileEntity[]; contextId: ContextId }) => {
-    const messageSubject = new Subject<UIMessagePart[]>();
+  const chat = ({ message, contextId }: { message: UIUserMessage; contextId: ContextId }) => {
+    const messageSubject = new Subject<{ parts: UIMessagePart[]; taskId: TaskId }>();
     let taskId: string | null = null;
 
     const iterateOverStream = async () => {
-      const stream = client.sendMessageStream({ message: createUserMessage({ text, files, contextId }) });
+      const stream = client.sendMessageStream({ message: createUserMessage({ message, contextId }) });
 
       for await (const event of stream) {
         match(event)
-          .with(
-            {
-              kind: 'task',
-            },
-            (task) => {
-              taskId = task.id;
-            },
-          )
+          .with({ kind: 'task' }, (task) => {
+            taskId = task.id;
+          })
           .with({ kind: 'status-update' }, (event) => {
+            taskId = event.taskId;
+
             const messageParts = handleStatusUpdate(event);
-            messageSubject.next(messageParts);
+
+            messageSubject.next({ parts: messageParts, taskId });
           });
       }
       messageSubject.complete();
diff --git a/apps/beeai-ui/src/api/a2a/part-processors.ts b/apps/beeai-ui/src/api/a2a/part-processors.ts
@@ -25,7 +25,7 @@ export function processMessageMetadata(message: Message): UIMessagePart[] {
   if (trajectory) {
     return [createTrajectoryPart(trajectory)];
   } else if (citation) {
-    const sourcePart = createSourcePart(citation, message.messageId);
+    const sourcePart = createSourcePart(citation, message.taskId);
 
     if (sourcePart) {
       return [sourcePart];
diff --git a/apps/beeai-ui/src/api/a2a/types.ts b/apps/beeai-ui/src/api/a2a/types.ts
@@ -4,9 +4,10 @@
  */
 
 import type { UIMessagePart } from '#modules/messages/types.ts';
+import type { TaskId } from '#modules/tasks/api/types.ts';
 
 export interface ChatRun {
   done: Promise<void>;
-  subscribe: (fn: (parts: UIMessagePart[]) => void) => () => void;
+  subscribe: (fn: (data: { parts: UIMessagePart[]; taskId: TaskId }) => void) => () => void;
   cancel: () => Promise<void>;
 }
diff --git a/apps/beeai-ui/src/api/a2a/utils.ts b/apps/beeai-ui/src/api/a2a/utils.ts
@@ -3,14 +3,20 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type { FilePart, FileWithUri, Message } from '@a2a-js/sdk';
+import type { FilePart, FileWithUri, Message, Part, TextPart } from '@a2a-js/sdk';
 import { v4 as uuid } from 'uuid';
 
-import type { FileEntity } from '#modules/files/types.ts';
 import { getFileContentUrl } from '#modules/files/utils.ts';
-import type { UISourcePart, UITextPart, UITrajectoryPart } from '#modules/messages/types.ts';
+import type {
+  UIMessagePart,
+  UISourcePart,
+  UITextPart,
+  UITrajectoryPart,
+  UIUserMessage,
+} from '#modules/messages/types.ts';
 import { UIMessagePartKind } from '#modules/messages/types.ts';
 import type { ContextId, TaskId } from '#modules/tasks/api/types.ts';
+import { isNotNull } from '#utils/helpers.ts';
 
 import type { CitationMetadata } from './extensions/citation';
 import { citationExtension } from './extensions/citation';
@@ -31,47 +37,51 @@ export function extractTextFromMessage(message: Message | undefined) {
   return text;
 }
 
-export function convertFileToFilePart(file: FileEntity): FilePart {
-  const { originalFile, uploadFile } = file;
-
-  if (!uploadFile) {
-    throw new Error('File upload file is not present');
-  }
-
-  return {
-    kind: 'file',
-    file: {
-      uri: getFileContentUrl({ id: uploadFile.id, addBase: true }),
-      name: uploadFile.filename,
-      mimeType: originalFile.type,
-    },
-  };
+export function convertMessageParts(uiParts: UIMessagePart[]): Part[] {
+  const parts: Part[] = uiParts
+    .map((part) => {
+      switch (part.kind) {
+        case UIMessagePartKind.Text:
+          const { text } = part;
+
+          return {
+            kind: 'text',
+            text,
+          } as TextPart;
+        case UIMessagePartKind.File:
+          const { id, filename, type } = part;
+
+          return {
+            kind: 'file',
+            file: {
+              uri: getFileContentUrl({ id, addBase: true }),
+              name: filename,
+              mimeType: type,
+            },
+          } as FilePart;
+      }
+    })
+    .filter(isNotNull);
+
+  return parts;
 }
 
 export function createUserMessage({
-  text,
-  files,
-  taskId,
+  message,
   contextId,
+  taskId,
 }: {
-  text: string;
-  files: FileEntity[];
+  message: UIUserMessage;
   contextId: ContextId;
   taskId?: TaskId;
 }): Message {
   return {
     kind: 'message',
-    messageId: uuid(),
+    role: 'user',
+    messageId: message.id,
     contextId,
     taskId,
-    parts: [
-      {
-        kind: 'text',
-        text,
-      },
-      ...files.map(convertFileToFilePart),
-    ],
-    role: 'user',
+    parts: convertMessageParts(message.parts),
   };
 }
 
@@ -91,10 +101,10 @@ export function getFileUri(file: FilePart['file']): string {
   return `data:${mimeType};base64,${bytes}`;
 }
 
-export function createSourcePart(metadata: CitationMetadata, messageId: string): UISourcePart | null {
+export function createSourcePart(metadata: CitationMetadata, messageId: string | undefined): UISourcePart | null {
   const { url, start_index, end_index, title, description } = metadata;
 
-  if (!url) {
+  if (!url || !messageId) {
     return null;
   }
 
diff --git a/apps/beeai-ui/src/modules/runs/contexts/agent-run/AgentRunProvider.tsx b/apps/beeai-ui/src/modules/runs/contexts/agent-run/AgentRunProvider.tsx
@@ -11,6 +11,7 @@ import { v4 as uuid } from 'uuid';
 
 import { buildA2AClient } from '#api/a2a/client.ts';
 import type { ChatRun } from '#api/a2a/types.ts';
+import { createTextPart } from '#api/a2a/utils.ts';
 import { getErrorCode } from '#api/utils.ts';
 import { useHandleError } from '#hooks/useHandleError.ts';
 import { useImmerWithGetter } from '#hooks/useImmerWithGetter.ts';
@@ -123,31 +124,36 @@ function AgentRunProvider({ agent, children }: PropsWithChildren<Props>) {
       setIsPending(true);
       setStats({ startTime: Date.now() });
 
-      try {
-        setMessages((messages) => {
-          const userMessage: UIUserMessage = {
-            id: uuid(),
-            role: Role.User,
-            parts: [{ kind: UIMessagePartKind.Text, id: uuid(), text: input }, ...convertFilesToUIFileParts(files)],
-          };
-          const agentMessage: UIAgentMessage = {
-            id: uuid(),
-            role: Role.Agent,
-            parts: [],
-            status: UIMessageStatus.InProgress,
-          };
-
-          messages.push(userMessage, agentMessage);
-        });
+      const userMessage: UIUserMessage = {
+        id: uuid(),
+        role: Role.User,
+        parts: [createTextPart(input), ...convertFilesToUIFileParts(files)],
+      };
+      const agentMessage: UIAgentMessage = {
+        id: uuid(),
+        role: Role.Agent,
+        parts: [],
+        status: UIMessageStatus.InProgress,
+      };
+
+      setMessages((messages) => {
+        messages.push(userMessage, agentMessage);
+      });
 
+      clearFiles();
+
+      try {
         const run = a2aAgentClient.chat({
-          text: input,
-          files,
+          message: userMessage,
           contextId: conversationId,
         });
         pendingRun.current = run;
 
-        pendingSubscription.current = run.subscribe((parts) => {
+        pendingSubscription.current = run.subscribe(({ parts, taskId }) => {
+          updateLastAgentMessage((message) => {
+            message.id = taskId;
+          });
+
           parts.forEach((part) => {
             updateLastAgentMessage((message) => {
               match(part)
@@ -188,7 +194,7 @@ function AgentRunProvider({ agent, children }: PropsWithChildren<Props>) {
         pendingSubscription.current = undefined;
       }
     },
-    [a2aAgentClient, files, conversationId, handleError, updateLastAgentMessage, setMessages],
+    [a2aAgentClient, files, conversationId, handleError, updateLastAgentMessage, setMessages, clearFiles],
   );
 
   const sources = useMemo(() => getMessageSourcesMap(messages), [messages]);
