fix(core): retain buffered messages until memory writes succeed

.changeset/memory-persist-retry.md

@@ -0,0 +1,5 @@
+---
+"@voltagent/core": patch
+---
+
+Keep conversation messages pending until their memory writes succeed, so failed writes can be retried without losing later messages.

packages/core/src/agent/conversation-buffer.ts

@@ -7,9 +7,9 @@ import { convertModelMessagesToUIMessages } from "../utils/message-converter";
 
 type MessageSource = "user" | "system" | "memory" | "response";
 
-interface PendingMessage {
-  id: string;
+export interface PendingMessage {
   message: UIMessage;
+  version: number;
 }
 
 export interface ConversationBufferCheckpoint {
@@ -62,6 +62,8 @@ const extractOpenAIItemId = (metadata: unknown): string => {
 export class ConversationBuffer {
   private messages: UIMessage[] = [];
   private pendingMessageIds = new Set<string>();
+  private pendingMessageVersions = new Map<string, number>();
+  private nextPendingMessageVersion = 0;
   private toolPartIndex = new Map<string, { messageIndex: number; partIndex: number }>();
   private activeAssistantMessageId?: string;
 
@@ -128,6 +130,9 @@ export class ConversationBuffer {
   restoreCheckpoint(checkpoint: ConversationBufferCheckpoint): void {
     this.messages = checkpoint.messages.map((message) => this.cloneMessage(message));
     this.pendingMessageIds = new Set(checkpoint.pendingMessageIds);
+    this.pendingMessageVersions = new Map<string, number>(
+      checkpoint.pendingMessageIds.map((id) => [id, ++this.nextPendingMessageVersion]),
+    );
     this.activeAssistantMessageId = checkpoint.activeAssistantMessageId;
     this.rebuildToolPartIndex();
     this.log("restore-checkpoint", {
@@ -137,32 +142,53 @@ export class ConversationBuffer {
   }
 
   drainPendingMessages(): UIMessage[] {
+    const drained = this.getPendingMessages();
+    this.pendingMessageIds.clear();
+    this.pendingMessageVersions.clear();
+    if (
+      this.activeAssistantMessageId &&
+      drained.some((item) => item.message.id === this.activeAssistantMessageId)
+    ) {
+      this.activeAssistantMessageId = undefined;
+    }
+    if (drained.length > 0) {
+      this.log("drain-pending", {
+        count: drained.length,
+        ids: drained.map((item) => item.message.id),
+      });
+    }
+    return drained.map((item) => item.message);
+  }
+
+  getPendingMessages(): PendingMessage[] {
     if (this.pendingMessageIds.size === 0) {
       return [];
     }
 
-    const drained: PendingMessage[] = [];
+    const pending: PendingMessage[] = [];
 
     this.messages.forEach((message) => {
       if (this.pendingMessageIds.has(message.id)) {
-        drained.push({ id: message.id, message: this.cloneMessage(message) });
+        pending.push({
+          message: this.cloneMessage(message),
+          version: this.pendingMessageVersions.get(message.id) ?? 0,
+        });
       }
     });
 
-    this.pendingMessageIds.clear();
+    return pending;
+  }
 
-    if (drained.length > 0) {
-      const drainedIds = new Set(drained.map((item) => item.id));
-      if (this.activeAssistantMessageId && drainedIds.has(this.activeAssistantMessageId)) {
-        this.activeAssistantMessageId = undefined;
-      }
+  markMessagePersisted(id: string, version: number): void {
+    if (this.pendingMessageVersions.get(id) !== version) {
+      return;
     }
 
-    if (drained.length > 0) {
-      this.log("drain-pending", { count: drained.length, ids: drained.map((item) => item.id) });
+    this.pendingMessageIds.delete(id);
+    this.pendingMessageVersions.delete(id);
+    if (this.activeAssistantMessageId === id) {
+      this.activeAssistantMessageId = undefined;
     }
-
-    return drained.map((item) => item.message);
   }
 
   getAllMessages(): UIMessage[] {
@@ -197,7 +223,7 @@ export class ConversationBuffer {
       ...(existing as Record<string, unknown>),
       ...metadata,
     } as UIMessage["metadata"];
-    this.pendingMessageIds.add(target.id);
+    this.markPending(target.id);
     return true;
   }
 
@@ -211,7 +237,7 @@ export class ConversationBuffer {
     this.registerToolParts(this.messages.length - 1);
 
     if (!options.markAsSaved) {
-      this.pendingMessageIds.add(hydrated.id);
+      this.markPending(hydrated.id);
     }
 
     this.log("append-existing", {
@@ -314,7 +340,7 @@ export class ConversationBuffer {
     }
 
     if (modified) {
-      this.pendingMessageIds.add(target.id);
+      this.markPending(target.id);
       this.registerToolParts(lastAssistantIndex);
     }
   }
@@ -370,19 +396,25 @@ export class ConversationBuffer {
     const cloned = this.cloneMessage(message);
     this.ensureMessageId(cloned);
     this.messages.push(cloned);
-    this.pendingMessageIds.add(cloned.id);
+    this.markPending(cloned.id);
     this.registerToolParts(this.messages.length - 1);
     this.log("append-message", { messageId: cloned.id, role: cloned.role, source });
 
     if (source === "memory") {
       this.pendingMessageIds.delete(cloned.id);
+      this.pendingMessageVersions.delete(cloned.id);
     }
 
     if (source === "response") {
       this.activeAssistantMessageId = cloned.id;
     }
   }
 
+  private markPending(id: string): void {
+    this.pendingMessageIds.add(id);
+    this.pendingMessageVersions.set(id, ++this.nextPendingMessageVersion);
+  }
+
   private registerToolParts(messageIndex: number): void {
     const message = this.messages[messageIndex];
     for (let index = 0; index < message.parts.length; index++) {

packages/core/src/agent/memory-persist-queue.spec.ts

@@ -94,6 +94,100 @@ describe("MemoryPersistQueue", () => {
     expect(memoryManager.saveMessage).toHaveBeenCalledTimes(1);
   });
 
+  it("keeps failed and unattempted messages pending for the next flush", async () => {
+    const messages = [createMessage("first"), createMessage("second"), createMessage("third")];
+    const memoryManager = {
+      saveMessage: vi
+        .fn()
+        .mockResolvedValue(undefined)
+        .mockRejectedValueOnce(new Error("storage unavailable")),
+    } as any;
+    const buffer = new ConversationBuffer();
+    buffer.ingestUIMessages(messages, false);
+    const oc = createOperationContext();
+    const queue = new MemoryPersistQueue(memoryManager, { logger: oc.logger });
+
+    await expect(queue.flush(buffer, oc as any)).rejects.toThrow("storage unavailable");
+    expect(memoryManager.saveMessage).toHaveBeenCalledWith(
+      oc,
+      expect.objectContaining({ id: messages[0].id }),
+      oc.userId,
+      oc.conversationId,
+      { throwOnError: true },
+    );
+    expect(buffer.getPendingMessages().map(({ message }) => message.id)).toEqual(
+      messages.map((message) => message.id),
+    );
+
+    await queue.flush(buffer, oc as any);
+    expect(memoryManager.saveMessage.mock.calls.map((call: any[]) => call[1].id)).toEqual([
+      messages[0].id,
+      ...messages.map((message) => message.id),
+    ]);
+    expect(buffer.getPendingMessages()).toHaveLength(0);
+  });
+
+  it("retries only messages that were not saved before a later failure", async () => {
+    const messages = [createMessage("first"), createMessage("second"), createMessage("third")];
+    const memoryManager = {
+      saveMessage: vi
+        .fn()
+        .mockResolvedValueOnce(undefined)
+        .mockRejectedValueOnce(new Error("storage unavailable"))
+        .mockResolvedValue(undefined),
+    } as any;
+    const buffer = new ConversationBuffer();
+    buffer.ingestUIMessages(messages, false);
+    const oc = createOperationContext();
+    const queue = new MemoryPersistQueue(memoryManager, { logger: oc.logger });
+
+    await expect(queue.flush(buffer, oc as any)).rejects.toThrow("storage unavailable");
+    expect(buffer.getPendingMessages().map(({ message }) => message.id)).toEqual([
+      messages[1].id,
+      messages[2].id,
+    ]);
+
+    await queue.flush(buffer, oc as any);
+    expect(memoryManager.saveMessage.mock.calls.map((call: any[]) => call[1].id)).toEqual([
+      messages[0].id,
+      messages[1].id,
+      messages[1].id,
+      messages[2].id,
+    ]);
+  });
+
+  it("keeps a message pending when it changes during an in-flight save", async () => {
+    let finishSave!: () => void;
+    const memoryManager = {
+      saveMessage: vi
+        .fn()
+        .mockImplementationOnce(
+          () =>
+            new Promise<void>((resolve) => {
+              finishSave = resolve;
+            }),
+        )
+        .mockResolvedValue(undefined),
+    } as any;
+    const buffer = new ConversationBuffer();
+    buffer.ingestUIMessages([createMessage("first")], false);
+    const oc = createOperationContext();
+    const queue = new MemoryPersistQueue(memoryManager, { logger: oc.logger });
+
+    const firstFlush = queue.flush(buffer, oc as any);
+    await Promise.resolve();
+    await Promise.resolve();
+    expect(memoryManager.saveMessage).toHaveBeenCalledTimes(1);
+    buffer.addMetadataToLastAssistantMessage({ updated: true });
+    finishSave();
+    await firstFlush;
+
+    expect(buffer.getPendingMessages()).toHaveLength(1);
+    await queue.flush(buffer, oc as any);
+    expect(memoryManager.saveMessage.mock.calls[1][1].metadata).toMatchObject({ updated: true });
+    expect(buffer.getPendingMessages()).toHaveLength(0);
+  });
+
   it("adds subagent metadata before persisting when parentAgentId is present", async () => {
     const memoryManager = {
       saveMessage: vi.fn().mockResolvedValue(undefined),

packages/core/src/agent/memory-persist-queue.ts

@@ -90,7 +90,7 @@ export class MemoryPersistQueue {
       return;
     }
 
-    const pending = buffer.drainPendingMessages();
+    const pending = buffer.getPendingMessages();
     if (pending.length === 0) {
       const payload = {
         conversationId: oc.conversationId,
@@ -104,7 +104,7 @@ export class MemoryPersistQueue {
       conversationId: oc.conversationId,
       userId: oc.userId,
       count: pending.length,
-      ids: pending.map((msg) => msg.id),
+      ids: pending.map(({ message }) => message.id),
     };
     this.logger?.debug?.("[MemoryPersistQueue] persisting", payload);
 
@@ -116,13 +116,22 @@ export class MemoryPersistQueue {
       | Map<string, AgentMetadataContextValue>
       | undefined;
 
-    for (const message of pending) {
+    for (const { message, version } of pending) {
       try {
         const messageWithMetadata = this.applySubAgentMetadata(message, {
           defaultMetadata: shouldApplySubAgentMetadata ? agentMetadata : undefined,
           toolCallMetadata,
         });
-        await this.memoryManager.saveMessage(oc, messageWithMetadata, oc.userId, oc.conversationId);
+        await this.memoryManager.saveMessage(
+          oc,
+          messageWithMetadata,
+          oc.userId,
+          oc.conversationId,
+          {
+            throwOnError: true,
+          },
+        );
+        buffer.markMessagePersisted(message.id, version);
       } catch (error) {
         this.logger?.error?.("Failed to save message", {
           conversationId: oc.conversationId,

packages/core/src/memory/manager/memory-manager.spec.ts

@@ -186,7 +186,7 @@ describe("MemoryManager", () => {
       );
     });
 
-    it("should handle errors gracefully", async () => {
+    it("should propagate storage errors when the caller needs to retry", async () => {
       // Create manager with mocked memory that throws error
       const errorMemory = new Memory({
         storage: new InMemoryStorageAdapter(),
@@ -204,10 +204,13 @@ describe("MemoryManager", () => {
 
       const message = createTestUIMessage();
 
-      // Should not throw even if save fails
+      await expect(
+        errorManager.saveMessage(context, message, "user-1", "conv-1", { throwOnError: true }),
+      ).rejects.toThrow("Save failed");
+
       await expect(
         errorManager.saveMessage(context, message, "user-1", "conv-1"),
-      ).resolves.not.toThrow();
+      ).resolves.toBeUndefined();
     });
   });
 

packages/core/src/memory/manager/memory-manager.ts

@@ -107,6 +107,7 @@ export class MemoryManager {
     message: UIMessage,
     userId?: string,
     conversationId?: string,
+    options?: { throwOnError?: boolean },
   ): Promise<void> {
     if (!this.conversationMemory || !userId) return;
 
@@ -202,6 +203,9 @@ export class MemoryManager {
           error: error instanceof Error ? { message: error.message, stack: error.stack } : error,
         },
       );
+      if (options?.throwOnError) {
+        throw error;
+      }
     }
   }