feat: LocalRuntime export / import (#609)

* refactor: Add export and import methods to LocalThreadRuntime and MessageRepository

* serialize messages as { message, parentId } pairs

* refactor: export `ExportedMessageRepository` type and add persisting head state

* pnpm changeset

.changeset/silver-queens-appear.md

@@ -0,0 +1,5 @@
+---
+"@assistant-ui/react": patch
+---
+
+feat: LocalRuntime export / import

packages/react/src/runtimes/local/LocalThreadRuntime.tsx

@@ -9,7 +9,7 @@ import type {
   Unsubscribe,
 } from "../../types";
 import { fromCoreMessages } from "../edge";
-import { MessageRepository } from "../utils/MessageRepository";
+import { ExportedMessageRepository, MessageRepository } from "../utils/MessageRepository";
 import type { ChatModelAdapter, ChatModelRunResult } from "./ChatModelAdapter";
 import { shouldContinue } from "./shouldContinue";
 import { LocalRuntimeOptions } from "./LocalRuntimeOptions";
@@ -239,4 +239,13 @@ export class LocalThreadRuntime implements ThreadRuntime {
       this.performRoundtrip(parentId, message);
     }
   }
+
+  export() {
+    return this.repository.export();
+  }
+
+  import(data: ExportedMessageRepository) {
+    this.repository.import(data);
+    this.notifySubscribers()
+  }
 }

packages/react/src/runtimes/utils/MessageRepository.tsx

@@ -13,6 +13,14 @@ type RepositoryMessage = RepositoryParent & {
   level: number;
 };
 
+export interface ExportedMessageRepository {
+  headId?: string | null;
+  messages: Array<{
+    message: ThreadMessage;
+    parentId: string | null;
+  }>;
+}
+
 const findHead = (message: RepositoryMessage): RepositoryMessage => {
   if (message.next) return findHead(message.next);
   return message;
@@ -224,4 +232,31 @@ export class MessageRepository {
       }
     }
   }
+
+  export(): ExportedMessageRepository {
+    const exportItems: ExportedMessageRepository["messages"] = [];
+
+    // hint: we are relying on the insertion order of the messages
+    // this is important for the import function to properly link the messages
+    for (const [, message] of this.messages) {
+      exportItems.push({
+        message: message.current,
+        parentId: message.prev?.current.id ?? null,
+      });
+    }
+
+    return {
+      headId: this.head?.current.id ?? null,
+      messages: exportItems,
+    };
+  }
+
+  import({ headId, messages }: ExportedMessageRepository) {
+    for (const { message, parentId } of messages) {
+      this.addOrUpdateMessage(parentId, message);
+    }
+
+    // switch to the saved head id if it is not the most recent message
+    this.resetHead(headId ?? messages.at(-1)?.message.id ?? null);
+  }
 }