diff --git a/.changeset/memory-persist-retry.md b/.changeset/memory-persist-retry.md new file mode 100644 index 000000000..c476f5d75 --- /dev/null +++ b/.changeset/memory-persist-retry.md @@ -0,0 +1,5 @@ +--- +"@voltagent/core": patch +--- + +Keep conversation messages pending until their memory writes succeed and retry failed writes on the next operation for the same conversation. Memory failures remain best effort for agent responses. diff --git a/packages/core/src/agent/agent.spec.ts b/packages/core/src/agent/agent.spec.ts index e2572f4fd..0f312ff2c 100644 --- a/packages/core/src/agent/agent.spec.ts +++ b/packages/core/src/agent/agent.spec.ts @@ -2532,6 +2532,67 @@ Use pandas and summarize findings.`.split("\n"), ); }); + it("returns the generated response after a memory write fails and saves it on the next turn", async () => { + const memory = new Memory({ storage: new InMemoryStorageAdapter() }); + const saveMessage = memory.saveMessageWithContext.bind(memory); + let failAssistantWrite = true; + vi.spyOn(memory, "saveMessageWithContext").mockImplementation(async (...args) => { + if (failAssistantWrite && args[0].role === "assistant") { + failAssistantWrite = false; + throw new Error("storage unavailable"); + } + return saveMessage(...args); + }); + const agent = new Agent({ + name: "TestAgent", + instructions: "Test", + model: mockModel as any, + memory, + }); + const response = (text: string) => ({ + text, + content: [{ type: "text", text }], + reasoning: [], + files: [], + sources: [], + toolCalls: [], + toolResults: [], + finishReason: "stop", + usage: providerUsage, + warnings: [], + request: {}, + response: { + id: text, + modelId: "test-model", + timestamp: new Date(), + messages: createAssistantResponseMessages(text), + }, + steps: [], + }); + vi.mocked(ai.generateText) + .mockResolvedValueOnce(response("First response") as any) + .mockResolvedValueOnce(response("Second response") as any); + const options = { + memory: { + userId: "retry-user", + conversationId: "retry-conversation", + options: { messageMetadataPersistence: true }, + }, + }; + + const first = await agent.generateText("First turn", options); + expect(first.text).toBe("First response"); + + const second = await agent.generateText("Second turn", options); + expect(second.text).toBe("Second response"); + + const saved = await memory.getMessages("retry-user", "retry-conversation"); + expect(saved.filter((message) => message.role === "assistant")).toEqual([ + expect.objectContaining({ parts: [{ type: "text", text: "First response" }] }), + expect.objectContaining({ parts: [{ type: "text", text: "Second response" }] }), + ]); + }); + it("should persist usage and finish reason in assistant message metadata for streamText", async () => { const memory = new Memory({ storage: new InMemoryStorageAdapter(), diff --git a/packages/core/src/agent/agent.ts b/packages/core/src/agent/agent.ts index 9bc2ed8b2..4a334253e 100644 --- a/packages/core/src/agent/agent.ts +++ b/packages/core/src/agent/agent.ts @@ -1485,7 +1485,7 @@ export class Agent { void this.recordStepResults(result.steps, oc); if (!shouldDeferPersist && shouldPersistMemory) { - await persistQueue.flush(buffer, oc); + await this.flushConversationMessagesBestEffort(persistQueue, buffer, oc); } const finalText = await executeOutputGuardrails({ @@ -1585,7 +1585,7 @@ export class Agent { } if (shouldDeferPersist && shouldPersistMemory) { - await persistQueue.flush(buffer, oc); + await this.flushConversationMessagesBestEffort(persistQueue, buffer, oc); } const feedbackValue = (() => { @@ -2217,7 +2217,7 @@ export class Agent { }); if (!shouldDeferPersist && shouldPersistMemory) { - await persistQueue.flush(buffer, oc); + await this.flushConversationMessagesBestEffort(persistQueue, buffer, oc); } // History update removed - using OpenTelemetry only @@ -4427,6 +4427,22 @@ export class Agent { } } + private async flushConversationMessagesBestEffort( + queue: MemoryPersistQueue, + buffer: ConversationBuffer, + oc: OperationContext, + ): Promise { + try { + await queue.flush(buffer, oc); + } catch (error) { + oc.logger.debug("Failed to persist conversation messages", { + error, + conversationId: oc.conversationId, + userId: oc.userId, + }); + } + } + /** * Get contextual logger with parent tracking */ diff --git a/packages/core/src/agent/conversation-buffer.ts b/packages/core/src/agent/conversation-buffer.ts index 02e428c6f..fa214e1dc 100644 --- a/packages/core/src/agent/conversation-buffer.ts +++ b/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(); + private pendingMessageVersions = new Map(); + private nextPendingMessageVersion = 0; private toolPartIndex = new Map(); 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( + 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), ...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,12 +396,13 @@ 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") { @@ -383,6 +410,11 @@ export class ConversationBuffer { } } + 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++) { diff --git a/packages/core/src/agent/memory-persist-queue.spec.ts b/packages/core/src/agent/memory-persist-queue.spec.ts index 881ae9971..25ae99d1b 100644 --- a/packages/core/src/agent/memory-persist-queue.spec.ts +++ b/packages/core/src/agent/memory-persist-queue.spec.ts @@ -94,6 +94,374 @@ 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 messages from a failed operation when the next operation flushes", async () => { + const oldMessage = createMessage("previous turn"); + const newMessage = createMessage("next turn"); + const memoryManager = { + saveMessage: vi + .fn() + .mockRejectedValueOnce(new Error("storage unavailable")) + .mockResolvedValue(undefined), + } as any; + const oldBuffer = new ConversationBuffer(); + oldBuffer.ingestUIMessages([oldMessage], false); + const oldContext = createOperationContext(); + + await expect( + new MemoryPersistQueue(memoryManager).flush(oldBuffer, oldContext as any), + ).rejects.toThrow("storage unavailable"); + + const newBuffer = new ConversationBuffer(); + newBuffer.ingestUIMessages([newMessage], false); + const newContext = createOperationContext(); + await new MemoryPersistQueue(memoryManager).flush(newBuffer, newContext as any); + + expect(memoryManager.saveMessage.mock.calls.map((call: any[]) => call[1].id)).toEqual([ + oldMessage.id, + oldMessage.id, + newMessage.id, + ]); + expect(memoryManager.saveMessage.mock.calls[1][0]).toBe(oldContext); + expect(oldBuffer.getPendingMessages()).toHaveLength(0); + expect(newBuffer.getPendingMessages()).toHaveLength(0); + }); + + it("handles a failed debounced save and retries it on a later operation", async () => { + const oldMessage = createMessage("previous turn"); + const newMessage = createMessage("next turn"); + const memoryManager = { + saveMessage: vi + .fn() + .mockRejectedValueOnce(new Error("storage unavailable")) + .mockResolvedValue(undefined), + } as any; + const oldBuffer = new ConversationBuffer(); + oldBuffer.ingestUIMessages([oldMessage], false); + const oldContext = createOperationContext(); + new MemoryPersistQueue(memoryManager, { debounceMs: 100 }).scheduleSave( + oldBuffer, + oldContext as any, + ); + + await vi.advanceTimersByTimeAsync(100); + expect(oldBuffer.getPendingMessages()).toHaveLength(1); + + const newBuffer = new ConversationBuffer(); + newBuffer.ingestUIMessages([newMessage], false); + await new MemoryPersistQueue(memoryManager).flush(newBuffer, createOperationContext() as any); + + expect(memoryManager.saveMessage.mock.calls.map((call: any[]) => call[1].id)).toEqual([ + oldMessage.id, + oldMessage.id, + newMessage.id, + ]); + expect(oldBuffer.getPendingMessages()).toHaveLength(0); + }); + + it("keeps both operation buffers when a later schedule resets the debounce timer", async () => { + const messages = [createMessage("previous turn"), createMessage("next turn")]; + const memoryManager = { saveMessage: vi.fn().mockResolvedValue(undefined) } as any; + const buffers = messages.map((message) => { + const buffer = new ConversationBuffer(); + buffer.ingestUIMessages([message], false); + return buffer; + }); + + new MemoryPersistQueue(memoryManager, { debounceMs: 100 }).scheduleSave( + buffers[0], + createOperationContext() as any, + ); + await vi.advanceTimersByTimeAsync(50); + new MemoryPersistQueue(memoryManager, { debounceMs: 100 }).scheduleSave( + buffers[1], + createOperationContext() as any, + ); + await vi.advanceTimersByTimeAsync(100); + + expect(memoryManager.saveMessage.mock.calls.map((call: any[]) => call[1].id)).toEqual( + messages.map((message) => message.id), + ); + expect(buffers.every((buffer) => buffer.getPendingMessages().length === 0)).toBe(true); + }); + + it("uses the flushing operation's trace when it cancels another queue's debounce", async () => { + const oldMessage = createMessage("scheduled turn"); + const newMessage = createMessage("flushing turn"); + const memoryManager = { saveMessage: vi.fn().mockResolvedValue(undefined) } as any; + const oldBuffer = new ConversationBuffer(); + oldBuffer.ingestUIMessages([oldMessage], false); + const newBuffer = new ConversationBuffer(); + newBuffer.ingestUIMessages([newMessage], false); + const oldContext = { + ...createOperationContext(), + operationId: "scheduled-operation", + traceContext: { id: "old-trace" }, + }; + const newContext = { + ...createOperationContext(), + operationId: "flushing-operation", + traceContext: { id: "new-trace" }, + }; + + new MemoryPersistQueue(memoryManager, { debounceMs: 100 }).scheduleSave( + oldBuffer, + oldContext as any, + ); + await new MemoryPersistQueue(memoryManager).flush(newBuffer, newContext as any); + await vi.advanceTimersByTimeAsync(100); + + expect(memoryManager.saveMessage).toHaveBeenCalledTimes(2); + expect(memoryManager.saveMessage.mock.calls[0][0]).toMatchObject({ + operationId: oldContext.operationId, + traceContext: newContext.traceContext, + }); + expect(memoryManager.saveMessage.mock.calls[1][0]).toBe(newContext); + }); + + it("attempts every retained buffer before rethrowing the first failure", async () => { + const firstMessage = createMessage("previous turn"); + const secondMessage = createMessage("next turn"); + const memoryManager = { + saveMessage: vi + .fn() + .mockRejectedValueOnce(new Error("first buffer unavailable")) + .mockResolvedValue(undefined), + } as any; + const firstBuffer = new ConversationBuffer(); + firstBuffer.ingestUIMessages([firstMessage], false); + const secondBuffer = new ConversationBuffer(); + secondBuffer.ingestUIMessages([secondMessage], false); + const queue = new MemoryPersistQueue(memoryManager, { debounceMs: 100 }); + const firstContext = { ...createOperationContext(), isActive: false }; + const secondContext = { ...createOperationContext(), isActive: false }; + + queue.scheduleSave(firstBuffer, firstContext as any); + queue.scheduleSave(secondBuffer, secondContext as any); + + await expect(queue.flush(secondBuffer, secondContext as any)).rejects.toThrow( + "first buffer unavailable", + ); + expect(memoryManager.saveMessage.mock.calls.map((call: any[]) => call[1].id)).toEqual([ + firstMessage.id, + secondMessage.id, + ]); + expect(firstBuffer.getPendingMessages()).toHaveLength(1); + expect(secondBuffer.getPendingMessages()).toHaveLength(0); + }); + + it("bounds retained failed buffers so old retries do not grow without limit", async () => { + const oldMessage = createMessage("old turn"); + const newMessage = createMessage("new turn"); + const memoryManager = { + saveMessage: vi + .fn() + .mockRejectedValueOnce(new Error("old buffer unavailable")) + .mockRejectedValueOnce(new Error("old buffer unavailable")) + .mockRejectedValueOnce(new Error("new buffer unavailable")) + .mockResolvedValue(undefined), + } as any; + const oldBuffer = new ConversationBuffer(); + oldBuffer.ingestUIMessages([oldMessage], false); + const newBuffer = new ConversationBuffer(); + newBuffer.ingestUIMessages([newMessage], false); + const oldContext = { ...createOperationContext(), isActive: true }; + const newContext = { ...createOperationContext(), isActive: true }; + const queue = new MemoryPersistQueue(memoryManager, { + maxRetryBuffers: 1, + retryRetentionMs: 60_000, + }); + + await expect(queue.flush(oldBuffer, oldContext as any)).rejects.toThrow( + "old buffer unavailable", + ); + await expect(queue.flush(newBuffer, newContext as any)).rejects.toThrow( + "old buffer unavailable", + ); + await queue.flush(newBuffer, newContext as any); + + expect(memoryManager.saveMessage.mock.calls.map((call: any[]) => call[1].id)).toEqual([ + oldMessage.id, + oldMessage.id, + newMessage.id, + newMessage.id, + ]); + expect(oldBuffer.getPendingMessages()).toHaveLength(1); + expect(newBuffer.getPendingMessages()).toHaveLength(0); + }); + + it("drops a failed buffer after its retry retention window while its context stays active", async () => { + const message = createMessage("expired turn"); + const memoryManager = { + saveMessage: vi + .fn() + .mockRejectedValueOnce(new Error("storage unavailable")) + .mockResolvedValue(undefined), + } as any; + const buffer = new ConversationBuffer(); + buffer.ingestUIMessages([message], false); + const context = { ...createOperationContext(), isActive: true }; + const queue = new MemoryPersistQueue(memoryManager, { + retryRetentionMs: 1_000, + }); + + await expect(queue.flush(buffer, context as any)).rejects.toThrow("storage unavailable"); + await vi.advanceTimersByTimeAsync(1_001); + + expect(memoryManager.saveMessage).toHaveBeenCalledTimes(1); + expect(buffer.getPendingMessages()).toHaveLength(1); + const entriesByManager = (MemoryPersistQueue as any).entriesByManager as WeakMap< + object, + { entries: Map } + >; + expect(entriesByManager.get(memoryManager)?.entries.size ?? 0).toBe(0); + }); + + it("uses one cleanup timer for failed buffers shared by queue instances", async () => { + const memoryManager = { + saveMessage: vi.fn().mockRejectedValue(new Error("storage down")), + } as any; + const firstBuffer = new ConversationBuffer(); + firstBuffer.ingestUIMessages([createMessage("first")], false); + const secondBuffer = new ConversationBuffer(); + secondBuffer.ingestUIMessages([createMessage("second")], false); + const options = { retryRetentionMs: 1_000 }; + const firstContext = { ...createOperationContext(), conversationId: "first", isActive: true }; + const secondContext = { ...createOperationContext(), conversationId: "second", isActive: true }; + + await expect( + new MemoryPersistQueue(memoryManager, options).flush(firstBuffer, firstContext as any), + ).rejects.toThrow("storage down"); + await expect( + new MemoryPersistQueue(memoryManager, options).flush(secondBuffer, secondContext as any), + ).rejects.toThrow("storage down"); + + expect(vi.getTimerCount()).toBe(1); + await vi.advanceTimersByTimeAsync(1_000); + expect(vi.getTimerCount()).toBe(0); + }); + + it("keeps retries separate for user and conversation IDs containing colons", async () => { + const firstMessage = createMessage("first conversation"); + const secondMessage = createMessage("second conversation"); + const memoryManager = { + saveMessage: vi + .fn() + .mockRejectedValueOnce(new Error("storage unavailable")) + .mockResolvedValue(undefined), + } as any; + const firstBuffer = new ConversationBuffer(); + firstBuffer.ingestUIMessages([firstMessage], false); + const firstContext = { ...createOperationContext(), userId: "a:b", conversationId: "c" }; + const secondBuffer = new ConversationBuffer(); + secondBuffer.ingestUIMessages([secondMessage], false); + const secondContext = { ...createOperationContext(), userId: "a", conversationId: "b:c" }; + + await expect( + new MemoryPersistQueue(memoryManager).flush(firstBuffer, firstContext as any), + ).rejects.toThrow("storage unavailable"); + await new MemoryPersistQueue(memoryManager).flush(secondBuffer, secondContext as any); + + expect(memoryManager.saveMessage.mock.calls.map((call: any[]) => call[1].id)).toEqual([ + firstMessage.id, + secondMessage.id, + ]); + expect(firstBuffer.getPendingMessages()).toHaveLength(1); + }); + + 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((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), diff --git a/packages/core/src/agent/memory-persist-queue.ts b/packages/core/src/agent/memory-persist-queue.ts index 7ebf81206..c1d19a656 100644 --- a/packages/core/src/agent/memory-persist-queue.ts +++ b/packages/core/src/agent/memory-persist-queue.ts @@ -8,11 +8,28 @@ import type { OperationContext } from "./types"; interface QueueEntry { timer?: NodeJS.Timeout; pendingPromise: Promise; + pendingTasks: number; + buffers: Map; +} + +interface QueueState { + entries: Map; + retentionCleanupTimer?: NodeJS.Timeout; + retentionCleanupAt?: number; +} + +interface QueueBuffer { + context: OperationContext; + retainedAt?: number; } export interface MemoryPersistQueueOptions { debounceMs?: number; logger?: Logger; + /** Maximum number of failed buffers retained for retry across this memory manager. */ + maxRetryBuffers?: number; + /** How long a failed buffer remains eligible for retry. */ + retryRetentionMs?: number; } export type MemoryPersistQueueMemoryManager = Pick; @@ -29,9 +46,16 @@ export interface AgentMetadataContextValue { * Debounced persistence manager responsible for writing buffered messages to memory. */ export class MemoryPersistQueue { + private static readonly entriesByManager = new WeakMap< + MemoryPersistQueueMemoryManager, + QueueState + >(); private readonly debounceMs: number; private readonly logger?: Logger; - private readonly entries = new Map(); + private readonly maxRetryBuffers: number; + private readonly retryRetentionMs: number; + private readonly state: QueueState; + private readonly entries: Map; constructor( private readonly memoryManager: MemoryPersistQueueMemoryManager, @@ -39,6 +63,15 @@ export class MemoryPersistQueue { ) { this.debounceMs = options.debounceMs ?? 200; this.logger = options.logger; + this.maxRetryBuffers = Math.max(0, options.maxRetryBuffers ?? 32); + this.retryRetentionMs = Math.max(0, options.retryRetentionMs ?? 5 * 60_000); + let state = MemoryPersistQueue.entriesByManager.get(memoryManager); + if (!state) { + state = { entries: new Map() }; + MemoryPersistQueue.entriesByManager.set(memoryManager, state); + } + this.state = state; + this.entries = state.entries; } scheduleSave(buffer: ConversationBuffer, oc: OperationContext): void { @@ -48,6 +81,13 @@ export class MemoryPersistQueue { const key = this.getKey(oc); const entry = this.getOrCreateEntry(key); + const existing = entry.buffers.get(buffer); + entry.buffers.set(buffer, { + context: oc, + retainedAt: existing?.retainedAt, + }); + this.pruneRetainedBuffers(); + this.scheduleRetentionCleanup(); if (entry.timer) { clearTimeout(entry.timer); @@ -55,7 +95,7 @@ export class MemoryPersistQueue { entry.timer = setTimeout(() => { entry.timer = undefined; - this.enqueuePersist(key, () => this.persist(buffer, oc)); + void this.enqueuePersist(key, () => this.persistAll(key, oc)).catch(() => {}); }, this.debounceMs); const logPayload = { @@ -70,6 +110,13 @@ export class MemoryPersistQueue { const key = this.getKey(oc); const entry = this.getOrCreateEntry(key); + const existing = entry.buffers.get(buffer); + entry.buffers.set(buffer, { + context: oc, + retainedAt: existing?.retainedAt, + }); + this.pruneRetainedBuffers(); + this.scheduleRetentionCleanup(); if (entry.timer) { clearTimeout(entry.timer); @@ -82,15 +129,47 @@ export class MemoryPersistQueue { }; this.logger?.debug?.("Flushing conversation persistence queue", flushPayload); - await this.enqueuePersist(key, () => this.persist(buffer, oc)); + await this.enqueuePersist(key, () => this.persistAll(key, oc)); + } + + private async persistAll(key: string, triggerContext: OperationContext): Promise { + const entry = this.entries.get(key); + if (!entry) return; + + this.pruneRetainedBuffers(); + let firstError: unknown; + let hasError = false; + + for (const [buffer, queued] of [...entry.buffers]) { + try { + await this.persist(buffer, queued.context, triggerContext); + if (buffer.getPendingMessages().length === 0) { + entry.buffers.delete(buffer); + } + } catch (error) { + queued.retainedAt ??= Date.now(); + if (!hasError) { + firstError = error; + hasError = true; + } + } + } + + this.pruneRetainedBuffers(); + this.scheduleRetentionCleanup(); + if (hasError) throw firstError; } - private async persist(buffer: ConversationBuffer, oc: OperationContext): Promise { + private async persist( + buffer: ConversationBuffer, + oc: OperationContext, + triggerContext: OperationContext, + ): Promise { if (!oc.userId || !oc.conversationId) { return; } - const pending = buffer.drainPendingMessages(); + const pending = buffer.getPendingMessages(); if (pending.length === 0) { const payload = { conversationId: oc.conversationId, @@ -104,7 +183,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 +195,27 @@ export class MemoryPersistQueue { | Map | 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); + // Keep the message's original operation metadata while tracing the actual write attempt. + const writeContext = + oc === triggerContext || !triggerContext.traceContext + ? oc + : { ...oc, logger: triggerContext.logger, traceContext: triggerContext.traceContext }; + await this.memoryManager.saveMessage( + writeContext, + messageWithMetadata, + oc.userId, + oc.conversationId, + { + throwOnError: true, + }, + ); + buffer.markMessagePersisted(message.id, version); } catch (error) { this.logger?.error?.("Failed to save message", { conversationId: oc.conversationId, @@ -136,6 +229,7 @@ export class MemoryPersistQueue { private enqueuePersist(key: string, task: () => Promise): Promise { const entry = this.getOrCreateEntry(key); + entry.pendingTasks++; entry.pendingPromise = entry.pendingPromise .catch(() => {}) @@ -147,8 +241,14 @@ export class MemoryPersistQueue { throw error; }) .finally(() => { + entry.pendingTasks--; const current = this.entries.get(key); - if (current === entry && !current?.timer) { + if ( + current === entry && + !current.timer && + current.buffers.size === 0 && + current.pendingTasks === 0 + ) { this.entries.delete(key); } }); @@ -159,14 +259,108 @@ export class MemoryPersistQueue { private getOrCreateEntry(key: string): QueueEntry { let entry = this.entries.get(key); if (!entry) { - entry = { pendingPromise: Promise.resolve() }; + entry = { pendingPromise: Promise.resolve(), pendingTasks: 0, buffers: new Map() }; this.entries.set(key, entry); } return entry; } + private pruneRetainedBuffers(): void { + const now = Date.now(); + const retained: Array<{ + buffer: ConversationBuffer; + entry: QueueEntry; + retainedAt: number; + }> = []; + + for (const entry of this.entries.values()) { + for (const [buffer, queued] of entry.buffers) { + if (queued.retainedAt === undefined) continue; + + if (now - queued.retainedAt >= this.retryRetentionMs) { + entry.buffers.delete(buffer); + this.logger?.warn?.("Dropping expired conversation persistence retry", { + conversationId: queued.context.conversationId, + userId: queued.context.userId, + pendingMessageCount: buffer.getPendingMessages().length, + retainedForMs: now - queued.retainedAt, + }); + continue; + } + + retained.push({ buffer, entry, retainedAt: queued.retainedAt }); + } + } + + retained.sort((a, b) => a.retainedAt - b.retainedAt); + while (retained.length > this.maxRetryBuffers) { + const oldest = retained.shift(); + if (!oldest) break; + const queued = oldest.entry.buffers.get(oldest.buffer); + if (queued?.retainedAt !== oldest.retainedAt) continue; + oldest.entry.buffers.delete(oldest.buffer); + this.logger?.warn?.("Dropping conversation persistence retry limit exceeded", { + conversationId: queued.context.conversationId, + userId: queued.context.userId, + pendingMessageCount: oldest.buffer.getPendingMessages().length, + maxRetryBuffers: this.maxRetryBuffers, + }); + } + + for (const [key, entry] of this.entries) { + if (!entry.timer && entry.buffers.size === 0 && entry.pendingTasks === 0) { + this.entries.delete(key); + } + } + } + + private scheduleRetentionCleanup(): void { + const nextExpiry = this.getNextRetentionExpiry(); + if (nextExpiry === undefined) { + if (this.state.retentionCleanupTimer) { + clearTimeout(this.state.retentionCleanupTimer); + this.state.retentionCleanupTimer = undefined; + this.state.retentionCleanupAt = undefined; + } + return; + } + if ( + this.state.retentionCleanupTimer && + this.state.retentionCleanupAt !== undefined && + this.state.retentionCleanupAt <= nextExpiry + ) { + return; + } + if (this.state.retentionCleanupTimer) clearTimeout(this.state.retentionCleanupTimer); + + this.state.retentionCleanupAt = nextExpiry; + this.state.retentionCleanupTimer = setTimeout( + () => { + this.state.retentionCleanupTimer = undefined; + this.state.retentionCleanupAt = undefined; + this.pruneRetainedBuffers(); + this.scheduleRetentionCleanup(); + }, + Math.max(1, nextExpiry - Date.now()), + ); + this.state.retentionCleanupTimer.unref?.(); + } + + private getNextRetentionExpiry(): number | undefined { + let nextExpiry: number | undefined; + for (const entry of this.entries.values()) { + for (const queued of entry.buffers.values()) { + if (queued.retainedAt !== undefined) { + const expiry = queued.retainedAt + this.retryRetentionMs; + nextExpiry = nextExpiry === undefined ? expiry : Math.min(nextExpiry, expiry); + } + } + } + return nextExpiry; + } + private getKey(oc: OperationContext): string { - return `${oc.userId ?? "unknown"}:${oc.conversationId ?? "unknown"}`; + return `${oc.userId?.length ?? 0}:${oc.userId ?? ""}:${oc.conversationId ?? ""}`; } private applySubAgentMetadata( diff --git a/packages/core/src/memory/manager/memory-manager.spec.ts b/packages/core/src/memory/manager/memory-manager.spec.ts index 6bfea0b99..4da866568 100644 --- a/packages/core/src/memory/manager/memory-manager.spec.ts +++ b/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(); }); }); diff --git a/packages/core/src/memory/manager/memory-manager.ts b/packages/core/src/memory/manager/memory-manager.ts index 348ba477d..fb152d465 100644 --- a/packages/core/src/memory/manager/memory-manager.ts +++ b/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 { 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; + } } }