@assistant-ui/react-a2a 0.2.22 → 0.2.24

package/src/A2AThreadRuntimeCore.ts

@@ -4,11 +4,13 @@ import { generateId, fromThreadMessageLike } from "@assistant-ui/core";
 import type {
   AppendMessage,
   AssistantRuntime,
+  ExportedMessageRepository,
   MessageStatus,
   ThreadAssistantMessage,
   ThreadHistoryAdapter,
   ThreadMessage,
 } from "@assistant-ui/core";
+import { MessageRepository } from "@assistant-ui/core/internal";
 import type { A2AClient } from "./A2AClient";
 import type {
   A2AArtifact,
@@ -23,6 +25,7 @@ import type {
 } from "./types";
 import {
   a2aMessageToContent,
+  contentPartsToA2AParts,
   isTerminalTaskState,
   taskStateToMessageStatus,
 } from "./conversions";
@@ -54,7 +57,8 @@ export class A2AThreadRuntimeCore {
   private readonly notifyUpdate: () => void;
 
   private runtime: AssistantRuntime | undefined;
-  private messages: ThreadMessage[] = [];
+  private readonly repository = new MessageRepository();
+  private exportedRepository: ExportedMessageRepository | undefined;
   private isRunningFlag = false;
   private abortController: AbortController | null = null;
   private pendingError: Error | null = null;
@@ -109,7 +113,69 @@ export class A2AThreadRuntimeCore {
   }
 
   getMessages(): readonly ThreadMessage[] {
-    return this.messages;
+    return this.repository.getMessages();
+  }
+
+  getMessageRepository(): ExportedMessageRepository {
+    this.exportedRepository ??= this.repository.export();
+    return this.exportedRepository;
+  }
+
+  private tryGetMessage(messageId: string) {
+    try {
+      return this.repository.getMessage(messageId);
+    } catch {
+      return undefined;
+    }
+  }
+
+  private tryGetMessages(
+    messageId: string,
+  ): readonly ThreadMessage[] | undefined {
+    try {
+      return this.repository.getMessages(messageId);
+    } catch {
+      return undefined;
+    }
+  }
+
+  private hasMessage(messageId: string): boolean {
+    return this.tryGetMessage(messageId) !== undefined;
+  }
+
+  private addOrUpdateMessage(
+    parentId: string | null,
+    message: ThreadMessage,
+  ): void {
+    this.repository.addOrUpdateMessage(parentId, message);
+    this.exportedRepository = undefined;
+  }
+
+  private switchToBranch(messageId: string): void {
+    this.repository.switchToBranch(messageId);
+    this.exportedRepository = undefined;
+  }
+
+  private resetRepositoryHead(messageId: string | null): void {
+    this.repository.resetHead(messageId);
+    this.exportedRepository = undefined;
+  }
+
+  private clearRepository(): void {
+    this.repository.clear();
+    this.exportedRepository = undefined;
+  }
+
+  private updateMessage(
+    messageId: string,
+    updater: (message: ThreadMessage) => ThreadMessage,
+  ): boolean {
+    const item = this.tryGetMessage(messageId);
+    if (!item) return false;
+    const message = updater(item.message);
+    if (message === item.message) return false;
+    this.addOrUpdateMessage(item.parentId, message);
+    return true;
   }
 
   getTask(): A2ATask | undefined {
@@ -146,8 +212,7 @@ export class A2AThreadRuntimeCore {
           this.agentCardValue = agentCard;
         }
         if (repo) {
-          const messages = repo.messages.map((item) => item.message);
-          this.applyExternalMessages(messages);
+          this.applyExternalMessageRepository(repo);
         }
       })
       .catch((error) => {
@@ -166,21 +231,22 @@ export class A2AThreadRuntimeCore {
 
   async append(message: AppendMessage): Promise<void> {
     const startRun = message.startRun ?? message.role === "user";
-    if (message.sourceId) {
-      this.messages = this.messages.filter(
-        (entry) => entry.id !== message.sourceId,
-      );
-    }
-    this.resetHead(message.parentId);
 
     const threadMessage = fromThreadMessageLike(
       message as any,
       generateId(),
       FALLBACK_USER_STATUS,
     );
-    this.messages = [...this.messages, threadMessage];
+    const parentId =
+      message.parentId === null
+        ? null
+        : message.parentId && this.hasMessage(message.parentId)
+          ? message.parentId
+          : this.repository.headId;
+    this.addOrUpdateMessage(parentId, threadMessage);
+    this.switchToBranch(threadMessage.id);
     this.notifyUpdate();
-    this.recordHistoryEntry(message.parentId ?? null, threadMessage);
+    this.recordHistoryEntry(parentId, threadMessage);
 
     if (!startRun) return;
     await this.startRun(threadMessage);
@@ -194,13 +260,13 @@ export class A2AThreadRuntimeCore {
     parentId: string | null,
     _config: { runConfig?: Record<string, unknown> } = {},
   ): Promise<void> {
-    this.resetHead(parentId);
-    this.notifyUpdate();
-
-    // Find the last user message to re-run
-    for (let i = this.messages.length - 1; i >= 0; i--) {
-      if (this.messages[i]!.role === "user") {
-        await this.startRun(this.messages[i]!);
+    const messages =
+      parentId === null
+        ? []
+        : (this.tryGetMessages(parentId) ?? this.getMessages());
+    for (let i = messages.length - 1; i >= 0; i--) {
+      if (messages[i]!.role === "user") {
+        await this.startRun(messages[i]!);
         return;
       }
     }
@@ -223,19 +289,125 @@ export class A2AThreadRuntimeCore {
     }
   }
 
-  applyExternalMessages(messages: readonly ThreadMessage[]): void {
+  private appendLinearChain(messages: readonly ThreadMessage[]): string | null {
+    let parentId: string | null = null;
+    let lastId: string | null = null;
+    const seen = new Set<string>();
+
+    for (const message of messages) {
+      if (seen.has(message.id)) continue;
+      seen.add(message.id);
+      this.addOrUpdateMessage(parentId, message);
+      parentId = message.id;
+      lastId = message.id;
+    }
+
+    return lastId;
+  }
+
+  private finalizeExternalApply(): void {
     this.assistantHistoryParents.clear();
-    this.messages = [...messages];
     this.recordedHistoryIds.clear();
-    for (const message of this.messages) {
+    for (const { message } of this.getMessageRepository().messages) {
       this.recordedHistoryIds.add(message.id);
     }
-    // Reset task-specific state to prevent leaking into new thread
     this.currentTask = undefined;
     this.currentArtifacts = [];
     this.notifyUpdate();
   }
 
+  applyExternalMessages(messages: readonly ThreadMessage[]): void {
+    if (messages.length === 0) {
+      this.clearRepository();
+    } else {
+      let expectedParentId: string | null = null;
+      let lastAppliedId: string | null = null;
+      let hardReplace = false;
+      const seen = new Set<string>();
+
+      for (const message of messages) {
+        if (seen.has(message.id)) continue;
+        seen.add(message.id);
+        const existing = this.tryGetMessage(message.id);
+        if (existing && existing.parentId !== expectedParentId) {
+          hardReplace = true;
+          break;
+        }
+        this.addOrUpdateMessage(expectedParentId, message);
+        expectedParentId = message.id;
+        lastAppliedId = message.id;
+      }
+
+      if (hardReplace) {
+        this.clearRepository();
+        lastAppliedId = this.appendLinearChain(messages);
+      }
+
+      this.resetRepositoryHead(lastAppliedId);
+    }
+
+    this.finalizeExternalApply();
+  }
+
+  private applyExternalMessageRepository(
+    loaded: ExportedMessageRepository,
+  ): void {
+    const headId = loaded.headId ?? loaded.messages.at(-1)?.message.id ?? null;
+    const ids = new Set<string>();
+    let degenerate = false;
+    for (const { message } of loaded.messages) {
+      if (ids.has(message.id)) {
+        degenerate = true;
+        break;
+      }
+      ids.add(message.id);
+    }
+    if (headId !== null && !ids.has(headId)) degenerate = true;
+
+    if (!degenerate) {
+      this.clearRepository();
+      let pending = [...loaded.messages];
+      const importedIds = new Set<string>();
+
+      while (pending.length > 0) {
+        const unresolved: typeof pending = [];
+        let progressed = false;
+        for (const item of pending) {
+          if (item.parentId !== null && !importedIds.has(item.parentId)) {
+            unresolved.push(item);
+            continue;
+          }
+          this.addOrUpdateMessage(item.parentId, item.message);
+          importedIds.add(item.message.id);
+          progressed = true;
+        }
+        if (!progressed) {
+          degenerate = true;
+          break;
+        }
+        pending = unresolved;
+      }
+    }
+
+    if (degenerate) {
+      this.clearRepository();
+      let previousId: string | null = null;
+      for (const { message } of loaded.messages) {
+        const existing = this.tryGetMessage(message.id);
+        this.addOrUpdateMessage(
+          existing ? existing.parentId : previousId,
+          message,
+        );
+        previousId = message.id;
+      }
+      this.resetRepositoryHead(previousId);
+    } else {
+      this.resetRepositoryHead(headId);
+    }
+
+    this.finalizeExternalApply();
+  }
+
   // --- Run logic ---
 
   private async startRun(userThreadMessage: ThreadMessage): Promise<void> {
@@ -258,7 +430,7 @@ export class A2AThreadRuntimeCore {
     this.currentArtifacts = [];
 
     const assistantParentId = userThreadMessage.id;
-    const assistantId = this.insertAssistantPlaceholder();
+    const assistantId = this.insertAssistantPlaceholder(assistantParentId);
     this.markPendingAssistantHistory(assistantId, assistantParentId);
 
     const abortController = new AbortController();
@@ -484,12 +656,14 @@ export class A2AThreadRuntimeCore {
     const parts: A2APart[] = [];
 
     if (message.role === "user") {
-      for (const part of message.content) {
-        if (part.type === "text") {
-          parts.push({ text: part.text });
-        } else if (part.type === "image") {
-          parts.push({ url: part.image, mediaType: "image/*" });
-        }
+      parts.push(...contentPartsToA2AParts(message.content));
+      for (const attachment of message.attachments ?? []) {
+        parts.push(
+          ...contentPartsToA2AParts(
+            attachment.content ?? [],
+            attachment.contentType,
+          ),
+        );
       }
     }
 
@@ -513,7 +687,7 @@ export class A2AThreadRuntimeCore {
     return a2aMsg;
   }
 
-  private insertAssistantPlaceholder(): string {
+  private insertAssistantPlaceholder(parentId: string): string {
     const id = generateId();
     const assistant: ThreadAssistantMessage = {
       id,
@@ -529,7 +703,8 @@ export class A2AThreadRuntimeCore {
         custom: {},
       },
     };
-    this.messages = [...this.messages, assistant];
+    this.addOrUpdateMessage(parentId, assistant);
+    this.switchToBranch(id);
     this.notifyUpdate();
     return id;
   }
@@ -538,19 +713,15 @@ export class A2AThreadRuntimeCore {
     messageId: string,
     content: ThreadAssistantMessage["content"],
   ) {
-    this.messages = this.messages.map((message) => {
-      if (message.id !== messageId || message.role !== "assistant")
-        return message;
+    this.updateMessage(messageId, (message) => {
+      if (message.role !== "assistant") return message;
       return { ...message, content };
     });
   }
 
   private updateAssistantStatus(messageId: string, status: MessageStatus) {
-    let touched = false;
-    this.messages = this.messages.map((message) => {
-      if (message.id !== messageId || message.role !== "assistant")
-        return message;
-      touched = true;
+    const touched = this.updateMessage(messageId, (message) => {
+      if (message.role !== "assistant") return message;
       return { ...message, status };
     });
     if (touched) {
@@ -562,10 +733,9 @@ export class A2AThreadRuntimeCore {
   }
 
   private getAssistantStatus(messageId: string): MessageStatus | undefined {
-    const msg = this.messages.find(
-      (m) => m.id === messageId && m.role === "assistant",
-    );
-    return msg?.status;
+    const msg = this.tryGetMessage(messageId)?.message;
+    if (msg?.role !== "assistant") return undefined;
+    return msg.status;
   }
 
   // --- Lifecycle helpers ---
@@ -582,18 +752,6 @@ export class A2AThreadRuntimeCore {
     this.setRunning(false);
   }
 
-  private resetHead(parentId: string | null | undefined) {
-    if (!parentId) {
-      if (this.messages.length) {
-        this.messages = [];
-      }
-      return;
-    }
-    const idx = this.messages.findIndex((message) => message.id === parentId);
-    if (idx === -1) return;
-    this.messages = this.messages.slice(0, idx + 1);
-  }
-
   // --- History persistence ---
 
   private recordHistoryEntry(parentId: string | null, message: ThreadMessage) {
@@ -612,7 +770,7 @@ export class A2AThreadRuntimeCore {
     if (!this.history) return;
     const parentId = this.assistantHistoryParents.get(messageId);
     if (parentId === undefined) return;
-    const message = this.messages.find((m) => m.id === messageId);
+    const message = this.tryGetMessage(messageId)?.message;
     if (!message || message.role !== "assistant") return;
     if (
       message.status?.type !== "complete" &&

package/src/conversions.test.ts

@@ -245,12 +245,61 @@ describe("contentPartsToA2AParts", () => {
     expect(result).toEqual([{ text: "hi" }]);
   });
 
-  it("converts image parts", () => {
+  it("converts image URL parts without stamping a mediaType", () => {
     const result = contentPartsToA2AParts([
       { type: "image", image: "https://img.com/a.png" },
     ]);
+    expect(result).toEqual([{ url: "https://img.com/a.png" }]);
+  });
+
+  it("applies the fallback MIME type to image URL parts", () => {
+    const result = contentPartsToA2AParts(
+      [{ type: "image", image: "https://img.com/a.png" }],
+      "image/png",
+    );
+    expect(result).toEqual([
+      { url: "https://img.com/a.png", mediaType: "image/png" },
+    ]);
+  });
+
+  it("applies the fallback MIME type and filename to image URL parts together", () => {
+    const result = contentPartsToA2AParts(
+      [{ type: "image", image: "https://img.com/a.png", filename: "a.png" }],
+      "image/png",
+    );
     expect(result).toEqual([
-      { url: "https://img.com/a.png", mediaType: "image/*" },
+      {
+        url: "https://img.com/a.png",
+        mediaType: "image/png",
+        filename: "a.png",
+      },
+    ]);
+  });
+
+  it("passes non-base64 data URLs through as URLs for image parts", () => {
+    const result = contentPartsToA2AParts([
+      { type: "image", image: "data:text/plain,hello" },
+    ]);
+    expect(result).toEqual([{ url: "data:text/plain,hello" }]);
+  });
+
+  it("converts image data URLs to raw bytes with the embedded MIME type", () => {
+    const result = contentPartsToA2AParts([
+      { type: "image", image: "data:image/png;base64,aGVsbG8=" },
+    ]);
+    expect(result).toEqual([{ raw: "aGVsbG8=", mediaType: "image/png" }]);
+  });
+
+  it("propagates image filenames", () => {
+    const result = contentPartsToA2AParts([
+      {
+        type: "image",
+        image: "data:image/png;base64,aGVsbG8=",
+        filename: "a.png",
+      },
+    ]);
+    expect(result).toEqual([
+      { raw: "aGVsbG8=", mediaType: "image/png", filename: "a.png" },
     ]);
   });
 
@@ -259,6 +308,64 @@ describe("contentPartsToA2AParts", () => {
     expect(result).toEqual([]);
   });
 
+  it("converts file parts with http URLs", () => {
+    const result = contentPartsToA2AParts([
+      {
+        type: "file",
+        data: "https://files.com/doc.pdf",
+        mimeType: "application/pdf",
+        filename: "doc.pdf",
+      },
+    ]);
+    expect(result).toEqual([
+      {
+        url: "https://files.com/doc.pdf",
+        mediaType: "application/pdf",
+        filename: "doc.pdf",
+      },
+    ]);
+  });
+
+  it("converts file parts with data URLs to raw bytes", () => {
+    const result = contentPartsToA2AParts([
+      {
+        type: "file",
+        data: "data:application/pdf;base64,ZmlsZQ==",
+        mimeType: "application/pdf",
+      },
+    ]);
+    expect(result).toEqual([{ raw: "ZmlsZQ==", mediaType: "application/pdf" }]);
+  });
+
+  it("converts file parts with raw base64 data", () => {
+    const result = contentPartsToA2AParts([
+      { type: "file", data: "ZmlsZQ==", mimeType: "text/csv" },
+    ]);
+    expect(result).toEqual([{ raw: "ZmlsZQ==", mediaType: "text/csv" }]);
+  });
+
+  it("falls back to the attachment MIME type when the file part MIME is empty", () => {
+    const result = contentPartsToA2AParts(
+      [{ type: "file", data: "ZmlsZQ==", mimeType: "" }],
+      "application/pdf",
+    );
+    expect(result).toEqual([{ raw: "ZmlsZQ==", mediaType: "application/pdf" }]);
+  });
+
+  it("omits mediaType when no MIME type is known", () => {
+    const result = contentPartsToA2AParts([
+      { type: "file", data: "ZmlsZQ==", mimeType: "" },
+    ]);
+    expect(result).toEqual([{ raw: "ZmlsZQ==" }]);
+  });
+
+  it("skips file parts with no data", () => {
+    const result = contentPartsToA2AParts([
+      { type: "file", mimeType: "application/pdf" },
+    ]);
+    expect(result).toEqual([]);
+  });
+
   it("skips unknown part types", () => {
     const result = contentPartsToA2AParts([
       { type: "text", text: "hi" },

package/src/conversions.ts

@@ -1,6 +1,7 @@
 "use client";
 
 import type { MessageStatus, ThreadAssistantMessage } from "@assistant-ui/core";
+import { httpUrlPattern, parseDataUrl } from "@assistant-ui/core/internal";
 import type { A2AMessage, A2APart, A2ATaskState } from "./types";
 
 function isImageMediaType(mediaType?: string): boolean {
@@ -89,18 +90,59 @@ export function taskStateToMessageStatus(state: A2ATaskState): MessageStatus {
 export function contentPartsToA2AParts(
   content: ReadonlyArray<{
     type: string;
-    text?: string;
-    image?: string;
+    text?: string | undefined;
+    image?: string | undefined;
+    data?: string | undefined;
+    mimeType?: string | undefined;
+    filename?: string | undefined;
   }>,
+  fallbackMimeType?: string,
 ): A2APart[] {
   return content
     .map((part): A2APart | null => {
       switch (part.type) {
         case "text":
           return { text: part.text ?? "" };
-        case "image":
+        case "image": {
           if (!part.image) return null;
-          return { url: part.image, mediaType: "image/*" };
+          const parsed = parseDataUrl(part.image);
+          if (parsed) {
+            return {
+              raw: parsed.data,
+              mediaType: parsed.mimeType,
+              ...(part.filename && { filename: part.filename }),
+            };
+          }
+          return {
+            url: part.image,
+            ...(fallbackMimeType && { mediaType: fallbackMimeType }),
+            ...(part.filename && { filename: part.filename }),
+          };
+        }
+        case "file": {
+          if (!part.data) return null;
+          const declaredMimeType = part.mimeType || fallbackMimeType;
+          if (httpUrlPattern.test(part.data)) {
+            return {
+              url: part.data,
+              ...(declaredMimeType && { mediaType: declaredMimeType }),
+              ...(part.filename && { filename: part.filename }),
+            };
+          }
+          const parsed = parseDataUrl(part.data);
+          if (parsed) {
+            return {
+              raw: parsed.data,
+              mediaType: parsed.mimeType,
+              ...(part.filename && { filename: part.filename }),
+            };
+          }
+          return {
+            raw: part.data,
+            ...(declaredMimeType && { mediaType: declaredMimeType }),
+            ...(part.filename && { filename: part.filename }),
+          };
+        }
         default:
           return null;
       }